pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
Authors: name: Hansheng Chen; name: Kai Zhang; name: Hao Tan; name: Leonidas Guibas; name: Gordon Wetzstein; name: Sai Bi
Published: 2025-10-16 17:59:51

Abstract: 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.