Mid-U Guidance: Fast Classifier Guidance for Latent Diffusion Models

https://wandb.ai/johnowhitaker/midu-guidance/reports/Mid-U-Guidance-Fast-Classifier-Guidance-for-Latent-Diffusion-Models--VmlldzozMjg0NzA1

Mid-U Guidance: Fast Classifier Guidance for Latent Diffusion Models

 

SD2.1과 중간-사용자 가이드라인을 사용하여 미적 손실이 반영된 예시 이미지

 

클래스 분류 가이던스는 확산 모델 샘플링에 추가적인 제어를 할 수 있도록 해줍니다. 예로는 CLIP을 이용하여 텍스트나 이미지 프롬프트에 맞는 샘플을 생성하거나, 사전 학습된 이미지 분류 모델을 사용해 특정 클래스의 이미지를 생성하는 것을 들 수 있습니다. 하지만 기존 방법들은 주로 이미지 입력에 대해 동작하여 잠재 확산 모델(latent diffusion model)에 적용할 때 성능 비용이 크게 증가하는 문제가 있었습니다.

 

이 글에서는 모델 가이던스의 새로운, 빠른 접근 방식을 소개하고 그 효과를 미적 가이던스 예시를 통해 시연합니다. 이를 통해 사람의 미적 선호에 맞춰 모델을 훈련하고, 이를 기존 Stable Diffusion 모델에 추론 시 적용함으로써 최소한의 계산 비용으로 더욱 만족스러운 출력을 생성하는 방법을 보여줍니다.

 

클래스 분류 가이던스 소개

모델 기반 가이던스(또는 클래스 분류 가이던스)는 추가적인 모델을 사용하여 확산 모델의 생성 과정을 조정하는 기술입니다. 여기서 분류기는 생성된 이미지가 텍스트 프롬프트나 이미지 클래스와 얼마나 잘 일치하는지를 측정하는 손실 신호를 계산하고, 이 손실 신호의 미분 값을 사용하여 추론 단계 사이에 노이즈가 있는 입력 x를 업데이트합니다.

 

대개 사용되는 모델들이 픽셀 공간에서 작동하기 때문에, 이 작업은 연산 및 메모리 소모가 큽니다. 이는 추론의 각 단계에서 분류 모델, 잠재 확산 모델의 경우 VAE 디코더, 그리고 확산 UNet 자체를 통해 기울기를 추적해야 하기 때문입니다.

 

잠재 확산 모델의 기존 가이던스 접근 방식 (a)과 제안된 mid-u 기법 (b)을 비교해보면, mid-u 기법은 가이던스 손실의 기울기를 추적해야 하는 경로가 훨씬 짧기 때문에 계산과 메모리 요구량이 훨씬 적습니다.

 

mid-u 가이던스의 핵심 통찰은 확산 UNet 자체가 내부적으로 입력에 대한 풍부한 표현을 생성한다는 점입니다. 이를 분류기의 시작점으로 활용할 수 있습니다. UNet은 여러 단계에서 시간 스텝과 프롬프트를 추가적인 조건으로 받아들이기 때문에, 이러한 내부 표현은 현재 노이즈가 있는 입력 x에 대한 풍부한 표현뿐만 아니라 프롬프트와 노이즈 수준에 대한 정보도 인코딩합니다. 이 정보는 가이던스 모델이 작업하는 데 유용한 추가 정보가 될 수 있습니다.

 

구현

이 글에서는 Stable Diffusion의 UNet에서 중간 블록(mid-block)의 출력 활성화를 사용하지만, 이론적으로는 어떤 내부 특징도 좋은 후보가 될 수 있습니다. 아래 다이어그램은 UNet 아키텍처와 내부 특징을 저장하는 지점을 보여줍니다:

Stable Diffusion의 UNet 아키텍처

 

UNet의 정상적인 순방향 패스 동안 이러한 출력을 저장하기 위해, 우리는 관련 모듈에 대해 다음과 같이 순방향 훅(forward hook)을 등록할 수 있습니다:

pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base").to(device)
def hook_fn(module, input, output):
    module.output = output 
pipe.unet.mid_block.register_forward_hook(hook_fn)

 

512px 정사각형 이미지의 경우, 이 중간 블록 출력은 (1280, 8, 8)의 형태를 가집니다. 우리는 이러한 특징을 바탕으로 여러 방법으로 분류 모델을 구축할 수 있습니다. 아래는 이 글의 예제에서 사용된 아키텍처로, 여러 개의 합성곱과 풀링 레이어를 사용해 입력을 512개의 특징 집합으로 줄이고, 이 집합을 몇 개의 선형 레이어로 통과시켜 최종 출력(예제에서는 10개의 클래스)을 생성합니다.

 

model = nn.Sequential(
    nn.Conv2d(1280, 256, kernel_size=3, padding=1), nn.ReLU(),
    nn.MaxPool2d(2, 2),
    nn.Conv2d(256, 128, kernel_size=3, padding=1), nn.ReLU(),
    nn.AdaptiveAvgPool2d(output_size=(2, 2)), nn.Flatten(),
    nn.Linear(128*4, 64), nn.ReLU(), nn.Linear(64, 10)
)

 

이 모델을 미적 분류기로 훈련하기 위해, 다음과 같은 과정을 수행할 수 있습니다:

1. 프롬프트나 캡션과 미적 평가 점수가 있는 이미지 데이터셋을 로드합니다.
2. 이미지 배치를 인코딩하고 잠재 변수(latents)에 노이즈를 추가합니다.
3. 프롬프트와 함께 UNet을 통해 이를 전달합니다.
4. 위에서 보여준 훅을 사용하여 중간 블록 출력을 저장합니다.
5. 이 중간 블록 출력을 점수 모델의 입력으로 사용하고, 평가 점수를 예측 목표로 설정합니다.
6. 이 과정을 여러 번 반복합니다.

"해변 위의 말"을 가이던스 없이 생성한 이미지(왼쪽)와 추가적인 미적 가이던스를 적용하여 생성한 이미지(오른쪽), 동일한 초기 시드에서 시작하였습니다.

 

추론 시에는 일반적인 샘플링 루프를 적용하되, 중간 블록 출력을 점수 모델에 전달하고 미적 손실의 기울기를 계산하는 추가적인 코드를 추가합니다. 이 기울기들을(가이던스 강도를 제어하기 위한 스케일링과 함께) 사용해 노이즈가 있는 잠재 변수를 수정하여 최종 결과의 미적 매력을 높이는 방향으로 조정합니다. 최소한의 추론 코드 노트북에서 이 코드를 확인할 수 있습니다.

 

훈련

훈련은 FastAI와 함께 진행 중인 강의의 일부로 개발된 miniai 라이브러리를 사용하여 수행되었습니다. (중간 블록 출력, 평가) 형태의 훈련 예제는 미리 생성되었으며, 128개의 배치로 불러왔습니다.

 

모델은 Kalming 초기화로 초기화되었고, Adam 옵티마이저와 one_cycle 학습률 스케줄을 사용하여 한 번의 에포크 동안 훈련되었습니다. 두 가지 버전의 모델이 훈련되었는데, 하나는 다양한 양의 노이즈가 포함된 이미지의 중간 블록 특징으로 훈련되었고, 다른 하나는 노이즈 추가 단계를 생략한 것입니다.

 

노이즈가 포함된 버전으로부터 미적 평가를 예측하는 것은 추론 시 모델이 노이즈가 있는 입력에서 작동하기 때문에 더욱 유용할 가능성이 있습니다.

 

데이터는 Simulacra Aesthetic Captions 데이터셋에서 가져온 128,000개의 이미지를 기반으로 합니다. 평가 점수는 매우 노이즈가 심한데, 이는 많은 다른 기여자들로부터 크라우드소싱되었고, 많은 이미지가 단일 평가(편향될 가능성이 있는)만을 받았기 때문입니다.

 

