One Step Diffusion via Shortcut Models
Authors: Kevin Frans
Danijar Hafner
Sergey Levine
Pieter Abbeel
Venue: ICLR 2025 Oral (2025)
Keywords: diffusion
flow-matching
fast inference
distillation
Primary area: generative models

Abstract: Diffusion models and flow matching models have enabled generating diverse and realistic images by learning to transfer noise to data. However, sampling from these models involves iterative denoising over many neural network passes, making generation slow and expensive. Previous approaches for speeding up sampling require complex training regimes, such as multiple training phases, multiple networks, or fragile scheduling. We introduce Shortcut Models, a family of generative models that use a single network and training phase to produce high-quality samples in a single or multiple sampling steps. Shortcut models condition the network not only on the current noise level but also on the desired step size, allowing the model to skip ahead in the generation process. Across a wide range of sampling step budgets, shortcut models consistently produce higher quality samples than previous approaches, such as consistency models and reflow. Compared to distillation, shortcut models reduce complexity to a single network and training phase and additionally allow varying step budgets at inference time.

Published as a conference paper at ICLR 2025
ONE STEP DIFFUSION VIA SHORTCUT MODELS
Kevin Frans
Danijar Hafner UC Berkeley
Sergey Levine UC Berkeley
Pieter Abbeel UC Berkeley
UC Berkeley
kvfrans@berkeley.een
ABSTRACT
Diffusion models and flow-matching models have enabled generating diverse and realistic images by learning to transfer noise to data. However, sampling from these models involves iterative denoising over many neural network passes, making generation slow and expensive. Previous approaches for speeding up sampling require complex training regimes, such as multiple training phases, multiple networks, or fragile scheduling. We introduce shortcut models, a family of generative models that high-quality samples in a single or multiple sampling steps. Shortcut models condition the network not only on the current noise level but also on the desired step size, allowing the model to skip ahead in the generation process. Across a wide range of sampling step budgets, shortcut models consistently produce higher quality samples than previous approaches, such as consistency models and reflow. Compared distillation, shorteut models reduce complexity to a single network and training phase and additionally allow varying step budgets at inference time.
single network and training phase to produce
INTRODUCTION
Iterative denoising methods such as diffusion (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020) and flow-mathhing (Lipman et al., 2022; Liu et al., 2022) have seen remarkable success in modelling diverse images (Rombach et al., 2022; Esser et al., 2024), video (Ho et al., 2022; Bar- Tal et al., 2024), audio (Kong et al., 2020), and proteins (Abramson et al., 2024). Yet, their weakness lies in expensive inference. Despite producing high-quality samples, these methods require an it- erative inference procedure-often requiring dozens hundreds of forward passes of the neural network-making generation slow and expensive. We posit that there exists a generative modelling objective which retains the benefits of diffusion training, yet can denoise in single step.
Flow Matching
Shortcut Models (ours)
128 Steps
Four Steps
One
Step
Figure 1: Generations of flow-matching models and shorteut models for different inference budgets. Shortcut models generate high-quality images across a wide range of inference budgets, including using single forward pass, drastically reducing sampling time by up to 128x compared diffusion and flow-matching models. In contrast, diffusion and fow-matching models rapidly deteriorate when queried in the few-step setting. The same starting noise used within each column and two models are trained on CelebA-HQ and Imagenet-256 (class conditioned).

---

