One Step Diffusion via Shortcut Models (2410.12557v3)

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.

• The paper introduces shortcut models that enable one-step sampling, reducing inference time by up to 128x compared to traditional iterative methods.
• It employs a single network trained across multiple step sizes, simplifying the complex multi-phase training typical in diffusion models.
• Empirical evaluations on CelebA-HQ and ImageNet-256 confirm that shortcut models maintain accuracy in many-step settings and excel in one-step generation.

One Step Diffusion via Shortcut Models

The paper "One Step Diffusion via Shortcut Models" presents an innovative approach to addressing the problem of time-consuming sampling in diffusion and flow-matching models used for image, video, audio, and protein modeling. This research proposes shortcut models that streamline the generative process by reducing the complexity associated with traditional iterative methods.

Overview of Diffusion and Flow-Matching Models

Diffusion and flow-matching models have gained prominence due to their capability to generate diverse and realistic data by transforming noise into meaningful information via learned ODEs. However, these models necessitate numerous neural network evaluations, resulting in slow and expensive generation processes. Traditional acceleration techniques employ complex multi-phase training strategies, requiring multiple models or delicate scheduling, which introduces additional layers of complexity and computational cost.

Introduction of Shortcut Models

Shortcut models serve as a novel solution by conditioning a single network on both the current noise level and the desired step size. This adaptability allows the model to effectively predict on multiple step sizes, including a single step, facilitating swift and high-quality data generation.

Key characteristics of shortcut models include:

• Single Network and Training Phase: These models bypass the complexities of multi-stage training typical in traditional distillation and consistency models.
• Flexibility in Inference: Shortcut models are versatile, accommodating various step budgets during inference. This adaptability contrasts with traditional models that deteriorate quickly when queried with fewer steps.
• Efficient Training: They require approximately 16% more compute than base diffusion models, making them computationally efficient.

Empirical Evaluation

The empirical results are compelling, demonstrating that shortcut models consistently outperform previous state-of-the-art methods such as consistency models and reflow in terms of sample quality across different step settings. Evaluations on benchmarks like CelebA-HQ and ImageNet-256 reveal that shortcut models preserve accuracy in many-step scenarios and significantly enhance performance in one-step settings compared to alternative approaches.

Claims and Contributions

• Superior Sampling Speed: Shortcut models can produce high-quality images in a single forward pass, reducing sampling time by up to 128x.
• Single Training Routine: Unlike the multi-phase procedures required by other models, shortcut models achieve end-to-end training in one go, eliminating scheduling complexities.
• Broad Applicability: Beyond image generation, the effectiveness of shortcut models extends to domains such as robotic control, showcasing their generalizability.

Implications and Future Directions

The introduction of shortcut models could significantly impact not only image synthesis but also various applications requiring rapid data generation without quality compromises. Theoretically, this research suggests a promising direction for achieving efficient generative modeling with minimal computational overhead. Future work could explore improvements in the integration of shortcut models with other AI domains, potentially enhancing versatility and performance.

In conclusion, this paper presents a significant advancement in generative modeling by addressing the limitations of traditional diffusion models. Shortcut models offer a streamlined, efficient approach to generating high-quality samples rapidly, setting a new benchmark in one-step generative modeling strategies. The release of model checkpoints and source code furthers the research community's ability to build upon and verify these findings.