Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

1. 문제 정의 (Problem)

• 배경: 확산 모델 (Diffusion Models, DMs) 은 고품질 이미지 생성에 탁월한 성능을 보이지만, 반복적인 디노이징 (denoising) 과정으로 인해 샘플링 비용이 매우 높습니다. 반면, 일관성 모델 (Consistency Models, CMs) 은 1~2 단계의 샘플링으로 고품질 이미지를 생성할 수 있어 속도가 빠릅니다.
• 제약 사항: DMs 에서 생성의 충실도 (fidelity) 와 다양성 (diversity) 을 조절하는 표준 기법인 클래스프리 가이드 (Classifier-Free Guidance, CFG) 는 CMs 에 적용하기 어렵습니다.
  • 기존 CMs 에 대한 가이드 방법은 별도의 DM Teacher 모델에서 지식을 증류 (Distillation) 받아야만 가능했습니다 (Consistency Distillation, CD 방식).
  • 이는 일관성 학습 (Consistency Training, CT) 으로만 학습된 모델 (Teacher DM 없이 데이터만으로 학습된 모델) 에는 적용할 수 없음을 의미합니다.
  • 즉, 기존에는 CMs 에 가이드를 적용하려면 반드시 Teacher DM 이 필요했고, 이를 통해 학습된 모델만 사용 가능했습니다.
• 핵심 질문: Teacher DM 없이, 이미 학습된 CM(특히 CT 모델) 에 사후적 (Post-Hoc) 으로 가이드 기능을 추가하여 CFG 와 유사한 효과를 얻을 수 있는가?

2. 제안 방법: JFDL (Joint Flow Distribution Learning)

저자들은 Teacher DM 없이 사전 학습된 CM 을 기반으로 가이드를 가능하게 하는 경량 정렬 방법인 JFDL을 제안합니다.

2.1 핵심 아이디어

• ODE 솔버로서의 CM: 사전 학습된 CM 을 확산 경로의 ODE 솔버로 간주합니다.
• 가상 무조건부 분포 (Pseudo-Unconditional Distribution): CM 이 무조건부 (unconditional, $\emptyset$) 경로를 학습하지 않았더라도, 조건부 (conditional) 경로와 무조건부 경로의 혼합 흐름 (Hybrid Flow) 을 통해 무조건부 분포를 근사할 수 있음을 이론적으로 증명합니다.
• 가이드 방향의 보간:
  1. 입력 노이즈 $x_t$에서 조건부 해 ($x_0^c$) 와 무조건부 해 ($x_0^\emptyset$) 를 추정합니다.
  2. 기존 CFG 와 유사하게 두 방향을 가중치 $\omega$로 보간하여 새로운 디노이징 방향을 생성합니다:
     $$u_\omega = \omega \cdot u_{cls} + (1-\omega) \cdot u_\emptyset$$
  3. 이 보간된 방향을 사용하여 일관성 손실 (Consistency Loss) 을 계산하고 모델을 미세 조정 (Fine-tuning) 합니다.

2.2 이론적 기반 및 검증

• 유사성 (Flow-based Models) 관점: JFDL 은 흐름 기반 모델 (Flow-based Models, FMs) 의 관점에서 분석됩니다.
• 가상 노이즈 (Pseudo-Noise) 의 정규성: JFDL 과정에서 생성되는 "혼합 신호 (mixed-signals)"가 실제로 가우시안 분포를 따르는지 검증했습니다.
  • 이론적 증명: 시간 $t$가 0 에 가까울 때, 혼합 신호는 데이터의 국소적 기울기 (gradient) 와 관련되지만, $t$가 커지거나 $\sigma_{max}$가 충분히 크면 가우시안 노이즈 항이 지배적이 되어 전체 분포가 가우시안에 수렴함을 증명했습니다 (Proposition 1, 2).
  • 실험적 검증: Shapiro-Wilk, Anderson-Darling, Kolmogorov-Smirnov 등 정규성 검정을 통해 생성된 가상 노이즈가 거의 모든 시간 단계에서 가우시안 분포임을 확인했습니다.

