논문 요약: Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-Based Models (F2D2)

이 논문은 Flow-based 생성 모델 (Flow Matching, Continuous Normalizing Flows 등) 에서 발생하는 계산적 병목 현상을 해결하기 위해 제안된 F2D2 (Fast Flow Joint Distillation) 프레임워크를 소개합니다. 기존 모델들은 고품질 샘플 생성은 가능하지만, 로그 가능도 (Log-likelihood) 평가에 수백수천 번의 신경망 함수 평가 (NFEs) 가 필요하다는 치명적인 단점이 있었습니다. F2D2 는 샘플링과 가능도 평가를 동시에 10 배 미만 (보통 28 단계) 의 NFE 로 정확하게 수행할 수 있게 합니다.

1. 문제 정의 (Problem)

• 로그 가능도 평가의 중요성: 생성 모델의 비교, 미세 조정 (Fine-tuning, 예: PPO, DPO), 하위 작업 적용 등에 로그 가능도 평가는 필수적입니다.
• 현재의 한계: Diffusion 모델이나 Flow Matching 모델은 현재 가장 우수한 생성 성능을 보이지만, 정확한 로그 가능도를 계산하려면 확률 흐름 ODE(Ordinary Differential Equation) 를 따라 발산 (Divergence) 항을 적분해야 합니다.
• 계산 비용: 이 적분 과정은 일반적으로 100~1000 번 이상의 NFE 를 요구하여, 실시간 적용이나 효율적인 미세 조정을 불가능하게 만듭니다.
• 기존 해결책의 부족: 최근 'Distillation' 기법들은 샘플링 단계를 1~10 단계로 줄이는 데 성공했으나, 이는 가능도 계산 능력을 포기하거나 여전히 전체 궤적에 대한 적분이 필요하게 만들어 빠른 가능도 평가는 여전히 해결되지 않은 문제였습니다.

2. 방법론 (Methodology)

저자들은 연속 정규화 흐름 (CNF) 에서 샘플링 궤적과 로그 밀도 진화가 공유된 속도장 (Velocity Field) 에서 파생된다는 통찰을 바탕으로 F2D2를 제안했습니다.

핵심 아이디어: 결합된 ODE 시스템의 동시 증류

기존 Flow Matching 은 다음 ODE 를 풉니다:
$$\frac{d}{dt} x_t = v_\theta(x_t, t)$$
로그 가능도 $z_t = \log p_t(x_t)$의 진화는 다음과 같은 결합된 ODE 시스템으로 표현됩니다:
$$\frac{d}{dt} \begin{bmatrix} x_t \\ z_t \end{bmatrix} = \begin{bmatrix} v_\theta(x_t, t) \\ -\text{div}(v_\theta(x_t, t)) \end{bmatrix}$$
여기서 $x_t$는 샘플링 경로, $z_t$는 로그 가능도이며, 두 식 모두 동일한 속도장 $v_\theta$와 그 발산 (Divergence) 항에 의존합니다.

F2D2 프레임워크

1. 결합 흐름 맵 (Joint Flow Map) 학습:
   • 단일 모델이 샘플링 경로 ($x_t \to x_s$) 와 누적 발산 ($z_t \to z_s$) 을 동시에 예측하도록 학습합니다.
   • 파라미터 공유를 통해 백본 네트워크 (Backbone) 는 공유하고, 발산 예측을 위한 별도의 헤드 (Head) 만 추가합니다.
2. 동시 증류 (Joint Distillation):
   • 샘플링 손실: 기존 Few-step 방법 (Shortcut, MeanFlow 등) 의 자기 일관성 (Self-consistency) 손실을 활용하여 빠른 샘플링을 학습합니다.
   • 가능도 손실: 발산 항에 대한 증류 손실을 추가하여, 적은 단계에서도 정확한 로그 밀도 적분을 학습합니다.
   • 학습 목표: $L_{F2D2} = L_{\text{sample}} + L_{\text{likelihood}}$를 최소화하여 두 과업을 동시에 최적화합니다.
