MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance

기술 요약: 학습이 필요 없는 분포 적응을 위한 MMD 가이던스

문제 정의
사전 학습된 확산 모델(Diffusion Models)은 무조건부 및 조건부 샘플 생성을 위한 강력한 생성적 사전 지식(Generative Priors) 역할을 합니다. 그러나 이러한 모델의 출력물은 사용자가 지정한 특정 타겟 분포와 일치하지 않는 경우가 많으며, 이는 참조 예시(예: 50~100개)가 제한적인 도메인 적응(Domain Adaptation) 작업에서 중요한 문제입니다. 이러한 모델을 재학습하거나 미세 조정(Fine-tuning)하는 것은 메모리 제약과 타겟 데이터셋의 부족으로 인해 계산적으로 불가능하거나 비실용적인 경우가 많습니다. 분류기 가이던스(Classifier Guidance)나 분류기 프리 가이던스(Classifier-free Guidance)와 같은 기존의 추론 시점 가이던스 방법들은 대리 목적 함수(예: 클래스 확률 또는 조건 강도)를 최적화함으로써 샘플링 궤적을 조정할 뿐, 생성된 분포를 타겟 참조 분포와 직접적으로 정렬하지는 못합니다. 결과적으로, 이러한 방법들은 적은 수의 참조 샘플로 정의된 임의의 타겟 분포를 명시적으로 일치시킬 수 없습니다.

방법론
본 논문은 생성된 샘플과 참조 데이터셋 사이의 최대 평균 편차(Maximum Mean Discrepancy, MMD)의 그래디언트를 사용하여 역방향 확산 과정(Reverse Diffusion Process)을 보강하는 학습이 필요 없는 메커니즘인 MMD 가이던스를 제안합니다.

1. 핵심 메커니즘: 이 방법은 각 타임스텝 $t$에서 현재 샘플 상태에 대한 제곱 경험적 MMD의 그래디언트를 차감함으로써 표준 역방향 확산 샘플링 업데이트를 수정합니다. 업데이트 규칙은 다음과 같이 공식화됩니다:
   $$z_{t-1} = \text{sampler}(z_t, t, \epsilon_\theta) - \lambda_t \nabla_{z_t} \widehat{\text{MMD}}^2(\hat{P}_t, \hat{Q})$$
   여기서 $\hat{P}_t$는 타임스텝 $t$에서의 생성된 배치(Batch)의 경험적 분포이고, $\hat{Q}$는 참조 샘플의 경험적 분포이며, $\lambda_t$는 가이던스 강도를 조절합니다.

2. 다이버전스(Divergence) 선택: 저자들은 MMD가 차원의 저주를 겪지 않으면서 제한된 데이터로부터 편향되지 않고 분산이 낮은 추정치를 제공하기 때문에, $f$-다이버전스(KL 등)나 와서스타인 거리(Wasserstein distance) 대신 MMD를 선택했습니다. 또한, 커널 기반 정식화 덕분에 효율적이고 폐쇄형(Closed-form)인 그래디언트 계산이 가능합니다.

3. 잠재 공간 구현: 계산 효율성을 보장하기 위해, 가이던스는 잠재 확산 모델(Latent Diffusion Models, LDMs)의 잠재 공간에서 직접 적용됩니다. 참조 이미지는 VAE 인코더를 사용하여 한 번만 잠재 표현으로 인코딩됩니다. 이는 MMD 계산의 차원을 낮추고 픽셀 수준의 노이즈를 필터링하여 신호 대 잡음비(Signal-to-noise ratio)를 개선합니다.

4. 프롬프트 인식 적응: 조건부 생성을 위해, 본 프레임워크는 **곱 커널(Product Kernel)**을 사용하는 프롬프트 인식 적응으로 확장됩니다. 이 커널 $k_{\otimes}([p, z], [p', z']) = k_{\text{text}}(p, p') \cdot k_{\text{image}}(z, z')$은 프롬프트(의미적)와 시각적(잠재적) 공간 모두에서 유사성을 공동으로 측정합니다. 이를 통해 디노이저(Denoiser)를 재학습하지 않고도 생성된 샘플을 텍스트 프롬프트 및 참조 쌍의 시각적 스타일과 정렬할 수 있습니다.

