보상 편향 분포 정합을 통한 소단계 생성기 강화

초록

최근 소수 단계 확산 증류의 발전으로 효율적인 이미지 생성이 가능해졌지만, 이러한 모델을 인간의 선호도와 일치시키는 것은 여전히 어려운 과제로 남아 있습니다. 우리는 소수 단계 플로우 생성기를 위한 분포 매칭 증류와 보상 기반 강화 학습을 통합하는 두 단계 프레임워크인 RTDMD(Reward-Tilted Distribution Matching Distillation)를 제안합니다. 우리는 보상 편향 교사 분포에 대한 KL 발산 최소화가 자연스럽게 분포 매칭 항과 보상 최대화 항으로 분해됨을 보입니다. 첫 번째 단계에서는 주변 일관성 분포 매칭 증류(AC-DMD)를 도입합니다. 이 방법은 하위 구간별 분포 매칭을 수행하고, 제한된 업데이트 하에서 변화하는 생성기 분포를 추적하도록 가짜 점수 모델을 돕기 위해 가짜 점수 목적 함수에 일관성 정규화기를 추가합니다. 두 번째 단계에서는 두 항을 공동으로 최적화합니다. 보상 최대화 항을 위해, 우리는 확률적 중간 전이에 대한 GRPO 스타일 추정기와 결정론적 최종 단계를 통한 직접 보상 역전파를 결합한 하이브리드 정책 그래디언트를 유도하고, 분산을 줄이기 위해 단계 부분집합 GRPO(SubGRPO)를 추가로 도입합니다. SD3, SD3.5, FLUX.2에 대한 실험은 RTDMD가 단 4번의 추론 단계만으로 선호도, 미적, 구성 메트릭 전반에 걸쳐 새로운 최첨단 결과를 달성하여 이전의 소수 단계 텍스트-이미지 생성 방법들을 능가함을 보여줍니다. 코드와 모델은 https://github.com/Harahan/RTDMD 에서 확인할 수 있습니다.

English

Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.