Shortcut Diffusion Training with Cumulative Consistency Loss: An Optimal Control View
Authors: Paribesh Regmi
Sandesh Ghimire
Rui Li
Venue: ICLR 2026 Poster (2026)
Keywords: One-Step Diffusion
Optimal Control
Shortcut Diffusion Models
Primary area: generative models

Abstract: Although iterative denoising (i.e., diffusion/flow) methods offer strong generative performance, they suffer from low generation efficiency, requiring hundreds of steps of network forward passes to simulate a single sample. Mitigating this requires taking larger step-sizes during simulation, thereby allowing one- or few-step generation. Recently proposed shortcut model learns larger step-sizes by enforcing alignment between its direction and the path defined by a base many-step flow-matching model through a self-consistency loss. However, its generation quality is significantly lower than the base model. In this paper, we formulate few-step generation as a controlled base generative process, and show that self-consistency loss can be understood through the lens of optimal control. This perspective naturally motivates its generalization to the proposed cumulative self-consistency loss that cumulatively penalizes misalignment along the entire trajectory. This encourages larger step-sizes that not only align with the base model at the current time step but also support alignment in the subsequent steps, facilitating high-quality generation. Furthermore, we draw a connection between our approach and reinforcement learning, potentially opening the door to a new set of approaches for few-step generation. Experiments show that we significantly improve one- and few-step generation quality under the same training budget. Implementation is available at: [https://github.com/paribeshregmi/Shortcut-CSL](https://github.com/paribeshregmi/Shortcut-CSL)

Published as a conference paper at ICLR 2026
SHORTCUT DIFFUSION TRAINING WITH CUMULATIVE CONSISTENCY Loss: AN OPTIMAL CONTROL VIEW
Paribesh Regmi'
Sandesh Ghimire
Rui Li
pr85370ris drsandes wet.mmirr@mmmml.m. com rxlics@rit.ed:
'Rochester Institute of Technology
ABSTRACT
Although iterative denoising (i.e., diffusion/flow) methods offer strong genera- tive performance, they suffer from low generation efficiency, requiring hundreds steps of network forward passes to simulate a single sample. Mitigating this requires taking larger step-sizes during simulation, thereby allowing one- few- step generation. Recently proposed shortcut model learns larger step-sizes by enforcing alignment between its direction and the path defined by base many- step flow-matching model through a self-consistency loss. However, its genera- tion quality is significantly lower than the base model. In this paper, we formu- late few-step generation as controlled base generative process, and show that self-consistency loss can be understood through the lens of optimal control. This perspective naturally motivates its generalization to the proposed cumulative self- consistency loss that cumulatively penalizes misalignment along the entire trajec- tory. This encourages larger step-sizes that not only align with the base model at the current time step but also support alignment in the subsequent steps, facili- tating high-quality generation. Furthermore, we draw connection between our approach and reinforcement learning, potentially opening the door to a new set of approaches for few-step generation. Experiments show that we significantly improve one- and few-step generation quality under the same training budget. Im- plementation is available
INTRODUCTION
Shortcut Model
Ours
Shortcut Model
Ours
Two Steps
One Step
Figure 1: Two- and one-step image generation using Shortcut model (Frans et al., 2025) and our Shortcut-CSI. method trained ImageNet-256 dataset. Shorteut model tends to develop artifacts in the generated image with one and two-step generation, and ours generates a smooth, clearer image.
Diffusion (Song & Ermon, 2019; Ho 2020: Song et al., 2021) and flow-matching (Lipman et al., 2023; Albergo et al., 2023) models have demonstrated remarkable capabilities in generating
"Corresponding Author

---

Published as a conference paper at ICLR 2026
high-quality images and video (Rombach et al., 2022; Luo et al., 2023a:b). audio (Huang et al., 2023; Liu et al., 2023a), and molecular graphs (Jo et al., 2022; Eijkelboom et al., 2024). The generation involves iteratively transforming random noise into structured data, typically requiring hundreds of steps of forward passes through a neural network. This renders the process inefficient and com- putationally expensive. It is key limitation in comparison to single-step generative models, such as VAE (Kingma& & Welling, 2014; Vahdat & Kautz, 2020; Regmi & Li, 2023), GAN (Goodfellow al., ())4,, and normalizing flows (Dinh al., 2017).
To achieve one-step generation, previous efforts train student model to learn a direct mapping from noise-image pairs generated by the pre-trained model (Luhman & Luhman, 2021; Zhao al., 2023; Zheng et al., 2023; Yin et al., 2024). A set of approaches avoid the cost of generating such pairs, and support flexible generation budgets by allowing few-step generation (Berthelot et al., 2023; Ghimire et al., 2023; Liu et al., 2023b; 2024; Salimans & Ho, 2022; Song et al., 2023). These methods either progressively distill knowledge by halving the number of generation steps at each stage (Salimans & Ho, 2022), or enforce straighter flow paths during distillation (Liu et al., 2023b; Lee et al., 2024). These distillation-basee approaches require two training phases: training a base diffusion model and distilling its knowledge into student model. Approaches like Consistency models (Song al., 2023; Song & Dhariwal, 2024) and Meanflow (Geng et al., 2025) propose a single-phase phase training for few-step generation by learning both the base model and the few-step model in single training pass. Recently, shortcut models, a family of models with step-sizes larger than that of the base models, show promising performance by conditioning the model's output not only on time-steps and noisy inputs but also on the simulation step-sizes (Frans et al., 2025). This simple but elegant design allows for generation under specified budget simply by conditioning the model on its corresponding step size. They jointly train base model and the shortcut models with a standard flow-matching loss and a celf-connistency loss (SL), which aligns the shorteut model's outputs with the base model each time step. Although the experiments show improved few-step generation performance, it still lags significantly behind the base flow-matching models, emphasizing the need for further improvement. Moreover, the introduction of the SL loss in the paper lacks theoretical foundation.
In this paper, we formulate few-step generation as a controlled base generative process and show that the SL loss admits a principled interpretation through the lens of the optimal control framework (Bryson & Ho, 1975). Through this viewpoint, we identify a critical limitation in the SL loss, and introduce a generalized version. Specifically, we show that the SL loss corresponds to special case of the optimal control objective, which penalizes misalignment between the shortcut and the base model only at the current time step, overlooking the downstream effects of this immediate deviation. By drawing on the notion of expected future cost in a controlled process, we propose a general cumu- lative self-consistency loss (CSL) that penalizes misalignment across the entire trajectory, enabling the few-step model to account for these downstream effects. Moreover, we establish connection between our algorithm and on-policy reinforcement learning methods (Mnih et al., 2013; Lillicrap et al., 2015) by interpreting the expected future cost as value function. Similar to on-policy al- gorithms aiming to maximize an agent's cumulative reward rather than optimizing for immediate gains, our approach optimizes the cumulative cost along the generation trajectory. Evaluating the generative performance on standard benchmarks shows that our method robustly outperforms the state-of afr single-phase training approaches with the same training budget. Our contributions can be summarized as follows:
We propose novel perspective that frames few-step generation as a controlled process, and interpret the SL loss as a special case of an optimal control objective.
We propose novel CSL loss that penalizes both immediate and future misalignments between the few-step model and the base generative model cumulatively over time.
Training with the proposed CSL loss significantly outperforms shorteut models and other few-step generation baselines on benchmark datasets.
2 RELATED WORKS
Few-Step Generation: Distillation- based approaches aim to transfer knowledge from a pretrained model to student model (Luhman & Luhman, 2021; Zhao et al., 2023; Zheng et al., 2023; Yin et al., 2024; Liu et al., 2023b; Salimans & Ho, 2022). Some methods allow only fixed-step generation or require generating large volumes of noise-image pairs (Luhman & Luhman, 2021; Zhao et al., 2023;
2

---

Published as a conference paper at ICLR 2026
Zheng et al., 2023; Yin et al., 2024; Liu al., 2023b), while others reduce steps via bootstrapping but demand retraining a new model for each step reduction (Salimans & Ho, 2022). Overall, these approaches are either computationally expensive or inefficient when scaling different step counts. Moreover, all these approaches follow a two-ppase pipeline: pretraining a base model followed by distilling knowledge into few-step models. In contrast, consistency models (Song et al., 2023; Song & Dhariwal, 2024; Lu & Song, 2025) introduce a single-phase training strategy that trains a model to map noisy inputs from any time step directly the final denoised version, essentially enforcing transformation consistency along the denoising trajectory. Remarkably, this class of models achieve one- or few-step performance close to the base diffusion or flow-matchns mawchingg model. However, they are complicated by design and naturally unstable while training, requiring substantial engineering efforts like adaptive weighing for training stability Lu & Song (2025). Recently, Frans et al. (2025) proposed shortcut models for few-step generation, featuring simple, intuitive design and straight- forward single-phase training recipe, making them accessible for real-world deployment. Shortcut models condition not only the noisy input and time step but also on the step size. By learning a family of denoising functions conditioned on the step size, they provide a flexible and efficient framework for few-step generation using single model across different computation budgets. The SL loss is employed to align the generation trajectory of the model at larger step sizes with that of the base model, yielding strong performance. Nevertheless, the performance of few-step models still lags considerably behind that of the base model.
We establish a principled theoretical foundation for few-step generation from a control-theoretic viewpoint, casting the self-consistency (SL) loss as a special case of an optimal control objective. Whereas prior work formulates training objectives that supervise the few-step model a single point in its trajectory treating each step in isolation from its downstream consequences, this viewpoint leads us to a trajectory-leve cumulative objective that explicitly accounts for these downstream ef- fects. In addition, our theoretical framework uncovers a previously unidentified connection between few-step generation and reinforcement learning.
Leveraging Control Theory to Guide Eertusion-Based Generation: A few works have adopted control-theoretic principles to steer or fine-tune base diffusion and flow-matching models (Wang et al., 2025; Sprague et al., 2024; Domingo-Enrich et al., 2025). OC-Flow (Wang et al., 2025) leverages an optimal control framework to guide pre-trained flow-matching models for controlled generation. Domingo-Ente Enrich et al. (2025) introduce fine-tuning method that interprets reward- based optimization through the lens of stochastic optimal control.
BACKGROUND
3.1
OPTIMAL CONTROL
Let X denote the state of a controlled dynamical system at time defined by:
(1) where b(X,' t) is the base drift, us (X is called the control vector field, (tt) is control coef- ficient, and Po is the initial distribution. These jointly define the evolution of the controlled process X Optimal control problem aims minimize the following objective, which is also called a value function:
d_x:" (b(X², f) + (t) us (x("",)) t)) dt, X² ~ Pu-
J(uo T.. 5) = (((m0(X),))))) + t() dt h(XT), 2. X
where f(ux(X (XF, () penalizes the magnitude of the control, g ((P, () is the intermediate state cost, and ((x)) is the terminal cost. A result which we will later use in our analysis, is the expression for the derivative of J(ug; v.,, 8) with respect the control parameters 0 is:
d.J (uo; Fz, 8) = S. 30 s(o (), dt dt+/ L. d0
(XX, 00
(t) 2J ( Iz, t)
(3)
The gradient in Equation (3) consists of an immediate and a propagated effect. The first term mea- sures how changes in control parameter 0 directly influence the instantancous control cost, while the second term captures how those same changes alter the future trajectory and thus the future costs. Togetheer, they quantify both the local and downstream impacts of the change in 0.
3

---

Published as a conference paper at ICLR 2026
3.2 FLOW MATCHING
Flow-matching models (Lipman et al., 2023; Albergo et al., 2023) formulate data generation as simulating an ordinary differential equation (ODE), iteratively converting random noise into realistic data samples. Given model with parameter 0 that outputs velocity field Vo, the denoising process represented as:
dr, (x), di, LE [0,1, xo 0(
The goal is to train the model such that the distribution of samples distribution. We consider the optimal transport formulation of flow-matching (Lipman et al., 2023) in which the model is trained to optimize the following objective:
1 matches the data
LPM (0) = Exp. [1vo(ZZ. vell] 2t = (1 (1-t)x + tr1, U/
21
Zg
Here, 21 sample from the data distribution, 21 D, and vy represents the direction from noise to the data sample. Multiple (æo. 21) pairs may correspond to the same 2,, making i, inherently stochastic. When trained with the objective in Equation (4), flow-matching models learn the ex- pected velocity conditioned on It, denoted as [xa)
3.3 SHORTCUT MODELS
The trajectory defined by the flow-matching ODE is typically curved, necessitating small step sizes for accurate simulation, resulting inefficient generation. Using larger step sizes to enable faster, few-step generation leads to deviations from the true trajectory and significantly degrades sample quality, often causing failure when generating with fewer than four steps. To overcome this limita- tion, Frans et al. (2025) propose shortcut models, which learn to take larger steps while accounting for the future curvature of the trajectory, thereby avoiding deviation from the original path. This ap- proach enables few-step generation with improved sample quality. Notably, both the flow-matching and shortcut models are parameterized by the same backbone network, with different variants con- ditioned the step size. step taken by the shortcut model is given by:
Te+d T. S(x,, d) dd
(5)
where d is the step size, and S(I,, t, d) denotes the normalized direction from 2, to the point 2++d on the trajectory. Shortcut models of different step sizes can be trained on the same backbone network. The flow-matching model is the one conditioned with d 0 i.c., 8(x,, f, 0). Shortcut models are trained to enforce consistency across step sizes: taking a single step of size 2d should be equivalent to taking two sequential steps of size d.
8 1( 8(x+ta, 1 d. d.d)
8 S t, 2d) Larget: 8 Larget
The target for the shortcut model S ((x,, I, 2d) is thus generated by bootstrapping. The joint training objective of the flow-matching and shorteut models is:
Lshorteut (0)
EEg. 2x. ..d [1 t, (z1 z0)l
+ || 2d) S uargel
2
Flow-Matching
Self Consistency
4 CONTROL THEORETIC FORMULATION OF FEW-STEP MODELS
Both the forward and backward processes in flow-matching and diffusion models are governed by the Fokker-Planck equation, which describes the evolution of probability density over time (Risken, 1996). From this perspective, one- or few-step generation can be viewed as training neural network to approximate the solution of an underlying ODE, Consequently, when training a shortcut model to approximate this solution, the resulting integration error is cumulative in nature. To formalize this intuition, we model the output of the shortcut model itself as the solution to an ODE, defined as follows:
dxx [b(X),t) + us (X( (x)x))) dt
(8)
4

---

Published as a conference paper at ICLR 2026
This formulation implies that the shortcut model's output incurs time that makes its trajectory deviate away from the base model. Consequently, the error at time t reflects the cumulative effect of all errors at earlier time points < Note that the model pre- dicts the shortcut velocity d.x²/ /dt, and the error uo is obtained by comparing this prediction with the hase velocity b( x t). Another advantage is that Equation (8) provides viewing the few-step model as a controlled process. Unlike the standard optimal control formu- lation, where control vector field Ug is introduced predefined objective (e.g., Equation (2)), here is an unknown error that is implicitly introduced along the generation trajectory of a shortcut model. Since our goal is solely to make the shortcut model's 's trajectory match the base model, we focus minimizing the error ug and do not consider intermediate costs g( (X,;", ,/, 1) or terminal cost h(X)). Accordingly, we define an objective by setting g h 0 in Equation (2).
uo uo(X;" t) at each point in
perspective of
steer the base process toward optimizing
J(ug Is, = S. f [uul(X) (),t)
The future cost, J. thus becomes the primary quantity of interest. It is important to note that ug = O for all and X, the controlled process matches exactly with the base process, and the path of the few-step model matches that of the base How-matching model.
Equations (8) and (9) reveal that changing the parameters 0 affects cumulative objective J in two ways: (i) by altering the immediate error ug ua(X (X.,. .s..s, s), and (ii) by changing the subsequent states X!", which in turn impacts the future errors Ug (X,",, t) for & > In the next section, we show that SL loss is a special case that accounts only for the immediate error, and propose its two generalizations.
4.1 SELF-CONSISTENCY Loss (SL) AS AN OPTIMAL CONTROL OBJECTIVE
Lemma 1 Let (uo (XF,
(Xr, 1),t) = I|ue (X2, - 1), where (() is the Dirac delta function.
The objective in Equation (9) becomes:
JSL (u0 Fz. 8) [ [Io((X) 5) t)dt (I.. s]12, F, X
(10)
From Equation (10), if we consider model with stepsize 2d as shorteut model dx- stepsize d as base drift 6(X;", t), then for the discrete case, ug is the difference between the resultant of the two steps taken by the base model and single step taken by the shorteut model: us uo(2,, So (I t, 2d)
di and that with
S unget-- The optimal control objective equals to the self-consistency loss term:
Jst. (uo; xr, t) (1 t, 2d)
S target
Equation (9) accumulates the cost ff over the full trajectory from t to but SL restricts this to only the current time f = via Dirac delta function. Thus, Jss. is suboptimal: it penalizes error ug only at t 8, ignoring how it propagates and affects future errors at >> along the trajectory.
5 CUMULATIVE SELF- COLL-CONSISTENCY Loss (CSL)
We relax the delta function constraint in JsL and propose the following two losses:
(U,X),,), (),1) = [[uo(wa, }) 44 > 8, where, x. = X the objective in Equation (9) becomes:
Lemma 2 Let, (u)(XV
JUSI (ust Es, 5) = / [|,0 (X),, t)] = (1- - s) 11u0 (Ze, s)|
(11)
where USL stands for Uniform Self- Seff-Consistency Loss. Just considers errors over the entire time
< t < 1; however, it makes a naive assumption that the error uo(X)" t) is independent of t, and is uniform over the entire trajectory, which is equal the value at t simply reduces to a weighted SL and suffers from the same limitation as JSL- We further relax the uniformity assumption and propose CSL loss.
Essentially, this
5

---

Published as a conference paper at ICLR 2026
shortcut model
immediate grad
model
cumulative grad
t=
Ug USL
1
CSL
(a)
(b)
Figure 2: (a)lllustration of the SL (left) and CSL (right) losses. The solid red arrow denotes error Ug. which is the difference between the direction taken by the base model and the shortcut model. For SL, we calculate the error only at the current time 8, while for CSL, we calculate and aggregate the at every point from / cumulative gradient in addition to the immediate gradient as denoted by dashed purple and green arrows respectively. (b) comparison between SL, USL, and CSL losses.
the final time / 1. CSL allows the model learn from
Lemma 3 Lei (X,, (),)) = [|un(XF, 1)12 Vi > 8, the objective in Equation (9) becomes:
[ |mo(( (X,;,)))200 d)) I, X
JCSI (uo 8)
(12)
The differences between the SL, USL, and CSL losses are illustrated in Figure 2(b). Starting from 8, SL penalizes the error up only that specific time. USL accounts for the error along the entire trajectory; however, it assumes the error uniform across time, effectively treating the error all points identical that & s. contrast, CSL captures the true error landscape along the entire trajectory, accumulating the егrоr at each time step.
5.1 ANALYSIS AND COMPARISON WITH SELF- EELF-CONSISTENCY LOss
From Equation (3), the gradient of JOSL with respect parameters is given as (for ott) 1):
dJest (ug; Fs
Oun a0
2uo + Va,Josi(uo;
dt
d8
Analyzing the integrand,
Oup
Oug
(2uo + V Vcsucuuu 1))
2uo
V 21
|lue (x;,
(13)
00
00
immediate grad
cumulative grad
In Equation (13), the gradient of JOSL comprises two components: an immediate gradient term corresponding to the eиror the current step, and cumulative gradient term corresponding to the cumulative error along the trajectory. In contrast, the gradient of the JsL and JUSL include only the immediate gradient. Note that the cumulative gradient is gradient of the future errors with respect to the state It- Optimizing Jost. therefore encourages the shortcut model not only to align with the base model at the current step. but also to guide the trajectory towards a x, that facilitates high-quality generation by supporting the alignment in the subsequent steps as well.
5.2 CSL ESTIMATION AND TRAINING
Since are working with discrete time steps, the loss JESI, in Equation (12) can be expressed as a summation of the nomn of at every discrete step along the trajectory:
R2'
1 Iluo dk, d)11,, R' 5 ne ...,...... R 1} (14)
JCSL (us Ends nd)-
....

---

Published as a conference paper at ICLR 2026
R R' denote the number of terms in the summation. Replacing R' with = (i.e. = 1) results in an SL objective. An illustration of the SL and CSL losses is presented in Figure 2. Estimating summation in Equation (14) requires simulation, which can be computationally demanding. However, we can approximate the summation by including only a few terms. Specifically.
find that using only two terms, i.e., R 2, significantly improves few- step generation performance with negligible computational overhead. The effectiveness of R = 2 estimation is because for few-step models with two four steps, even two-step simulation covers substantial portion of the trajectory, covering the full trajectory for 2-step model and 50% of the trajectory for 4-step shortcut model. We have presented the Algorithm, and details of how gradient backpropagates from future (k >> n) steps to the current (k = 1) step in Appendix Section E.
5.3 CONNECTION TO REINFORCEMENT LEARNING
In our setting, the reinforcement learning analogy be drawn by viewing the agent attempting to transform noise into a data sample, where the states correspond to intermediate noisy samples along the generation path, the action is the direction of the next step, and the reward is defined as the negative magnitude role similar to a value function: it quantifies the cumulative cost of following policy along the trajectory. Moreover, since few-step generation involves a small, fixed number of steps, we do not require a separate network approximate JCSL: instead, it can be estimated directly by rolling out the actions. A more detailed discussion on this connection is provided in the Appendix Section F.
ug. Under this view, the objective JOSL in Equation (12) plays a
Table 1: Comparison of FID-50K scores of the baselines and our method on CelebA-256 and CIFAR1O, (* indicates that results are taken from Frans et al. (2025)
CelebA-256
CIFARIO
128-Step
Four-Step
Two-Step
One-Step
128-Step
Four-Step
Twa-Step
One-Step
Two Phase Training
Reflow
12.80+0.08
13.77:0005
14.48
16.07=000
13.93=005
14.92=0.05
15.59=0.14
16.98=0006
PD
7.96+007
14.4940.07
16.7310002
20.40:..15
7.89
10.75
11.80
13.26::0.09
Single Phase Training
FM
7.92.10.4
62.810.04
112.1
321.2.10.4
7.95
65.03
177.910.12
385.110.11
CT
53.7
19.0
33.2
ST
7.83110.44
9.36110.5
12.561002
20.4611002
7.37
9.15
11.79
19.80
ST-USL
7.95
9.18
12.00:10.5
19.41
7.37
9.35
11.6540.08
19.57
ST-CSL
7.88+0.04
8.98+00.2
10.96::002
18.37=0002
7.13#0003
8.10=0.09
9.24=0.09
17.76=0002
EXPERIMENTS
In this section, we empirically evaluate the effect of proposed CSL on performance. In Section 6.1 compare the generative performance of method with the baselines on the unconditional gen- eration task on CelebA256 (Liu et al., 2015) and CIFAR 10 (Krizhevsky, 2009). In Sections 6.2 and 6.3, assess if the improvement brought by CSL consistent across varying backbone net- work size and varying ratio of now-matching to bootstrap targets (B K) (see Algorithm 1) along with the training time comparison of the methods. In Section 6.4 assess the effect of increas- ing the number of terms R on the performance. Finally, in Section 6.5 we evaluate the baselines and or model on class-conditionn generation in ImageNet256 benchmark. For evaluations, we use FID-50K and Fl scores the comparison metrics (metric details are provided in the Appendix Section H).
6.1 PERFORMANCE COMPARISON
We consider two categories of baselines for performance comparison. 1. Two-Phase distillation approaches: Progressive Distillation (PD) (Salimans & Ho, 2022), Reflow (Liu et al., 2023b).
7

---

Published as a conference paper at ICLR 2026
DIT-S-2 (33 M)
DIT-B-2 (130 M)
DIT-L-2 (458 M)
20
24
18
FID-50K (1)
20
16
16-
12
12
Steps
Steps
Steps
CelebA
Shortcut
CelebA + Shortcut-CSL
CIFAR10
Shortcut
CIFAR10 + Shortcut-CS,
Figure 3: FID score comparison between shortcut and shortcut-CSL (ours) across different backbone network sizes (parameter counts). The networks are diffusion transformers (DiT) of varying sizes. Our model consistently outperforms the shortcut model across all network sizes and generation steps on both CelebA-256 and CIFARIO datasets.
2. Single-Phase training approaches: Training (CT) (Song et al., 2023), Shortcut Models (FT)ss et al., 2025)
(FM) (Lipman et al., 2023), Consistency
For fair comparison with the shortcut model, make sure the training budget is the same for both methods. We use the same number of flow-matching and bootstrap targets per batch. To achieve this, since create two bootstrap targets per datapoint for R datapoints for bootstrap targets in our case, while for the shortcut model we use datapoints. We require additional compute for the additional network forward pass step during bootstrap targets generation (step 18 in Algorithm 1), which adds only 10% of extra computation and training time in practice. Similar Frans al. (2025), we use the medium-scale diffusion transformer model DiT-B-2 (Peebles & Xie, 2023) as a backbone network. All the other hyperparameters are the same, and the details are provided in the Appendix, The FID-50K scores for the baselines and our method are reported in Table 1, the base flow-matching model is the one with 128-steps. Our method, trained with JUSL and JOSL. are referred to as shortcut-USL (ST-USL) and shortcut-CSL (ST-CSL), respectively. The results show that ST-USL consistently outperforms the baseline ST in one-, two-, and four-step generation, except for the four-step case on CIFARIO. Notably, performance of ST-CSL surpasses both ST and ST-USL on both datasets across all the steps, with the exception the base 128-step model on CelebA, where it achieves performance comparable ST. This shows the effectiveness of using CSL improving the generation quality of few-step models. Additionally, while PD achieves superior one-step generation performance on CIFAR-10, and Reflow does so on CelebA, both methods require two-phase training procedure. Furthermore, they offer less favorable trade-off between computational cost and performance, underperforming on 2- and 4-step generation tasks compared our method. PD also requires training and deploying separate models for each generation budget, which is time-inefficient.
we only make of K/2
Next, we empirically asses the analysis in Sec- tion 5.1 that our method trains the model to not only minimize current but also future errors. We compare trained Shortcut and Shortcut-CSI. models for CI- FARIO by generating samples from random noise t = 0 and measuring the misalignment uy between the two-step and it's base four-step trajectories at in- termediate time steps t 0.5 and t 1.0. For two-step trajectory, the value of ug at / = 0.5 and t= 1.0 corresponds to immediate and future emror respectively. Table 2 shows that while both methods have similar value of up 0.5. Shortcut-CSI. has a lower value at / 1.0, demonstrating its advantage in reducing overall error.
Table 2: Squared misalignment (uî) for
CIFAR-10 10, averaged over 100 samples.
Method
= 0.5
== 1.0
ST
0.5 x 10-3
2.5 x 10-3
ST-CSL 0.5 x 10-3 1.4 x 10-3
6.2 PERFORMANCE WITH INCREASING BACKBONE NETWORK SIZE
We evaluate whether the improvement of ST-CSL ST is consistent across different backbone network sizes. We train both models using small (DiT-S-2, 33M parameters), medium (DiT-B-2, 130M parameters), and large (DiT-L--2, 458M parameters) diffusion transformer networks.
8

---

Published as a conference paper at ICLR 2026
The FID scores for models trained on CelebA and CIFARIO are shown in Figure 3. We observe that as the model size increases, the overall few-step generation quality the shortcut model improves.
Importantly. ST-CSL consistently outperforms ST across all model sizes and generation steps. This shows that the improvement of CSL over SL robust across varying network sizes.
Table 3: Wall-clock time (in hours) for training (100 epochs) an FM, ST, and ST-CSL methods with different sizes Diffusion Transformers on the CelebA dataset.
Training Time Comparison: In Table 3, we compare the time taken to train the flow-matching (FM), Shortcut(ST),, and methods. Training FM is efficient of all but it doesn't t support few- step generation. Compared to ST, our method only consumes 6 10% more training time on average. Please refer to the Appendix Section for detailed time and memory consumption anal- ysis.
FM
ST
ST-CSL
DÝT-S-2 (33M)
5.6
6.2
DiT-B-2 (130M)
6.4
7.4
7.8
DiT-L-2 (458M)
11.4
12.2
6.3 ANALYSING THE EFFECT OF THE RATIO B K
We further evaluate whether the improvement of ST-CSL is consistent we vary the ratio flow- matching to bootstrap targets (B K) while training. We report the FID scores with varying ratio in Table 4. Although increasing the percentage of bootstrap targets introduces additional computational cost, results show that it significantly enhances the generation performance. Importantly, our ST- CSL consistently outperforms the ST baseline in few-step generation across all ratios. Moreover, observe that at the ratio 1 1, our method achieves two-step performance within 2.0 points of ST's 128-step performance on CelebA, and within 1.5 points on CIFAR-10.
Table 4: FID score comparison between shorteut and shortcut-CSL (ours) for different ratios of now-matching to bootstrap targets (B : K). Our model consistently outperforms shortcut model across all ratios and generation steps on both CelebA-256 and CIFAR10 datasets.
CelebA-256
CIFARIO
B:K
Method
128-Step
Four-Step
Two-Step
One-Step
128-Step
Four-Step
Two-Step
One-Step
ST
7.83
9.36
12.56
20.46
7.37
9.15
11.79
19.80
4:1
ST-CSL
7.88
8.98
10.96
18.37
7.13
8.10
9.24
17.76
ST
7.73
9.00
11.51
18.50
6.98
8.53
11.00
18.33
2:1
ST-CSL
7.75
8.58
10.09
16.37
6.95
7.98
9.08
16.51
ST
8.01
9.00
11.14
16.51
6.56
8.17
10.54
16.43
1:1
ST-CSL
7.56
8.54
9.91
15.51
6.67
6.95
7.94
13.96
6.4 EFFECT OF R
We investigate the effect of increasing the number of terms in the CSL objective in Equation (14). Specifically, we experiment with R 1,2, where larger values of Rt yield more accurate esti- corresponds
Table 5: FID-50K scores for one- and few-
step generation for different values of R CI- FARIO.
mations of the CSL. Note that to the SL loss used in the shortcut models. For all settings of R, the number of flow-matching and bootstrap targets per batch is kept constant. We observe that, compared to R 1, R incurs approximately 5% more training time. while R 4 incurs about 30% more time. The FID scores for few-step generation under differ- ent values of R are reported in Table 5. The results demonstrate that increasing R consistently improves generation performance, highlighting the benefit of more accurate CSL estimation.
CIFAR10
R
Four-Step
Two-Step
One-Step
1
8.174 0.10
10,544 0.08
16.43+ 0.06
2
6.95+ 0.07
7.94 0.07
13.96+ 0.10
4
6.66- 0.11
7.11 0.05
13.10= 0.06

---

Published as a conference paper at ICLR 2026
6.5 CLASS CONDITIONAL GENERATION ON IMAGENET
We would like to evaluate further the efficacy of the proposed CSL on the class-conditional genera- tion task. For this, we train the baseline ST and our ST-CSL methods the ImageNet dataset, and calculate the FID and the FI scores of the generated samples. For this analysis, we also compare with the recent MeanFlow (Geng al., 2025) bascline different percentages of bootstrap target, which corresponds to the B K ratio for ST and ST-CSL. The FID for base 128-step generation 15.21. The results in Table show significant gains of our ST-CSL. method when compared to Meanflow and ST. Additionally, although Meanflow performs better than ST in one-step generation, offers better performance vs efficiency tradeoff as shown by its superior performance at two- and four-step generations. Moreover, compared to Table 1, the gain of ST-CSL over ST is larger in advanced and challenging datasets like ImageNet-222 2566: for one-step generation, our method improves FID by nearly 20 points and F1 score by 4%.
Furthermore, in Figure 1, images generated by ST with one and two steps have noticeable artifacts; for instance, in the flying-bint example, despite the plain background, distinct distortions appear around the object. In contrast, our (ST-CSL) method yields smoother and cleaner generations.
Table 6: FID and Fl score comparison between baselines our method on the class-conditional gen- eration task on ImageNet-256, trained DiT-B-2 backbone. We compare the performance across different ratios B K for ST and ST-CSL, and different percentage of bootstrap targets for Mean- flow. ST-CSL consistently outperforms ST and Meanflow across all ratios and generation steps.
FID
Fi-score t
Method
Four-Step
Two-Step
One-Step
Four-Step
Two-Step
One-Step
Reflow
40.86
43.68
50.40
0.56
0.54
0.51
FM
104.96
210.39
325.78
0.26
0.09
0.00
Mcanflow (15%)
48.22
50.13
59.63
0.55
0.55
0.54
Meanflow (10%)
39.65
41.10
50.01
0.57
0.57
0.55
Mcanflow (5%)
34.09
36.61
45.12
0.58
0.57
0.56
ST (4:1)
24.70
35.73
64.12
0.61
0.56
0.46
ST (2:1)
23.47
32.54
55.55
0.63
0,58
0.50
ST (1:1)
24.17
32.22
51.78
0.63
0,59
0.51
ST-CSL (4:1)
16.98
21.77
45.84
0.63
0.60
0.50
ST-CSL (2:1)
16.21
18.71
37.60
0.64
0.62
0.53
ST-CSL (1:1)
15.71
17.35
31.66
0.64
0.63
0.56
7 CONCLUSION AND FUTURE WORK
We formulate few-step generation controlled process, using the flow-matching model as its base process. This perspective provides the theoretical foundation for the self-consistency loss and motivates our proposed cumulative self-consistency loss. By training with this objective, we achieve significant improvement in one- and few-step generation performance over the baselines.
In this work, consider the objective J without any intermediate state cost, i.e., ((X] An interesting direction for future work is to explore ways of incorporating this cost to influence the trajectory by penalizing undesirable intermediate states. Furthermore, the connection of our formu- lation to reinforcement learning offers opportunities to leverage techniques from the reinforcement learning literature further enhance few-step generation quality.
0.
ACKNOWLEDGEMENT
This work is supported by the National Science Foundation under NSF Award No. 2045804 and the National Institute of General Medical Sciences of the National Institutes of Health under Award No. R35GM156653. We acknowledge Research Computing at the Rochester Institute of Technology (RC) for providing computational resources.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Research computing at rochester institute of technology. https://oi.or//03477///////0000
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unity- ing framework for flows and diffusions. CoRR, 2023.
David Berthelot, Arnaud Autef. Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.04222 2023.
Arthur E. Bryson and Yu-Chi Ho. Applied Optimal Control: Optimization, Estimation, and Control. Taylor & Francis, New York, 1975.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference Learning Representations (ICLR), 2017. URL https://
openreview net/ Eorum?id=4Hkpnntt5
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In International Conference Learning Representations (ICLR), 2025.
Floor Eijkelboom, Grigory Bartosh, Christian A. Naesseth, Max Welling, and Jan- ail-Willem van de Meent. Variational flow matching for graph generation. In Neural Information Processing Systems (NeurIPS), 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. In International Conference on Learning Representations (ICLR), 2025.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505.13447, 2025.
Sandesh Ghimire, Jinyang Liu, Armand Comas, Davin Hill, Aria Masoomi, Octavia Camps, and Jennifer Dy. Geometry of score based generative models. arXiv preprint 13X 2023.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in Neural Informa- fion Processing Systems (NeurIPS), 2014.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Rongjie Huang, Jiawei Huang, Dongchao Yang, Yi Ren, Luping Liu, Mingze Li, Zhenhui Ye, Jinglin Liu, Xiang Yin, and Zhou Zhao. Make-an-audiol Tuxdioo generation with prompt-enhanced diffusion models. In International Conference on Machine Learning (ICML), 2023.
Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International Conference on Machine Learning (ICML), 2022.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representutions (ICLR), 2014.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. Technical Report.
Sangyun Lee, Zinan Lin, and Giulia Fanti. Improving the training of rectified flows. In Neural Information Processing Systems (NeurIPS), 2024.

---

Published as a conference paper at ICLR 2026
Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.0200
2015.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In Conference on Learning Representations (ICLR), 2023.
Haohe Liu, Zehua Chen, Yi Yuan, Xinhao Mei, Xubo Liu, Danilo Mandic, Wenwu Wang, and Mark D Plumbley. AudioLDM: Text-to-audio generation with latent diffusion models. In Inter- national Conference on Machine Learning (ICML), 2023a.
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Kleventh International Conference Learning Repre- sentations (ICLR), 2023b.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and qiang liu, Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. In The Twelfih International Conference on Learning Representations (ICLR), 2024.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Cheng Lu and Yang Song. Simplifying. stabilizing and scaling continuous-time consistency models. In International Conference on Learning Representations (ICLR), 2025.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2 101.0230100828 2021.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthe- sizing higl-resolution images with few-step inference. arXiv preprint arXiv:2310.04379, 2023a.
Zhengxiong Luo, Dayou Chen, Yingya Zhang, Yan Huang, Liangsheng Wang, Yujun Shen, Deli Zhao, Jinren Zhou, and Tien-Ping Tan. Videofusion: Decomposed diffusion models for high- quality video generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2023b.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602. 2013.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEFF/CVF International Conference on Computer Vision (ICCV), 2023.
Paribesh Regmi and Rui Li. Adavae: Bayesian structural adaptation for variational autoencoders. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Hannes Risken. The Bokker-Plancke Equation: Methods of Solution and Applications. Springer Series in Synergetics. Springer, 2nd edition, 1996. ISBN 5---5400--653-66.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations (ICLR), 2022.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In The Twelfth International Conference on Learning Representations (ICLR), 2024.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems (NeurIPS). 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. based generative modeling through stochastic differential equations. In Interna- fional Conference on Learning Representations (ICLR), 2021.
12

---

Published as a conference paper at ICLR 2026
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning (ICML), 2023.
Christopher Iliffe Sprague, Arne Elofsson, and Hossein Azizpour. Stable autonomous flow match- ing. arXivv preprint urXiv:2402.05774, 2024.
Arash Vahdat and Jan Kautz. NVAE: deep hierarchical variational autoencoder. In Neural Infor- mation P Processing Systems (NeurIPS), 2020.
Luran Wang, Chaoran Cheng, Yizhen Liao, Yanru Qu and Ge Liu. Training free guided flow- matching with optimal control. In The Thirteenth International Conference on Learning Repre- sentations (ICLR), 2025.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2024.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. UniPC: unified predictor- corrector framework for fast sampling of diffusion models. In Neural Information Processing Systems (NeurIPS), 2023.
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference Machine Learning (ICML), 2023.
13

---

Published as a conference paper at ICLR 2026
APPENDIX
A DERIVATION OF EQUATION (3)
We have,
Jus; Fz. 8) = S. (f(mo(X), ().0) (x(, t)) dt + h(X)),
= X
Step 1: An expression for the gradient of J:
do /. [(uo(Xi' 1110) + ((x 1) di + +(x)))
d
[. 20 (X), t)dt t), [
(X] + 9(X),, .)] 0x;
(h(X
dt + 00
(15)
ox.
00
Step 2: Expression for gradients
dxy dx
and then
dx," d0
The state evolves according to
dxT
6(X)) ux(X),
di
The state-to-state sensitivity is given by:
dXX.¹
A(L, 8) =
126
dXuu
For a small time increment such that s + < t, using the chain rule along the trajectory,
dx a dX
dXy
dx
dxw
dXy
stc
In terms of A(t,
A(t,s) +e), ((++))
(16)
Using Taylor expansion for X
8451
X = X + + (0())o((X))))) s]] + +(=)))
Thus,
2 A(s 8) = =1
[(XX) (x)) 0(6))
(17)
Ox
By definition of the derivative,
i
A(I, + + A(1, lim c-
A(t,
Øs
8
Using Equations (16) and (17),
A +tt,s+)) + [b s) + (x))
s)] O(s2) -((2)
i
A(t, 8) = lim
Os
e-001
i
=-((t,
[(X" 8) (X)
(X
0x.
14

---

Published as a conference paper at ICLR 2026
Therefore,
① dxp
dxy
b(x 8
(x)))
Ös dx
dx
0x
dxy."
J.
dx dx ax
dXy
((X + (X)) ds
(18)
dXy
dx
s=
Equation (18) can be generalized to obtain gradient with respect to the parameter 0 (Chen et al., 2018), which is given by the following integral:
dx:
dxx 0
b(X ((S) us (X),
d0
dx 00
dXF
Ouo (xx, 8)
((s)
00 ds
(19)
dX
Step 3: Final Expression
Substituting the gradient representation from Equation (19) into Equation (15), we get:
dJ (ug; g2.
(uoo(Xx)) 1),))
d0
00
(X,", t), + +((xx.t)
d.y." dxy
Qus t')
+
a(t))
dt'
de
0X0
00
(h()
dxi axp
ol1) Oug de
0xp
00
Interchanging the order of integration gives
d.J (uj; F2, s)
[:¹
00 (uo(XX dt
d0
f.
(X + g(X,, dXp
dt + oh(xF) dxp
Ous(XF dt'
xxe dxy (+))
0X.
dxp
00
Recognizing the value gradient, (x, = x
["' O[f(uo(X] t)tt + g(XY, dXT
ah(XF) dXF
Va, J te, (us
+ OxP dxy.
OXP²
dxy
we arrive at the final expression
d.J [uo; Fz,
': (uo(XF) 1), di S² 0g
Dus (X]'")))
(()) Vz, [uo; x4, dt
(20)
d8
00
B LATENT SPACE INTERPOLATABILITY
To analyze whether the one-step model learned by approach captures interpolatable latent space, we examine the model's outputs when fed with interpolated noise samples between two independently drawn Gaussian noise vectors. Specifically, given 20,2% 1(0, T), generate interpolated inputs using
Za VI - azo + Valg E 10.
We then apply one-step denoising each interpolated sample using our shortcut-CSL model. The resulting images shown Figure exhibit smooth transitions across the interpolation path, suggest- ing that the model has an interpolatable latent space.
15

---

Published as a conference paper at ICLR 2026
Figure 4: Image generated using the shorteut- CSL (ours) method with one-step denoising applied to variance-presse preeeeeiireeriigg interpolation between two Gaussian noise samples. The leftmost and right- most images correspond to independently drawn noise samples, while the intermediate images were produced from interpolated samples.
128 Steps
0
Four Steps
Two Steps
Step
Figure 5: Image generated by the shortcut- CSL (ours) method by four-, two-, and one-step denois- ing. 128 steps denoising corresponds to the base flow-matching model.
C SAMPLES FROM SHORTCUT-CSL
In Figure 5, provide some samples from the proposed shorteut-CSL method for different gener- ation budgets.
D HYPERPARAMETER DETAILS
For experiments in Tables 1 and 6, we use the & Xie, 2023) the backbone network for both the shorteut baseline and our method. For ImageNet- 256, due to computational limitations, we sample a subset of 30 classes with ~40K total image samples for class-conditional training. For this and also CelebA-256. downsample images from
diffusion transformer DiT-B-2 (Peebles
16

---

Published as a conference paper at ICLR 2026
3x256x256 to 4x32x32 using the VAE encoder from the Stable Diffusion framework (Rombach et al., 2022), specifically the sd-vae- ft -ema variant. We then train the diffusion model this compressed latent space. Detailed hyperparameters for our (ST-CSL) approach are provided in Table 7. Here, M denotes the total number of denoising steps in the base flow-matching model. The baseline shortcut method (ST) uses the same hyperparameter configuration, except with R 1. The experiments were run on NVIDIA A100 GPU.
Table 7: Hyperparameter setting for the reported results in Tables and 6.
Hyperparameter
Value
epochs
300 (CelebA), 500 (CIFARIO, ImageNet)
M
128
64 (CelebA, CIFARIO), 128 (ImageNet)
K
16 (CelebA, CIFARIO), 32 (ImageNet)
R ema
2
0.9999
optimizer
A
learning rate
0.0001
weight decay
0.1
E
ALGORITHM
The algorithm for training shortcut models with the proposed CSL loss is detailed in Algorithm 1,
Algorithm Training Shortcut Models with cumulative Seff-Consistency Loss (Shortcut-CSL)
1: K < #bootstrap targets
2: <- #flow-matching targets
3: R #terms in JESL estimation
4: while not converged do
5: To 11 1), 31 D, (d,t) t(d,t)
6:
Ie - (1 t)zo + tx
7:
for B batch elements do
8:
Starget d-
20
9:
10:
end for
11:
for
K/R batch elements do t, 2 2t, 8 target
12:
13:
for R iterations do
14:
81 So(2', s.
15:
2 + x + S
8, So(I'g
d)
17:
APPEND S arger, stopgrad((8; + 82)/2))
21 + 1 + So d,2d) 2d
19:
s = s + 2d
end for end for
21:
22:
1 f, o 5- Z₁
23:
for
to R do
24:
+ V, Vol[S((2 (2' s, 2d) S 3r]]ll
[r]l
25:
x' = s + 2d
_ So (z s, 2d) 2d
26:
27:
end for end while
28:
17

---

Published as a conference paper at ICLR 2026
The cumulative grad term in 2 Equation (13)) term is realized in practice by line 24 of Algo- rithm 1. For R = 2, lines 23-26 operate as follows: Starting from an initial state z at time t, we compute a shorteut step with stepsize 2d as:,
So (0 t,2).
and evaluate the squared-loss objective against the fixed target produced by rolling out steps with stepsize d (lines 13-19). In practice, PyTorch autograd computes this gradient. We then move along the shorteut step obtain the "next state" and updated time (line 25 and 26). Now, in the second iteration, we again compute So(z', 2d) at this "next state", and evaluate the squared loss with the corresponding target. Crucially, during this second loss evaluation, the "next state" depends on the shortcut step taken in the first iteration (line 25). As result, the gradient from the second loss flows backward into the computation of the first shortcut step. This ensures that the first shortcut step is also updated as to reduce the loss incurred at the second step. This is precisely how the gradient backpropagates from the second (future) step the first (current) step and realizes the cumulative gradient term in Equation (13).
F CONNECTION TO REINFORCEMENT LEARNING
The optimal control (OC) viewpoint of few-step models naturally connects to reinforcement learn- ing. In OC literature, the following OC objective (Equation (2)) at optimality is known as value function, which mirrors the RL value function that represents the cumulative future cost along a trajectory:
Nuggz = [." (F(ua(X2,0)0) + (xx", M(T));
X = Te
Since our proposed JCSL has the same structure structure-accumullee accumulating the consistency error || 1 along the trajectory- can formulate training as an RI. problem: the shorteut model acts as the actor, while a value network serves as the critic estimating JOSI. via temporal-differee formulate our method as reinforcement problem as follows:
learning. We can
Starting with random noise at = O, an agent aims to transform it into meaningful image in (1/d)- steps, taking a each step of size d going from = 0 to = The agent is trained to take a direction, moving along which generates the best image (i.e., same as the direction taken by the base 1 128-step model in our case). For this, we define the reward as the negative of the CSL (-JOsL), which we aim to maximize.
-JCSL acts as our value function, and we can train a value network to estimate it. We can use the following temporal difference learning learn value network Vg:
Vo 1, d) = -luo( d]l1 + V (o dd,d)
(21)
where It+dd is calculated using Equation (5). Here, us ( t, d) is the immediate misalignment where 0 is the parameter of the few-step model, and Vo represents the future misalignment from time + d onwards, predicted by the value network.
In reinforcement learning terms, the (1 st/d)-step model serves as actor, and the value model serves as a critic, where the actor tries to maximize the estimated value. Instead of training a separate value network to estimate CSL, opted for R-rollout (especially. R = 2) it offered good balance of performance and efficiency. We leave further exploration with techniques from the reinforcement learning literature for future work.
G EMPIRICAL TIME AND MEMORY CONSUMPTION COMPARISON
Table reports the wall-clock time (in hours) and GPU memory usage (in GB) for training Flow- Matching (FM), shorteut model (ST), and our method (ST-CSL) the CelebA dataset for 100 epochs, using small, medium, and large Diffusion Transformer (DiT) models. All experiments were conducted on single NVIDIA A 1100 GPU.
18

---

Published as a conference paper at ICLR 2026
Table 8: Wall-clock time (in hours) and GPU memory usage (in GB) for training Flow-Matching (FM), shortcut (ST) and Ours (ST-CSL) methods for Diffusion Transformers (parameter counts) of small, medium, and large sizes.
Time (hours)
Memory (GB)
Model
FM
ST
ST-CSL
FM
ST
ST-CSL
DIT-S-2 (33M)
5.6
6.2
21.35
21.35
DiT-B-2 (130M)
6.4
7.4
7.8
28.66
31.67
28.96
DiT-L--2 (458M)
11.4
12.2
45.47
38.52
FM is the most efficient of all in time and memory, but does support few-step generation.
Our method requires slightly more time than the shortcut model but uses less memory in larger models. This trade-off arises from how each method processes batches during training. The shortcut model processes the entire batch in a single step from the current state, which minimizes time but leads high memory consumption. In contrast, our method uses two-step process: it first com- putes intermediate step using part of the batch, then computes the next step using the intermediate outputs (as described in lines 23-26 in Algorithm 1). Splitting the batch this way reduces memory usage but requires additional time for the extra computation.
For two-stage methods like Progressive Distillation (PD) and Reflow, we first need to train the FM model and distill the knowledge to the student model. Reflow requires twice the time taken by FM. For PD, we train a student model that learns to generate samples in half as many steps as the teacher model. Therefore, separate model needs to be trained for 128, 64, 32, .... 2 and I step generation successively, consuming significantly large amount of time.
H METRIC DETAILS
Fréchet Inception Distance (FID). FID measures the discrepancy between the real and generated data distributions by comparing their feature-level means and covariances. The features are obtained from an Inception v3 network. Lower FID indicates that the generated distribution matches the
real distribution more closely. Given real features Fr (f7 Nyr and generated features F9
Na let (J4r, E,,) and (Mg. Egg) denote their empirical means and covariances.
FID ||, - 1gll2 Tr E, Egg 2(E,E,,))) -
(22)
Precision Precision measures the quality of generated samples by quantifying the fraction of gen- erated points that lie within the local support of the real data manifold. For each real feature f"" define the --nearesttneighhbr radius
kNN.
N,
(23)
A generated point gj is considered valid if it lies within the radius of at least one real point. Precision is:
N Σ
Precision
1 1153
-fll 512200
(24)
Recall Recall measures the diversity of generated samples by quantifying how much of the real data manifold is covered by the generated distribution. For each generated feature 13. define
5
kNN dist(Jf, F*),
j= Ng-
(25)
A real point f is considered covered if a generated sample lies within its radius. Recall is:
N.
Recall
Σ
min l1 - -19112 s 8(0)
(26)
N
19

---

Published as a conference paper at ICLR 2026
F1 Score. The F1 score summarizes precision and recall into single measure, balancing sample quality and diversity. is defined as the harmonic mean:
2 Precision Recall
(27)
Precision + Recall
I CONVERGENCE WITH TRAINING STEPS
DIT-5-2
DIT-B-2
DIT-L-2
RD-50K (4)
100K
150K
200K
250K
300K
Training Steps
Figure 6: We evaluate the convergence of shortcut-CSL. by tracking the FID score at regular training intervals across different backbone DiT network sizes on the CelebA-256 dataset. Our results show that performance consistently improves with increasing model size at every stage of training.
LLM USAGE
LLM is lightly used to polish writing in the paper. The used prompt Please polish the following sentence/ sentences
20


Reviews:
[{"review_id": "QFwILQB8KY", "reviewer": "Reviewer_kzb9", "summary": "The paper proposes an improvement over shortcut models to distill multi-step diffusion pipelines, thereby improving inference speed. The central idea of the paper is that the shortcut models only consider immediate errors, which can lead to misalignment with the model's true trajectory in severe step compression. The authors propose to penalise not just the immediate misalignment but also further misalignments down the trajectory of the teacher model.", "questions": "The authors discuss the increasing training cost with increasing R. Would be good to know if the authors experimented with R>4 or if the authors have an idea of how training scales with R.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- By framing the shortcut problem as an optimal control one, the authors provide mathematical grounding for their work. They model self consistency as an optimal control objective and then relax constraints to incorporate errors further along the ODE trajectory\n\n- The authors include fair comparisons that ensures equal training budgets with baselines\n\n- Training algorithm is included to aid reproducibility", "weaknesses": "- The paper lacks sufficient clarity, specifically in the connection of the work to Reinforcement Learning. Linking J_csl as a value function can be better explained, given the authors have mentioned this in the abstract as an important section of the paper.", "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": "ffsnnT3iO4", "reviewer": "Reviewer_z676", "summary": "This paper proposes Cumulative Self-Consistency Loss (CSL) as an extension of the shortcut training loss (SL) used in one-step and few-step diffusion models. The authors reinterpret shortcut training through an optimal-control lens. CSL generalizes this by integrating consistency along the remaining trajectory, encouraging alignment not just locally but cumulatively over time. The paper derives a continuous formulation of this loss, motivates it with a “cumulative gradient” analysis, and implements a discrete estimator involving R rollouts (typically R=2). Empirically, CSL yields improved FID scores on CIFAR-10 and CelebA-256 for 1-, 2-, and 4-step sampling, with minimal additional compute cost.", "questions": "1. How is the ∇ₓJ term in Eq. 13 realized in practice? Are gradients ever propagated through rollout steps, or is CSL purely a multi-point supervision objective?  \n2. Why does Algorithm 1 advance the 2d step from the midpoint state rather than from the starting point? Was this empirically motivated?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Conceptually clear reinterpretation of shortcut loss using an optimal-control framework, identifying why standard self-consistency only optimizes instantaneous alignment.  \n- The cumulative loss is simple to implement and shows consistent, measurable FID improvements with small computational overhead (~5–10%).  \n- Thorough experimental evaluation with ablations (number of rollout terms, backbone scale, training cost) and comparisons across DiT backbones.", "weaknesses": "- The “cumulative gradient” claimed in Eq. 13 does not seem to match the implementation. Algorithm 1 detaches rollout targets with `stopgrad` and does not backpropagate through future steps, so it seems that the optimization reduces to multiple local MSEs rather than true cumulative assignment. \n- The midpoint update rule used during 2d rollouts is unconventional and unexplained. \n- Evaluation scope is a bit narrow, authos only evaluate on CIFAR-10 and CelebA-256, with large baselines borrowed from prior work under different setups.", "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": "HnWhESUPbG", "reviewer": "Reviewer_CsHB", "summary": "This paper improves upon the original shortcut model paper by introducing a cumulative self-consistency loss. In the paper, the original shortcut model loss is first framed as an optimal control problem, where the objective is to minimize *total* accumulated error. However, the standard loss only penalizes error at the *current* timestep, and thus does not account for future error. This paper derives the optimal objective that maximizes for future error minimization, analogous to maximizing future return rather than reward in an RL framework. Practically, it is found that using a discrete number of future steps (R=2 or R=4) already leads to sizable improvement. Experiments show that the proposed cumulative SL reliably improves upon the base SL, and ablations show that the improvement can be steadily improved with additional compute allocated to R.", "questions": "See section above for further questions associated with each topic.\n\nAs a side note, the discussion in the Appendix on the connection to TD-based value learning methods is quite interesting, and could be a worthy direction for future work.\n\nIn general, the ideas in this paper can be strengthened by a more thorough clarity in writing. The introduction and background sections can be condensed, and the relation between optimal control and the shortcut loss can benefit from additional textual commentary explaining the intuition behind the equations. With clarifications to the questions in the weaknesses section, I would consider a raise to the review score.\n\nCan the cumulative self-consistency objective be applied to other few-step techniques such as consistency models or meanflow?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper highlights a well-posed analysis of the shortcut modelling loss, showing that is an instantaneous approximation to an objective that should be expanded to account for future errors. Using this viewpoint, the paper derives a clear true objective, and shows that even a discrete approximation to this true objective results in performance improvement. In this way, the paper clearly ties together empirical improvement towards a theoretically-motivated insight.\n\nThe clarity in the paper is passable. Figure 2 provides a solid intuition as to why the cumulative objective is a more precise target. (See below for weaknesses).\n\nThe significance of this paper is solid, as it provides a new viewpoint which can be applied to distillation methods in general. The experiments are conducted on standard benchmarks at a reasonable network size, and wall-clock experiments are presented to take computational requirements into account. Error bars are included.", "weaknesses": "The empirical performance of this method could be better framed by comparing to recent prior works, specifically Meanflow. \n\nIn terms of clarity, the notation introduced in Section 3/4/5 can be refined. A full explanation of optimal control as given in section 3 may not be necessary, especially terms such as b, g, and h which are not used in the final relation to shortcut model loss. In equation 4, it will help to make it clear that u(X) is the *error* term which has an optimality at zero, but the network itself does not predict u(X) but rather a velocity (that must be compared to the teacher velocity). \n\nFor equation 14, is the gradient with relation to future errors taken through multiple evaluations of the network? If so, is the error term for k>1 backpropgated to all previous steps including the current step? An algorithmic explanation would strengthen this section.\n\nWhile the text argues that equations 11 and 12 should be different, the resulting integrals appear identical. What is the precise difference between these two objectives? Commentary would strengthen this section.", "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": "a3nvJCpj85", "reviewer": "Reviewer_sxLT", "summary": "The paper addresses the low sampling efficiency of diffusion/flow models by enabling one- or few-step generation. Building on “shortcut” samplers that learn large step sizes via a self-consistency loss against a many-step base flow-matching model, the authors reinterpret this loss through optimal control, treating few-step sampling as a controlled base process. From this view, they propose a cumulative self-consistency loss (CSL) that penalizes misalignment not only at the current step but also along future steps of the trajectory, encouraging large yet reliable steps that maintain downstream sample quality. They further draw connections to reinforcement learning. Experiments indicate improved one-/few-step generation quality under the same training budget, and the proposed shortcut-CSL variant outperforms baselines on certain tasks.", "questions": "My primary concern is the limited practical performance gains demonstrated (see Weaknesses).", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "• The paper is clearly written with precise notation; figures and tables are informative and easy to follow.\n\n• The proposed shortcut-CSL shows better performance than prior shortcut variants on some tasks, supporting the value of the cumulative alignment objective.", "weaknesses": "1. Metrics are too narrow. Evaluating primarily with FID is insufficient, and the observed FID gains are minor in some cases. Please broaden the evaluation metrics like F1, CLIP-score, and qualitative comparisons, and, if feasible, a small human preference study. This multifaceted assessment would make the empirical case more solid.\n\n2. The teaser (Fig. 1) visualizations look only slightly improved over baselines. \n\n3. The dataset scope is limited. Results are reported only on CelebA and CIFAR-10. While broader evaluation is costly, for fair comparisons with other methods, you should include additional datasets—e.g., ImageNet (64×64 and/or 256×256), FFHQ, or another modern benchmark—while keeping ablations on a single dataset to manage workload. This would better test generality and strengthen claims about scalability and robustness.", "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)