pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
Authors: Hansheng Chen
Kai Zhang
Hao Tan
Leonidas Guibas
Gordon Wetzstein
Sai Bi
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion models
flow models
few-step generation
distillation
Primary area: generative models
TL;DR: pi-Flow distills a pre-trained flow model into a policy-based flow model using simple imitation learning.

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 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\times 256$, 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, $\pi$-Flow achieves substantially better diversity than state-of-the-art DMD models, while maintaining teacher-level quality.

Published as a conference paper at ICLR 2026
PI-FLOW: POLICY-BASED FEW-STEP GENERATION VIA IMITATION DISTILLATION
Hansheng Chen'
Kai Zhang² Hao Tan Leonidas Guibas'
Gordon Wetzstein'
Sai Bi?
Stanford University 2 Adobe Research
T--Flow: Few-Step Distillation Made Easy!
Predict fust flow velacity policy not shorlao!!
The polcy rols out carved DCE tngechey with mismal overhaad Apply initaaiion learning match palce
T--Flow #5..
VAE
": # 2200 baaybayyy Miae, 423642
PiT
2: Brior EMM Anisonis
Policy (Fow))
I-Fow
netwwork evalustions Many staps
BIFF FLLM
Figure High quality 4-NFE text-to-image generations by T-Flow, distilled from FLUX.1-12B (top-right three images) and Qwen-Image-200 (all remaining images). r-Flow preserves the teacher's coherent structures, fine details (e.g., skin and hair), and accurate text rendering, while avoiding diversity collapse (see Fig, for sample diversity).
ABSTRACT
Few-step diffusion or flow-based generative models typically distill velocity- predicting teacher into student that predicts a shortcut towards denoised data. This format mismatch has led complex distillation procedures that often suffer from a quality-diversity trade-off. To address this, we propose policy-based flow models (r-FFlow). m-Flow modifies the output layer of student flow model predict network-free policy at one timestep. The policy then produces dynamic flow velocities at future substeps with negligible overhead, enabling fast and ac- curate 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 ss along the policy's ' trajectory using a standard f2 flow matching loss. By simply mimick- ing the teacher's behavior, 7-Flow enables stable and scalable training and avoids the quality-diversity trade-off. On ImageNet 256 it attains a 1-NFE FID of 2,85, outperforming previous 1-NFE models of the same DiT architecture. On FLUX.1- 12B and
at NFEs. π --Flow achieves substantially better diver- sity than state-of-the-art DMD models, while maintaining teacher-level quality.

---

Published as a conference paper at ICLR 2026
INTRODUCTION
Diffusion and flow matching models (Sohl-Dickstein al., 2015 Ho al. 2020: Song & Ermon 2019 Lipman et al. [202] & Vanden-Fijndeny (0033) have become the dominant method for visual generation, delivering compelling image quality and diversity. However, these models rely on costly denoising process for inference, which integrates a probability flow ODE (Song et al. 2021) over multiple timesteps, each step requiring a neural network evaluation. Commonly, the inference cost of diffusion models is quantified by the number of function (network) evaluations (NFEs).
To reduce the inference cost, diffusion distillation methods compress a pre-trained multi-step model (the teacher) into student that requires only one few network evaluation steps. Existing distilla- tion approaches avoid ODE integration by taking one or a few shortcut steps that map noise to data, where each shorteut path predicted by the student network, referred to shoricut- prrecccttn model. Learning these shortcuts is a significant challenge because they cannot be directly inferred from the teacher model. This necessitates the use of complex training methods, such as progres- sive distillation (Salimans & Ho 2022, Liu et al. 2023 2024), consistency distillation (Song et al., 2023). and distribution matching (Sauer et al., 2024a] Yin et al., 2024b a, Salimans et al., 2024). In turn, the sophisticated training often lead to degraded image quality from eπrоr accumulation or compromised diversity due mode collapse.
To sidestep the difficulties in shortcut- preedictngg distillation, we propose a novel policy-based flow model (r-- or pi-Flow) paradigm: given noisy data one timestep, the student network predicts a network-free policy, which maps new noisy states to their corresponding flow velocities with negligible overhead, allowing fast and accurate ODE integration using multiple substeps of policy velocities instead of network evaluations.
To train the student network, we introduce policy-based imitation distillation (I-ID), a DAgger- style (Ross et al. (011]) on-policy imitation learning (IL) method. I-ID trains the policy its own trajectory: at visited states, we query the teacher velocity and match the policy's 's output to it, using the teacher's corrective signal to teach the policy recover from its mistakes and reduce error accumulation. Specifically, the matching employs a standard fg loss aligned with the teacher's flow matching objective, thus naturally preserving its quality and diversity.
We validate our paradigm with two types of policies: a simple dynamic-40 (t) (DX) policy and an advanced GMFlow policy based on Chen et al. (2025). Experiments show that GMFlow policy outperforms DX policy and delivers strong ImtgeNet 256² FIDs at 1- and 2-NFE generation. To demonstrate its scalability, we distill FLUX.1-125 (Black Forest Labs] (022b) and Dowen-mmah 20B (Wu et al. 2025) text-to-image models into 4-NFE Flllow students, which achieve state-of- the-art diversity, while maintaining teacher-level quality.
We summarize the contributions of this work as follows:
We propose --Flow, a new paradigm that decouples ODE integration substeps from network evaluation steps, enabling both fast generation and straightforward distillation.
We introduce --ID, novel on-policy IL method for few-step -FFlow distillation, which reduces the training objective to a simple &g flow matching loss.
We demonstrate strong performance and scalability of 7-Flow. particularly, its superior diversity and teacher alignment compared to other sate-oo-theeaad 4-NFE text-to-image models.
PRELIMINARIES
In this section, we briefly introduce flow matching models (Lipman al. 2023 Liu et al. 2023 and the notations used in this paper.
Let p(æo) denote the (latent) data probability density, where æo E RD is data point. A stan- dard flow model defines an interpolation between a data sample and a random Gaussian noise (O, I), yielding the diffused noisy data æ; ayZy + o,e, where t E (0, 1] denotes the diffusion time, and a₂ = t,o, = t are the linear flow noise schedule. The optimal transport map across all marginal densities p(æ.)) fg Nax,; 0x20, of )p(oo))))) day be described by
2

---

