Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning

1. 문제 정의 (Problem)

확산 모델 (Diffusion Models) 은 이미지 생성의 품질과 안정성 측면에서 획기적인 발전을 이루었으나, 반복적인 샘플링 과정으로 인해 추론 속도가 느리고 계산 비용이 높다는 단점이 있습니다. 이를 해결하기 위해 **분포 매칭 증류 (Distribution Matching Distillation, DMD)**와 같은 증류 기법이 등장하여 적은 단계 (Few-step) 로 고품질 이미지를 생성할 수 있게 했습니다.

그러나 기존 DMD 기반 방법론들은 다음과 같은 한계를 가집니다:

• 품질 저하: 증류 과정에서 제이슨-샨논 (Jensen-Shannon) 또는 KL 발산 (KL divergence) 을 최소화하는 데만 집중하여 인간의 미적 선호도나 실제 데이터의 뉘앙스를 반영하지 못합니다.
• 강화학습 (RL) 통합의 비효율성: 최근 DMD 에 RL 을 결합하려는 시도 (예: DMDR) 가 있었으나, 생성된 원시 샘플 (Raw Sample, $x_0$) 을 기반으로 보상 (Reward) 을 평가하는 방식은 초기 학습 단계에서 노이즈가 심하고 불안정합니다.
• 최적화 충돌: RL 의 보상 최적화 방향과 DMD 의 증류 방향이 서로 상충되어 (Gradient Conflict), 학습이 불안정해지거나 '보상 해킹 (Reward Hacking)'과 같은 현상이 발생하여 최종 품질이 제한됩니다.

2. 제안 방법론: GDMD (Methodology)

저자들은 원시 픽셀이 아닌 '증류 그라디언트 (Distillation Gradients)'를 RL 의 주요 보상 신호로 활용하는 새로운 패러다임인 GDMD를 제안합니다.

핵심 아이디어

기존의 샘플 기반 점수 매기는 방식 대신, DMD 목적함수에서 유도된 그라디언트 자체를 RL 의 평가 대상으로 삼습니다. 이를 통해 증류 과정의 업데이트 방향을 RL 이 직접적으로 가이드하도록 설계했습니다.

주요 구성 요소

1. 증류 인식 그라디언트 수집 (Distillation-aware Gradient Collection, DaGC):
   • 기존 Flow-GRPO 등에서 사용하는 SDE(Stochastic Differential Equation) 샘플링 대신, DMD 고유의 무작위성을 활용합니다.
   • 다양한 노이즈 강도 ($t$) 와 가짜 점수 추정기 ($\mu_{fake}$) 의 업데이트 버전을 활용하여 그룹화된 그라디언트를 수집합니다.
2. 암시적 그라디언트 점수 매기기 (Implicit Gradient Scoring, IGS):
   • 기존 보상 모델 (CLIP, HPS 등) 은 이미지와 텍스트를 입력받으므로, 그라디언트 정보를 직접 평가할 수 없습니다.
   • 저자들은 DMD 의 손실 함수를 변형하여, 그라디언트 정보가 포함된 **가상 타겟 텐서 ($x_{tar}$)**를 생성합니다.
   • 이 $x_{tar}$를 VAE 디코더를 통해 이미지로 변환한 후 기존 보상 모델에 입력하여 그라디언트의 품질을 간접적으로 점수화합니다.
3. 부정적 인식 선호도 최적화 (Negative-aware Preference Optimization, NaPO):
   • DiffusionNFT 에서 영감을 받아, 양의 정책 (Implicit positive policy) 과 음의 정책 (Implicit negative policy) 을 정의합니다.
   • 수집된 그라디언트 기반 보상을 활용하여 학생 모델의 속도장 (Velocity field) 을 업데이트하며, 증류 방향과 RL 방향의 정렬을 유지합니다.

학습 프로세스

• Joint Training: DMD 손실 ($\mathcal{L}_{DMD}$) 과 RL 손실 ($\mathcal{L}_{RL}$) 을 가중치 합으로 결합하여 동시 학습합니다.
• Cold-start 불필요: 그라디언트 기반 평가는 초기 학습 단계의 불안정한 샘플 생성에 덜 민감하므로, 별도의 Cold-start 단계를 거치지 않고도 안정적인 학습이 가능합니다.

3. 주요 기여 (Key Contributions)

1. 그라디언트 기반 점수 매기기 전략의 선구적 도입: DMD-RL 결합 증류에서 샘플 기반이 아닌 그라디언트 기반 평가 전략을 처음 제안하고, 그 이점을 분석했습니다.
2. 전용 파이프라인 설계: DaGC, IGS, NaPO 등 그라디언트 수준의 가이드를 효과적으로 활용하기 위한 종합적인 아키텍처를 설계했습니다.
3. SOTA 성능 달성: 기존 증류 방법론 및 DMDR 보다 우수한 품질을 달성하여 Few-step Diffusion 모델의 새로운 State-of-the-Art(SOTA) 를 확립했습니다.

4. 실험 결과 (Results)

GDMD 는 SDXL-Base 와 SD3-Medium 모델을 기반으로 4 단계 (4-step) 생성 모델로 증류되었으며, 다음과 같은 결과를 보였습니다.

• 정량적 평가 (Quantitative):
  • GenEval: 텍스트 - 이미지 정합성 (Text-Image Alignment) 평가에서 기존 DMD, DMDR 및 Teacher 모델을 모두 상회했습니다.
  • 이미지 품질: CLIP Score, HP(Human Preference) Score, Aesthetic Score, Pick Score, ImageReward 등 모든 지표에서 기존 최첨단 방법론 (DMD2, Flash Diffusion 등) 보다 높은 점수를 기록했습니다. 특히 훈련에 사용되지 않은 'Unseen Rewards'에서도 뛰어난 일반화 성능을 보였습니다.
• 정성적 평가 (Qualitative):
  • 텍스트 프롬프트에 대한 정합성이 뛰어나며, 손가락 왜곡, 과포화 (Oversaturation), 아티팩트 등의 문제가 현저히 감소했습니다.
  • Teacher 모델 (100 단계) 보다도 더 높은 품질을 보여주는 경우가 많았습니다.
• Human Preference:
  • 인간 평가에서 Teacher 모델, DMD, DMDR, DMD+NFT 등 모든 경쟁 모델 대비 **55%~64%**의 승률을 기록하며 압도적인 선호도를 받았습니다.

5. 의의 및 결론 (Significance)

이 논문은 증류 (Distillation) 와 강화학습 (RL) 의 시너지를 재정의했다는 점에서 중요한 의미를 가집니다.

• 학습 안정성: 초기 학습 단계의 노이즈가 많은 샘플 평가를 피하고, 그라디언트 정보를 직접 활용함으로써 최적화 불안정성을 해결했습니다.
• 효율성: 별도의 SDE 샘플링 과정이나 Cold-start 단계를 제거하여 학습 파이프라인을 간소화하고 효율성을 높였습니다.
• 품질의 한계 돌파: 증류된 모델이 Teacher 모델의 품질 한계를 넘어서는 것을 증명하며, 적은 단계로 고품질 이미지를 생성하는 실용적인 솔루션을 제시했습니다.

결론적으로 GDMD 는 Few-step Diffusion 모델의 품질과 속도를 동시에 개선하는 강력한 프레임워크로, 향후 이미지 생성 분야의 표준적인 접근법으로 자리 잡을 가능성이 높습니다.