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

TL;DR

pi-Flow reframes few-step diffusion/flow distillation as imitation learning on a policy: the student no longer predicts a denoised shortcut, it predicts a tiny network-free policy at each outer timestep that rolls out several ODE substeps at negligible cost. A single L2 loss matching the student’s policy velocity to the teacher’s velocity along the student’s own trajectory replaces JVPs, GANs, and auxiliary networks. Result is 4-step T2I that matches the teacher’s quality and keeps its diversity — historically the hard part of DMD- style methods. Accepted to ICLR 2026, with distilled FLUX.1, FLUX.2, Qwen-Image, and ImageNet DiT releases.

Key claims

The “predict shortcut to clean data” format used by standard few-step distillation is structurally mismatched with the velocity- predicting teacher, and this mismatch is what drives the quality– diversity trade-off in DMD/consistency/rectified-flow students [§1, §3].
A policy-based output head — the student emits a network-free policy at one outer timestep that produces dynamic velocities at future ODE substeps — closes the format gap without extra forward passes per substep [§3, Fig. 1].
Imitation distillation (pi-ID): a single L2 loss between the policy velocity and the teacher velocity, evaluated along the policy’s own ODE trajectory, suffices for stable training — no JVPs, no auxiliary student/teacher discriminator, no GAN, no real-data anchor [§3].
4-step T2I distillates achieve “state-of-the-art diversity and teacher-aligned quality” on FLUX.1-dev and Qwen-Image backbones [§4, abstract].
Method generalizes across teachers (FLUX.1, FLUX.2, Qwen-Image, ImageNet DiT) and across step budgets (1-/2-step on ImageNet, 4-/8-step on FLUX) — confirmed by the released checkpoint family [§5, model cards].

Method

The student keeps the same backbone as the teacher but replaces the final velocity head with a policy head that emits a small, fixed parametric object (a “network-free policy”) at each outer timestep. That policy is cheap to evaluate at intermediate substeps, so the student can integrate the ODE accurately over a few substeps without running the full network again — this is where the few-step speedup comes from while keeping per-substep velocities dynamic instead of collapsing to a linear shortcut.

Training is imitation distillation. For a sampled outer timestep, the student rolls out its policy to obtain a trajectory; at points along that trajectory the teacher’s velocity is queried, and the loss is L2 between policy velocity and teacher velocity. There is no JVP, no adversarial loss, no second network. Because the supervision is along the student’s own rollout, the train/inference distribution matches by construction — this is a structural fix to the same “DMD supervises a single sampled timestep but inference rolls out N steps” gap that TMD attacked via flow-head rollout inside DMD.

Results

Headline result is 4-step T2I on FLUX.1-dev and Qwen-Image that holds teacher-class diversity (where DMD-style students typically collapse) while matching teacher FID/aesthetic. Released checkpoints span 4- and 8-NFE FLUX.1, 4-NFE FLUX.2, Qwen-Image (pi-Qwen), and 1-/2-step ImageNet DiTs, with HF demo spaces for each. Detailed numerical tables (HPS, AeS, CLIP, diversity metrics on OneIG-Bench) are in §4 of the paper; the wiki notes pi-Flow as the baseline that subsequent work (ArcFlow: Unleashing 2-Step Text-to-Image Generation via High-Precision Non-Linear Flow Distillation) had to beat at 2 NFE.

Why it’s interesting

This is the original pi-Flow paper — the wiki already discusses pi-Flow as a baseline in ArcFlow: Unleashing 2-Step Text-to-Image Generation via High-Precision Non-Linear Flow Distillation (which claims to beat pi-Flow / TwinFlow / SenseFlow / Qwen-Image-Lightning at 2 NFE with <5% of parameters trained), so filing the underlying paper closes that citation chain. The connection to Diffusion Distillation is the central one: pi-Flow occupies a distinct axis from the rest of the wiki’s cluster — TMD attacks the train/inference gap via flow-head rollout inside DMD, DP-DMD attacks DMD’s mode-seeking via role- separated steps, ArcFlow attacks linear-shortcut underfitting via a non-linear closed-form ODE, while pi-Flow drops the shortcut format entirely and replaces it with a per-step policy + pure imitation loss. Same first author and Stanford/Adobe team also produced Asymmetric Flow Models (AsymFlow), filed last week — the LakonLab codebase ships pi-Flow, AsymFlow, and GMFlow together, suggesting a coherent research program around non-shortcut flow parameterizations.