用策略模仿蒸馏，让生成模型一步到位，又快又准。
pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
- 学生模型预测动态策略，替代传统速度预测，提升推理效率。
- 在ImageNet上仅1次网络评估即达FID 2.85，超越同类模型。
- 适合追求高效高质生成的用户，尤其适用于大模型图像生成。
少步扩散或流模型通常将预测速度的教师模型蒸馏到学生模型中，但格式不匹配导致蒸馏复杂且常面临质量与多样性权衡。为此，我们提出策略型流模型（π-Flow）。π-Flow 修改学生流模型输出层，使其在单个时间步预测一个无网络依赖的策略；该策略随后在后续子步生成动态流速，几乎无额外开销，实现快速准确的常微分方程（ODE）积分。为使策略轨迹与教师一致，我们引入新型模仿蒸馏方法，使用标准ℓ₂流匹配损失，在策略轨迹上对齐策略速度与教师速度。通过直接模仿教师行为，π-Flow 实现稳定可扩展训练，避免质量-多样性权衡。在 ImageNet 256² 上，其 1-NFE FID 达 2.85，优于同架构先前 1-NFE 模型。在 FLUX.1-12B 与 Qwen-Image-20B 上，4 NFE 下显著提升多样性，同时保持教师级质量。
原文摘要 · Abstract (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 ($π$-Flow). $π$-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, $π$-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 previous 1-NFE models of the same DiT architecture. On FLUX.1-12B and Qwen-Image-20B at 4 NFEs, $π$-Flow achieves substantially better diversity than state-of-the-art DMD models, while maintaining teacher-level quality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。