pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
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.
68f19b086e0bef323a68fc1a
Hansheng Chen
claimed_verified
2025-10-17T04:11:30.593Z
638067fcb334960c987fbeda
/avatars/87f1eaaf6b3a9c0d47d6f406261ccc18.svg
Hansheng Chen
user
Lakonik
68f19b086e0bef323a68fc1b
Kai Zhang
68f19b086e0bef323a68fc1c
Hao Tan
68f19b086e0bef323a68fc1d
Leonidas Guibas
68f19b086e0bef323a68fc1e
Gordon Wetzstein
68f19b086e0bef323a68fc1f
Sai Bi
pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
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.
Policy-based flow models enable efficient and high-quality image generation by distilling teacher models into student models with dynamic flow velocities, improving diversity and quality.
velocity-predicting teacher student flow model policy-based flow models $\pi$-Flow dynamic flow velocities ODE integration imitation distillation $\ell_2$ flow matching loss ImageNet 256$^2$ FID MeanFlow DiT architecture FLUX.1-12B Qwen-Image-20B NFE