One Step Diffusion via Shortcut Models

TL;DR

Shortcut models offer a simplified, single-network approach to fast, high-quality image generation by conditioning on step size, outperforming previous methods across various sampling budgets.

Contribution

Introduction of shortcut models that enable high-quality, flexible, and fast sampling with a single network and training phase, eliminating complex training procedures.

Findings

Shortcut models outperform previous methods like consistency models and reflow.

They produce higher quality samples across various sampling steps.

They enable variable step budgets at inference without retraining.

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.