2.3 두 가지 변형

1. Naive JFDL: 사전 학습된 CM 이 무조건부 ($\emptyset$) 클래스 ODE 솔버로 학습된 경우를 가정합니다.
2. Random JFDL: 대부분의 CM 이 무조건부 ODE 를 학습하지 않았다는 점을 고려하여, 무조건부 앵커 ($\emptyset$) 대신 데이터 분포에서 무작위로 선택된 다른 클래스를 사용하여 무조건부 경로를 근사합니다. 이는 Teacher DM 없이도 CT 모델에 가이드를 적용할 수 있게 해줍니다.

3. 주요 기여 (Key Contributions)

1. DM-Free Post-Hoc Guidance: DM Teacher 없이 사전 학습된 CM 에만 가이드 기능을 추가할 수 있는 첫 번째 방법론 (JFDL) 을 제안했습니다.
2. 이론적 통찰: 흐름 기반 모델 (Flow-based Models) 의 관점에서 JFDL 의 유효성을 설명하고, 생성된 가상 노이즈가 가우시안 분포를 따른다는 이론적 증명과 실험적 검증을 제공했습니다.
3. CT 모델의 가이드 가능화: 무조건부 샘플링을 명시적으로 설계하지 않은 Consistency Training (CT) 모델에도 JFDL 을 적용하여 효과적인 가이드 튜닝이 가능함을 보였습니다.
4. 성능 향상: CIFAR-10 및 ImageNet 64x64 데이터셋에서 기존 무가이드 (Unguided) CM 대비 FID(Fréchet Inception Distance) 를 크게 개선했습니다.

4. 실험 결과 (Results)

• 데이터셋: CIFAR-10, ImageNet 64x64.
• 메트릭: FID (1-step 및 2-step 샘플링 기준).
• 주요 결과:
  • CIFAR-10: 1-step 샘플링에서 Naive JFDL 과 Random JFDL 모두 Baseline (ECT) 대비 FID 를 개선했습니다 (예: 3.40 $\to$ 3.29/3.24).
  • ImageNet 64x64: 1-step 및 2-step 모두에서 FID 가 유의미하게 감소했습니다 (예: 1-step 5.84 $\to$ 4.38/4.68).
  • 가이드 스케일 ($\omega$) 효과: CFG 와 유사하게, 작은 $\omega$ (1~2) 에서 FID 가 개선되다가 너무 큰 $\omega$에서는 다양성 손실로 인해 성능이 저하되는 경향을 보였습니다.
  • 효율성: JFDL 미세 조정은 ECT 전체 학습 데이터의 약 7.5% 만 사용하고, GPU 시간은 약 30% 수준으로 매우 경량화되었습니다.
• 시각적 결과: 가이드 스케일이 증가함에 따라 이미지의 대비 (contrast) 가 높아지고 세부 사항이 선명해지는 CFG 고유의 특성이 재현되었습니다.

5. 의의 및 결론 (Significance)

• 기술적 격차 해소: 현재 CMs 의 가이드 방법은 Teacher DM 에 의존하는 CD 방식에 국한되어 있었으나, JFDL 은 이를 극복하고 CT 모델에도 직접 가이드를 적용할 수 있는 길을 열었습니다.
• 실용성 증대: 별도의 Teacher 모델 학습 없이 기존에 학습된 CM 을 빠르게 가이드 가능하게 만들어, 생성 모델의 배포 및 활용 비용을 크게 낮춥니다.
• 이론적 확장: 확산 모델, 흐름 기반 모델, 일관성 모델 간의 관계를 통합적으로 분석하여, ODE 솔버 관점에서의 가이드 메커니즘에 대한 새로운 통찰을 제공했습니다.

요약하자면, 이 논문은 Teacher DM 없이도 사전 학습된 Consistency Model 에 CFG 와 유사한 가이드 기능을 부여하는 JFDL을 제안함으로써, 고속 생성 모델의 제어 가능성과 품질을 동시에 향상시키는 획기적인 방법을 제시했습니다.