5. 그래디언트 역학: MMD 제곱의 그래디언트는 두 가지 항으로 구성됩니다: 생성된 샘플 간의 반발력을 만들어내는 배치 내(Intra-batch) 항(다양성 증진)과 참조 샘플을 향한 인력을 만드는 교차(Cross-term) 항(분포 일치 유도)입니다. 이론적 분석은 교차 항 그래디언트에 대한 집중 보증(Concentration Guarantees)을 제공하여, 제한된 참조 샘플이 있더라도 모집단 그래디언트를 신뢰성 있게 추정할 수 있음을 보장합니다.

주요 기여

• MMD 가이던스: 제한된 참조 세트로 지정된 타겟 분포에 확산 모델을 적응시키기 위한 학습이 필요 없는 추론 시점 방법을 도입했습니다.
• 프롬프트 인식 확장: 조건부 생성에서 텍스트 프롬프트와 시각적 스타일을 공동으로 정렬하기 위해 곱 커널을 사용하는 실용적인 확장을 제시했습니다.
• 효율적인 구현: LDM의 잠재 공간에서 직접 작동하는 확장 가능한 구현을 통해, 픽셀 공간 가이던스에 비해 계산 오버헤드를 크게 줄였습니다.
• 경험적 검증: MMD 가이던스가 샘플 충실도(Fidelity)를 유지하면서 우수한 분포적 정렬을 달enc함을 보여주는 종합적인 실험을 수행했습니다.

실험 결과
저자들은 합성 가우시안 혼합 모델(GMM)과 실제 이미지 벤치마크(FFHQ, CelebA-HQ, 그리고 Stable Diffusion 출력의 스타일별 서브셋)에서 MMD 가이던스를 평가했습니다.

• 합성 데이터: 8-모달 및 25-모달 GMM에서 MMD 가이던스는 타겟 혼합 비율을 성공적으로 복구하고 불균형한 분포를 교정하였으며, Fréchet Distance(FD) 및 Kernel Distance(KD) 측면에서 베이스라인을 능가했습니다.
• 실제 이미지: FFHQ(예: 특정 안경 스타일을 가진 얼굴 생성) 및 CelebA-HQ 실험에서, MMD 가이던스는 가이드가 없는 샘플링, 분류기 가이던스, 심지어 미세 조정을 요구하는 사용자 훈련 베이스라인보다 높은 충실도와 분포 커버리지 점수를 달성했습니다.
• 조건부 생성: 텍스트 조건부 설정(Stable Diffusion XL 및 PixArt 사용)에서, 이 방법은 텍스트 프롬프트를 준수하면서 참조 세트와 일치하도록 생성된 이미지의 시각적 스타일(예: "애니메이션 스타일의 고양이" 또는 "Pixar 스타일의 자동차")을 성공적으로 적응시켰습니다.
• 효율성: 런타임 분석 결과, LDM의 잠재 공간에서 수행되는 MMD 가이던스의 추가적인 계산 비용은 SDXL과 같은 대규모 모델에서도 가이드가 없는 샘플링에 비해 무시할 수 있는 수준임을 나타냈습니다.
• 강건성: 이 방법은 다양한 참조 샘플 수(50~100개만으로도 효과적으로 수렴)와 서로 다른 도메인 간극(Domain gap)에 대해 강건함을 입증했습니다.

의의 및 주장
본 논문은 M-MD 가이던스가 계산 자원이 재학습에 제한적이거나 사용자 특정 데이터가 부족한 시나리오에서 생성 모델을 배포하는 데 실질적인 해결책을 제공한다고 주장합니다. 모델 파라미터를 수정하지 않고 추론 시점에 완전히 작동함으로써, 이 방법은 즉각적인 배포를 가능하게 하고 개인정보를 보호합니다(사용자 데이터가 학습에 사용되지 않으므로). 저자들은 이 접근 방식이 사전 학습된 파운데이션 모델의 능력과 특정 사용자 요구 사항 사이의 간극을 효과적으로 메우며, 생성된 샘플의 다양성과 품질을 유지하면서 분포적 정렬을 달성한다고 강조합니다. 이 연구는 MMD 가이던스를 기존의 미세 조정 및 프롬프트 엔지니어링 기술을 보완하는 도구로 위치시키며, 특히 개인정보 보호가 중요하거나 자원이 제한된 환경에서 매우 가치 있는 기술입니다.