Shortcut Models Enable Single-Step Image Generation for Diffusion Models

One-Step Diffusion via Shortcut Models

By Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel (UC Berkeley)

Diffusion models have revolutionized image generation by producing diverse and realistic images. However, their Achilles' heel is the slow and computationally expensive sampling process, which typically involves multiple denoising steps. In a recent paper, researchers from UC Berkeley introduce shortcut models-a novel approach that can generate high-quality samples in a single step, significantly speeding up the process.

Key Technical Changes

• Single Network, Single Training Phase: Unlike previous methods that required complex training regimes (e.g., multiple networks or fragile scheduling), shortcut models use a single network and a single training phase.
• Conditioning on Step Size: The model conditions not only on the current noise level but also on the desired step size. This allows it to skip ahead in the generation process, making it highly flexible for varying inference step budgets.

Why It Matters

For practitioners, this means:

• Faster Inference: Single-step generation can drastically reduce latency and computational costs.
• Simplified Training: No need for complex multi-stage training procedures or delicate scheduling.
• Flexibility: The ability to adjust the number of steps at inference time without retraining.

Technical Details

Naive Diffusion Models and Limitations

Naive diffusion models generate images by iteratively denoising a noisy input. However, they struggle with few-step generation due to inherent uncertainties in the noise-to-data paths. Specifically:

• Training Paths: These are created by randomly pairing data and noise, leading to overlapping paths.
• Direction Uncertainty: Given only the current state ( x_t ), there's uncertainty about the direction ( v_t ) towards the data point.
• Flow-Matching Models: While they learn a deterministic ODE, their paths are not straight and must be closely followed. Fewer inference steps lead to bias towards the dataset mean, causing mode collapse and off-track generations.

Overview of Shortcut Model Training

Shortcut models address these issues by:

• Conditioning on Step Size: The model is conditioned on both the current noise level and the desired step size ( d ).
• Training Objectives:
  • At ( d \approx 0 ), the shortcut objective aligns with the flow-matching objective, trained by regressing onto empirical ( E[v_t | x_t] ) samples.
  • For larger ( d ), targets are constructed by concatenating a sequence of two ( d/2 ) shortcuts. Both objectives can be trained jointly without requiring a two-stage procedure or discretization schedule.

Behavior Over Fewer Denoising Steps

• Flow-Matching Models: As the number of denoising steps decreases, these models suffer from degradation and mode collapse.
• Shortcut Models: They maintain high-quality sample generation even with fewer steps. This is achieved by effectively skipping ahead in the generation process while maintaining directionality towards diverse data points.

Benchmarks and Results

Across a wide range of sampling step budgets, shortcut models consistently outperform previous approaches such as consistency models and reflow. Compared to distillation methods, they offer:

• Reduced Complexity: A single network and training phase.
• Flexibility: Varying step budgets at inference time without retraining.

Conclusion

Shortcut models represent a significant advancement in the field of generative modeling. By enabling high-quality image generation in a single step, they not only speed up the process but also simplify training and inference. This could have profound implications for applications requiring real-time or resource-constrained image generation, such as augmented reality, video processing, and more.