분포 매칭을 보상으로 재구성한 Diffusion Distillation

본 논문은 Diffusion 모델의 고품질 이미지 생성에 있어 가장 큰 병목인 반복적인 샘플링 과정을 단축하고, 동시에 학생 모델의 성능을 교사 모델에만 의존하지 않도록 하는 새로운 프레임워크를 제시한다. 기존의 Distribution Matching Distillation(DMD)은 교사와 학생 사이의 역 KL 다이버전스를 최소화하는 손실 \(L_{DMD}\) 을 사용해, 몇 단계만으로도 높은 충실도의 이미지를 생성하도록 학습한다. 그러나 DMD는 교사 모델에만 기반을 두기 때문에 학생 모델이 교사의 한계를 그대로 물려받는 문제가 있다. 최근 연구들은 RL과 결합해 보상을 추가함으로써 이 한계를 극복하려 했지만, 대부분은 단순히 두 손실을 가중합하는 방식으로, 최적화 충돌과 불안정성을 야기한다. 이에 저자들은 “분포 매칭을 보상으로 재구성”한다는 근본적인 전환을 시도한다. 구체적으로, DML의 그래디언트를 정책 그래디언트 형태로 변형해 보상 \(R_{dm}\) 을 정의한다. \(R_{dm}\)은 실제 스코어와 가짜 스코어의 차이 \(s_{real} - s_{fake}\)에 현재 타임스텝의 노이즈와 모델 출력 차이를 곱한 형태이며, 이는 RL에서 흔히 보는 보상 함수와 동일한 역할을 한다. 이렇게 하면 DMD와 RL을 별개의 목표가 아니라 하나의 보상 최적화 문제로 통합할 수 있다. 하지만 \(R_{dm}\)은 타임스텝이 커질수록 분산이 급격히 증가한다. 이는 고노이즈 단계에서 스코어 차이가 불안정해져, 학습이 발산하거나 수렴이 느려지는 원인이 된다. 이를 해결하기 위해 저자들은 RL에서 널리 쓰이는 Group Normalization(GN)을 차용한다. 그룹 단위로 보상의 평균과 표준편차를 이용해 정규화된 어드밴티지 \(A_{i,t}^{dm}\) 를 계산하고, 클리핑 및 중요도 가중치 \(r_i(\theta)\)와 결합해 안정적인 정책 업데이트를 수행한다. 이 방식을 Group Normalized Distribution Matching(GNDM)이라 명명한다. GNDM은 두 가지 주요 장점을 제공한다. 첫째, 보상의 스케일이 그룹 내에서 평균 0, 분산 1로 정규화되므로, 학습 초기에 발생하던 급격한 파라미터 변동을 크게 완화한다. 둘째, 그룹 내에서 동일한 타임스텝 \(t\)와 노이즈 타임스텝 \(t'\)를 공유함으로써 샘플 간 다양성을 유지하면서도 일관된 업데이트를 가능하게 한다. 또한, 논문은 Importance Sampling(IS)을 손실에 자연스럽게 통합한다. 기존 DDPO에서 제시된 IS 추정량을 차용해, 과거 파라미터 \(\theta_{old}\) 로 생성된 트래젝터리를 현재 파라미터 \(\theta\) 에 재가중치함으로써, 한 번의 샘플링으로 다중 스텝 업데이트를 수행한다. 이는 샘플링 비용을 크게 절감하면서도 성능 저하를 방지한다. 다중 보상 확장을 위해, 저자들은 기존 DMDR에서 제안된 외부 보상 \(R_o\) (예: CLIP Score, 인간 선호도 모델)을 GN 기반 어드밴티지 \(A_i^{o}\) 와 결합한다. 최종 어드밴티지는 \