보상 기반 모멘트 매칭 증류(Rewarded Moment Matching Distillation)를 이용한 확산 모델 미세 조정

증류 (Distillation) 및 강화학습 (Reinforcement Learning, RL) 미세 조정 (fine-tuning)은 확산 (diffusion) 사후 학습의 주요 핵심 축입니다. 전통적으로 이들은 개별적으로 연구되어 왔으나, 이러한 단계들 사이의 상호작용, 특히 미세 조정이 증류된 모델의 생성 품질에 어떤 영향을 미치는지에 대해서는 여전히 이해가 부족한 상태입니다. 본 논문에서는 확산 모델을 동시에 증류하면서 보상 함수 (reward function)를 최대화하는 새로운 프레임워크인 보상 기반 모멘트 매칭 증류 (Rewarded Moment Matching Distillation, RMMD)를 소개합니다. RMMD는 샘플링 루프를 온폴리시 (on-policy) 학습에 맞게 조정하고, 증류 손실 (distillation loss)을 적분 KL 정규화 (integral KL regularization)의 대리 지표로 재용도화함으로써, 고급 증류 기술(예: 8단계 모멘트 매칭)의 고충실도 "자연스러움 (naturalness)" 특성을 보존합니다. ImageNet에서 FID-Reward 파레토 프런트 (Pareto fronts)를 평가함으로써, RMMD가 단일 단계 베이스라인 (DI++) 및 다단계 경쟁 모델 (DRaFT, HyperNoise)에 비해 우수한 트레이드오프 (trade-offs)를 달성함을 입증합니다. 마지막으로, 우리는 RMMD를 최첨단 기상 예측 모델인 GenCast에 적용하여, 연속 순위 확률 점수 (Continuous Ranked Probability Score, CRPS) 지표를 최적화하는 동시에 모델을 증류합니다. 그 결과, 증류된 모델은 7.5배의 속도 향상을 달성하는 동시에 타겟 기상 변수의 93%에서 교사 모델 (teacher model)보다 뛰어난 성능을 보였으며, 더 나은 보정 (calibrated) 성능을 나타냈습니다. 이는 RMMD가 복잡하고 고차원적인 과학적 영역으로 확장 가능하다는 것을 증명합니다.