PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hansheng Chen, Kai Zhang, Hao Tan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

arXiv:2510.14974v3cs.LGcs.AIcs.CV

TL;DR

Few-step distillation must reduce network evaluations without losing the quality and diversity of multi-step flow models, but shortcut-based training is complex and vulnerable to these trade-offs. π-Flow predicts a network-free policy and trains it on-policy against teacher velocities, achieving teacher-level quality with stronger diversity across ImageNet and large text-to-image models. The DX policy remains less robust than GMFlow, marking a scope boundary for the simpler policy class.

  • Problem

    Shortcut-based few-step distillation requires complex training and can degrade image quality or diversity through error accumulation and mode collapse.

  • Method

    π-Flow predicts a network-free policy whose future velocities are matched to teacher velocities along the policy trajectory with a standard ℓ2 flow matching loss.

  • Results

    Across ImageNet, FLUX.1-12B, and Qwen-Image-20B, π-Flow attains teacher-level quality while outperforming competitors in diversity and teacher alignment.

  • Takeaways & Limitations

    π-Flow provides a scalable framework for efficient few-step generation that avoids the reported quality–diversity trade-off.

  • Takeaways & Limitations

    The simpler DX policy has limited robustness because its predicted trajectory is not adaptive to perturbations in the current state.

Abstract

from arXiv · show

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 $\ell_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 256$^2$, 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.

1 INTRODUCTION

π-Flow replaces shortcut prediction with a network-free policy that integrates dense future substeps, while π-ID trains this policy on its own trajectory against teacher velocities. The approach targets efficient generation without the quality–diversity trade-off and scales from ImageNet to large text-to-image models.

  • 1 INTRODUCTION: Existing shortcut-based distillation requires complex training because shortcut paths cannot be directly inferred from the teacher, risking quality degradation or mode collapse.These methods compress multi-step teachers into students requiring only one or a few network evaluations.
  • 1 INTRODUCTION: π-Flow predicts a network-free policy at one timestep, which generates future flow velocities for dense ODE integration with negligible overhead.This decouples network evaluations from integration substeps.
  • 1 INTRODUCTION: π-ID matches policy velocities to teacher velocities along the policy’s own trajectory using a standard ℓ2 flow matching loss.The on-policy corrective signal is intended to reduce error accumulation while preserving teacher quality and diversity.
  • 1 INTRODUCTION: GMFlow outperforms DX, achieves strong ImageNet 256^2 FIDs at 1- and 2-NFE generation, and scales to 4-NFE FLUX.1-12B and Qwen-Image-20B students with teacher-level quality and state-of-the-art diversity.The supplied passage reports these outcomes without giving the corresponding FID values.

2 PRELIMINARIES

Flow matching generates samples by integrating a neural velocity field, but each velocity query requires a network evaluation and high-quality sampling typically needs more than 10 steps. π-Flow instead predicts a network-free policy once and integrates its velocity over multiple substeps.

  • 2 PRELIMINARIES: A standard flow model interpolates data with Gaussian noise over diffusion time and uses a neural network to approximate the probability-flow ODE velocity.Sampling starts from noise and solves the ODE toward the data state.
  • 2 PRELIMINARIES: Flow matching models couple sampling efficiency to solver precision because every velocity query evaluates a neural network, while high-quality sampling typically requires over 10 steps.The cost is commonly measured by the number of network function evaluations.
  • 2 PRELIMINARIES: π-Flow defines a network-free policy that maps each state and time to a flow velocity, with the policy conditioned on the trajectory’s initial state.The student predicts an entire dynamic policy rather than a single velocity.
  • 2 PRELIMINARIES: Generation first predicts the policy from the initial state, then integrates its velocity across multiple substeps to obtain a less noisy state.This separates one network evaluation from subsequent policy queries.
  • 2 PRELIMINARIES: By decoupling network evaluations from ODE integration substeps, π-Flow combines few network evaluations with dense teacher-like integration.The policy must provide efficient closed-form velocities while remaining compatible, expressive, and robust.

0 POLICY

