Flash-DMD：効率的な蒸留と連携強化学習による高忠実度数ステップ画像生成に向けて

要旨

拡散モデルは生成モデルの主要なクラスとして台頭しているが、その反復的なサンプリング過程は計算コストが高いままである。タイムステップ蒸留は生成を加速する有望な技術であるが、大規模な訓練を必要とし、画像品質の劣化を招くことが多い。さらに、美的魅力やユーザー嗜好といった特定の目的に対して、強化学習を用いてこれらの蒸留モデルをファインチューニングすることは、極めて不安定で報酬ハッキングに陥りやすいことで知られている。本研究では、蒸留とRLベースの精緻化を統合した高速収束を実現する新しいフレームワーク、Flash-DMDを提案する。具体的には、まず、訓練コストを大幅に削減しつつリアリズムを向上させる効率的なタイムステップ認識蒸留戦略を提案する。これはDMD2のわずか2.1%の訓練コストでそれを上回る性能を示す。次に、RL目標によるファインチューニングとタイムステップ蒸留訓練を同時に継続する共同訓練スキームを導入する。継続中の蒸留から得られる安定した明確な損失が強力な正則化器として機能し、RL訓練過程を効果的に安定化させ、ポリシー崩壊を防止することを実証する。スコアベースモデルとフローマッチングモデルを用いた大規模な実験により、提案するFlash-DMDが、数ステップのサンプリング体制において、大幅に高速に収束するだけでなく、視覚的品質、人間の嗜好、テキスト-画像整合性の指標において既存手法を上回る、最先端の生成品質を達成することを示す。本研究は、効率的で高忠実度、かつ安定した生成モデルを訓練するための効果的なパラダイムを提示する。コードは近日公開予定である。

English

Diffusion Models have emerged as a leading class of generative models, yet their iterative sampling process remains computationally expensive. Timestep distillation is a promising technique to accelerate generation, but it often requires extensive training and leads to image quality degradation. Furthermore, fine-tuning these distilled models for specific objectives, such as aesthetic appeal or user preference, using Reinforcement Learning (RL) is notoriously unstable and easily falls into reward hacking. In this work, we introduce Flash-DMD, a novel framework that enables fast convergence with distillation and joint RL-based refinement. Specifically, we first propose an efficient timestep-aware distillation strategy that significantly reduces training cost with enhanced realism, outperforming DMD2 with only 2.1% its training cost. Second, we introduce a joint training scheme where the model is fine-tuned with an RL objective while the timestep distillation training continues simultaneously. We demonstrate that the stable, well-defined loss from the ongoing distillation acts as a powerful regularizer, effectively stabilizing the RL training process and preventing policy collapse. Extensive experiments on score-based and flow matching models show that our proposed Flash-DMD not only converges significantly faster but also achieves state-of-the-art generation quality in the few-step sampling regime, outperforming existing methods in visual quality, human preference, and text-image alignment metrics. Our work presents an effective paradigm for training efficient, high-fidelity, and stable generative models. Codes are coming soon.

Flash-DMD：効率的な蒸留と連携強化学習による高忠実度数ステップ画像生成に向けて

Flash-DMD: Towards High-Fidelity Few-Step Image Generation with Efficient Distillation and Joint Reinforcement Learning

要旨

Support