Source-linked AI summary
One Step Diffusion via Shortcut Models
Kevin Frans, Danijar Hafner, Sergey Levine, Pieter Abbeel
TL;DR
Diffusion and flow-matching models produce high-quality samples but are slow because inference requires many denoising passes, while earlier acceleration methods add training complexity. Shortcut models use one network trained in one phase, conditioned on noise level and desired step size, and achieve high-quality generation across inference budgets while matching or outperforming prior methods. Their main limitation is a remaining quality gap between many-step and one-step generation.
Problem
Diffusion and flow-matching models require expensive iterative inference, while few-step flow matching can fail because predicted directions average over ambiguous data possibilities.
Method
Shortcut models train one end-to-end network conditioned on timestep and desired step size, using small-step flow matching and self-bootstrapping at larger step sizes.
Results
Shortcut models support high-quality many-step, few-step, and one-step generation, matching baseline diffusion quality at many steps and matching or outperforming two-stage distillation at fewer steps.
Takeaways & Limitations
Shortcut models provide a single-run, schedule-free approach that supports variable inference budgets and extends beyond image generation to robotic control.
Takeaways & Limitations
A remaining gap exists between many-step and one-step generation quality.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Diffusion and flow-matching models generate high-quality diverse data but require expensive iterative inference. Shortcut models address this by conditioning one end-to-end network on desired step size, supporting high-quality generation across inference budgets and matching or outperforming prior approaches.
- Diffusion and flow-matching methods often require dozens to hundreds of neural-network forward passes, making generation slow and expensive.
- Shortcut models condition a single network on both noise level and desired step size, enabling it to jump ahead during denoising.They require one training run, no schedule or careful warmup, and approximately 16% more training compute than a base diffusion model.
- A single shortcut model supports many-step, few-step, and one-step generation on CelebA-HQ and ImageNet-256.
- Shortcut models match baseline diffusion quality for many-step generation and consistently match or outperform two-stage distillation in few- and one-step settings.
- The paper compares shortcut models with prior diffusion, flow-matching, distillation, and end-to-end methods under fixed architecture and compute, and extends them to robotic control.In robotic control, shortcut policies maintain comparable performance at an order-of-magnitude lower inference cost.
2 BACKGROUND
Flow-matching models learn an ODE from noise to data by predicting expected velocities along interpolated paths. With finite or very large sampling steps, averaging ambiguous directions and following curved paths can cause few-step generations to leave the intended trajectory and collapse toward dataset averages.
- Flow matching linearly interpolates noise x0 and data x1 as xt = (1 − t)x0 + tx1, with velocity vt = x1 − x0.
- Because multiple noise-data pairs can produce the same xt, flow models predict the expected velocity E[vt | xt] over plausible directions.
- Sampling follows the learned denoising ODE from noise to data, typically approximated by Euler updates over small discrete intervals.
- Few-step ambiguity: With large finite steps, curved flow paths create discretization error, while averaging ambiguous directions can move generations toward the dataset mean.
- Few-step ambiguity: At t = 0, one-step flow matching points toward the dataset mean and therefore fails for multimodal data distributions.
3 SHORTCUT MODELS FOR FEW STEP GENERATION
Shortcut models learn step-size-conditioned denoising shortcuts in one jointly optimized training run, combining flow matching at d=0 with self-consistency targets for larger steps. This lets one model support varied sampling budgets, including one-step generation, without schedules or two-stage training.
- Shortcut model formulation: Shortcut models condition on timestep t and desired step size d to learn denoising jumps across different sampling budgets.They generalize flow matching beyond instantaneous velocity prediction to larger jumps.
- Sampling behavior: Conditioning on d helps account for curved ODE trajectories, avoiding the large discretization errors caused by naively taking large sampling steps.The model can jump toward the correct next point instead of following an inaccurate straight approximation.
- Training objectives: The flow-matching and self-consistency objectives are jointly optimized in a single training run rather than through separate stages.The procedure avoids a discretization schedule and careful warmups, while using a mixture of empirical and self-consistency targets.
- Training objectives: At d →0, the shortcut objective uses empirical flow-matching targets to ground small-step generation in the denoising ODE.Training samples random noise-data pairs and regresses the expected velocity.
- Training objectives: Larger-step targets are constructed by concatenating two consecutive shortcuts of half the step size, enforcing self-consistency recursively.The second shortcut is evaluated from the model-predicted intermediate point.
- Training details: Classifier-free guidance is used at d=0 but omitted at larger steps because its linear approximation becomes error-prone there.Exponential moving average weights and weight decay are used to improve stability and sample quality.
4 RELATED WORK
Prior acceleration methods distill or train one-step models using synthetic simulations, bootstrapping, or discretization-dependent consistency procedures. Shortcut models reduce bootstrap requirements and retain a nonbootstrapped many-step operating mode.
- Distillation of diffusion models: Diffusion distillation and rectified-flow methods often generate synthetic data by fully simulating the denoising ODE, while bootstrapping can warm-start that simulation.These approaches target one-step or few-step sampling from pretrained diffusion models.
- Consistency modelling: Consistency models can be trained end-to-end or used as students for distillation, but shortcut models instead construct intermediate points from the learned ODE.This distinguishes the shortcut training target from the consistency-model setup described here.
- Comparison with consistency models: Shortcut models require log2(T) bootstraps rather than T and support many-step generation as a nonbootstrapped base case.They also bypass several consistency-model practices, including strict discretization schedules and perceptual loss.
5 EXPERIMENTS
Experiments evaluate shortcut models across image-generation step budgets, scale, interpolation, and robotic control. They find that shortcut models retain quality with few or one step while preserving many-step performance and extending to multimodal control policies.
- Inference budgets: Shortcut models degrade less than naïve flow-matching models as denoising steps decrease, while maintaining baseline flow-model performance at high step counts.Naïve flow matching shows stronger low-step degradation, including blurriness and mode collapse; shortcut models preserve similar sample distributions at one and few steps.
- Inference budgets: Shortcut models provide high-quality generation at 128, 4, and 1 denoising steps within a single training run.They reduce sampling steps relative to diffusion and flow matching while allowing the inference budget to be selected after training.
- Scaling: One-step generation quality continues to improve as shortcut-model parameter count increases, despite the method being bootstrap-based.This scaling behavior is examined because bootstrap-based approaches can lose performance improvements with model scale.
- Latent interpolation: Variance-preserving interpolations between Gaussian noise samples produce qualitatively smooth transitions of semantically plausible, model-generated images.The interpolation visualization uses one-step denoising for each interpolated noise point, without explicit regularization of the noise-to-data mapping.
- Robotic control: Shortcut-model policies achieve strong performance on Push-T and Transport while limiting inference to a single function call.These tasks were selected because baseline methods struggle; one-step diffusion policies catastrophically fail in the comparison.
6 DISCUSSION
Shortcut models support few-step and one-step generation through step-size conditioning, with a single unscheduled training run. Evaluations show they outperform prior single-phase methods and remain competitive with two-stage distillation, while a gap persists between many-step and one-step quality.
- Shortcut models learn a step-size-conditioned model grounded at small step sizes and trained by self-bootstrapping at larger step sizes.
- The training recipe is end-to-end, uses one run without scheduling, and avoids separate pretraining and distillation phases.
- Evaluations on CelebA-HQ and Imagenet-256 show shortcut models outperform prior single-phase methods and remain competitive with two-stage distillation methods.
- Practitioners are advised to normalize datasets to unit variance, reduce self-consistency targets if training is erratic, use non-zero weight decay, and evaluate with EMA parameters.
- The authors provide an open-source implementation and model checkpoints through the shortcut-models GitHub repository.
- A remaining limitation is a quality gap between many-step and one-step generation, despite the framework’s simple one-step training objective.
A APPENDIX
The appendix presents image examples and an Imagenet-256 comparison showing shortcut models across sampling budgets and model sizes. The reported setup spans CelebA-HQ and class-conditional Imagenet generation at 256x256 resolution.
- CelebA-HQ examples use unconditional 256x256 generation from a DiT-B model trained for 400k iterations.Each triplet uses the same noise at 128, 4, and 1 denoising steps.
- Imagenet examples use class-conditional 256x256 generation from a DiT-XL model trained for 800k iterations.The examples use the same noise at 128, 4, and 1 denoising steps, with CFG 1.5 during training.
- The Imagenet-256 Shortcut Model (XL) reaches substantially lower FID than the DiT-B architecture in Table 1 and is competitive with previous state-of-the-art models.The table notes that compute constraints prevent matching the compute of the best previously reported generative models.
B TRAINING DETAILS
The training-details appendix specifies that the model architecture follows the DiT design described by Peebles and Xie, using DiT-B unless otherwise stated.
- Training hyperparameters use the DiT-B architecture unless another model is explicitly specified.
B.1 COMPUTATION
Experiments run on TPUv3 nodes with JAX implementations, and each training run typically takes one to two days.
- All experiments use TPUv3 nodes and JAX, with typical training-run durations of 1-2 days.Runtime varies by method.
B.2 METHOD DETAILS
The comparison uses a shared implementation framework and describes training and sampling procedures for flow-matching, reflow, progressive distillation, consistency methods, and live reflow.
- Implementation setup: All comparison methods use the same codebase and architecture, with training budgets selected for roughly equal compute.Some methods receive more compute when necessary to produce reasonable results.
- Flow-Matching: The flow-matching base model is trained for 800k iterations and sampled with deterministic Euler integration.A 400k-iteration checkpoint serves as the teacher for two-stage distillation methods.
- Reflow: Reflow generates 50K synthetic targets for CelebA-HQ and 1M for ImageNet before 400k additional distillation iterations.Classifier-free guidance is used for synthetic-data generation when applicable, but not in the resulting distilled model.
- Progressive Distillation: Progressive distillation runs eight equal phases over 400k iterations, with each student learning from two sequential bootstrap steps before becoming the next teacher.This allocates 50k training steps per phase.
- Consistency Distillation: Consistency distillation enforces velocity consistency between x_t and x_{t+d}, where d = (1/128), using the consistency model to construct target velocities.The target velocity is computed after estimating x_1 from the point x_{t+d}.
- Consistency Training: Consistency training samples interpolated noise–data pairs and applies a binary-time discretization schedule that doubles its intervals across phases.The schedule begins with one interval, then uses two, four, and so on.
- Live Reflow: Live reflow trains a d-conditioned model with 75% of each batch using flow matching and the remainder using self-generated bootstrap targets.Target generation is limited to eight denoising steps, yet live reflow requires over four times the computation of any other method.