π-Flow policies generate dynamic velocities for dense ODE integration after one student evaluation, with DX and GMFlow variants trading simplicity against robustness. π-ID trains these policies on their own trajectories by matching teacher velocities, while GMFlow provides expressiveness and perturbation robustness.

  • DX policy: A single student evaluation predicts a grid of DX denoising moments, and linear interpolation supplies the policy at arbitrary intermediate times.The grid uses N evenly spaced times between tdst and tsrc.
  • DX policy: DX policy is fast and expressive but lacks robustness because its predicted denoising moments do not adapt to perturbations in the initial state.The limitation concerns perturbations in xt.
  • GMFlow policy: GMFlow predicts a factorized Gaussian mixture velocity distribution whose parameters define closed-form velocities at future states and times.The output factorizes the data dimension into sequence length L and channel size C, with K mixture components.
  • GMFlow policy: GMFlow policies can approximate any N-step trajectory when K = N ·C mixture components are available.Theorem 1 states that the trajectory velocities can be approximated arbitrarily well at N distinct times.
  • Policy-based imitation distillation: π-ID samples intermediate states from a detached policy rollout, then matches the policy velocity to the teacher using an on-policy imitation objective.The detached rollout uses high-accuracy ODE integration, while on-policy training provides corrective signals against error accumulation.
  • Policy-based imitation distillation: On-policy imitation bounds performance degradation by O(nε), improving over the O(n2ε) compounding-error behavior of off-policy cloning.Here n is the number of substeps and ε is the average imitation error per step.

5 EXPERIMENTS

Experiments evaluate π-Flow across ImageNet and large text-to-image models, comparing policy variants, prior few-step students, and practical efficiency. Across these settings, π-Flow shows strong quality, diversity, teacher alignment, robustness, and near-shortcut inference speed.

  • 5.2 IMAGENET DIT: GMFlow consistently outperforms DX on 1-NFE ImageNet FID and remains consistent across K values, whereas DX is sensitive to N.The comparison is conducted after 40k training iterations; a 0.05 GM dropout rate further improves FID and Recall over no dropout.
  • 5.2 IMAGENET DIT: π-Flow outperforms the original MeanFlow DiT at both 1-NFE and 2-NFE generation without JVP operations or adaptive loss scalings.FACM attains a leading 2-NFE FID but still relies on JVP, while π-Flow uses a minimal training framework.
  • 5.3 FLUX.1-12B AND QWEN-IMAGE-20B: π-Flow outperforms other few-step students on roughly 70% of reported text-to-image metrics without obvious weaknesses in any specific area.This result covers the COCO-10k, HPSv2, and OneIG-Bench evaluations.
  • 5.3 FLUX.1-12B AND QWEN-IMAGE-20B: π-Flow achieves the highest diversity scores and best teacher-referenced FIDs by clear margins, especially at 4 NFEs, while maintaining teacher-like structure.The same-initial-noise comparison shows structurally similar outputs to the teacher, unlike VSD students that repeat structures.
  • 5.3 FLUX.1-12B AND QWEN-IMAGE-20B: Data-dependent and data-free π-Flow models achieve nearly identical results, supporting use when high-quality data is unavailable.The study also reports that π-Flow outperforms alternative students on detail rendering and avoids their noted texture or text-rendering weaknesses.
  • 5.3 FLUX.1-12B AND QWEN-IMAGE-20B: Thirty-two policy integration substeps cost around 15 ms, only 3% of network time, keeping overall 4-NFE inference speed on par with shortcut-predicting models.π-Flow averages 32 policy substeps after each policy-generation network evaluation.

6 RELATED WORK

Few-step distillation methods generally predict shortcuts toward less-noisy states, using objectives from direct regression to distribution matching. These approaches face quality, velocity-construction, diversity, and auxiliary-network challenges.

  • Shortcut-predicting distillation methods compress multi-step teachers by predicting paths toward less-noisy states.
  • Direct x_0 regression can degrade quality by producing blurry results, while progressive distillation reduces NFEs through successive multi-step output regression.
  • Consistency methods improve over x-based regression but require implicitly constructed velocities and remain limited by accumulated velocity errors.
  • Distribution matching can improve quality but risks mode collapse, style drift, and reliance on auxiliary networks.

7 CONCLUSION

The paper presents π-Flow as a scalable few-step generation framework that combines network-free policy integration with simple on-policy imitation. Across ImageNet, FLUX.1-12B, and Qwen-Image-20B, it reports teacher-level quality and stronger diversity and teacher alignment than competitors.

  • π-Flow consistently attains teacher-level image quality while outperforming competitors in diversity and teacher alignment across three distilled model families.The evaluated models are ImageNet DiT, FLUX.1-12B, and Qwen-Image-20B.
  • The framework uses a fast policy for dense ODE substeps and reduces distillation to a single ℓ2 loss to mitigate error accumulation and quality–diversity trade-offs.
  • The authors identify robust policy families, improved distillation objectives, and extensions such as video generation as future directions.
  • The paper provides detailed training procedures and hyperparameters to facilitate reproduction.

B.1 GM TEMPERATURE