3. 구현 변형:
   • Shortcut-F2D2: 반군 (Semigroup) 속성을 기반으로 한 자기 일관성 손실을 확장.
   • MeanFlow-F2D2: 오일러 방정식 (Eulerian equation) 기반의 MeanFlow 식을 발산 항에 적용.
   • LSD-F2D2: 라그랑주 (Lagrangian) 자기 증류 방식에 발산 예측 추가.

3. 주요 기여 (Key Contributions)

1. 최초의 Few-step 가능도 평가: Diffusion/CNF 기반 모델에서 정확한 로그 가능도 평가를 소수의 NFE(1~10 단계) 로 가능하게 한 최초의 프레임워크입니다.
2. 모듈형 설계: 기존 Few-step 샘플링 모델 (Shortcut, MeanFlow 등) 에 쉽게 적용 가능하며, 발산 예측 헤드 하나만 추가하면 됩니다.
3. 새로운 응용: 최대 가능도 자기 유도 (Maximum Likelihood Self-Guidance):
   • 학습된 모델이 자체적으로 예측한 로그 가능도를 '리워드'로 사용하여, 샘플링 시작 전 초기 잡음 ($x_0$) 을 최적화하는 기법을 제안했습니다.
   • 이는 외부 보상 모델 없이도 샘플 품질을 극대화할 수 있게 합니다.

4. 실험 결과 (Results)

CIFAR-10, ImageNet 64x64, CelebA-64 데이터셋에서 실험을 수행했습니다.

• 정확한 가능도 평가:
  • 기존 Few-step 모델들은 NFE 가 줄어들면 로그 가능도 (NLL/BPD) 가 무의미하게 변하거나 발산했으나, F2D2 는 2~8 단계에서도 Teacher 모델 (1024 단계) 과 유사한 보정된 (Calibrated) 가능도를 제공합니다.
  • 예시: CIFAR-10 에서 2 단계 MeanFlow-F2D2 는 Teacher 모델의 BPD(3.12) 에 근접한 1.63 의 NLL 을 기록했습니다.
• 샘플 품질 유지:
  • 가능도 평가를 추가하더라도 샘플 품질 (FID) 이 크게 저하되지 않았습니다. 오히려 MeanFlow-F2D2 는 기존 MeanFlow 보다 더 낮은 FID 를 기록하기도 했습니다.
• 자기 유도 (Self-Guidance) 의 효과:
  • F2D2 를 활용한 자기 유도 기법을 적용한 2 단계 MeanFlow는 동일한 크기의 1024 단계 Flow Matching 모델보다 CIFAR-10 에서 더 낮은 FID를 달성했습니다. 이는 빠른 가능도 평가가 생성 품질을 높이는 강력한 신호로 작용함을 입증했습니다.
• 2D 체커보드 실험:
  • 분석적으로 정확한 가능도를 알 수 있는 2D 합성 데이터에서, F2D2 는 1 단계 NFE 만으로도 정확한 밀도 분포를 복원하는 반면, 기존 방법은 완전히 실패함을 보였습니다.

5. 의의 및 결론 (Significance)

• 장기적인 병목 현상 해결: Flow-based 모델의 가장 큰 약점이었던 "빠른 샘플링 vs 정확한 가능도 평가" 간의 트레이드오프를 해결했습니다.
• 새로운 알고리즘적 가능성: 빠른 가능도 평가는 강화 학습 (RL), 선호도 최적화 (DPO/GRPO), 모델 비교 등 다양한 후처리 (Post-training) 기법과 하위 작업에 Flow 기반 모델을 적용할 수 있는 길을 열었습니다.
• 효율성: 추가적인 계산 비용 없이 (단순히 발산 예측 헤드 추가 및 미세 조정) 기존 모델의 성능을 극대화할 수 있어, 생성 모델의 실용성과 확장성을 높이는 중요한 기여를 했습니다.

요약하자면, F2D2는 생성 모델이 "빠르게" 그리고 "정확하게" 확률을 계산할 수 있게 함으로써, 생성 AI 의 다음 단계인 효율적인 미세 조정과 고급 응용 분야로의 확장을 가능하게 하는 핵심 기술입니다.