우리는 이를 회귀 문제로 다룰 수도 있었지만, 대신 이를 분류 작업으로 훈련하기로 결정했습니다 (10개의 클래스 중 하나로 평가를 예측). 이후에 모델 출력을 단일 점수로 해석했습니다. 이 접근 방식은 클래스 기반의 가이던스를 위해 적응될 수 있으며, 목표와 손실 함수를 다른 목표(예: CLIP 텍스트 임베딩과의 정렬 등)로 변경하는 것도 가능합니다.

 

샘플 결과

아래는 W&B 테이블에 기록된 다양한 가이던스 스케일로 생성된 예시 이미지들입니다. 각 행마다 고정된 시드를 사용하여 비교를 더 용이하게 했습니다.

 

또한, 부정적인 가이던스 스케일을 사용하여 '미적이지 않은' 이미지를 생성할 수도 있습니다. 이는 예상보다 더 재미있는 결과를 만들어냅니다:

 

"제임스 본드"를 가이던스 없이 생성한 이미지(왼쪽)와 부정적 스케일을 사용해 미적이지 않게 만든 이미지(오른쪽)

이 모든 예시들은 Stable Diffusion 2.1 베이스 모델을 기반으로 하고 있습니다.

 

평가

최소한의 평가 인터페이스는 사용자가 선호하는 이미지를 선택하도록 되어 있으며, Colab 노트북에서 ipywidgets로 구현되었습니다.

 

같은 시드와 프롬프트로 생성된 두 이미지를 대상으로, 하나는 가이던스 기법을 사용한 것과 하나는 사용하지 않은 것을 인간 자원봉사자들에게 비교 선택하도록 요청해 선호도를 추정할 수 있습니다. 작은 규모의 테스트에서 가족들이 선호하는 이미지를 선택한 결과, 사람들은 약 70%의 확률로 미적 가이던스를 적용한 출력을 선호했습니다. 특히 가이던스 스케일이 높을수록 더 강한 선호도를 보였고, (스케일=10 이하에서는 효과가 거의 구분되지 않았으며, 선호도 점수는 약 50%로 떨어졌습니다). 이 결과를 너무 심각하게 받아들이기 전에, 더 나은 모델과 가이던스 스케일의 효과를 철저히 탐구한 대규모 연구가 필요합니다.

 

미적 가이던스를 넘어서

미적 가이던스는 이러한 유형의 가이던스가 가진 잠재력을 보여주는 훌륭한 방법이지만, 이게 유일한 사용 사례는 아닙니다. 이 접근 방식을 활용할 수 있는 추가적인 아이디어는 다음과 같습니다:

• 특정 클래스의 이미지를 생성하는 것 (예: '독일 셰퍼드')을 분류 모델을 이용해 수행
• CLIP 가이던스: CLIP 텍스트 또는 이미지 인코더와 일치하는 CLIP 임베딩을 생성하도록 모델을 훈련
• 디스크리미네이터 가이던스: 분류기를 훈련하여 중간 블록 특징을 기반으로 실제 이미지와 합성된 이미지를 구별하도록 하는 방법

특히 디스크리미네이터 가이던스는 현재 세대의 모델들이 겪고 있는 눈에 띄는 아티팩트를 줄일 수 있는 잠재력을 가지고 있어 추가 탐색에 있어 매우 유망한 방향으로 보입니다.

 

결론

잠재 확산 모델의 내부 표현에 대한 가이던스를 위해 모델을 훈련하는 것은 탐구할 가치가 있는 유망한 방법으로 보입니다. 이 글을 통해 커뮤니티가 이 개념을 더 실험해 보도록 장려하는 것이 제 희망입니다. 저 자신은 현재 이 프로젝트에 더 시간을 쓸 수 없게 되었지만요.