pi-Flow: 模倣蒸留によるポリシーベースの数ステップ生成

pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation

October 16, 2025

著者: Hansheng Chen, Kai Zhang, Hao Tan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

cs.AI

要旨

数ステップの拡散またはフローベースの生成モデルは、通常、速度を予測する教師モデルを、ノイズ除去されたデータへのショートカットを予測する学生モデルに蒸留します。この形式の不一致は、品質と多様性のトレードオフに悩まされる複雑な蒸留手順を引き起こしてきました。この問題に対処するため、我々はポリシーベースのフローモデル（pi-Flow）を提案します。pi-Flowは、学生フローモデルの出力層を修正し、1つのタイムステップでネットワークフリーのポリシーを予測します。このポリシーは、将来のサブステップで動的なフロー速度を生成し、追加のネットワーク評価なしに高速かつ正確なODE積分を可能にします。ポリシーのODE軌跡を教師モデルの軌跡に一致させるために、我々は新しい模倣蒸留アプローチを導入します。このアプローチでは、標準的なℓ₂フローマッチング損失を使用して、ポリシーの軌跡に沿ってポリシーの速度を教師モデルの速度に一致させます。教師モデルの動作を単純に模倣することで、pi-Flowは安定したスケーラブルなトレーニングを可能にし、品質と多様性のトレードオフを回避します。ImageNet 256^2では、1-NFE FIDが2.85を達成し、同じDiTアーキテクチャのMeanFlowを上回ります。FLUX.1-12BおよびQwen-Image-20Bでは、4 NFEsで、最先端の数ステップ手法よりも大幅に優れた多様性を実現しつつ、教師モデルレベルの品質を維持します。

English

Few-step diffusion or flow-based generative models typically distill a velocity-predicting teacher into a student that predicts a shortcut towards denoised data. This format mismatch has led to complex distillation procedures that often suffer from a quality-diversity trade-off. To address this, we propose policy-based flow models (pi-Flow). pi-Flow modifies the output layer of a student flow model to predict a network-free policy at one timestep. The policy then produces dynamic flow velocities at future substeps with negligible overhead, enabling fast and accurate ODE integration on these substeps without extra network evaluations. To match the policy's ODE trajectory to the teacher's, we introduce a novel imitation distillation approach, which matches the policy's velocity to the teacher's along the policy's trajectory using a standard ell_2 flow matching loss. By simply mimicking the teacher's behavior, pi-Flow enables stable and scalable training and avoids the quality-diversity trade-off. On ImageNet 256^2, it attains a 1-NFE FID of 2.85, outperforming MeanFlow of the same DiT architecture. On FLUX.1-12B and Qwen-Image-20B at 4 NFEs, pi-Flow achieves substantially better diversity than state-of-the-art few-step methods, while maintaining teacher-level quality.