報酬傾斜分布マッチングによる少数ステップ生成器の強化

要旨

近年の数ステップ拡散蒸留の進展により、効率的な画像生成が可能となったが、これらのモデルを人間の選好に合わせることは依然として課題である。本研究では、報酬傾斜分布マッチング蒸留（RTDMD）を提案する。これは2段階フレームワークであり、分布マッチング蒸留と報酬誘導型強化学習を数ステップフロー生成器に統合するものである。報酬傾斜教師分布に対するKL divergenceを最小化することが、分布マッチング項と報酬最大化項に自然に分解されることを示す。第1段階では、環境整合分布マッチング蒸留（AC-DMD）を導入する。これは部分区間ごとの分布マッチングを実行し、偽スコア目的関数に一貫性正則化項を追加することで、限られた更新回数下で偽スコアモデルが変化する生成器分布を追跡できるようにする。第2段階では、両方の項を共同最適化する。報酬最大化項については、確率的な中間遷移に対するGRPOスタイルの推定器と、決定論的な最終ステップを通じた直接的な報酬逆伝播を組み合わせたハイブリッド方策勾配を導出し、さらに分散低減のためのステップサブセットGRPO（SubGRPO）を導入する。SD3、SD3.5、FLUX.2を用いた実験により、RTDMDがわずか4回の推論ステップで、選好、美観、構成性の各指標において新たな最先端結果を達成し、従来の数ステップテキスト-to-画像生成手法を凌駕することを示す。コードとモデルは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.