The appendix introduces inference temperature scaling for GMFlow and a scheduled trajectory mixing strategy for reducing teacher-query cost and exposure to out-of-distribution states. The latter transitions training from teacher trajectories to on-policy imitation.

  • Training transitions from fully off-policy teacher behavior cloning to mixed teacher-policy segments and finally on-policy π-ID as the teacher ratio decays to zero.
  • The strategy addresses FLUX.1 dev’s weaker robustness to out-of-distribution states because it lacks true classifier-free guidance.
  • Scheduled trajectory mixing keeps teacher-query compute constant by covering coarse intervals initially and progressively filling them with the fast policy.
  • Teacher step boundaries are randomly sampled under the current teacher-ratio constraint, varying step sizes and locations while preserving the scheduled teacher-covered length.

B.3 MICRO-WINDOW VELOCITY MATCHING

Micro-window velocity matching replaces instantaneous matching with an averaged signal that improves training stability and reduces effective sampling steps. The appendix also reports implementation choices for temperature, time shifting, and teacher-step schedules.

  • Micro-window velocity matching smooths training, stabilizes the DX policy, and lets Δt = 3/128 mimic 128-step teacher sampling with approximately 43 steps.
  • Removing the micro-window trick degrades teacher alignment for 4-NFE GM-FLUX.
  • The experiments use a fixed time-shift parameter m = 3.2 and halve the final raw-time step size for FLUX.1 and Qwen-Image.

C ADDITIONAL IMPLEMENTATION DETAILS AND HYPERPARAMETERS

The implementation uses BF16 mixed precision and 8-bit Adam, while GMFlow hyperparameters trade off expressiveness, compatibility, and robustness. Toy experiments compare pixel-wise, element-wise, and patch-wise factorizations for their effects on generated appearance.

  • All models use BF16 mixed precision and 8-bit Adam without weight decay, with EMA weights and a dynamic moment schedule at inference.
  • Pixel-wise factorization produces neutral colors and detailed textures, whereas element-wise factorization yields over-saturated colors and over-smoothed textures.
  • Patch-wise factorization results in “confetti” artifacts in the toy-model comparison.
  • K = 8 and C = VAE latent channel size generally provide a useful balance among GMFlow compatibility, expressiveness, and robustness.Larger K improves expressiveness but can complicate training, while larger C improves robustness but reduces expressiveness.

D DISCUSSION ON FLUX.1 SCHNELL

The FLUX.1 schnell comparison is inequitable because schnell is distilled from the closed-source FLUX.1 pro rather than FLUX.1 dev. Nevertheless, π-Flow receives higher human preference and maintains more coherent structures in the reported comparison.

  • FLUX.1 schnell is an inequitable direct baseline because it is distilled from closed-source FLUX.1 pro, unlike the students distilled from publicly available FLUX.1 dev.
  • FLUX.1 schnell has better data and prompt alignment than FLUX.1 dev but weaker preference alignment than FLUX.1 dev and all its students.
  • π-Flow was preferred 59.5% of the time over FLUX.1 schnell in a 200-image human preference study.The study used HPSv2 prompts, and the result aligned with the HPSv2.1 preference metric.
  • Qualitative comparisons report frequent structural errors for FLUX.1 schnell, while π-Flow maintains coherent structures.

E PROOF OF THEOREM 1

The proof establishes that a Gaussian mixture with at most N · C components can approximate trajectories at N specified times. It derives this through moment equations, convex-hull arguments, and the Richter–Tchakaloff theorem, then connects the result to closed-form GMFlow velocities.

  • A Gaussian mixture with at most N · C components suffices to approximate any N-step trajectory in R^C.The proof first obtains an N · C-atomic probability measure and replaces Dirac deltas with narrow Gaussians.
  • Theorem 2 guarantees a probability measure whose ODE moment equations hold at every one of N distinct times.
  • The proof reduces the multivariate Fredholm equation to a univariate factorized case and shows that the zero vector lies in the relevant convex hull.
  • A separating-hyperplane contradiction establishes the convex-hull condition, after which Carathéodory’s theorem yields a finite-support solution.
  • The Richter–Tchakaloff theorem bounds the required atomic support by dim V ≤ N · C, completing the approximation argument.
  • The resulting closed-form GMFlow velocity is represented by a policy function π and can be extended to discrete support through a limiting construction.

G ADDITIONAL QUALITATIVE RESULTS.

Additional figures provide uncurated qualitative samples from the OneIG-Bench prompt set. The samples are presented in two batches without a stated quantitative comparison.

  • Figures 12 and 13 show uncurated random batches from the OneIG-Bench prompt set.
Loading 2510.14974v3…