1. 연구 배경 및 문제 정의 (Problem)

확산 모델 (Diffusion Models, DMs) 은 고품질 이미지 생성 분야에서 지배적인 패러다임으로 자리 잡았으나, 두 가지 주요 한계가 존재합니다.

1. 높은 추론 비용: DM 은 샘플링 시 많은 시간 단계 (time-steps) 를 거치는 반복적인 탈노이즈 (denoising) 과정이 필요하여 생성 속도가 느립니다.
2. 소량 데이터 (Few-shot) 적응의 어려움: 새로운 도메인이나 주제에 모델을 적응 (Adaptation) 시킬 때, 기존에 학습된 지식과 새로운 데이터 사이의 균형을 맞추기 어렵습니다. 특히 소량의 데이터 (Few-shot) 로 적응할 경우 과적합 (Overfitting) 이나 다양성 저하가 발생합니다.

기존의 해결책은 2 단계 파이프라인 (Distill-then-Adapt 또는 Adapt-then-Distill) 을 사용했습니다.

• Distill-then-Adapt: 먼저 모델을 빠르게 만드는 과정 (Distillation) 을 거친 후 적응시킵니다. 하지만 학생 모델 (Student) 이 적응 용량에 한계를 느껴, 소량 데이터에서 과하게 평탄화된 (Over-smoothed) 결과가 나오거나, 원본 확산 손실 (Diffusion loss) 로 미세 조정 시 증류의 이점이 사라지는 문제가 있습니다.
• Adapt-then-Distill: 먼저 적응시킨 후 증류합니다. 이는 품질을 높일 수 있으나, 증류 과정이 소량 데이터에서 과적합되기 쉽고, 증류된 모델이 적응된 교사의 성능에 종속되어 다양성이 떨어질 수 있습니다.

이러한 2 단계 방식은 설계가 복잡하고, 소량 데이터 환경에서 품질과 다양성을 동시에 확보하기 어렵다는 문제가 있습니다.

2. 제안 방법: Uni-DAD (Methodology)

저자들은 Uni-DAD (Unified Distillation and Adaptation of Diffusion Models) 를 제안합니다. 이는 증류 (Distillation) 와 적응 (Adaptation) 을 단일 단계 (Single-stage) 파이프라인으로 통합하여, 소량 데이터 (Few-shot) 환경에서 빠르고 (Few-step), 고품질이며 다양한 이미지를 생성하는 방법입니다.

핵심 구성 요소

Uni-DAD 는 두 가지 상보적인 학습 신호를 결합합니다:

1. 이중 도메인 분포 매칭 증류 (Dual-domain Distribution-Matching Distillation, DMD):

   • 목적: 학생 모델이 소스 도메인 (Source, 예: FFHQ 얼굴) 의 다양성을 유지하면서 타겟 도메인 (Target, 예: 아기 사진) 의 구조를 학습하도록 유도합니다.
   • 구현: 고정된 소스 교사 (Source Teacher, $\epsilon_{src}$) 와 선택적 온라인 타겟 교사 (Target Teacher, $\epsilon_{trg}$) 의 점수 (Score) 를 활용합니다.
   • 수식: 학생 모델의 생성 분포가 소스 분포와 타겟 분포 모두에 근접하도록 KL 발산을 최소화하는 그래디언트를 계산합니다.
     • $L_{DMD}^{trg+src} = (1-a) L_{DMD}^{src} + a L_{DMD}^{trg}$
     • 여기서 $a$는 가중치로, 도메인 간 거리가 멀수록 타겟 교사의 비중을 높여 구조적 적응을 돕습니다.

2. 멀티헤드 생성적 적대 신경망 (Multi-head GAN) 손실:

   • 목적: 소량 데이터 환경에서의 과적합을 방지하고 타겟 도메인의 사실성 (Realism) 을 강화합니다.
   • 구현: 생성된 이미지와 실제 타겟 이미지를 구별하는 판별자 (Discriminator) 를 사용합니다.
   • 특징: 단일 헤드가 아닌 멀티헤드 구조를 사용하여, 디코더의 여러 블록 (Feature scales) 에서 로컬 및 글로벌 스케일의 특징을 동시에 판별합니다. 이는 소량 데이터에서 모드 붕괴 (Mode collapse) 를 방지하고 훈련을 안정화합니다.
   • 가짜 교사 (Fake Teacher): 학생 모델의 생성 분포를 추적하는 $\epsilon_{fk}$를 사용하여 GAN 프레임워크 내의 음수 (Negatives) 를 최신 상태로 유지합니다.