Published as a conference paper at ICLR 2026
(a) Standard flow model Many network steps
(b) Shortcut-pon phoricut-predicting model
(c) Policy-based model Few network steps
Few network steps
A Many integration su aubeteps
Many integration steps
Few integration steps
det
Pollcy at (X, E)
X
dext
Figure 2: Comparison between (a) standard flow model (teacher), (b) Chortcut-predictio model, and (c) our policy-based model. The shorteut- prreecccttpriiiiigggg model skips all intermediate states, whereas the our-based model retains all intermediate substeps with minimal overhead.
the following probability flow ODE (Song et al. dx
2021 Liu 2022): :
2t E Paoorpii planppppaaalaee p(
Tol
JRP Top (zol xx,) dxp
dt
with the denoising posterior p(Zrola:)) N Marsomemo ( p(as)
At test time, the model can generate samples by first initializing the noise 21 +- e and then solving the ODE to obtain lim;-so Pt-
dæz In practice, flow matching models approximate the ODE velocity de using a neural network (Z(,),)) () with learnable parameters 0, trained using the & flow matching loss:
Lo Bg
(ar, 1)] with sample velocity
Et Zo
2
Since each velocity query requires evaluating the network (Fig. 2 (a)), flow matching models couple sampling efficiency with solver precision. Despite the progress in advanced solvers (Karras et al.,
2022 Zhang Chen 2023 Lu et 2022 2023 Zhao et al., 2023). high-quality sampling typically
requires over 10 steps due to inherent ODE truncation error, making it computationally expensive.
3 --FL:: POLICY-BASED FEW-STEP GENERATION
In T-Flow, we define the policy as a network-free function :: RD X R + RDD that maps a state (æ,, to flow velocity. A policy be network-free if it only needs describe single ODE trajectory, which fully determined by its initial state (It... fare) with fw 2 t. In this case, the policy for each trajectory must be dynamically predicted by neural network conditioned on that initial state (Itt,. lar). ). We therefore adapt a flow model to output not a single velocity, but an entire dynamic policy that governs the full trajectory. Formally, define the policy function space F { <: IDD x R RD }. Then, our goal is to distill policy generator network GΦ: R x R with learnable parameters Φ, such that π (2e, t) = Gø (Ii (arc) (x.,t). t).
As shown in Fig. (c), -FFlow performs ODE-based denoising from twe to tass via two stages:
A policy generation step, which feeds the initial state (I).., twx)) to the student network Gø to produce the policy π, i.e., t- Go( ).
Multiple policy integration substeps, which integrates the ODE by querying the policy velocity over multiple substeps, obtaining a less noisy state by Btt. <- Dtg + plse (T.,, dr.
Unlike previous few-step distillation methods, r-Flow decouples network evaluation steps from ODE integration substeps. This allows to combine the key advantages of two paradigms: per- forms only a few network evaluations for efficient generation, similar to a shortcut-predd poretcctttreiicting model, while also executing dense integration substeps, just like a standard flow matching teacher. Thanks to its teacher-like ODE integration process, a m-Flow student offers unprecedented advantage in training. as we can now follow well-established imitation learning (IL) approaches to directly match the policy velocity (Tx t) the teacher velocity Ge(2,, t), as discussed later in $ 4
To identify the appropriate function classes of student policies for fast image generation, we need to consider the following requirements:
Efficiency. The policy should provide closed-fome velocities with minimal overhead, so that
rolling out dense (e.g., 100+) substeps incurs negligible cost compared to network evaluation.
3

---

Published as a conference paper at ICLR 2026
Compatibility. The policy should have compact set of parameters that can be easily predicted
by the student G6 with standard backbones (e.g.. DiT (Peebles & Xie 2023 )-
Expressiveness. The policy should be able to approximate complicated ODE trajectory starting from certain initial state Dt
Robustness. The policy should be able to handle trajectory variations that arise from perturba- tions to the initial state Pt.... For instance, a suboptimal student network will produce an erroneous mapping from at π. This introduces the randomness that the policy needs accommodate throughout the rollout. Consequently, the policy function should adapt its velocity output to vari- ations in its state input 24, which is challenging requirement for network-free free functions.
3.1 DYNAMIC- 20 POLICY
We introduce simple baseline policy called dynamic-ap policy (DX policy). DX policy defines
π (#; !) == -.-a (t) where a/ approximates the posterior moment Exo-p(so)(ee) [2)] in Eq. d Along a fixed trajectory starting from an initial state (Itu (see, the posterior moment is only depen- dent I. Therefore,
first predict grid of ae(t)) N evenly spaced times 41, LN E [Laa (sre] by a single evaluation of the student network (T (sse). This is achieved by expanding the out- put channels of the student network and performing u-to-a7 reparameterization. Then, for arbitrary t E [daa: fae). we obtain the approximated moment (1) by a linear interpolation over the grid.
Apparently, DX policy is fast, compatible, and expressive enough that any N-step teacher tra- jectory can be matched with N grid points. However, robustness limited because a() is not adaptive to perturbations in 2t.
3.2 GMFLOW POLICY
For stronger robustness, we incorporate an advanced GMFlow policy based on the closed- form GM velocity field in Chen et al. (2025). GMFlow policy expands the network out- put channels to predict a factorized Gaussian mixture (GM) velocity distribution ((())x;..)) = I Ek= AaNN(u;; pikk 2²T), where A₁ E Rt, Mux E RC, E R2 are GM parameters pre- dicted by the network, T. C factorizes the data dimension D into sequence length L. and channel size C, and K is a hyperparameter specifying the number of mixture components. Intuitively, the student network G maps the initial state Tt,.. to multiple denoising modes that parameterize the GMFlow policy. The policy then enables closed-form velocity expression at future state (æ,, t) for any 0 < t < tue (see For details). The speed and compatibility of GMFlow has already been discussed in Chen et al. (2025), thus we focus on amalyzing its expresiveness and robustness.
Expressiveness. With the x C factorization, each individual C-dimensional GM needs to be expressive enough to approximate C-dimensional chunk of the teacher trajectory. In we rigorously prove the following theorem, demonstrating GMFlow's 's expressiveness.
Theorem (A GMFlow policy with K = N C' can accurately approximate any N-step trajectory). Given pairwise distinct times .,..., In E (0, 1] and vectors æt,, tt,, E R for n I,..., N, there exists GM parameterization of p( with N C' components, such that æ.., can be approxi- mated arbitrarily well using Eq (I) at t, for every n I....,N..
In practice, we can K < N C (c.g., K 8) since the teacher trajectory is mostly smooth. More analysis of GMFlow hyperparameters are presented in C.I
Robustness. GMFlow is highly robust against trajectory perturbation due to its probabilistic origin. Unlike DX policy, GMFlow models a fully dynamic denoising posterior (Eq. (23 dependent on both 2t and t. Leveraging its robustness, the policy be flexibly altered via GM dropout in
training ($ and GM temperature inference ($ B.I. both improving generalization performance.
4

---

Published as a conference paper at ICLR 2026
Algorithm 1: On-policy T-ID.
Input: NFE, teacher Go, student Gg. condition
Go
Sample tue from NPE: NPE....,
Policy
2 Initialize 2im (data-free data-dependent)
stopgrad
Datached policy Fo
Gø (2trun lxc,C)
dropout)
Teacher valocity (x..))
ATD +- stopgrad(T)
Detached policy HD
Palicy valocitt (X;
L +-
Figure 3: On-policy flow imitation dis- tillation. Intermediate states are sampled along the detached policy rollout, where the loss matches the policy to the teacher.
for finite samples Ufrrr NFB. f(rc) do
Bles + Dime TD(as, 1)
Lo - Lo + [1o (,)) nt (2s, t) 9 - Adam( VoLo)
II optimizer step
4 r-ID: POLICY-BASED IMITATION DISTILLATION
With the policy rollout sharing the same format as the teacher's ODE integration, it is straightfor- ward to adopt imitation learning to learn the policy by directly matching the policy's velocity the teacher's velocity. In this section, we introduce simple policy-based imitation distillation (7-ID) algorithm based on DAgger-style (Ross et al. 2011) on-policy imitation.
On-policy imitation learning is robust to error accumulation since it trains the policy on its own trajectory, allowing the teacher's s corrective signal to steer a deviating trajectory back on track. As shown in Fig. 3 and Algorithm [ for a time interval from tate ta (i.e., I-NFE segment), feed the initial state (I (xx)) to the student network G& to obtain the policy π. We then sample an intermediate time E (tas, fax] and roll out detached policy TD from fax using high- accuracy ODE integration (with a small step size of 1/128), yielding an intermediate state æ on the policy trajectory. This is fed both the learner policy and the frozen teacher Ga. which produce their respective velocities. Finally, we compute a standard f2 flow matching loss between the two velocities, and backpropagate its gradients through the policy π to the student network G Because the student forward/backward pass dominates compute while policy and teacher queries are relatively cheap, we may repeat the rollout-and aaattcind-mm step multiple times for additional teacher supervisions. In practice, we sample two intermediate states per student forward pass.
first
Data-dependent and data-free --ID The initial state Tt... can be obtained via forward diffusion from real data To (data-dependent Algorithm 2) or via x-Flow's reverse denoising from random noise 21 (data-free Algorithm 3). Both methods have roughly the same computational cost, and comparable performance, demonstrated in the experiments ($
Error bounds and convergence. As discussed by Ross al. (011), on-policy imitation learning guarantees that the performance of the learned policy is bounded by the teacher's performance plus an error term that scales as Onn)). where n is the number of substeps and E the average imitation error per step, which is strictly better than the O(n2s) compounding- erroo behavior of off-policy behavior cloning. Moreover, the sequence of on-policy iterates converges in performance to the best policy in the function class, under the student's is capacity constraint.
5 EXPERIMENTS
To demonstrate the versatility T-Flow, we evaluate it with three distinct image generation mod- els of different scales and architectures: DiT(SiT)-XL/2 (675M) (Peebles & Xie 2023 Ma al.) 2024 Vaswani et al. 2017) for ImageNet 2562 (Deng et al. 2009) class-conditioned generation, FLUX.1 I2BB (Black Forest Labs 2024b) and Qwen-Image-200 [Wu al. 2025) for text-to-image generation.
5.1 IMPLEMENTATION DETAILS
In this subsection, we discuss key implementation details essential to model performance. More training details and hyperparameter choices presented in C
5

---

Published as a conference paper at ICLR 2026
Table 1: 1-NFE generation results of π-Flow with DX and GMFlow policies on ImageNet. Tested after 40K training it- erations. FM stands for standard flow matching.
Table 2: Comparison with previ-
ous few-step DiTs on ImageNet.
Model
NFE
FIDJ
Policy
Teacher
FIDL
IS1
Precision
Recall
iCT
20.30
DX (N 10)
REPA
4.73
327.6
0.781
0.514
iMM
1x2
7.77
DX (N -20)
REPA
4.44
329.8
0.786
0.531
McanFlow
2.20
DX (N- =40)
REPA
4.90
321.8
0.778
0.537
FACM (REPA)
1.52
GM (Æ
REPA
3.07
336.9
0.789
0,572
T-Flow (GM-REPA)
1.97
GM (K 32)
REPA
3.08
341.7
0.791
0,562
iCT
34.24
GM (K- 32)
FM
3.65
282.0
0.797
0,533
Shortcut
10.60
GM (K 32) w/o dropout
FM
4.14
279.6
0.799
0,525
McanFlow
3,43
T-Flow (GM-FM)
3.34
5-Flow (GM-REPA)
2,85
GM dropout. Dropout is a widely adopted technique in supervised/imitation learning and rein- forcement learning to improve generalization (Srivastava et al. [2014] Cobbe et al. 2019). For the GMFlow policy,
introduce GM dropout in training to stochastically perturb and diversify x-ID rollouts to make the policy more robust to potential trajectory variations. Given the GM mixture weights Asa of the detached policy TD, we sample a binary mask for each component k == 1,.... and multiply it into A. synchronously across all i = 1, L. The masked weights are then renormalized and used for the detached rollout. By exploring alternative GM modes, this simple technique improves the policy's robustness, yielding better FID on ImageNet 256² 5.2).
Handling FLUX.1 dev teacher. On-policy imitation learning assumes the teacher is robust to out- of-distribution (OOD) intermediate states and can steer trajectories back on track. This generally holds for standard flow models with classifier-fres guidance (CFG) (Ho Salimans, 202T). which exhibit error- coro-correttion behavior (Chidambaram et al. 2024). However, FLUX.1 dev (Black Forest Labs 2024b) is a guidance-distilled model without true CFG and less robust to OOD inputs. To mitigate OOD exposure, we adopt a scheduled trajectory mixing strategy, which rolls out the trajectory using a mixture teacher and student with a linearly decaying teacher ratio (see $ B.2 for details).
5.2 IMAGENET DIT
Our study utilizes two pretrained teachers with the same DiT architecture: standard flow match- ing (FM) DiT (the baseline in Chen et al.] (2025)), and the REPA DiT YYu et al.] 2025). Interval CFG (Kynkäänniemi al. 2024) is applied both teachers to maximize their performance. Each Flllow student is initialized with the teacher weights and then fully finetuned using the -- loss.
Evaluation metries. We adopt the standard evaluation protocol ADM (Dhariwal Nichol 2021) with the following metrics: Fréchet Inception Distance (FID) (Heusel et al.) 2017), Inception Score (IS), and Precision-Recall (Kynkäänniemi et al. 2019).
Comparison of DX and GMFlow policies. As shown in Table | both policies yield strong 1- NFE FIDs after 40k training iterations, with the GMFlow policy consistently outperforming the DX policy by a clear margin. Notably, the DX policy exhibits sensitivity to the hyperparameter N (number of grid points), whereas the GMFlow policy produces consistent results across different values of K (number of Gaussians).
Comparison with prior few-step DiTs. In Table 2 we compare Flow (GM policy with K 32) to prior few-step DiT's on ImageNet 2562: iCT (Song & Dhariwal (2024), Shorteut models (Frans al., 2025), iMM (Zhou et al., (025), MeanFlow (Geng et al. (025)) and the concurrent work FACM (Peng et al.) 2025). FACM (distilled from REPA) improves MeanFlow with auxiliary loss and attains a leading 2-NFE FID, though it still relies on the inefficient JVP operation. In contrast, FFloww uses minimal training framework with no JVP and adaptive loss scalings, yet still outperforms the original MeanFlow DiT across both 1-NFE and 2-NFE generation.
Ablation study GM dropout. From the two bottom rows in Table I we conclude that our standard implementation with a 0.05 GM dropout rate yields better FID and Recall compared to the setting without dropout, confirming the effectiveness of our GM dropout technique.

---

Published as a conference paper at ICLR 2026
Table Quantitative comparisons on COCO-10K dataset and HPSv2 prompt set.
COCO-100 prompts
HPSV2 prompts
Model
Distill method
NFF
Data align.
Pronpt align.
Prof. align.
Teacher align.
Prompt align.
Pref. align.
FID1
pFID
CLIPt
VQAT
HPSv2.1f
FIDL pFID
CLIPt
VQAT
HPSv2.11
FLUX.I dev
50
27.8
34.9
0.268
0.900
0.309
0.284
0.805
0.314
Turbo
GAN
26.7
0.267
0.900
18.5
0.313
Hyper-FLUX
CD+Re
8
29.8
33.3
0.268
0,894
0.309
15.6
22.2
0,285
0.807
0.315
a-Flow (GM-FLUX)
x-ID
8
29.0
35.4
0.268
0.901
0.311
12.6
15.9
0.285
0.810
0.316
Semsellow (FLUX)
VSD+CD+GAN
4
34.1
44.2
0.266
0,879
0.308
23.3
28.2
0.283
0.806
0.318
a-Flow (GM-FLUX)
29.8
0.269
0.903
19.2
0.313
3-Flow (GM-FLUX)
I (data-frec)
29.7
36.2
0.269
0.905
0.310
14.4
19.7
0,287
0.813
0.314
Owen-Image
50x2
34.1
0.282
0.936
0.872
0.309
Owen-Image Lighting
VSD
4
37.5
51.6
0.280
0.935
0.322
15.6
19.7
0.299
0.867
0.328
a-Flow (GM-Qwen)
D-ID
36.0
0.281
0.934
16.6
0.860
0.310
3-Flow (GM-Qwen)
a-ID (data-free)
4
36.0
45.7
0.282
0.936
0,315
12.9
16.8
0.301
0.862
0.312
Table 4: Quantitative comparisons on OnelG-Bench CCang et al.
2025
Model
Distill Method
NFE
Aligmment
Text
Diversity!
Stylet
Reasoning
FLUX. dev
50
0.790
0.556
0.238
0.370
0.257
FLUX Turbo
GAN
8
0.791
0.334
0.234
0.370
0.239
Hyper-FLUX
CD+Re
8
0.790
0.530
0.198
0.369
0.254
J-Flow (GM-FLUX)
-t-DD
0.792 8
0.517
0.234
0.369
0.256
SenseFlow (FLUX)
VSD+CD+GAN a-ID
4
0.776
0,384
0.151
0.343
0.238
3-Flow (GM-FLUX)
0.799
0.437
0.229
0.360
0.251
7-Flow (GM-FLUX)
I-ID (data-free)
0.799
0.460
0,224
0,363
0.249
Owen-Image
50x2
0.880
0.888
0.194
0.427
0.306
Owen-Image Lightnin
VSD
4
0.885
0,923
0.116
0.417
0.311
F-Flow (GM-Qwen)
a-ID (data-free)
0.875
0.892
0.180
0.434
0.298
3-Flow (GM-Qwen)
0.881
0.890
0.176
0.433
0.300
5.3 FLUX. AND QWEN-IMAGE-20B
For rext-to-image generation, we distill the 12B FLUX.1 dev (Black Forest Labs, 2024b) and 20B Qmene (Wu et al. 2025) models into x-Flow students. During student training, we freeze the base parameters inherited from the teacher and finetune only the expanded output layer along with 256-rank LoRA adapters (Hu et al.] 2022) the feed-forward layers. For data-dependent distillation, we prepare 2.3M one-megapixel (IMP) images captioned with Qwen2.5-VL. (Bai et al.) 2025). In the data-free setting, we use only the generated captions as conditioning inputs while keeping the same IMP resolution when initializing the noise.
Evaluation protocol. We conduct a comprehensive evaluation 1024 high-resolution image
generation from three distinct prompt sets: (a) 10K captions from the COCO 2014 validation set (Lin) ltt al. 2014), (b) 3200 prompts from the HPSv2 benchmark (Wu et al. 2023). and (c) 1120 prompts from OneIG-Bence (Chang et al. 2025). For the COCO and HPSv2 sets, we report common metrics including FID (Heusel al. |007 7). patch FID (pFID) (in al. (2024). CLIP similarity (Radford cet al. 2021). VQAScore (Lin et al. 2024b), and HPSv2.1 1 (Wu et al. 2023). On COCO prompts, FIDs are computed against real images, reflecting data alignment. On HPSv2, FIDs are computed against the 50-step teacher generations, reflecting teacher alignment. CLIP and VQAScore measure prompt alignment, while HPSv2 captures human preference alignment. For OneIG-Bened we adopt its official evaluation protocol and metrics. All quantitative results are presented in Table 3 and 4
Competitor models. We compare π- Flow against other few-step student models distilled from the same teacher. For FLUX, we compare against: 4-NFE SenseFlow (Ge et al. 2026), primarily lever- aging variational score distillation (VSD) (Wang et al., 2023), also known as distribution matching distillation (DMD) (Yin et al., 2024b); 8-NFE Hyper-FLUx (Ren et al. 2024), trained with consis- tency distillation (CD) (Song et al. (023) and reward models (Re) (Xu et al. (023); 8-NFE FLUX Turbo, based on GAN-like adversarial distillation (Goodfellow et al. 2014|| Sauer al. 2024b). For Dwen-Image we compare with the N-EEE Qwen-Imaee Lighting based on VSD (ModelTC] (0225)
Note that the 4-NFE FLUXX.1 schnell is distilled from the closed-source FL.UX.1 pro instead of the
7

---

Published as a conference paper at ICLR 2026
Siver Apdin
Fritow Sedtes
With a clear bax hoverkig above showing 9ystem Update: Now gunderviu genereete pliei cheenful rotorvarenng plaststed
.T-Foow (GM-FLUXX AFNFE
FLUX:1 1 dov SD-NFF
SonseFiow Jo (FLUX) N-NFE
PYSTEM UP3ATT
garfeaing gr
futuratic garden.
F-ooow wen) 4-NFE
Owen-Image 50k2-NFE
Qwer-Inagge Lightning 4-NFE
Young woman in dlyuemi neco strecnvear pases on bussinn todyo street at night, surounddd illaminueed Coptored in photograph Mallie individurrs.
(GM-FLUXX
BD-NNN
SenseFlow (FLUX)
(GM-Qwen) 4-NFE
Qwer-Image 50x2-NFE
Qwen-Image Lightniin 4-NFE
greoenslty prodiieeed
Mn golf swllin bright purmmr day. Mor Morionghar stttrthwssly streame betund fic. complementing hwr kriwe phyatipan
TT-loow (GM-FLUX) 4-NFE
FLUX.1 50-NFF
SenseF jow 4-NFE
Sawccenttuatee by fiting liggl adconed' striking gromemo
butted staatosss
she enecutte her mutigwen firesse.
w-Flow (GM-Qwen) 6-NFE
Qwen-Image 50x2-NFE
Owen-Image Lightning 4-NFE
Figure 4: Images generated from the same batch of initial noise by --Flows, teachers, and VSD stu- dents (SenseFlow, Doeen-Image Lightning). --Flow models produce diverse structures that closely mirror the teacher's. In contrast, VSD students tend to repeat the same structure. Notably, Sense- Flow mostly generates symmetric images.
publicly available FLUX.1 dev (Black Forest Labs 2024a). so we do not compare with it directly, but include further discussion in $
Strong all-around performance. As shown in Table and Table 4 π Flow demonstrates strong all- around performance, outperforming other few-step students on roughly 70% of all metrics, without exhibiting obvious weaknesses in any specific area.
Superior diversity and teacher alignment. r-Flow consistently achieves the highest diversity scores and the best teachee-referencee FIDs by clear margins, especially in the 4-NFE setting. These results strongly suggest that π -Fllow effecttvvly avoids both diversity collapse and style drift. As a result, most of its scores closely match those of the teacher, with some even slightly surpass- ing the teacher scores (e.g., prompt alignment and several Qwen-Image OnelG scores). Its strong teacher alignment is also evident in Fig. where -lolow generates structurally similar images to the teacher's 's from the same initial noise.
Comparison with VSD (DMD) students. VSD models are notable for high visual quality, some- times surpassing teachers in quality and preference metries. However, they are widely known to
8

---

Published as a conference paper at ICLR 2026
VELOCITY
VELOCITY
GRANO PHEA
STOP
STOP
---low (GM-FLUX) 8-NFE
FLUX Turbo 8-NFE
Figure 5: Images generated from the same initial noise by x-Flow and FLUX Turbo. T-Flow renders coherent texts, whereas FLUX Turbo underperforms in text rendering.
IT-Flow (GM-FLUX) 8-NFE
Hyper-FLUX 8-NFE
Aoshima's & mesterplece depicts a forest illuminated by morning light.
A painting
destination by
Figure 6: Images generated from the same initial noise by --Flow and Hyper-FLUX, -FFlow pro- duces notably finer details, as highlighted in the zoomed-in patches.
suffer from mode collapse, reflected in our experiments: both SenseFlow and Qwen-Image Light- ning show significant drops in diversity and FIDs. Visual examples in Fig. | further highlight the collapse, where different initial noises produce visually similar images with only minor variations. In contrast, --Flow maintains high quality and diversity without sacrificing either aspect.
Comparison with other students. FLUX Turbo achieves better data alignment FIDs than the teacher due to GAN training, yet its text rendering performance is significantly weaker, as shown in Fig. 5 Meanwhile, Hyper-FLUX often produces undesirable texture artifacts and fuzzy details, whereas 7- Floow achieves superior detail rendering, as shown in Fig. 6
Data-dependent data-free. As shown in Table 3 and Table data-dependent and data-free --Flow models achieve nearly identical results. This demonstrates the practicality of π --Flow in scenarios where high-quality data unavailable.
GMFlow vs. DX policy. Consistent with prior ImageNet findings, the DX policy slightly underper- forms comapred the GMFlow policy (Table 5 highlighting the latter's superior robustness.
Convergence. Figure 7 illlssrrtee the convergence of --Flow (GM-Qwen) over training iterations. Both FID and Patch FID scores initially improve rapidly, outperforming Image Lightning within the first 400 iterations, and continue improve steadily thereafter. This contrasts with previ- ous GAN or VSD-basee methods that often require frequent checkpointing and cherry-picking (Ge al. 2026), demonstrating the scalability and robustness of our approach.
Inference time. To validate that the policies are indeed fast enough so that the overhead is negligible compared to shorteut-predicting models, compare the inference times of 4-NFE --Flow models and Qwen-wweeaee Image Lightning in Table 6 For F-loow, each policy generation step (with one network evaluation) is followed by 32 policy integration substeps average. The results in Table 6 show that 32 policy substeps cost around 15 ms in total, which is only 3% of the network time. Therefore, the overall speed of -Fllow is par with predicit models.

---

Published as a conference paper at ICLR 2026
Table 5: Comparisons between DX and GMFlow policies on
FID
texa-to-image generation.
20
Qwentightring
Patch FID
IPPS prompts
OnelG-Bench
Policy
Teacher
FID1
pFID)
HPSv2.11
Text!
Diversityt
16
Qwen Lightning
DX (N= 10) FLUX
14.9
20.9
0.313
0.397
0.225
GM (K- FLUX
14.3
19.2
0.313
0.437
0.229
DX (N 10) Owen-Image
12.7
17.0
0.306
0.869
0.185
2K
4K
6K
0K
GM (K Qwen-Image
12.8
16.6
0.310
0.892
0.180
Teacher-referenced FID and Patch FID of GM-Qwen evaluated on HPSv2 prompts.
Figure 7:
Table 6: Per-NFE inference time of π-Flow models and the shortcut-preet shortcut-predicting model (Qwen-Image Lightning). Tested on an A 100 GPU with 12 CPU cores (3.0 GHz).
Model
Network time (sec)
Policy time (sec)
Qwen-Image Lightning
0.465
a r-Flow (DX-Que (DX-Qwen)
0.464
0.015
a-Flow (GM-Qwen)
0.465
0.014
RELATED WORK
Prior work on diffusion model distillation primarily focuses on predicting shortcuts towards less noisy states, with training objectives ranging from direct regression to distribution matching.
Early work (Luhman & Luhman 2021) directly regresses the teacher's s ODE integral in a single step, but suffers from degraded quality, since regressing To with Eg loss tends to produce blurry results.
Progressive distillation methods (Salimans & Ho 2022 Liu et al. 2023 2024 Frans et al., 2025) make further improvements via multi-stage process that progressively increases the student's step size and reduces its NFE by regressing the previous stage's multi-step outputs with fewer steps, yet this introduces error accumulation.
Consistency-based models (Song et al. 2023 Gu et al. 2023 Kim et al. 2024 Song & Dhari- wal 2024 Geng et al. 2025) Boffi al.
2025 implicitly impose velocity- based regression loss, which improves quality compared to x based regression. However, the flow velocity of a shortcut- predicting student must be constructed implicitly using either inaccurate finite differences expen- sive Jacobian-vector products (JVPs). Moreover, their quality is still limited due to accumulation of velocity errors into the integrated shorteut. Therefore, in practice, consistency distillation of- ten augmented with additional objectives to improve quality (Ren et al. 2024 further complicating training.
Zheng et al.) 2025
Conversely, distribution matching approaches Yin et al. 2024b a Sauer et al. 2024b Zhou et al.) 2024 Luo et al. 2024) Salimans et al. 2024 Zhou et al. 2025) adopt score matching and adversar- ial training align the student's output distribution with the teacher's. ss. The VSD objective achieves superior quality but risks diversity loss due to mode collapse; GAN and SiD objectives balance qual- ity and diversity but cause style drift. Their common reliance auxiliary networks introduces additional tuning complexity and may lead to stability issues at scale (Ge et al. 2026).
CONCLUSION
We introduced policy-based flow models (r-FFow), novel framework for few-step generation in which the network outputs a fast policy that enables accurate ODE integration via dense substeps reach the denoised state. To distill --Flow models, we proposed simple on-policy imitation learn- ing approach that reduces the training objective to single & loss, mitigating error accumulation and quality-diversity trade-offs. Extensive experiments distilling ImageNet DIT, FLUX. I-2B, and models show that few-step m-Floww consistently attain teacher-level image qual- ity while significantly outperforming competitors in diversity and teacher alignment. π- F-Flow offers scalable, principled paradigm for efficient, high-quality generation and opens new directions for future research, such as exploring more robust policy families, improved distillation objectives, and extensions to other applications (e.g., video generation).
10

---

Published as a conference paper at ICLR 2026
Reproducibility statement. To facilitate reproduction, describe the detailed training procedures in Algorithms 2 and S
and list all important hyperparameters in $
Acknowledgements. This project Qualcomm Innovation Fellowship and partially done while Hansheng Chen was an intern at Adobe Research. We would like to thank Jianming Zhang and Hailin Jin for their great support throughout the internship, and Xingtong Ge for the help evaluating SenseFlow.
partially done while Hansheng Chen supported by the
REFERENCES
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In ICLR, 2023.
Shuai Bai, Keqin Chen, Xuejing Liu, Jialin Wang. Wenbin Ge, Sibo Song, Kai Dang, Peng Wang, Shijie Wang, Jun Tang. Humen Zhong, Yuanzhi Zhu, Mingkun Yang. Zhaohai Li, Jianqiang Wan, Pengfei Wang, Wei Ding. Zheren Fu, Yiheng Xu, Jiabo Ye, Xi Zhang, Tianbao Xie, Zesen Cheng, Hang Zhang, Zhibo Yang. Haiyang Xu, and Junyang Lin. Qwen2.5-vl technical report, 2025. URL
wtt.///.iorrroooooooooo hbtp://...222222/2 13923
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In NeurlPS, NIPS' pp. 1171-1179, Cambridge, MA, USA, 2015. MIT Press.
Black Forest
Labs. Flux.1 [schnell].
https///huggangacccccc//ppppps 2024a.
Black Forest Labs. Flux.
2024b.
Nicholas Matthew Boffi, Michael Samuel Albergo, and Eric Vanden-Eijnden. Flow map matching with stochastic interpolants: A mathematical framework for consistency models. TMLR, 2025. ISSN 2835-8856. URL httL::t:://// httttpss///ppnrevveew net
Jingjing Chang, Yixiao Fang, Peng Xing, Shuhan Wu, Wei Cheng, Rui Wang, Xianfang Zeng. Gang Yu, and Hai-Bao Chen. Oneig-bench: Omni- Omni-dimensional nuanced evaluation for image generation. In NeurlPS, 2025.
Hansheng Chen, Kai Zhang, Hao Tan, Zexiang Xu, Fujun Luan, Leonidas Guibas, Gordon Wet- zstein, and Sai Bi. Gaussian mixture flow matching models. In /CML, 2025.
Muthu Chidambaram, Khashayar Gatmiry, Sitan Chen, Holden Lee, and Jianfeng Lu. What does guidance do? a fine-grained analysis in a simple setting. In NeurlPS, 2024. URL https:
7/ //openreveew. www.:ee/foorne www.:n/oooommmddddeeeaaaaa
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generaliza- tion in reinforcement learning. In ICML, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, pp. 248-255, 2009. doi: 10. 101109999PP2220092226646
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In ICLR, 2022.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan (eds.), NeurlPS, 2021. URL net forum? d=AANuCvzaVE..
htttss://ppenret
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer. Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion En- glish, Kyle Lacey. Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. ICMI, 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. In ICLR, 2025. URL htt.:htttt/ppeeeeeeeeeee net/forum oof.forrmmmii0000866....... 2B61nXc5

---

Published as a conference paper at ICLR 2026
Xingtong Ge, Xin Zhang, Tongda Xu, Yi Zhang, Xinjie Zhang, Yan Wang, and Jun Zhang. Sense- flow: Scaling distribution matching for flow-based text-to-image distillation. In ICLR, 2026.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, J. Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. In NeurIPS, 2025.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger (eds.), NeurIPS, volume 27. Cur- ran Associates, Inc., 2014. URL/attss://ppoeeeeeee www.:i/gssseurrii..cccooo.. ccppppe files/
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free distillation of denoising diffusion models with bootstrapping. In ICML Workshop, 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In NeurlPS, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS Workshop, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurlPS, 2020.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In ICLR, 2022. URL
https::/cppeneeviett htt.://wwwwweeeeeteeeeeeeeeeeeee dew.emenddnneeeeeeeeeeee
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. NeuriPS, 2022.
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In CVPR, 2024.
Dongjun Kim, Chich Misih--sinnha Lai, Wei- Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning proba- bility flow ODE trajectory of diffusion. In ICLR, 2024. URL Forum? 1d=ym feDDD
net.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2014.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. In NeurIPS, 2019.
Tuomas Kynkäääniemii, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in limited interval improves sample and distribution quality in diffusion models. In NewrlPS, 2024.
Shanchuan Lin, Anran Wang, and Xiao Yang. Sdxl-lightning: Progressive adversarial diffusion distillation, 2024a. URL
13929
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and Lawrence Zitiick. Microsoft coco: Common objects in context. In David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars (eds.), ECCV, pp. 740-755, Cham, 2014. Springer International Publishing. ISBN 978-3-319-10602-11
Zhiqiu Lin, Deepak Pathak, Baiqi Li, Jiayao Li, Xide Xia, Graham Neubig, Pengchuan Zhang, and Deva Ramanan. Evaluating text-to-visual generation with image-to-text generation. In ECCV, 2024b.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In ICLR, 2023. URL
12

---

Published as a conference paper at ICLR 2026
Qiang Liu. Rectified flow: A marginal preserving approach optimal transport, 2022. URL
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL htttps://pperreiieeeee -
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and qiang liu. Instaflow: One step is enough for hugh-quality diffusion-bused text-to-image generation. In ICLR, 2024. URL https:/ openreview. net/ http:/ooummmmm wet.:e/ooooommdd.4/yy.. (irrrummb2dd114222bb23
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), NeurIPS, 2022. htLhttss:
//openzeview.c www.:eet/ooore http://vwweee..eeeoomooooeeoom. /ew.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models, 2023. URL https://arxiv: org/abs/223 /222//2211 01095
Eric Luhman and Troy Luhman. Knowledge distillatioon in iterative generative models for improved sampling speed, 2021. URL
02388
Weijian Luo, Zemin Huang. Zhengyang Geng, J. Zico Kolter, and Guo jun Qi. One-step diffusion distillation through score implicit matching. In NeurIPS, 2024.
Nanye Ma, Mark Goldstein, Michael Albergo, Nicholas M. Boffi, Eric "anden-Eijndens and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In ECCV, 2024.
ModelTC.
Owen-image-lightinng Owen Iraan-Imaae-aae Lightning 2025.
httppa///gttuub.cco//oodlell
William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, 2023.
Yansong Peng, Kai Zhu, Yu Liu, Pingyu Wu, Hebei Li, Xiaoyan Sun, and Feng Wu. Flow-anchored consistency models, 2025. URL htttt:://aaxiii// org/ abs/ 03738
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML, pp. 8748 8063-776,,488 2021.
Yuxi Ren, Xin Xia, Yanzuo Lu, Jiacheng Zhang, Jie Wu, Pan Xie, XING WANG, and Xuefeng Xiao. Hyper-SD: Trajectory segmented consistency model for efficient image synthesis. In NeurIPS, 2024.
Hans Richter, Parameterfreie abschätzung und realisierung von erwartungswerten. Blätter der DGVFM, 3(2):147-162, Apr 1957. ISSN 1864-0303. doi: 10.1007/BF02808864 URL
https://dof. 0011
Stephane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction no-regret online learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudík (eds.), AISTATS, volume 15 of Proceedings of Machine Learning Research, pp. 627-635, Fort Lauderdale, FL, USA, 11-13 Apr 2011. PMLR. URL URLL/tttss//pprceeding press/vi5/ ross llos.lla... html
Tim Salimans and Jonathan Ho. Progressive distillation fast sampling of diffusion models. In ICL,R, 2022.
Tim Salimans, Thomas Mensink, Jonathan Heek, and Emiel Hoogeboom. Multistep distillation of diffusion models via moment matching. NewrIPS, 37:36046-36070, 2024.
13

---

Published as a conference paper at ICLR 2026
Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rom- bach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. In SIGGRAPH Asia, SA '2, New York, NY. USA, 2024a. Association for Computing Machin- ery. ISBN 9798400711312. doi: 10. 10.11457368052886668777 3687625. URL httpt:t///o....ooooo wtw.odiooog//... 1145/3680 3687625
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. In ECCV, pp. 87-103, 2024b.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML. pp. 2256-2265, 2015.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In ICLR, 2024. URL http://swppppeeeeeee.....e///rrr
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NeurIPS, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In JCML, 2023.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. JMLR, 15(1):1929-199 1958, January 2014. ISSN 1532-4435.
Vladimir Tchakaloff.
Formules de cubatures mécaniques à coefficients non négatifs. Bulletin des Sciences Mathématiques, (1(2):123-134, 1957.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), NeurIPS, volume 30. Curran Associates, Inc., 2017.
Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Prolific- dreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. In NeurlPS, 2023.
Chenfei Wu, Jiahao Li, Jingren Zhou, Junyang Lin, Kaiyuan Gao, Kun Yan, Sheng ming Yin, Shuai Bai, Xiao Xu, Yilei Chen, Yuxiang Chen, Zecheng Tang, Zekai Zhang, Zhengyi Wang, An Yang, Bowen Yu, Chen Cheng, Dayiheng Liu, Deqing Li, Hang Zhang, Hao Meng, Hu Wei, Jingyuan Ni, Kai Chen, Kuan Cao, Liang Peng, Lin Qu, Minggang Wu, Peng Wang, Shuting Yu, Tingkun Wen, Wensen Feng, Xiaoxiao Xu, Yi Wang, Yichang Zhang, Yonggiang Zhu, Yujia Wu, Yuxuan Cai, and Zenan Liu. Qwen-image technical report, 2025. URL httpt:tt/:///at http:::s///xxvv.orrog. abs/
2508. 02324
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis, 2023. URL httph:ttt///aiiiio wwteeiiiv.rrgaaaae wtwp//a//////aaoeaeeeo.oeeeeeeee 2306. 09341
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: learning and evaluating human preferences for text-to-image generation. In NeurIPS, pp. 15903-15935, 2023.
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William Freeman. Improved distribution matching distillation for fast image synthesis. In NewrIPS, 2024a. URL
n ttt.://ooormmmooooooocccaaaaaa
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Frédo Durand, William Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In CVPR, 2024b.
14

---

Published as a conference paper at ICLR 2026
Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers is casier than you think. In ICLR, 2025.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator: In ICLR, 2023. URL
net forst//orrmmiii www.:nne/frrmmmi--cccooo. wdd 7hfb4 6P
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. In NeurlPS, 2023.
Kaiwen Zheng, Yuji Wang, Qianli Ma, Huayu Chen, Jintao Zhang, Yogesh Balaji, Jianfei Chen, Ming-Yu Yu Liu, Jun Zhu, and Qinsheng Zhang. Large scale diffusion distillation via score- regularized continuous-time consistency, 2025. URL 08431
Linqi Zhou, Stefano Ermon, and Jiaming Song. Inductive moment matching. In ICML, 2025.
Mingyuan Zhou, Huangjie Zheng, Zhendong Wang, Mingzhang Yin, and Hai Huang. Score identity distillation: exponentially fast distillation of pretrained diffusion models for one-step generation. In ICMI,, 2024.
15

---

Published as a conference paper at ICLR 2026
Algorithm 2: Data-dependent on-policy x-ID training loop with time shifting.
Input: NFE, teacher Gs. data cona-connition distribution p(ao, c). shift m
Output: Student G Initialize student params Ø
2 &- NFE NFE
2 ............ }
Il can be adjusted to reduce final step size
for finite samples wo; c. pl a0 c). NO,, F), U(O,1) do
Fire min Tyx Tux E and Tac 27'
Tdst f max Tást
Tha E SU (0) and Táx < Tec
Lve <- I-(m 1)
I1 time shifting
Esser et al. 2024
Fte - Ole 2o + T
Co lue.
ND - stopgrad(T)) or Tn dropout(stopgrad(())
10
Ls +-
11
for finite samples T ~ U UTTo, Twe) do
12
t 1=)--1)-
13
x: - Dia + I To(as, t)
14
L$ t Ls + GG(,t n mxv, t)
II can be replaced with Eq. Il ooptimizer step
6
15
+- Adam( voLo)
Algorithm 3: Data-free on-policy 7-ID training loop with time shifting.
Input:
NFE, teacher Co. condition distribution p(c), shift
Output: Student G Initialize student params
2 for finite samples e ~ p(c). æ1 ((0, do
Tue 4 1, ta +
L.
while Tax do
Tiu - Tse NTE
Il can be adjusted to reduce final step size
fas
I/ time shifting Esser al. 2024
D)
(ælatu lane,
TD <- stopgrad or TD +- dropout(stopgraa(())
10
for finite samples 7 U(Ta, Twa) do
11
t t 1+(m-1)r
12
2i f- Dim + St TD aD
13
L &- Lx
(a L, <(x,, 1)l
Il can be replaced with Eq.
Wiss &- Bie + h To( dl
14
15
Tue +- Teas Luo f Laa
16
+ Adam( VoLo)
Il optimizer step
A USE OF LARGE LANGUAGE MODELS
In preparing this manuscript, used large language models (LLMs) general-purpose writing assistants for grammar corrections, rephrasing, and clarity/coccision edits. All LLM LLM-suggested edits were reviewed and verified by the authors, who take full responsibility for the final manuscript.
B ADDITIONAL TECHNICAL DETAILS
B. GM TEMPERATURE
Inspired by the temperature parameter in language models, we introduce similar temperature pa- rameter for the GMFlow policy during inference. Let T' >> 0 be the temperature parameter. Given
C-dimensional GM velocity distribution q(ul.
Et=1 AuN(u; pk, s2 s²I), the new GM prob-
16

---

Published as a conference paper at ICLR 2026
(a) Off-policy:
decay
(b) Mixture of teacher and detached policy
decay
(c) On-policy: teacher ratio=0%
teacher ratio-100%
Xtre 88
X
(see Fig. 3)
Xtaet
Polley If
Policy x
stopgrad dropout)
Datachee polisy rollout
TeacherGy Gg
Detached pollcy Ri
Folicyxt
Figure 8: Three stages of scheduled trajectory mixing. (a) Off -of-p behavior cloning with a teacher ratio of 1. (b) Mixed teacher and detached- polliy segments with decaying teacher ratio. (c) On- policy imitation learning with a teacher ratio of 0 (Fig.
ability with temperature T is defined
a (2) JRe q (ul du
qr (u (
(3)
Although @r ( does not have general closed-form expression, it can be approximated by the following expression, which works very well practical implementation:
K
qr (ul
N $
(4)
AT
For the distilled FLUX and Owen-Image models, we set T = 0.3 for 4-NFE generation and T = 0.7 for 8-NFE generation. An exception is that we do not apply temperature scaling the final step, as we found this can impair texture details. As shown in Table 7 ablating GM temperature from the 4-NFE GM-FLUX leads to degraded teacher alignment.
B.2 SCHEDULED TRAJECTORY MIXING FOR -DISTILEDD TEACHERS
To reduce -iitr---istribttion exposure in imitation learning, scheduled sampling (Bengio et al., 2013) stochastically alternates between expert (teacher) and learner policy during trajectory integra- tion, decaying the expert probability from I to O. However, naively applying it to r-ID is impractical because the teacher flow model Go is much slower than the network-free policy TD.
To maintain constant compute throughout training, we introduce a scheduled trajectory mixing strat- egy. Since the teacher is slow, we fix the total number of teacher queries, allow each query to cover a coarse, longer step initially, and gradually shrink the teacher step size while filling the gaps with the fast policy Tp. As shown in Fig. (a), training initially adopts fully off-policy teacher trajectory (behavior cloning). At the beginning time la of each teacher step, we roll in the learner policy m, integrate it over the same interval from i, to lg. and match its average velocity to the teacher velocity with the &g loss:
ts x (1 ( t) dt
Ls = E
Go ( ta)
(5)
2
to to Lo
& (b)), we then mix teacher and detached-policy segments while using the same loss, and linearly decay the teacher ratio-the the sum of teacher step lengths divided by the total interval length tare fast. Finally, when the teacher ratio reaches 0, training reduces to on- policy IDD All teacher step boundaries (starts and ends) are randomly sampled within the interval [tass twee] under the teacher ratio constraint, that step sizes and locations vary while the total teacher-covered length follows the current ratio schedule.
As training progresses (Fig.
We apply scheduled trajectory mixing exclusively when distilling the FLUX.1 dev model, as it lacks real CFG. Since omitting CFG doubles the teacher's speed, we increase the number of intermediate samples (teacher steps) 4 accordingly.
B.3 MICRO-WINDOW VELOCITY MATCHING
For on-policy T-ID, in practice we found that replacing the instantancous velocity matching loss in Algorithm [ with modified average velocity loss over miero time window generally benefits
17

---

foreground, with Published as a conference paper at ICLR 2026
Table Ablation study on 4-NFE --Flow (GM-FLUX), evaluated on the HPSv2 prompt set using teacher-referenced FID metrics (reflecting teacher alignment).
GM
Micro
FIDL
PFID
temperature window
14.3
19.2
14.9
20.1
14.6
20.3
FLUX.1 dev
FLUX.1 dev
43-NFE
128-NFE
with aerodynamic design shown motion against backdrop of skyline The car features sharp angles, large wheels with orange accents, and prominent front grille. The cityscape includes tall buildings with illuminated windows, and the sky is painted with hues of orange and blue as the sun sets. The lighting is warm and golden, the sun setting behind city, casting a glow over scene. The car is positioned the skyline in the background, creating sense depth and movement.
futuristic, sleek
Figure 9: The 128-NFE FLUX.1 dev often generates blurry images, whereas the 43-NFE FLUX.1 dev reduces the blur and produces sharper edges.
training. The modified loss is defined as:
27
[" A-Ar Ls = E 2 Go(
1) de
(6)
Al
where At is the window size. We At 3/128 (three policy integration steps) for all FLUX.1 and Qwen-Image experiments.
The benefits micro-window velocity matching are threefold:
It generally smooths the training signal, reducing sensitivity to sharp local variations in the teacher trajectory.
It stabilizes the less robust DX policy. In the ImageNet experiments, we observe that training with the DX policy diverges without this modification.
With At 3/128, the policy effectively mimics teacher sampling with 128 R 43 steps instead of 128 steps. For the guidance-distilled FLUX.1 dev model, we observe that the teacher often generates blurry images using 128-step sampling, while 43-step sampling yields sharper results (see Fig. (9). This behavior is inherited by the student, so micro-window, velocity matching helps reduce blur.
3
18

---

Published as a conference paper at ICLR 2026
Table 8: Hyperparameters used in the ImageNet experiments.
I-NFE
2-NFE
GM-FM (K = 32)
GM-REPA
GM-REPA (K = 32)
DX-REPA
DX-REPA
DX-REPA (N = 40)
GM-REPA (K 32)
(R = 8)
(N = 10)
(N 20)
GM dropout
0.05
0.05
0.05
0.05
intermediate states
2
2
2
Window size (raw)
10/128
5/128
3/128
Shift
1.0
1.0
1.0
1.0
Teacher CFG
2.7
3.2
3.2
3.2
3.2
3.2
2.8
Teacher CFG interval
e [..0.6]
(0,
E [0,0.7]
[0.
[0.
[0,0.7]
[0.
Learning rate
5c-5
5c-5
5c-5
5c-5
Sc-5
5c-5
5c-5
Batch size
4096
4096
4096
4096
4096
4096
4096
# of training iterations
in Table 2
140K
140K
24K
EMA param
7.0
7.0
7.0
7.0
7.0
7.0
7.0
Karras et al. 2024
Table 9: Hyperparameters used in FLUX and Qwen-Image experiments.
4-NFE
8-NFE
GM-FLUX
GM-Qwen (K 8)
DX-FLUX
DX-Qwen (N= 10)
GM-FLUX (K 8)
(K' 8)
(N 10)
GM dropout
0.1
0.1
GM temperature
0.3
0.3
0.7
of internediate states
4
Window size (raw) Ar
3/128
3/128
3/128
3/128
3/128
m
3.2
Final step size scale
0.5
0.5
0.5
0.5
Teacher CFG
3.5
4.0
3.5
Leaming rate
lc-4
1o-1
1c-4
1c-4
Batch size
256
256
256
256
256
of training itcrations
3K
3K
of decay itcrations
B.2
2K
2K
EMA
7.0
7.0
7.0
7.0
7.0
Karras et al. 2024
As shown in Table ablating the micro window trick from the 4-NFE GM-FLUX leads to degraded teacher alignment.
B.4 TIME SAMPLING
For high resolution image generation, Esser et al. (2024) proposed time shifting mechanism to rescale the noise strength. Let 7 the pre-shift raw time and be the shift hyperparameter, the shifted time is defined as t == (+m-)))
mr
Following this idea, ID samples times uniformly in raw-time space and then applies the shift to remap those samples. Detailed time sampling routines are given in Algorithms 2 and 3
For FLUX.1 and Qwen-Image, we use fixed shift of FLUX.1's official dynamic shift at IMP resolution. In addition, several diffusion/foww models
3.2, which is rounded approximation
reduce the noise strength the final step to improve detail (Karras et al., 2022 Wu et al.) 2025). Accordingly, for FLUX.1 and lwenn-Imae we halve the final step size (relative to previous steps) in raw-time space.
19

---

Published as a conference paper at ICLR 2026
Table 10: FLUX.1 schnell evaluation results on COCO-10k dataset and HPSv2 prompt set.
COCO- 10k prompts
HPSv2 prompts
Model
Distill method NFE
Data align.
Prompi align.
Pref. align.
Prompi align.
Pref. align.
FIDI pFID!
CLIPT
VQAt
HPSv2.11
CLIPT VQAT
HPSV2.11
FLUX.1 schnell
GAN
21.8
29.1
0.274
0.913
0.297
0.297
0.843
0.301
FLUX.1 schnell
FLUX.1 dev
IK-Flow (GM-FLUX)
4-NFE
50-NFE
4-NFE
The president being abducted by aliens.
puppy driving
Spidermsn
portrait a molting skull with intricate abstract dotails, crnated by Toath Whop Beople, and Den Mumtoord using Octane Render.
The image depicts beautiful goddess of spring wearing wreath and flawy green skirt, created by artist wiop.
Figure 10: Typical failure cases of FLUX.1 schnell. For reference, we also show the corresponding FLUX.1 dev and F-Flow results from the same initial noise.
C ADDITIONAL IMPLEMENTATION DETAILS AND HYPERPARAMETERS
All models are trained with BF16 mixed precision, using the 8-bit Adam optimizer (Kingma & Ba 2014] Dettmers et al. 2022) without weight decay. For inference, EMA weights with a dynamic moment schedule (Karras et al. 2024). Detailed hyperparameter choices are listed in Table andd
20

---

Published as a conference paper at ICLR 2026
Element-wise factorization (Hf x x 16) x
Pixel-wise factorization (# x W) 16
Patch-wise factorization (#/2 x W22 x
Fuzzy olay crestures playing a game in hesvily bilted green pasture.
Figure II: Comparison of different L x C' factorization schemes for the GMFlow policy, evaluated on the toy models after 4000 optimization iterations. The default pixel-wise factorization VAE latent channels) proposed by Chen et al. (2025) produces neutral colors and detailed textures. In contrast, element-wise factorization (C = 1) leads to over-saturated colors, high contrast, and over-smoothed textures, while patch-wise factorization (C 2 x 2 x VAE latent channels) results "confetti" artifacts and noisy textures.
C.I DISCUSSION ON GMFLOW POLICY HYPERPARAMETERS
For the GMFlow policy, we observed that the hyperparameters suggested by Chen et al. (2025)
(K = C VAE latent channel size) generally work well. These parameters play important roles in balancing compatibility, expresiveness, and robustness. A larger K improves expresiveness but impairs compatibility as if may complicate network training. A larger C: improves robustness (since GMFlow models correlations within each C--dimensional chunk) but impairs expresiveness (raises the theoretical K N C' bound), In addition, improving expresiveness may generally compromise robustness, due to the increased chance of encountering outlier trajectories during inference.
To further justify the design choice of pixel-wise factorization in GMFlow (where the latent grid is factorized by L H x W, C VAE latent channel size), we conduct toy model experiments to test alternative factorizations. In each experiment, we initialize set of GMFlow parameters Aik, Mile with K 32 components, and directly optimize them to overfit the FLUU teacher's behavior over the entire time domain (0, 1], using simple T-ID (Algorithm I a fixed initial noise I1 and a fixed text prompt. This setup isolates the inductive bias the policy itself, without any influence from the student network. As shown in Fig. [ pixel-wise factorization achieves the best results within 4000 optimization iterations, producing neutral colors and rich textures, and is therefore the most suitable choice for the GMFlow policy.
D DISCUSSION ON FLUX.1 SCHNELL
The official 4-NFE FLUX.1 schnell model (Black Forest Labs
2024a) (based on adversarial distil- lation (Sauer et al. 2024a)) is distilled from the closed-source FLUX.1 1 pro instead of the publicly available FLUX.1 dev. This makes a direct comparison to the student models in Table iiequitable.
For reference, nevertheless, we include the COCO-10k and HPSv2 metrics for FLUX.1 schnell in Table IO, These metrics reveal a trade-off: while FLUX.1 I schnell achieves significantly better data and prompt alignment than FLUX.1 dev, its preference alignment is substantially weaker than FLUX.1 dev and all of its students.
To validate this observation, FLUX) was compared against FLUX.1 schnell on 200 images generated from HPSv2 prompts. --Flow was preferred by users 59.5% of the time, aligning with the HPSv2.1 preference metric. Furthermore, qualitative comparisons in Fig. 10 reveals that FLUX.1 schnell is prone to frequent structural errors g., miissinggexttrr/disttoteeed limbs), whereas Flloow maintains coherent structures.
conducted human preference study. Our 4-NFE --Flow (GM-
E PROOF OF THEOREM 1
We will prove that a GM with N- -CC components suffices for approximating any N N-step trajectory in RC by first establishing Theorem 2 and then applying the Tchakal-Th theorem show
21

---

Published as a conference paper at ICLR 2026
that mixture of N C Dirac deltas satisfy all ODE moment equations, which finally leads to N C Gaussian components.
Theorem 2. Given pairwise distinct times iN,...,ty (0, 1] and vectors 2i Lt E RC for n = 1,...,N, .N. there exists a probability measure p(dx)) on I, such that Eq holds at t = tn for every N.
E.1 MOMENT EQUATION
For every / E (0, 1], the ODE moment equation has the following equivalent forms:
Et
2o
3t R
p
p(dxoo/z))
t
p(dxoo (14)
{
2o p(dlzol(as)
R
RC
20 2e + tæ;
4
p pld2oo/
RRC
1.
(Io
2t + txt W (Irr) a, p((())(
RRC
tp(x+)
(00 + tă, W (x ag of p(dxo)) = 0.
RC
Let (t, (xo 2t tab ; W M(x; Do ofI) be kernel function. The above equation can be written as multivariate homogeneous Fredholm integral equation of the first kind:
/ g(t zo)p(dxo)
(8)
To prove Theorem 2
we need to show that there exists a probability measure p(dxo) on IRC that solves the Fredholm equation f tu for every = 1, , N.
E.2 UNIVARIATE MOMENT EQUATION
To prove the existence of a solution to the multivariate Fredholm equation, we can simplify the proof into univariate case by showing that element-wise probability factorization p(dæz))
e p(dr.o) exists that solves the Fredholm equation. In this case, Eq. 7 can written as:
= 1,2, CC
| ( Jst + + (a a0 p()p()r;))) II N (1 10 of p()
Vi = 1,2, ,, [ In t (2,r ar p(da0)))) =
(9)
To see this, we need to prove that there exists a probability measure p(zo) on E that solves the following univariate Fredholm equation t t,, for every n. I.... N:
g(t, g(t,zo)p(dzo) = O, R
(10)
where z0) (Do x, (N of) is the univariate kernel function.
E.3 CONVEX COMBINATION
Lemma 1, Define the vector function:
7: R + N
7(2)) = (g(t), æo), g(), .o), zo)).
(1I)
Then, the zero vector lies in the convex hull in R i.e.:
0 E conv 7z00) 2o R RRjRR
(12)
22

---

Published as a conference paper at ICLR 2026
Proof. Define S = conv| ((zo) E R }. Assume for the sake of contradiction that 0 & S. By the supporting and separating hyperplane theorem, there exists w * E RN,, such that:
AX E S, (w, x) S 0.
(13)
In particular, this implies that:
VrO E R, (w, ((o)) SO.
(14)
Define h(mo) = (w, ((xo)) En= (nn 20). Recall the definition of g(t 2o) that:
x(t,xo)
(Io
It W apto. of)
20
3t + thk
(It
exp
(15)
V2Tf2
20
Let n² be an index with wn. # for which the exponential term above decays the slowest, i.e.:
of
of
min
Wn- + #0%
(16)
20
20
Note that since is monotonic, for every n # n² with Wn # O, we have aig 20
o
Therefore,
207
20
as |zo| 20, h(xo) is dominated by the n"-th component, iee.:
30 It,. + In- 2t. exp (Nig- -tt, 2o)² + 0(1)-
h(x0) Wn-
(17)
V2Tt2.
20
Because the term xo - It., + +ba it,,, changes sign between -00 and +o0, h(x)) takes both positive and negative values. This contradicts the hyperplane implication that hxxo) S 0. Therefore, conclude that 0 E S.
By T Carathéodory's theorem, the zero vector can expressed as a convex combination of at most N +1 points on (ato). Therefore, there exists finite-support probability measure pl (d2z) consisting of N + Dirac delta components that solves the univariate Fredholm equation at tn for every 1, ... .N, completing the proof of Theorem 2
E.4 N C COMPONENTS SUFFICE
Richter's extension to Tchakaloff's theorem states as follows.
Theorem (Richter (1957); Tehakaloff (1957)) Let be finite-dimensional space of measurable functions on RC For some probability measure p(dæo) on R², define the moment functional:
Agg] [ g(zn)p(ddro)
R,
(18)
Then there exists a K -attmii measure p" (dæo) = Ek Agder (dao) with A > and < dim V
such that:
Vg E V, A[g] = / ((zo))p((dx))
= Σ
Ang(uz)
(19)
k=1
with the scalar function gi (fn, zu) (Iio - 245 + t W(=; 0y I) there exists probability measure p(dma) such that f2p 9: (b,, 3g p()))))))))))))) = for every i, n. Then, by the Richao--Tchakalff theo- rem, there also exists K-atomic measure with K S dim V < N C that satisfies all the moment equations. By taking the upper bound, this implies the existence
By Theorem we know that for >> = spant gi (tns xo) 1i= 1,... C, In
an C-atomic probability
measure (dxx)) = W.C A (dæa) with A > O, NNC A = 1 that solves the Fredholm equation (Eq. (8) & Ln for every n 1, N.
23

---

Published as a conference paper at ICLR 2026
Finally, since N (x,; an of is continuous, the N C' Dirac deltas in p" (dæo)) can be replaced by a mixture of C' narrow Gaussians, such that tn is approximated arbitrarily well for every n, Le.:
Vn 1, N, Jgo top(ddan) at..))
Ft,
lim a-+0
tn
p(dzo)= (d
fgo mop(dxolaa..)
Lt.,
p(dea)=Eis An5ug (dzo)
it..
(20)
This completes the proof of Theorem 1
F DERIVATION OF CLOSED-FORM GMFLOW VELOCITY
In this section, provide details regarding the derivation of closed-form GMFlow velocity, which was originally presented by Chen et al. (2025) but not covered detail.
Given the u-based GM prediction ((.)).....) = Σ A&N(u: uu; μk, 2) with A E f. p4 E TRC E RR,, we first convert into the wo-based parameterization by substituting u density function, which yields:
2t. 20
into the o
( xtne) = Σ k=1
AzNi Mxi sT
(21)
with the new parameters Mxz Btys Ee Mag and 8x on 8. Then, for any t bwo and any Rt E RC, the denoising posterior at (o:r, 1) is given by:
p(x, xo)
g(<olxx)) Z Z.p(an p( (00)
of (20 (200)))))))
(22)
where Z is normalization factor dependent on ata. Dz, tw t. Using the definition of forward diffusion p(æ://zo) = N (ax; 0xa0, I),,, we have:
N (1 an00,
q( xx()
q( (mm))
Z NN (t ar No. oi I)
N xpi 1 xt. c
N Z'
Foi ceem
af
1
.
9(200) (x Øoi
2"
Olbtr Pta
where
o?
2
Ç
of
o.
The result can be further simplified into a new GM:
K Σ
((xolo1)
ALW wo:
(r: sk
(23)
1=1
24

---

Published as a conference paper at ICLR 2026
with the following parameters:
s²²
(24)
$2C +
11 8 V + Mskk
(25)
22 +
expox
A
(26)
Σ k= expo
where the new logit dj. is given by:
d log Ak + Mak (v {(HA
(27)
22 +1
Finally, the closed-form GMFlow velocity at (a 1) is given by function n:
g (x [00
T: RC x R RR x(x,,t)
K 21 Ek=1 Alp41
(28)
Extension discrete support. The closed-form GMFlow velocity can also be generalized dis- crete support by taking lim,, -0 ()), which yields the simplified parameters:
(29)
(v V 2 2(pere)
a = log Av + pxx
(30)
G ADDITIONAL QUALITATIVE RESULTS.
We show additional uncurated results of FLUX-based models in Fig. and 13
25

---

Published as a conference paper at ICLR 2026
I-Floo (GM-FLUX) 4-NFE
SenseFlow (FLUX)
It-Flow (GM-FLUX) 8-NFE
Hyper-FLUX 8-NFE
FLUX Turbo 8-NFE
FLUXX1 1 dev 50-NFE
4-NFE
This pittlow falitoos athiinent eepressstorimm wisue style, clanetter schor pist mail syans sowrring neck glslor, athrnss, smaaiiers,
alamnda buss strwet. Vilowet aloo sights lanterru, cuemo,, fystrunt and sunlight creaale dynamiic attrosptiien.
bamee hardenes male wanter stands tall, his Niood drenchen face and determmned syes reflecting flene resorve Fie grlos limitatell deconeted sword, its niT greaming mhe din lint. The Baraque fis she acons, with diamate righting casting doop ahasoows hi worn feutures and kimizare armor rexuves.
through
hegeedrr Cuctussily comdor Apatstrasse
restaaibin fabee the
eicluser noathesit anw.oottoor www.urrttioooooooo www.koorrrdeitioeee....
it ytaay wanh no io gov betivinn and the bame: Stort curt-oabens lig actrty contrux wet the uphobdaryy Eucktig corobia nealy amanged whis Nizd sterd, far aboob, vnd pilnxc A
ampli agtieco wetloseee us www.srs/aaccoooooooooooooooooo enterng the nemanno Amggra Nowt.
Stiretg woroon 2T mayy sut2t PaI oarr loxe rote gola cuz Cosing e1 Dreare, piciig over sheatee A&7 actrmed wit sundawscs misigge Higigghiiggg Ser mabectey ICE, zomowau eupressIon realoo
alvemavig beveu owwrsst kndl
CotaTce, not 1338 wws.chooost dow.dd/wo p.w.ooo chou Cawra. a abrast www.ea/..ooaer........... pn CONTEE.
www.hwwwiinggggoomm wwr.ssig................ prosriaa wanch soneenst noman ooo www.ghwwng..oo poumanet
Figure 12: An uncurated random batch from the OnelG-Bendi prompt set, part A.
26

---

Published as a conference paper at ICLR 2026
IT-Flow (GM-FLUX) 4-NFE
SenseFlow (FLUX)
IT-Foow (GM-FLUX) 8-NFE
Hyper FLUX 8-NFE
FLUX Turbo 8-NFE
FLUX.1 1 dev 50-NFE
4-NFE
SOWERA
A row blue buses kilea deneath a mosal carcay marked "GATE 12 INTEROITY EXPRESS at busy bus temminal An LED screen scrotts "NEXT DEPARTURE: 15.30. Travelers 6D coffee fromo poper cups, whie neartyy newestard laboled "C1TY NEVS shaplays headkines and postoards of local londiaanss.
Fusien
FUSION
Fipils
Faslen
Couli the difeeoore betowan nuclew Ansion and fossion chart//p reattion somic nuche
showe udicationull chagram canoonataicall strociurea? Tipic nandenaterilt alruccure
Greas ice osdes Doar to valk, IN00 with OUSUNST vnd porTIAL The Suntio Frceers wncoo, and grss extercong IhS INNNSS DeT and swengg DOONS experence The SA5 scere-s depistor 30
with dipisson ofseeer, EAA-Ereee
hanza
attoncsiingg Shae
Austchop of sidterre Tokycobaccps.
WONDERS Stargeiing Nights
Celestial Wanders WONDERS SEO Stop Nidets
Celastaa Windors Stangaring Night PALE
Celestial Wonders'
Celestial Wonders
Celestial Wonders
Keogaynn Hight
exparttatee Footer informationnnincoooo thery
avakible seline"
Young vomua 51 weoding ottie boy abos whlte boct axtcita dalcce twn, stig-Oitg garreTE ovettyy aowoy aWocoutts whtu rust rop with Seadhort and iihcaw maple Sur patur, raslnne
Figure 13: An uncurated random batch from the OnelG-Bench prompt set, part B.
27


Reviews:
[{"review_id": "gP04gxVFSy", "reviewer": "Reviewer_3qdq", "summary": "This paper proposes pi-Flow, a policy-based paradigm for distilling multi-step flow-matching teachers into few-NFE students. Instead of predicting a single shortcut velocity, the student network outputs a network-free policy that can be rolled out with many cheap sub-steps to integrate the probability-flow ODE. Training uses on-policy imitation distillation (pi-ID), a velocity matching loss applied along the student’s own trajectory. Two policy families are introduced—DX and GMFlow. Experiments demonstrate state-of-the-art 1-NFE FID on ImageNet-256 with DiT-XL/2 and high-quality 4-NFE text-to-image generation from 12B/20B teachers (FLUX.1, Qwen-Image), with notably better diversity than prior few-step methods.", "questions": "1. How does pi-Flow compare to MeanFlow in terms of training stability and wall-clock time? Does avoiding JVP give measurable speedups during training?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. By predicting a policy rather than a shortcut, pi-Flow retains the dense ODE trajectory of the teacher while requiring only a few network evaluations. This insight is both conceptually clean and practically powerful.\n2. The velocity matching on the student’s own rollout avoids the complex progressive/consistency schedules of prior work and directly inherits teacher quality/diversity.\n3. Strong empirical results across scales. 1-NFE FID 2.90 on ImageNet-256 (DiT-XL/2) beats all prior DiT-based few-step models. 4-NFE distillation of 12B/20B T2I models preserves fine details and text rendering while achieving substantially better diversity than SOTA few-step baselines.", "weaknesses": "1. Table 1 shows K=8 vs. K=32 gives similar FID, but no ablation on $L\\times C$ factorization, mixture renormalization, or temperature scaling in inference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fusR042A3C", "reviewer": "Reviewer_KQWK", "summary": "This paper proposes pi-Flow, a new framework for few-step generative modeling that reframes flow-based distillation as policy learning. Instead of directly predicting the terminal point of a probability flow (as in velocity or consistency distillation), pi-Flow trains a network to output a policy that governs the velocity field across multiple integration substeps without additional network evaluations. The authors introduce a training algorithm, Policy-Based Imitation Distillation (pi-ID), inspired by on-policy imitation learning (DAgger), which lets the student model imitate the teacher’s flow dynamics along its own trajectory. Two specific parameterizations are explored:DX and GMFlow. Empirical results on ImageNet-256 and large text-to-image backbones (FLUX, Qwen-Image) show that pi-Flow achieves superior quality–diversity trade-offs and state-of-the-art one-step and few-step sampling performance.", "questions": "1. How sensitive is pi-Flow performance to the policy horizon (number of substeps)? Could adaptive or learned horizons improve efficiency?\n2. What are the computational trade-offs between DX and GMFlow in practice—does GMFlow’s expressivity justify its added complexity?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The policy-based perspective is a conceptually elegant and novel reinterpretation of few-step distillation, bridging diffusion/flow modeling and imitation learning.\n2. The pi-ID training procedure is well-motivated and effectively stabilizes learning by mitigating compounding errors—an issue that plagues prior consistency and flow distillation methods.\n3. The proposed network-free policy formulation enables fine-grained ODE integration with very few network calls, yielding strong efficiency gains.", "weaknesses": "1. While the policy formulation is innovative, the theoretical justification (why and when a policy yields better approximation of teacher trajectories) is mostly intuitive; a more formal analysis of error bounds or convergence would strengthen the contribution.\n2. The DX and GMFlow variants differ significantly in nature, but the paper does not fully disentangle whether gains arise from the pi-Flow framework itself or the added modeling flexibility of GMFlow.\n3. The approach introduces additional hyperparameters (policy length, Gaussian components) whose tuning cost is not well quantified.\n4. Although pi-Flow is claimed to preserve diversity, a quantitative comparison using standard mode-collapse or coverage metrics (beyond FID) would help substantiate this claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vSJcC2gBr4", "reviewer": "Reviewer_EmE8", "summary": "This paper introduces Pi-Flow, a novel framework for distilling fast samplers by training a policy model rather than a traditional trajectory-based generator. Instead of predicting instantaneous velocities along a diffusion trajectory, Pi-Flow outputs an entire policy function that governs the trajectory’s evolution. This approach allows the distilled model to represent the generative process as a global functional mapping, rather than a sequence of stepwise updates.", "questions": "1. **Clarify Multistep Training Scheme**  \n   The paper mentions both single-step and multistep training setups, but the procedure for multistep training is not clearly described. Additional details on how the multistep policy is trained, would improve clarity.  \n\n2. **Expand Description of GM Flow**  \n   The explanation of GM Flow is incomplete. The paper lacks a formal definition of the training objective for each policy and does not give sufficient details on how the architecture is adapted to enable policy generation. Providing explicit loss formulations, architectural modifications, and training details would make the contribution clearer and more reproducible.\n\n---\n\n## References  \n[1] **Fast Sampling of Diffusion Models via Operator Learning.**  \n[2] **Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation.**  \n[3] **Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion.**  \n[4] **Guided Score identity Distillation for Data-Free One-Step Text-to-Image Generation.**   \n[5] **Phased Consistency Models.**", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Novel Policy-Based Distillation Approach**  \n   Distilling a policy function that predicts an entire trajectory, while conceptually related to operator learning in DFNO (Deep Fourier Neural Operator) [1], is novel and interesting.\n\n2. **Comprehensive Large-Scale Experiments**  \n   The authors conduct extensive experiments on large-scale text-to-image models such as **Flux** and **Qwen**, demonstrating that the method scales effectively to state-of-the-art architectures.  \n\n3. **Strong Text-to-Image Performance**  \n   Empirical results show that Pi-Flow outperforms prior distillation methods in the text-to-image domain, highlighting its potential as a scalable and efficient alternative to conventional diffusion distillation techniques.", "weaknesses": "1. **Unfair Baseline Comparisons**  \n   In Table 2, the paper primarily compares Pi-Flow with training-from-scratch methods such as iMM and Mean Flow, rather than distillation-based methods like Score Identity Distillation (SiD) [2] or Consistency Training Models (CTM) [3]. This makes it difficult to assess Pi-Flow’s relative effectiveness as a distillation framework.  \n\n2. **Limited Evaluation on Standard Diffusion Backbones**  \n   While large-scale experiments on Flux demonstrate the method’s scalability, evaluations on more established diffusion models such as **Stable Diffusion 1.5** and **SDXL** would strengthen the paper. These benchmarks are widely used in prior distillation works (e.g., Long and Short Classifier-Free Guidance [4], Phased Consistency Models [5]), and results on them would allow for a fairer comparison.  \n\n3. **Inconsistent Efficiency Reporting**  \n   Since Pi-Flow predicts the full trajectory at once (i.e., the entire policy function), the architecture requires increased channel capacity and computational cost. Consequently, NFE (number of function evaluations) is no longer a fair efficiency metric. Metrics such as FLOPs or latency should be reported instead to provide a meaningful assessment of inference efficiency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AvPlsh1tf3", "reviewer": "Reviewer_1tXL", "summary": "This paper describes a method for the distillation of diffusion/flow networks. It assumes a pretrained flow model, and aims to learn a student model, initialized from the same weights, that can produce outputs with a small number of function evaluations. The key novelty with relation to prior distillation methods is the idea to predict a *policy* rather than a direct velocity. This policy can then be used to perform ODE denoising over many iterations, without the need for further network calls. The two policy classes are a DX policy, which predicts a sequence of (x0) predictions according to various t, and a GMFlow policy, which predicts a Gaussian mixture over the data. The GMFlow policy can be solved analytically to predict the velocity at a given t. To train the student model, on-policy distillation is used to fit the policy's predictions to the teacher. The proposed method outperforms prior works on a range of benchmarks such as DiT, FLUX, and Qwen-Image distillation.", "questions": "- I am a bit confused how this distillation procedure avoids the mean-collapse issue with naive distillation. For example, if the fixed trajectory is at t=1 (i.e. pure noise), then the DX policy head will attempt to predict E[x0], which is just the mean under the dataset. Thus, how is it that the student network can be used to generate images in a single NFE (which by definition should map pure noise to an image?)\n- How are t_src and t_dst chosen (page 5)? Does NFE stand for NFE under the teacher model, or under the student model?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper provides a novel contribution in the idea to predict a \"policy\" from the student, rather than directly predicting velocity. This innovation allows a single network call to define a series of ODE steps, which is key to retaining performance while reducing the number of function calls. In terms of quality and clarity, the paper is well-written, with a clear explanation of the method and thorough experimental results. This paper has a chance to be significant in the sense of a practical, effective flow distillation method.", "weaknesses": "- The description of the GMFlow policy class is hard to understand, and could benefit from a more intuitive explanation. Specifically, the paper would benefit from a more explicit instantiation of the mapping from (student network) -> (policy parameters) -> (velocity prediction at x_t, t).\n- The proposed method requires denoising the policy via an ODE *during the training process*, which adds a considerable overhead to the training pipeline. This stands in contrast to prior works which reduce the need for ODE rollouts during training (e.g. consistency models, progressive distillation, etc).\n- The comparison in Table 2 is not necessarily fair, as it compares the proposed pi-Flow method, which is a distillation method that requires a pretrained teacher, to iCT/Shortcut/Meanflow which are *from-scratch* methods that do not require a teacher.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)