Published as a conference paper at ICLR 2025
We consider the end-to-end setting, in which one-step denoising is acquired by single model over a single training run. Closely related are previous twoo-stage methods which take existing diffusion models and later distill one-step capabilities into them. These stages introduce complexity and require either generating a large synthetic dataset (Luhman & Luhman, 2021; Liu et al., 2022) or propagating through series of teacher and student networks (Ho et al., 2020; Meng et al., 2023). Consistency models (Song et al., 2023) are step closer to the end-to-end setting, but their dependency on large amounts of bootstrapping requires careful learning schedule throughout training, Two- stage or hightly-cheddlee procedures suffer from a need to specify when to end training and begin distillation. In contrast, end-to-end methods can be trained indefinitely to continually improve.
We present shortcut models, class of end-to-end generative models that produce high-quality gen- erations under any inference budget, including in single sampling step. Our key insight to condition the neural network not only on the noise level but also the desired step size, enabling it to accurately jump ahead in the denoising process. Shortcut models can be seen as performing self- distillation during training time, and thus do not require separate distillation step and are trained over a single run. No schedules or careful warmups are necessary. Shortcut models are efficient to train, requiring only 16% more compute than that of base diffusion model.
Empirical evaluations display that shortcut models satisfy a number of useful desiderata. On the commonly used CelebA-HO and Imagenet-256 benchmarks, single shorteut model can handle many-step, few-step. and one-step generation. Accuracy is not sacrificed - in fact, many-step gen- eration quality matches those of baseline diffusion models. At the same time, shorteut models can consistently match or outperform two-stage distillation methods in the few- and one-step settings.
The key contributions of this paper are summarized as follows:
We introduce shortcut models, class of generative models that generate high-quality samples in a single forward pass, by conditioning the model on the desired step size. Unlike distillation or consistency models, shorteut models are trained in a single training run without schedule.
We perform comprehensive comparison of shorteut models to previous diffusion and flow- matching approaches on CelebAHQ-256 and ImageNet-256 under fixed architecture and com- pute. Shortcut models match or exceed the distillation methods that require multiple training phases and significantly outperform previous end-to-end methods across inference budgets.
To demonstrate the generality of shortcut models beyond image generation, we apply them to robotic control and replace diffusion policies with shortcut policies. We observe that shortcut models maintain comparable performance under order-of-magnitude lower inference cost.
We release model checkpoints and the full training code for replicating our experimental results: http:///jttbbbboookvvvira////oot
2 BACKGROUND
Diffusion and flow-matching. A recent family of models, including diffusion (Sohl-Dickstein al., 2015; Ho al., 2020; Song et al., 2020) and flow-matching' (Lipman al., 2022; Liu et al., 2022) models, approach the generative modelling problem by learning an ordinary differ- ential equation (ODE) that transforms noise into data. In this work, we adopt the optimal transport flow-mater now-matching objective (Liu et al., 2022) for simplicity. We define 2; as a linear interpolation be- tween data point V1 ~ DD and noise point Zo ~ ~ I) of the same dimensionality. The velocity Vt is the direction from the noise to the data point: and
24 = (1 - t) Fo + t
112 = 21
To-
(1)
Given xo and T1, the velocity i; is fully determined. But given only Te. there are multiple plausible pairs (zo, 21) and thus different values the velocity can take on, rendering it random variable. Flow models learn neural network to estimate the expected value it Elv, 2t] that averages over all plausible velocities at 2₁. The flow model can be optimized by regressing the empirical velocity of randomly sampled pairings of noise To and data 1 pairs:
Uo t) R E..DD [ete x:]
f (0) = EZ [I1 (x1 - xo)
(2)
We consider flow-matching as special case of diffusion modelling (Kingma & Gao, 2024), and use the terms interchangeably.
2

---

Published as a conference paper at ICLR 2025
Training pairings
Learned ODE
4 Steps
2 Steps
Step
æ0 noise
221 data
2/a Dz
Flow Matching Training
Flow Matching Inference
Naive diffusion and flw-mmthhhng models fail at few-step generation. Left: Training paths are created by randomly pairing data and noise. Note that the paths overlap; there is inherent uncertainty about the direction fy to the data point, given only 24- Right: While flow-matching models learn a deterministic ODE, its paths are not straight and have to be followed closely. The predicted directions 2t point towards the average of plausible data points. The fewer inference steps, the more the generations are biased towards the dataset mean, causing them to go off track. At the first sampling step, the model points towards the dataset mean and thus cannot generate multi-modal data in single step (see red circles).
Figure 2:
To sample from flow model, noise point To is first sampled from the normal distribution. This point is then iteratively updated from xi to æ following the the denoising ODE defined as following the learned flow model (i l). In practice, this process is approximated using Euler sampling over small discrete time intervals.
Few-step ambiguity. While perfectly trained ODE deterministically maps the noise distribution to the data distribution in continuous time, this guarantee is lost under finite step sizes. As illustrated in Figure 2, flow-matching learns to predict the average direction from z; towards the data, so following the prediction with a large step size will jump to an average of multiple data points. At I the model receives pure noise as input and (zo; are randomly paired during training, so the predicted velocity at = 0 points towards the dataset mean. Thus, even at the optimum of the flow matching objective, one step generation will fail for any molti- data distribution.
SHORTCUT MODELS FOR FEW STEP GENERATION
We introduce shortcut models, new family of denoising generative models that overcomes the large number of sampling steps required by diffusion and flo-mmtchhing models. Our key intuition is that train single model that supports different sampling budgets, by conditioning the model not only on the timestep t but also on a desired step size d.
As shown in Figure 2, How-matching learns an ODE that maps noise data along curved paths. Naively taking large sampling steps leads to large discretization error and in the single-step case, to catastrophic failure. Conditioning on allows shorteut models account for the future curvature, and jump to the correct next point rather than going off track. We refer to the normalized direction from I, towards the correct next point a' as the shortcut s(x,, t, d)d:
2'4d
+ s(xe, d)
Our aim is to train a shorteut model (2,,,, t, d) to learn the shorteut for all combinations of 21, t. and d. Shortcut models can thus be seen as a generalization of flow-matching models to larger step sizes: whereas flow-matching models only learn the instantaneous velocity, shorteut models additionally learn to make larger jumps. At + O, the shorteut equivalent to the flow.
naive way compute targets for training se (æt, t, would be fully simulate the ODE forward with a small enough step size (Luhman & Luhman, 2021; Liu et al., 2022). However, this approach computationally expensive, especially for end-to-end training. Instead, leverage inherent self-consistency property of shortcut models, namely that one shortcut step equals two consecutive shortcut steps of half the size:
4, =
s(1/=a 4 d, d2/2
This allows us to train shortcut models using self-consistency targets for d > 0 and using the flow- matching loss (Equation base case for d = 0. In principle, train the model any

---

Published as a conference paper at ICLR 2025
(2 ()
Regress E[u, [z,]
a) Diffusion Flow Matching
Train
towards
8o (Tte fo
two smaller
steps
Shortcut Models
Overview of shorteut model training. At d 0, the shortcut objective is equivalent to the flow-matching objective, and can be trained by regressing onto empirical E[0, [z,] samples. Targets for larger shortcuts are constructed by concatenating sequence of two d/2 shortcuts. Both objectives can be trained jointly; shortcut models do not require a two-stage procedure or discretization schedule.
Figure 3:
distribution of d ~ p(d). In practice, we split the batch into a fraction that is trained with d = and another fraction with randomly sampled d >> O targets. We thus arrive at the combined shortcut model loss function:
_ (0) Exgn ai (1.)))p(..))
[150 0) - (x) xo)1 ( L, 2d) - Stargee 11?
Flow Matching
Self- Consistency
(5)
where
Starget
s0( L, d)/2 so s0( t d, d)/2
and 2/ = Izz sø (Ze, d)d.
Intuitively, the above objective learns a mapping from noise to data which is consistent when queried under any sequence of step sizes, including directly in single step. The flow-matching portion of the objective grounds the shortcut model at small step size to match empirical velocity samples. This ensures that the shortcut model develops base generation capability when queried with many steps, exactly as an equivalent flow-matching model does. In the self-consistency portion, appropriate targets for larger step-sizes are constructed by concatenating sequence of two smaller shortcuts. This propagates the generation capability from mulli-step to few-step to one-step. The combined objective can be trained jointly, using a single model and over single end-to-end training run.
3.1 TRAINING DETAILS
We now present simple framework for training shortcut models via the objective described above. At each stage, we opt for design decisions which encourage training stability and simplicity.
Regressing onto empirical samples. As + 0, the shortcut is equivalent to instantaneous flow. Thus,
train the shortcut model at d = 0 using the loss given by Equation i.e. by sampling random (Io. I1) pairs and fitting the expectation over This term can be seen as grounding the small-step shortcuts to match the data denoising ODE. We find that sampling (0, 1) uniformly is the simplest and works as well as any other sampling scheme.
Enforcing self-consistency. Given that the shortcut model accurate small step-size, our next goal is to ensure that the shortcut model maintains this behavior at larger step-size. We rely on self-generated bootstrap targets for this purpose. To limit compounding approximation error, desirable to limit the total length of the bootstrap paths. We therefore opt for binary recursive formulation in which two shortcuts are used to construct a twice-as-large shortcut (Figure 3).
We must decide on number of steps M to represent the smallest unit of time for approximating the ODE; we use 128 in our experiments. This creates (128) + 1 = 8 possible shortcut lengths according to d E (1/128, 1/64 ... 1/2, 1). During each training step, sample Iı, t. and random < 1, then take two sequential steps with the shortcut model. The concatenation of these two steps then used as the target to train the model 2d.
Note that the second step is queried at 2'44 under the denoising ODE and not the empirical data pairing, Í.e. it is constructed by adding the predicted first shorteut Ft, and not by interpolating
4

---

Published as a conference paper at ICLR 2025
Algorithm Shortcut Model Training
Algorithm 2 Sampling a ~ (\0,1)
while not converged do
~ (0, I), ~ D, (d,t) p(d,t) Zy f- (1 - t) Zo + tx
d 4 1/M +
Noise data point
for first batch elements do
for
[0. [0..... M do
Starget f- 21 20
Flow-matching target
T f- x + (0 L, d)
d +- 0
|+ d
for other batch elements do
return Σ
t, d) -- so sa [2,
First small step
Fitd t 2t + 8 s
Follow ODE
St+d + sa(Ietd L + d, d)
Second small step
Stargtt +- stopgrad(:; 0t + S240) /2
Self-consistency target
< Vollso(Ze,1 2d) Starger ||
towards 21 from the dataset. When is at the smallest value (e.g. 1/128), we instead query the model d =
Joint optimization. Equation 5 consists of an empirical flow-matching objective and self- consistency objective, which are jointly optimized during training. The variance of the empirical term is much higher, as it regresses onto random noise pairings with inherent uncertainty, whereas the self-consistency term uses deterministic bootstrap targets. We found it helpful construct batch with significantly more empirical targets than self-consistency targets.
The above behavior also gives us room for computational efficiency. consistency targets than empirical targets, and seff-consistency targets are also more expensive to generate (requiring two additional forward passes). We can therefore construct training batch by combining ratio of 1 k empirical targets with k self-consistency targets. We find k (1/4) to be reasonable. In this way, we can reduce the training cost shortcut model be roughly only 16% more than that of an equivalent diffusion model?
Training requires less self-
Guidance. Classifier-free guidance (CFG; Ho & Salimans, 2022) has proven to be an essential tool for diffusion models to reach high generation fidelity. CFG provides a linear approximation of tradeoff between the class-conditional and -uncnniiiioall denoising ODE. We find that CFG helps small step sizes but is error- error-prone at larger steps when linear approximation is not appropriate. We therefore use CFG when evaluating the shortcut model at but forgo elsewhere. A limitation of CFG in shorteut models is that the CFG scale must be specified before training.
Exponential moving average weights. Many recent diffusion models use an exponential moving average (EMA) over weight parameters to improve sample quality. EMA induces smoothing effect on the generations, which is especially helpful in the in diffusion modelling since the objective has inherent variance. We find that similarly shortcut models, variance from loss the d level can result in large oscillations in the output at d = 1. Utilizing EMA parameters for generating self-consistency targets alleviates this issue.
Weight decay. We find that weight decay is crucial for enabling stability. especially early on in training. When the shortcut model is near initialization, the self-consistency targets generates are largely noise. The model can latch on to these incoherent targets, resulting in artifacting and bad feature learning. We find that proper weight decay causes these issues to disappear, and enables us to bypass the need for discretization schedules careful warmups.
Discrete time sampling. practice, we can reduce the burden of the shortcut network by only training on relevant timesteps. During training, we first sample d, then sample only at the discrete points for which the shortcut model will be queried, i.c. multiples of d. We train the self-consistency objective only at these timesteps.
Approximatiig a hackward pass twice the compute of forward pass. Each shortcut update uses (forward) + 2 (hackward) (1/4)*2 (self-consistency targets) compute units, units diffusion update.
5

---

Published as a conference paper at ICLR 2025
4 RELATED WORK
Distillation of diffusion models. A number of prior works have explored the distillation of pre- trained diffusion models into a one-step or ftw-step model (Luo, 2023). Knowledge distillation (Luhman Luhman, 2021) and rectified flows (Liu et al., 2022) generate synthetic dataset by fully simulating the denoising ODE. As full simulation is expensive, number of methods have been proposed that utilize bootstrapping to warm-start the ODE simulation (Gu et al., 2023; Xie et al., 2024). Alternatives to L2 distance have been proposed for distillation targets, such as ad- versarial (Sauer et al., 2023) or most closely relates to techniques using binary ditellistillttion (Salimans & Ho, 2022; Meng al,, 2023; Berthelot et al., 2023), which divides distillation into log2(T) stages of increasing step-size, shortening the required bootstrap paths. Unlike these prior works, we focus on learning a one-step generative model end-to-end, without a separate pretraining and distillation phase. Our method is computationally cheaper than full simulation methods (e.g. rectified flows, knowledge distillation) and avoids the multiple teacher-student phases of progressive distillation methods.
mattuiinnn (Yin et al., 2024b;a) objectives. Our work
Consistency modelling. Consistency models (Song et al., 2023) are a family of one-step gener- ative models that learn to map from any partially-noisd data point to the final data point. While such models can be used as a student for distillation (Luo et al., 2023; Geng et al., 2024), consis- tency training has also been proposed train consistency models end-to-end from seratch (Song et al., 2023; Song & Dhariwal, 2023). Our work tackles the same problem setting, but approaches it differently consistency training enforces consistency among empirical Ty and Fr+d samples, which accumulates irreducible bias at each discretization step due to ambiguity. We instead sample rnda from the learned ODE, avoiding this issue. Shortcut models require only log2(2") bootstraps, whereas consistency models require T. Shortcut models naturally support many-step generation as a con-boosttrappe base case, whereas consistency models require bootstrapping in all cases. In ad- dition, shorteut models are practically simpler, many of the consistency model tricks e.g. using a strict discretization schedule, using perceptual loss rather than L2 loss can be bypassed entirely.
EXPERIMENTS
We present a series of experiments evaluating quality, scalability, and robustness of shortcut models. In terms of quality, shortcut models show few- and one-shot generation capability competetive with prior two-stage methods, and outperform alternative end-to-end methods. At the same time, shortcut models maintain the performance of base diffusion models on many-step generation. We show that under increasing model scale, generation capability continues to improve. Shortcut models learn an interpolatable latent noise space, and are robust to applications such as robotic control.
5.1 How DO SHORTCUT MODELS COMPARE TO PRIOR ONE-STEP GENERATION METHODS?
In this section, carefully compare approach to number of prior methods, by training each objective from scratch with the same model architecture and codebase. We utilize the DłT-B diffu- sion transformer architecture Peebles & Xie (2023). We consider CelebAHQ-256 for unconditional generation, and Imagenet-256 for class-conditional generation. For all runs, we utilize the Adam W optimizer (Loshchilov, 2017) with a constant learning rate of 1 10-4 and weight decay of 0. All runs use the latent space of the sd-vae ft -mae autoencoder (Rombach et al., 2022).
Comparison to prior work. We compare against two categories of diffusion-basee prior meth- ods: two-stage methods, which pretrain diffusion model then distill it separately, and end-to-end methods, which train a one-step model from scratch over single training run.
Diffusion (DiT-B) represents standard diffusion model, faithfully following the setup in (Pee- bles & Xie, 2023).
Flow Matching replaces the diffusion objective with an optimal transport optimal-transport objective, following Liu et al. (2022). Together with diffusion, these two methods provide baseline for the perfor- mance of an iterative many-step denoising model. The rest of the methods all adopt the flow- matching objective as base, using standard flow-matching model teacher if applicable.

---

Published as a conference paper at ICLR 2025
Reflow provides comparison to standard two-stage distillation approach which fully evalu- ates a teacher model to generate synthetic (To, r1) pairs. We follow methodology from Liu et al. (2022), and generate 50k synthetic examples for CelebAHQ, and 1M synthetic examples for Ima- genet. Rach example requires 128 forward passes to generate. Reflow is comparable to knowledge distillation (Luhman & Luhman, 2021), except the student model is trained among all E (0,
rather at just t= O.
Progressive Distillation provides comparison to two-stage binary time-distillation approach, which aligns with our proposed method. Following Salimans Ho (2022), starting from pre- trained teacher model, series of student models are distilled, each with 2x larger step-size than the previous. To match our methodology, classifier-free guidance is used during the first distilla- tion phase.
Consistency Distillation compares a consistency-mont model based two-stage distillation strategy. Pairs of (Zt.Zz+dd Zz+aa) are generated via teacher diffusion model, and a separate student consistency model trained enforce consistent #1 prediction among these pairs.
Consistency Training provides a comparison to a prior end-to-end method which trains a one- step model from scratch, most closely matching our setting. A consistency model is trained over empirical (ar, It++-) d) samples from the dataset. Following Song et al. (2023), the time discretiza- tion bins scheduled to increase over the training run.
Live Reflow is an additional end-to-end method we propose, in which a model is trained on both flw-matching and Reflow-distilled targets simultaneously. The model is conditioned on each tar- get type separately. Distillation targets are self self-generate every training step using full denoising, so this method is considerably expensive computationally. We include for comparison purposes.
Together, these works span the space of prior two-stage distillation and end-to-end training methods. To ensure rigorous comparison, we run all prior work comparisons on the same codebase, using an equal or greater amount of total computation than our method.
Evaluation. We evaluate models on samples generated with 128, 4, and 1 diffusion step(s), using the standard Frechet Inception Distance (FID) metric. We report the FID-50k metric, as is standard in prior work. Following standard practice, FID is calculated with respect to statistics over the entire dataset, compression is applied the generated images, and images resized to 299x299 with bilinear upscaling and clipped to (- 1). We use the EMA model parameters during evaluation.
Table highlights the abilities of shorteut models to retain accurate generation under few- and one-step sampling. With the exception of two-stage progressive distillation, shortcut models out- perform all prior methods, without the need for multiple training stages. Note that progressive distillation models lose the ability for many-step sampling, whereas shorteut models retain this abil- ity. Unsurprisingly, diffusion and flow-matching models display poor performance on 4- and 1-step generation. Interestingly, in our experiments shortcut models display a slightly better FID than a comparable mot-hia model. A blind hypothesis that the self-consistency loss acts as form of implicit regularization of the model, although we leave this investigation to future work. Further visual examples are provided in Appendix A.
5.2 WHAT IS THE BEHAVIOR OF SHORTCUT MODELS UNDER VARYING INFERENCE BUDGETS?
We now analyze the behavior of shortcut models when images are generated with varying step counts. While naïve flow-matching models approximate the denoising ODE with a tangent velocity, shortcut models are trained to approximately integrate the denoising ODE, providing more accurate behavior low step counts. Figure 4 shows that, while greater numbers of steps are always helpful, the degradation at low step counts is much more pronounced in naïve flow-matching models than in shortcut models. In our experiments, we find that many-step performance not sacrificed, and shortcut models maintain the performance of bascline flow models even at high step counts.
Few-step artifacts in flow models resemble blurriness and mode collapse (Figure 1, left). Short- cut models generally do not suffer from such issues, and globally match the equivalent many-step images generated from the same initial noise. Artifacts in shortcut models resemble errors in high- precision details. One-step shorteut models therefore be used proxy: if a downstream user wishes to refine a one-step generated image, they may simply regenerate the image from the same initial noise, but with more generation steps.
7

---

Published as a conference paper at ICLR 2025
CelebAHQ-256 (unconditioned)
Imagenet-256 (class conditioned)
128-Step 4-Step
1-Step
128-Step
4-Step
1-Step
Two phase training
Progressive Distillation
(302.9)
(251.3)
14.8
(201.9)
(142.5)
35.6
Consistency Distillation
59.5
39.6
38.2
132.8
98.01
136.5
Reflow
16.1
18.4
23.2
16.9
32.8
44.8
End-to-end (single training run)
Diffusion
23.0
(123.4)
(132.2)
39.7
(464.5)
(467.2)
Flow Matching
7.3
(63.3)
(280.5)
17.3
(108.2)
(324.8)
Consistency Training
53.7
19.0
33.2
42.8
43.0
69.7
Live Reflow (ours)
6.3
27.2
43.3
46.3
95.8
58.1
Shorteut Models (ours)
6.9
13.8
20.5
15.5
28.3
40.3
Table I: Comparison of training objectives under equivalent architecture (DiT-B) and com- pute. FID-50k scores (lower is better) are shown over 128, 4, and 1-step denoising. Shortcut models provide high-quality samples under any inference budget, within single training run. Compared to diffusion and flow-matching, shortcut models drastically reduce needed sampling steps. Com- pared to distillation approaches, short models simplify the training pipeline and provides flexibility to choose inference budget after training. Parentheses represent evaluation under conditions that the objective is not intended to support.
FID under decreasing inference budget (CelebAHQ)
300
Flow-Matching
280.0
250
Model
Figure 4: Behavior of flow-matching and short- cut models over decreasing numbers of denois- ing steps. While naïve flow-matching models incur degradation and mode collapse, shortcut models able to maintain similar sample distribution at few- and one-step generation. This capability comes no expense to generation quality under large infer- ence budget. See Figure for qualatative examples.
200
FID-S0k
150
117.8
63.3
11.9
6.9
7.2
7.0
7.
9.0
13.8
16.4
of Denoising steps
FID model (Imagenet-22)
140
Small (5)
S
Base
120
Large (L.
Extra Large (XL)
FDD4096
100
B
L
50K
150K
250K
XL
Training Steps
Figure One-step generation quality continues to improve as model parameter count in- creases. While generative models tend to display continual improvement with model scale, bootstrap-based methods such as Q-learning have been shown to lose this property (Kumar al., 2020). We show that shortcut models, while bootstrap-based method, retains the ability to scale accuracy with model size.
8

---

Published as a conference paper at ICLR 2025
Figure 6: Interpolations between two sampled noise points. All displayed images are model generated. Each horizontal set represents images generated by one-step denoising of variance- preserving interpolation between two Gaussian noise samples.
5.3 HOw DOES SHORTCUT MODEL PERFORMANCE INCREASE WITH MODEL SCALE?
A powerful trend in deep learning is that large werrpaaaaeeerrize. models increase in capabilities when scaled up. This trend has been less clear in bootstrap- methods such learning Q- functions for reinforcement learning (Kumar et al., 2020; Sokar et al., 2023; Obando-Ceron et al., 2024). Common issues resemble form of rank-collapse: by re-training on self-generated outputs, model expressivity may become limited. As shortcut models are inherently amine whether shortcut models can extract benefit from model scale. As highlighted Figure 5, shortcut model architecture can achieve increas-
pansstra-bbseede we ex-
even in the one-step setting, our Transformer-bant ingly accurate generation as model size increased. These results hint that even though shorteut models rely on bootstrapping, they nevertheless avoid drastic collapse and continue to scale with model parameter count. A shortcut model trained with DiT-XL size is able to achieve a one-step FID of 10.6 and 128-step FID of 3.8 on Imagenet-256, as shown in Table 2 in the Appendix.
5.4 CAN SHORTCUT MODELS GIVE US AN INTERPOLATABLE LATENT SPACE?
Shortcut models represent a deterministic mapping from noise to data. To examine the structure of this mapping, we interpolate between initial noise samples and view the corresponding genera- tions. Figure 6 displays examples of interpolations generated in this fashion. pair of noise points (=g: xb) are sampled, then interpolated parienie- manner:
z noy v n²z!
Despite no explicit regularization the noise-to-data mapping, the resulting interpolated genera- tions display a qualitatively smooth transition. Intermediate images appear to be semantically plau- sible. Note that in the visualized generations, all images are generated. While do not explore interpolation between existing images, it may be possible to add noise to an existing image and interpolate those intermediate points, is done by Ho et al. (2020).
5.5
Do SHORTCUT MODELS WORK IN NON- IMAGE DOMAINS?
The generative modeling objectives described in this work are all domain- agnostic. yet common benchmarks traditionally involve image generation. We now evaluate the generality of the shortcut models formulation by training shortcut model policies on robotic control tasks.
Specifically, build off the methodology presented by Chi et al. (2023). This diffusson policy framework trains an cearvation-condition
model to predict robot actions in an iterative manner. The original work uses 100 denoising steps, and we aim to reduce inference to one step by training a shortcut model instead. We utilize the same network structure and hyperparameters the original work, except for changing AdamW weight decay from 0.001 to 0.1, and including conditioning flow-matching target rather than epsilon prediction
term for d. We additionally
epsilon-predictioon target.
Figure highlights the robotic control tasks, Push-T and Transport, which were selected from Chi et al. (2023) as tasks where baseline methods struggle. We compare to iterative denoising meth- ods IBC (Florence et al., 2022) and diffusion policy (Chi et al., 2023), as well as one-step methods LSTM-GMM (Mandlekar al., 2021) and BET (Shafiullah al., 2022). Results display the pability of shortcut model policies to achieve strong performance on robotic control, while limiting inference cost single function call. In contrast, one-step diffusion policies catastrophically fail.

---

Published as a conference paper at ICLR 2025
# Steps
Push-T
Transport
IBC
100
0.90
0.00
Diffusion Policy
100
0.95
1.00
LSTM-GMM
1
0.67
0.76
BET
1
0.79
0.38
Diffusion Policy
1
0.12
0.00
Shortcut Policy (ours)
0.87
0.80
Figure 7: Shortcut models can represent multimodal policies similar manner as diffusion policies, while reducing the number of denoising steps to 1. Shown on the left are trajectories from the Push-T (top) and Transport (bottom) tasks. each case, generative model trained on human demonstrations, and the model is queried to produce actions given set of past observations.
6 DISCUSSION
This work introduces shoricut models, a new type of diffusion- generative model that supports few-step and one-step generation. They key idea behind shortcut models is learn a step-xize conditioned model that grounded to the data at small step-size, and trained via self- seff-bootstrapping at larger step-sizes.
Shortcut models provide simple recipe that requires scheduling and can be trained end-to-end in a single run. Unlike prior distillation methods, shortcut models do not require separate pretraining and distillation phases. In comparison to consistency training, shortcut models do not need a training schedule, use only standard L2 regression, and require less bootstrapping steps.
Empirical evaluations CelebA HQ and Imagenet-256 show that shortcut models outperform prior single-phase methods, and are competetive to two-stage distillation methods. Despite a reliance on bootstrapping, shortcut models stable to train and scale with model parameter size. We showcase how shortcut models can be used in non-image domains, such as robotic control.
Best practices. For practitioner wishing utilize shorteut models, recommend to keep the following in mind. Standard diffusion model practice should be followed e.g. normalizing the dataset unit variance. the self-consistency loss displays erratic behavior, one can reduce the ratio of self-consistency targets in the training batch. A non-zero weight decay is highly recommended. Evaluating with EMA parameters provides a consistent improvement, but is not strictly neccessary.
Implementation. We provide clean open-source implementation of shortcut models, along with model checkpoints, available at: https://github. http://iihuuuooovvviieeeee comm/vvrran//shoreeett. models
Limitations. While shortcut models provide a simple framework for training one-step generative models, a key inherent limitation that the mapping between noise and data entirely dependent on an expectation over the dataset. In other generative model flavors (e.g.. GANs (Goodfellow et al., 2020), VAEs (Kingma, 2013)) possible to adjust this mapping, which has the potential to simply the learning problem. Secondly, in our shortcut model implementation there remains gap between many-step generation quality and one-step generation quality.
Future work. Shortcut models open up range of future research directions. In Table 1, shortcut models display slightly better many-step generation than a base flow model-is is there a way for one- step generation improve many-step generation (rather than typically the opposite)? Extensions to the shortcut model formulation may be possible that iteratively adjust the noise-to-data mapping (e.g., with a Reflow-like procedure (Liu et al., 2022)), which reduce the gap between many- step and one-step generation. This work provides stepping stone towards developing the ideal generative modelling objective-a simple recipe which satisfies the trifecta of fast sampling, mode coverage, and high-quality generations.
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGEMENTS
This work was supported in part by an National Science Foundation Fellowship for KF, under grant No. DGE 2146752. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the NSF. PA holds concurrent appointments as Professor at UC Berkeley and as an Amazon Scholar. This paper describes work performed at UC Berkeley and is not associated with Amazon. We thank Google TPU Research Cloud (TRC) for granting us access to TPUs for research.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard. Joshua Bambrick, al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Omer Bar-Tal Hila Chefer. Omer Tov, Charles Herrmann, Roni Paiss, Shiran Zada, Ariel Ephrat, Junhwa Hur, Yuanzhen Li, Tomer Michaeli, et al. Lumiere: A space-time diffusion model for video generation. arXiv preprint 2tfl:::4001229444 2024.
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.042488 2023.
Andrew Brock. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:/809. 1096, 2018.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint 23Xiv:2003.0433,, 2023.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning, In Conference Robot Learning. pp. 168-168.10 PMLR, 2022.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. arXiv preprint 24X.v:2406.144: 14548, 2024.
Ian Goodfellow, Jean Pouget-Abadie, Mchdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):133 144, 2020.
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free dis- tillation of denoising diffusion models with bootstrapping. In ICML 2023 Workshop on Structured Probabilistic Inference &{& Generative Modeling. 2023.
Jonathan Ho and Tim Salimans.
Classifier-free
diffusion guidance.
arXiv preprint
arXiv:2207. 10.07.12550 12598, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi. and David J Fleet. Video diffusion models. Advances Neural Information Processing Systems, 35:8633- 8646, 2022.
11

---

Published as a conference paper at ICLR 2025
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
Diederik P Kingma. Auto-encoding variational bayes. arXiv preprint arXiv: 2013.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint wwXv:2200900976,, 2020.
Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit under--araametttiizaio inhibits data-efficient deep reinforcement learning. arXiv preprint arXiv:2010.1449) 2020.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint wrxiv:2209.03000 2022.
2017. Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2 20101001.... 2021.
Loshchilov. Decoupled weight decay regularization. arXiv preprint
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthe- sizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04378, 2023.
Weijian Luo. A comprehensive survey knowledge distillation of diffusion models. arXiv preprint arXiv:2334.04262, 2023.
Ajay Mandlekar, Danfei Xu, Josiah Wong. Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Li Fei- Fei, Silvio Savarese, Yuke Zhu, and Roberto Martin-Martin. What matters in learning from offline 2021.
human demonstrations for robot manipulation. arXiv preprint
Chenlin Meng. Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Johan Obando-Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures experts unlock parameter scaling for deep rl. arXiv preprint M4iv:2402.086099 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205, 2023.
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 confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005//2 2022.
Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings, pp. 1-10, 2022.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. arXiv preprint arXiv:2311.17044 2023.
Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing sys- tems, 35:22955-22994 2022.
Jascha Sohl- Sohk-Diicst,i,, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2256-2265. PMLR, 2015.
12

---

Published as a conference paper at ICLR 2025
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phe- nomenon in deep reinforcement learning. International Conference on Machine Learning, pp- 32145-32168. PMLR, 2023.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310.1411 2023.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
Sirui Xie, Zhisheng Xiao, Diederik P. Kingma, Tingho Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, and Ruiqi Gao. Em distillation for one-step diffusion models. ArXiv, abs/2405.16852, 2024. URL https:/ hesss//appl /api..p...// www.nticscholoooc. CorpusID: : 270062581.
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. arXiv preprint 1484777405.14667 2024a.
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 IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 1613-6663,, 2024b.
13

---

Published as a conference paper at ICLR 2025
A
APPENDIX
Figure 8: Examples generated from the unconditional CelebA-HQ dataset at 256x256 resolution. Each triplet represents an image generated from the same noise, at (128, 4, 1) denoising steps from top to bottom. Results are generated from DiT-B size model trained for 400k iterations.
14

---

Published as a conference paper at ICLR 2025
Figure 9: Examples generated from the class-conditional Imagenet dataset at 256x256 resolution. Each triplet represents an image generated from the same noise, at (128, 4, 1) denoising steps from top to bottom. Results are generated from a DiT-XL size model trained for 800k iterations. CFG of 1.5 is used during the training process (not during generation).
15

---

Published as a conference paper at ICLR 2025
FID
Sampling Steps
Param Count
Epochs Trained
Adversarial approaches
BigGAN-deep (Brock, 2018)
6.96
StyleGAN-XL (Sauer et al., 2022)
2.3
Denoising approaches
DīT-XL (Peebles & Xie, 2023)
2.27
500
675M
640
ADM-G (Dhariwal & Nichol, 2021)
4.59
250
426
LDM-4-G (Rombach al., 2022)
3.6
500
400M
106
Shorteut Model (XL)
3.8
128
676M
250
Shorteut Model (XL)
7.8
676M
250
Shortcut Model (XL)
10.6
676M
250
Table 2: Comparison to ofte-oof-hheart generative models Imagenet-256. Due compute straints, we cannot train models with the same compute as the best previously reported generative models. However, results demonstrate that shortcut models quality with model size. The Shortcut Model (XI.) using the DiT-XI. architecture reaches substan- tially lower FID that the DiT-B architecture in Table and is competetive with previous state-of- the-art models. Note that the compared models use different architectures and compute budgets. Moreover, parameter count is not completely proportional to training compute due to weight shar- ing, e.g. in convolutional and attention layers. Horizontal dashes indicate unreported quantities.
able to competitively scale up
B TRAINING DETAILS
Batch Size
(CelebA-HQ), 256 (Imagenet)
Training Steps
400,000 (CelebA-HQ) 800,000 (Imagenet)
Latent Encoder
sd-vae-mse-ft (256x256x3 to 32x32x4)
Latent Downsampling
Ratio of Empirical to Bootstrap Targets
0.75
Number of Total Denoising Steps (M)
128
Classifier Free Guidance
0 (CelebA-HQ), 1.5 (Imagenet)
Class Dropout Probability
0.1
EMA Parameters Used For Bootstrap Targets?
Yes
EMA Parameters Used For Evaluation?
Yes
EMA Ratio
0.999
Optimizer
Adam
Learning Rate
0.0001
Weight Decay
0.1
Hidden Size
768
Patch Size
Number of Layers
12
Attention Heads
12
ML.P Hidden Size Ratio
Table Hyperparameters used during training. Model architecture follows that described in Peebles & Xie (2023), specifically DiT-B unless mentioned otherwise.
B.1 COMPUTATION
All experiments are run on TPUv3 nodes, and methods are implemented in JAX. While runtimes vary per method, each training typically takes 1-2 days complete.
16

---

Published as a conference paper at ICLR 2025
B.2 METHOD DETAILS
Followed are description of details behind the comparison methods. All methods are implemented in the same codebase, using the same architecture. Training budgets chosen that comparisons utilize roughly equal compute (or more compute if neccessary for the method to produce reasonable results).
Flow-Matching, We use the standard linear interpolation for generating velocity targets. We train the base model for 800k iterations, and denoise images using deterministic Euler sampling. A checkpoint at 400k iterations is used as teacher model for two-stage distillation methods.
Reflow. We use a base flw-matcching model create the synthetic targets used in reflow. For CelebA-HQ. 50K targets are generated. 1M targets are generated for Imagenet. The distillation process is then run for 400k additional iterations on the synthetic dataset. CFG is used to generate synthetic data when applicable, and no CFG is used on the resulting distilled model.
Progressive Distillation. We run progressive distillation for the standard distillation time (400k iterations), split into equal sections for each phase of distillation. As there are distinct phases, this results 50k training steps per phase. At each phase, a teacher model is used to sample two sequential bootstrap steps, which student model is trained to mimic. At the end of training phase, the student becomes the new teacher. CFG is used only for the first phase (distilling 128-step model into a 64-step model).
Consistency Distillation. We train a consistency model in velocity space, to match the base frame- work of flow-matching. Thus each iteration, consistency in velocity prediction is enforced be- tween two points (Iz. Tt+a). The two points are computed by taking a small = (1/128) step from T2 in the direction of the teacher flow model. Target velocities are computed from 2x4-d by querying the consistency model, using this prediction to estimate æ1, then computing target velocity as (x1 21)/(1 t).
Consistency Training, A consistency model is trained in the same format as described above. In- stead of sampling (It, Zt+a) from teacher model, the pairs are sampled by interpolating noise and empirical data samples at different strength. Following the methodology in Song et al. (2023), a dis- cretization schedule is used. We use a binary-time schedule similar that progressive distillation. Specifically, the first phase has single discretization interval, then two, then four, etc.
Live Reflow, In this method, we convert Reflow procedure that can run in a single training run. Similar to a shortcut model, we train d-conditioned model. A percentage of the batch will train the model at d 0 towards the flow-matching loss. We use 0.75 as this ratio, the same as the shortcut model. The other part of the batch consists of self- eelf-generated bootstrap targets generated by fully denoising a set random noise. As this is a considerably expensive prodecure, we limit the number of denoising steps for target generation to 8. Even so, live reflow takes over 4x the computation of any other method.
17


Reviews:
[{"review_id": "0sZUactcQL", "reviewer": "Reviewer_JHdH", "summary": "The paper introduces shortcut models, a modification to the diffusion approach to be conditioned on noise level and step size. The approach is able to train end-to-end flexible models that simultaneously allow many-steps and one-step sampling.", "questions": "- Comparing the various approach using the same compute budget could not be the best way as different methods may require different compute budget to converge to their optimal performance, did you take this consideration into account? It would be helpful to detail how the compute budget for each method has been fixed, and if you expect different convergence rates for each method. Moreover a plot of performance versus compute would be a nice addition if it's feasible.\n\n- How does this method relates in any way to the ICLM 24 paper \" Optimizing DDPM Sampling with Shortcut Fine-Tuning\", I think the paper should be discussed as very relevant to the proposed method. \n\n- It would be nice to introduce a comparison on additional compute required by post training distillation methods compared to the proposed solution which just increase the required compute by 16% over a standard diffusion training. A table highlighting the total compute overdue required by each method should be added.\n\n- The method presents a lot of hyperparameters to be tuned for training, as the selection of M (total number of sampling steps) and k, have you ablated how varying this affect the final results? do you have estimated the impact of choosing a bigger/smaller k? Can you show any ablation study varying some of these parameters, and provide generale guidelines for their selection to pratictioners?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper is well written and easy to follow\n\n- The proposed idea is novel and promising, this work pioneers a new family of diffusion/flow based generative models.\n\n- The method is experimentally evaluated in a controlled setting, an extensive comparison with alternative methods is present.\n\n- An open source codebase to replicate all the experiments and pretrained models will be release to falicitate further research", "weaknesses": "- The performance of 1 step generated samples still lags behind multi-step methods; 1-step generated images presents some artifacts so the usability of the proposed solution in a pratical setting is limited to generate proxy images with 1 step and resample them with multiple steps\n\n- end-to-end training to reach a few sample diffusion model is an interesting solution as it removes the burden of performing two separate training phases and leads to a flexible model. However the best results in 1 step regime are obtained by Progressive Distillation, thus demonstrating that the best option for a one step model is to distill it in a separate training phase. \n\n- It may be possible that learning few step models in a joint way introduces instabilities to the learning process especially when using much more complex data distribution than ImageNet or CelebA dataset. There are no guarantees on how the proposed approach would scale to large data regime training.", "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": "5yg86IHZ6Q", "reviewer": "Reviewer_w2mF", "summary": "This paper proposes shortcut modes, an end-to-end training method that could produce high-quality images in either single or multiple sampling steps. The shortcut models take an additional input of the desired step size, compared to diffusion or consistency models.", "questions": "* Why the CFG scale must be specified before training?\n* Why use 128 as the number of steps? What about something as large as 1024? How would this design choice affect the quality? \n* Why not consider a two-stage distillation training, so that a pretrained diffusion model could be converted to a shortcut model?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The paper is well written in the context of reducing sampling steps of diffusion models.\n* The proposed method is novel and clearly presented.\n* The experiments are comprehensive and convincing.", "weaknesses": "* Lacks explicit comparison of training compute against other methods.\n* The issue with CFG could hinder future practical use.", "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": "OTSwXtQ5wo", "reviewer": "Reviewer_PitZ", "summary": "The paper proposes denoising generative model called shortcut model. These models are trained end-to-end and can generate samples in both single-step and few-step setting. The core idea is to condition the model on step size in addition to time step, which is typically done in conventional diffusion models. The hypothesis is that the additional conditioning on step size allows shortcut models to account for the future curvature. The model is trained via a loss objective that combines the optimal-transport flow matching loss with a self-consistency loss. The self-consistency loss is enforced by ensuring that the prediction of model at a given skip size matches the target obtained by two sequential smaller skips (for the same skip distance) with the shortcut model. The high computational cost as well as variance during training is dealt with by computing self-consistency loss on a small subset of samples in the batch. The proposed method has been tested on CelebAHQ-256 and ImageNet-256, and has good single-step image generation quality.", "questions": "1. For the results of Progressive Distillation in Table 1, a reasonable 128-step model would be the obtained after taking the initial (say, T=1024) step model and then distilling it progressively until we get 128-step model. Was this model used to report metrics in the table or, did the authors use the final 1-step model to generate samples with 128 steps? \n2. Appendix B.1 - could the authors specify the amount of TPUv3 nodes needed to run these experiments in 1-2 days? \n3. For consistency training, is there a reason why the authors chose to deviate from the proposed discretization schedule by Song et al. (2023) during training? From my experience with these models, the choice of discretization schedule matters a lot, and deviations affect the performance. Further, does this paper implement consistency training (which includes perceptual loss)[1] or the follow up paper [2] which uses pseudo-Huber loss and proposes several improvements over their previous paper.\n4. Could the authors include architectural details/positional encoding details of both of the conditioning (time/noise level and step size)?\n5. Could the authors include some qualitative results of one-step generation of some other methods they have trained on in the appendix? FID score doesn’t always correlate with human perception of images. \n6. Minor: Related work: for consistency models, the line needs revision - “shortcut models are practically simpler, as many consistency model tricks - e.g. using a strict discretization schedule, using perceptual loss rather than L1 loss - can be bypassed entirely. ” In their latest paper, Improved Consistency Training [2], Song et al. already bypass the need of perceptual loss; they use only pseudo-Huber loss.\n7. Minor: some typos: competetive -> competitive (Section 5, line 2), qualatative -> qualitative (Figure 4 caption)\n\n[1] Song, Yang, et al. \"Consistency models.\" arXiv preprint arXiv:2303.01469 (2023).\n[2] Song, Yang, and Prafulla Dhariwal. \"Improved techniques for training consistency models.\" arXiv preprint arXiv:2310.14189 (2023).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. _Simplicity_: The proposed method is simple, easy to understand and simplifies some challenges previously encountered in other methods to obtain one-step or few-step models. For eg. It gets rid of two-stage training required for diffusion distillation as well as complex training scheduling required for consistency models. The loss objective is intuitive and builds upon flow matching loss and introduces an additional self-consistency condition. \n2. _Writing_: The paper is well-written. The core method has been explained well. The paper talks about challenges that can be encountered while training shortcut models and ways to address these challenges. \n3. _Computational efficiency_: The method seems computationally efficient as with an additional (marginal) cost of training diffusion models, shortcut models can be trained to generate samples in single-step as well as few steps. Further, even on larger datasets such as ImageNet-256 and CelebAHQ-256, these models can be trained in 1-2 days (though the exact number of TPUv3 nodes is unclear).  \n4. _Applications_: The authors demonstrate advantages of this method on both image-generation as well as on robotic control tasks showing ways to generate an iterative sequence of actions in a single denoising step per action.", "weaknesses": "Note: The template of this paper seems to be different and doesn’t have line numbers present in the standard template for papers under review. \n\n1. The proposed method for shortcut sampling is general and should work for any gaussian probability paths ($x_t = \\alpha_t x_0 + \\sigma_t \\epsilon$) however this hasn’t been explored or described in the paper. The paper considers only one choice of optimal transport probability path ($x_t = (1-t) x_0 + t \\epsilon$). However, nothing constraints the method to be restricted to this particular choice. (One motivation for this choice of parametrization could be that OT paths are said to have less curvature but in this case, these are conditional OT paths and are not exactly straight lines paths between $x_1$ and $x_0$.) It is worth exploring ablations on how shortcut models perform on other parametrization of gaussian probability paths/sampling paths. In addition, the text should also be updated to have a more general parametrization.\n2. The paper should also include quantitative performance on more commonly used datasets such as ImageNet-64 and CIFAR-10. While reporting results on larger datasets such as CelebAHQ-256 and ImageNet-256 is impressive, reporting metrics on these smaller datasets is also important. Most of the prior papers report metrics on these datasets which makes it easier to compare the performance of shortcut models to a larger number of previously proposed methods/models which are not currently reported in the paper (Note: I understand that it is impossible to replicate all the previously proposed methods due to the vastness of literature in the space).\n3. Minor: I personally feel that images such as in Figure 1 are misleading as the original flow-matching objective was not specifically meant for single step generation, and the performance is expected to deteriorate because these models don’t learn an ODE integrator. There are some advantages of the original flow-matching models as well over shortcut models. These models are continuous time models which can be evaluated at any time between [0, 1] whereas shortcut models are discrete time and can only be evaluated at specific discrete time steps.", "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": "BWskyFRmVk", "reviewer": "Reviewer_NSZo", "summary": "The paper introduces a shortcut design (self-consistency objective) in flow models for efficient and high-quality sampling. Unlike complex methods, this shortcut leverages existing networks and requires a single training phase.\nThe shortcut design involves incorporating an additional step size condition into the generative models, enabling them to learn the objective while considering the step size. This approach showcases improved performance in image and policy generation domains.", "questions": "* In Equation 5, Flow-Matching and Self-Consistency are intended to correspond to different units of time. Flow-Matching is represented by d=0, whereas Self-Consistency is associated with d > 0. Does it affect the training data samples or procedures?  \n\n* In Algorithm1, why stopgrad is applied in self-consistency target? Will it disable the gradient update?  \n\n* Table1 suggests two-phase training can give better results than all end-to-end methods in one step sampling: progressive distillation in 1-step:14.8 and 35.6 for Celeb and ImageNet. Does the two-phase method has higher performance ceiling than end-to-end? \n\n* In Figure 7, does increasing the sampling steps lead to a higher success rate for Short-cut Policy, similar to the diffusion policy? In policy planning, the primary focus is on achieving multi-mode and high precision performance.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a shortcut design (self-consistency objective) in flow models to maintain consistency in intermediate sampling trajectories. This shortcut involves adding a self-consistency objective during training without significant additional burden. Empirically, it consistently outperforms other end-to-end methods in terms of generation quality.", "weaknesses": "Limited comparison with one-step generator GAN is provided. In Appendix Table 2, StyleGAN-X demonstrates the best performance (2.3). Does this suggest that GAN is currently the optimal choice for one-step sampling?", "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 (Oral)