훈련 프로세스

훈련은 다음 세 가지 업데이트를 번갈아 수행합니다:

1. 학생 모델 (Student, $G$) 업데이트: DMD 손실과 GAN 생성자 손실을 최소화.
2. 가짜 교사 및 판별자 ($\epsilon_{fk} + D$) 업데이트: MSE 와 GAN 판별자 손실을 최소화.
3. 타겟 교사 (\epsilon_{trg) 업데이트 (선택적): 타겟 데이터로 미세 조정하여 적응 성능 향상.

3. 주요 기여 (Key Contributions)

1. 최초의 단일 단계 파이프라인: 확산 모델의 증류와 적응을 동시에 수행하여, 새로운 도메인에서 빠르고 고품질이며 다양한 생성을 가능하게 합니다.
2. 새로운 손실 함수 설계:
   • Dual-domain DMD: 소량의 타겟 데이터에서도 소스 도메인의 다양성 (포즈, 배경 등) 을 보존하면서 타겟 구조를 학습합니다.
   • Multi-head GAN: 소량 데이터 환경에서의 과적합을 줄이고 사실성을 높입니다.
   • Optional Target Teacher: 구조적으로 먼 도메인 적응을 돕습니다.
3. 성능 입증: FSIG (Few-shot Image Generation) 와 SDP (Subject-driven Personalization) 벤치마크에서 기존 최첨단 (SoTA) 방법들보다 적은 샘플링 단계 (≤4 단계) 로 더 나은 품질과 다양성을 달성했습니다. 특히 기존 2 단계 파이프라인보다 품질과 다양성 면에서 종종 우위를 보였습니다.

4. 실험 결과 (Results)

연구는 FSIG (FFHQ 기반, guided-DDPM) 와 SDP (Stable Diffusion v1.5 기반, DreamBooth) 벤치마크에서 평가되었습니다.

• 품질 (Quality):
  • FSIG: 10-shot 설정에서 3 단계 (NFE=3) 샘플링으로 기존 비증류 방법 (CRDI, DDPM-PA) 보다 낮은 FID (더 좋은 품질) 를 기록했습니다. 특히 구조적으로 먼 도메인 (MetFaces, Cats) 에서도 우수한 적응 능력을 보였습니다.
  • SDP: 1 단계 (NFE=1) 샘플링으로 DreamBooth (2x50 단계) 와 유사한 주체 일관성 (DINO, CLIP-I 점수) 을 유지하면서도 텍스트 프롬프트 정합성 (CLIP-T) 을 잘 따랐습니다.
• 다양성 (Diversity):
  • 기존 증류 방법들은 다양성이 급격히 떨어지는 경향이 있었으나, Uni-DAD 는 Intra-LPIPS 점수를 통해 기존 비증류 방법과 유사한 수준의 다양성을 유지하면서도 증류의 이점 (속도) 을 얻음을 보였습니다.
• 계산 비용 (Computational Cost):
  • 추론: 기존 방법 (25 단계) 대비 4.2 분 (5K 이미지 기준) 으로 추론 시간을 획기적으로 단축했습니다.
  • 훈련: 2 단계 파이프라인 (3 GPU·h) 보다 Uni-DAD 는 더 적은 훈련 비용 (2.2~2.8 GPU·h) 으로 더 나은 성능을 냈습니다.

5. 의의 및 결론 (Significance)

Uni-DAD 는 확산 모델의 증류와 적응을 분리된 단계가 아닌 통합된 과정으로 재정의했습니다. 이 방법은 다음과 같은 의의를 가집니다:

• 실시간 개인화 응용 가능성: 소량의 데이터로 새로운 주제나 스타일에 빠르게 적응하면서도, 몇 단계의 샘플링만으로 고품질 이미지를 생성할 수 있어 실시간 개인화 애플리케이션에 실용적입니다.
• 효율성과 품질의 균형: 기존 2 단계 방식의 복잡성과 성능 저하 문제를 해결하며, 소량 데이터 환경에서의 과적합과 다양성 손실 문제를 효과적으로 완화합니다.
• 유연성: 사전 증류된 모델이나 사전 적응된 모델을 초기값으로 사용할 수 있어 (Checkpoint-agnostic), 다양한 시나리오에 유연하게 적용 가능합니다.

결론적으로, Uni-DAD 는 제한된 데이터와 계산 자원 하에서도 고품질의 확산 모델 생성을 가능하게 하는 강력한 프레임워크를 제시합니다.