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

少步数扩散模型或基于流的生成模型通常会将一个预测速度的教师模型蒸馏到学生模型中，由后者直接预测通向去噪数据的捷径。这种师生模型输出格式不一致的问题导致了复杂的蒸馏流程，且往往面临生成质量与多样性之间的权衡。为解决这一问题，我们提出了基于策略的流模型（$\pi$-Flow）。$\pi$-Flow 修改了学生流模型的输出层，使其在单个时间步上预测一个无需神经网络参与的策略。该策略随后可在后续子步骤中动态生成流速度，计算开销极小，从而在这些子步骤上实现快速而精确的常微分方程（ODE）积分，且无需额外的网络前向计算。为了使策略所生成的ODE轨迹与教师模型对齐，我们提出了一种新颖的模仿蒸馏方法，即沿策略自身的轨迹，使用标准的 $\ell_2$ 流匹配损失，使策略输出的速度逼近教师模型对应的速度。通过简单地模仿教师行为，$\pi$-Flow 实现了稳定且可扩展的训练，并避免了质量与多样性之间的权衡。在 ImageNet 256$^2$ 数据集上，$\pi$-Flow 达到了 1-NFE（单次网络函数求值）下 2.85 的 FID 分数，优于相同 DiT 架构下的 MeanFlow 方法。在 FLUX.1-12B 和 Qwen-Image-20B 模型上，仅需 4 次 NFE 时，$\pi$-Flow 在保持教师模型级别生成质量的同时，显著提升了生成结果的多样性，超越了当前最先进的少步生成方法。

论文旨在解决少步扩散或流式生成模型中的质量-多样性权衡问题。传统方法通过蒸馏一个预测速度的教师模型来训练学生模型，但学生模型通常直接预测从噪声到数据的捷径，导致输出格式不匹配，从而需要复杂的蒸馏流程，并难以同时保证生成质量和样本多样性。这是一个在高效生成模型中日益受到关注的问题，尤其在追求快速采样（如1~4步生成）时尤为突出。

提出Policy-based Flow Models（π-Flow），其核心思想是修改学生流模型的输出层，使其不再直接预测最终流动方向，而是预测一个无需神经网络参与的策略（network-free policy）。该策略能够在后续子步骤中动态生成流速度，实现对ODE路径的精细建模，而无需额外的网络评估。通过引入模仿蒸馏机制，使策略轨迹上的速度匹配教师模型的速度，使用标准的ℓ2流匹配损失，从而实现更稳定、可扩展的训练，并自然规避了质量与多样性的冲突。

π-Flow在ImageNet 256²上实现了1-NFE（一步采样）FID为2.85，优于相同DiT架构下的MeanFlow；在FLUX.1-12B和Qwen-Image-20B模型上，仅用4次函数评估（NFE=4）即展现出显著更优的多样性，同时保持教师级别的生成质量。实验设计清晰，覆盖多种主流大模型和高分辨率图像生成任务，验证了方法的通用性和可扩展性。目前尚未提及开源代码，但其方法可集成到现有流匹配框架中，未来值得探索在视频生成、跨模态生成等动态系统中的应用。

1. Flow Matching for Generative Modeling (Lipman et al., 2023) 2. Analytic-DPM: An Analytic Perspective on Training Diffusion Probabilistic Models (Zhang et al., 2023) 3. MeanFlow: Faster and Higher-Quality Diffusion Models via Expected Mean Flow Estimation (Chen et al., 2024) 4. Rectified Flow: Accelerating Diffusion Models with Straight Flows (Liu et al., 2022) 5. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow-Based Diffusion Probabilistic Models (Liu et al., 2023)