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

arXiv: https://arxiv.org/abs/2510.14974 · HTML: https://mlsanity.com/paper/2510.14974/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 66 references that are in the map.
- First posted on arXiv 16 Oct 2025.

## Builds on (most cited first, 30 of 56)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation](https://mlsanity.com/paper/2305.16213/) (2023)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models](https://mlsanity.com/paper/2211.01095/) (2022)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks](https://mlsanity.com/paper/1506.03099/) (2015)
- [Adversarial Diffusion Distillation](https://mlsanity.com/paper/2311.17042/) (2023)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [Improved Distribution Matching Distillation for Fast Image Synthesis](https://mlsanity.com/paper/2405.14867/) (2024)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Fast Sampling of Diffusion Models with Exponential Integrator](https://mlsanity.com/paper/2204.13902/) (2022)

## Built on by (most cited first, 29 of 29)

- [Improved Mean Flows: On the Challenges of Fastforward Generative Models](https://mlsanity.com/paper/2512.02012/) (2025)
- [Distribution Matching Distillation Meets Reinforcement Learning](https://mlsanity.com/paper/2511.13649/) (2025)
- [Filter, Then Reweight: Rethinking Optimization Granularity in On-Policy Distillation](https://mlsanity.com/paper/2606.02684/) (2026)
- [Transition Matching Distillation for Fast Video Generation](https://mlsanity.com/paper/2601.09881/) (2026)
- [Causal-rCM: A Unified Teacher-Forcing and Self-Forcing Open Recipe for Autoregressive Diffusion Distillation in Streaming Video Generation and Interactive World Models](https://mlsanity.com/paper/2606.25473/) (2026)
- [D-OPSD: On-Policy Self-Distillation for Continuously Tuning Step-Distilled Diffusion Models](https://mlsanity.com/paper/2605.05204/) (2026)
- [TwinFlow: Realizing One-step Generation on Large Models with Self-adversarial Flows](https://mlsanity.com/paper/2512.05150/) (2025)
- [Flow Map Distillation Without Data](https://mlsanity.com/paper/2511.19428/) (2025)
- [Parallel Decoding Distillation for Fast Image and Video Generation](https://mlsanity.com/paper/2607.26004/) (2026)
- [Bidirectional Normalizing Flow: From Data to Noise and Back](https://mlsanity.com/paper/2512.10953/) (2025)
- [DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment](https://mlsanity.com/paper/2603.22125/) (2026)
- [SubFlow: Sub-mode Conditioned Flow Matching for Diverse One-Step Generation](https://mlsanity.com/paper/2604.12273/) (2026)
- [Bi-Anchor Interpolation Solver for Accelerating Generative Modeling](https://mlsanity.com/paper/2601.21542/) (2026)
- [Duality Models: An Embarrassingly Simple One-step Generation Paradigm](https://mlsanity.com/paper/2602.17682/) (2026)
- [Distilling Drifting Transformers with Representation Autoencoders](https://mlsanity.com/paper/2606.15553/) (2026)
- [Glance: Accelerating Diffusion Models with 1 Sample](https://mlsanity.com/paper/2512.02899/) (2025)
- [Multi-Resolution Flow Matching: Training-Free Diffusion Acceleration via Staged Sampling](https://mlsanity.com/paper/2607.01642/) (2026)
- [Stabilizing, Scaling & Enhancing MeanFlow for Large-scale Diffusion Distillation](https://mlsanity.com/paper/2605.17834/) (2026)
- [FourTune: Towards Fully 4-Bit Efficient Post-Training for Diffusion Models](https://mlsanity.com/paper/2607.05711/) (2026)
- [FlowConsist: Make Your Flow Consistent with Real Trajectory](https://mlsanity.com/paper/2602.06346/) (2026)
- [SODA: Sensitivity-Oriented Dynamic Acceleration for Diffusion Transformer](https://mlsanity.com/paper/2603.07057/) (2026)
- [IB-Flow: Information Bottleneck-Guided CFG Distillation for Few-Step Text-to-Image Generation](https://mlsanity.com/paper/2607.09133/) (2026)
- [A Latent Distribution Perspective on Evaluating and Improving Latent Generative Models](https://mlsanity.com/paper/2609.24088/) (2026)
- [WorldPlay2: Extending Real-Time Interactive World Models in Control and Horizon](https://mlsanity.com/paper/2609.35560/) (2026)
- [TACD: Distilling Efficient Text-to-Motion Models via Terminal Amplification Control](https://mlsanity.com/paper/2610.02867/) (2026)
- [Kandinsky 6.0 Video: Foundation Models for Synchronized Video and Audio Generation](https://mlsanity.com/paper/2610.05608/) (2026)
- [Consistent Distribution Matching for Data-Free Diffusion Distillation](https://mlsanity.com/paper/2610.09221/) (2026)
- [Parametric Trajectory Distillation for Few-Step Video Generation](https://mlsanity.com/paper/2610.11498/) (2026)
- [Kandinsky S2V 1.0: Audio-Driven Avatar Video Generation](https://mlsanity.com/paper/2610.19290/) (2026)

## Abstract (arXiv)

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 ℓ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 2562, 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.