Fast Diffusion Model (2306.06991v2)

Abstract: Diffusion models (DMs) have been adopted across diverse fields with its remarkable abilities in capturing intricate data distributions. In this paper, we propose a Fast Diffusion Model (FDM) to significantly speed up DMs from a stochastic optimization perspective for both faster training and sampling. We first find that the diffusion process of DMs accords with the stochastic optimization process of stochastic gradient descent (SGD) on a stochastic time-variant problem. Then, inspired by momentum SGD that uses both gradient and an extra momentum to achieve faster and more stable convergence than SGD, we integrate momentum into the diffusion process of DMs. This comes with a unique challenge of deriving the noise perturbation kernel from the momentum-based diffusion process. To this end, we frame the process as a Damped Oscillation system whose critically damped state -- the kernel solution -- avoids oscillation and yields a faster convergence speed of the diffusion process. Empirical results show that our FDM can be applied to several popular DM frameworks, e.g., VP, VE, and EDM, and reduces their training cost by about 50% with comparable image synthesis performance on CIFAR-10, FFHQ, and AFHQv2 datasets. Moreover, FDM decreases their sampling steps by about 3x to achieve similar performance under the same samplers. The code is available at https://github.com/sail-sg/FDM.

Overview of Fast Diffusion Models for Accelerated Training and Sampling

The paper presents a novel perspective on accelerating Diffusion Models (DMs) by drawing a theoretical connection between their diffusion processes and stochastic optimization methods, particularly focusing on the use of momentum techniques akin to those employed in Momentum Stochastic Gradient Descent (SGD). By leveraging these connections, the authors propose a Fast Diffusion Model (FDM) to improve both the training and sampling efficiency of existing DMs, demonstrating significant cost reductions without sacrificing performance.

Theoretical Foundations

Diffusion Models (DMs) are a class of generative models renowned for their capability in capturing intricate data distributions through sequential noise addition and removal. However, they are often hindered by high computational costs, primarily during the training and sampling phases. This paper proposes an accelerated approach by framing the diffusion process within the context of stochastic optimization. Specifically, it establishes a relationship with the optimization procedure of SGD, where the forward diffusion step can be viewed as a stochastic approximation of an optimization process. This foundation allows utilizing techniques from accelerated gradient descent, particularly the incorporation of momentum.

Momentum in optimization accelerates convergence by using an extrapolated step computed from past gradients, which provides a more robust search direction by mitigating oscillations. In the context of DMs, this idea manifests as a momentum-enhanced diffusion process, designed to expedite convergence to the target distribution while maintaining stability.

Methodological Innovations

The paper introduces the Fast Diffusion Model (FDM), which integrates momentum into the forward diffusion process. The essence of FDM lies in augmenting the vanilla diffusion process by including a momentum term, effectively transforming it into a system resembling a Damped Oscillation. This integration involves careful calibration to achieve a critically damped state, thereby eliminating overshoot and achieving faster convergence characteristics.

The authors derive continuous formulations of the diffusion process for computational tractability, which allows the efficient computation of the necessary perturbation kernels inherently aligned with momentum-based dynamics. This advance enables the seamless incorporation of momentum-based processes into existing DM architectures.

Empirical Results

Extensive empirical evaluations demonstrate the effectiveness of FDM across several datasets, including CIFAR-10, FFHQ, and AFHQv2. The results indicate that FDM notably accelerates training, reducing the number of training samples required by approximately 50% compared to standard DMs while achieving comparable synthesis performance. Similarly, FDM also significantly reduces the number of sampling steps by about threefold, facilitating faster sample generation processes. These enhancements are demonstrated on multiple diffusion frameworks such as VP, VE, and EDM, underscoring the versatility and effectiveness of the proposed method.

Practical and Theoretical Implications

The introduction of FDM presents substantial implications for practical applications where computational resources are constrained, opening up possibilities for more efficient real-time applications and wider deployment of diffusion-based generative models. Theoretically, this work enriches the understanding of the diffusion process by revealing its deep-rooted connection with stochastic optimization, providing a conceptual bridge that can inspire further advancements in generative modeling.

Future Directions

Future research could explore further refinements in the application of momentum to other variants of DMs and investigate adaptive strategies for dynamically adjusting momentum parameters in response to model and data characteristics. Extending the framework to higher-dimensional data and investigating its impact on stability and performance can also yield valuable insights. Furthermore, integrating FDM with more sophisticated samplers could lead to additional performance gains and broaden the range of practical applications in AI. The groundwork laid by this paper thus sets the stage for ongoing exploration to realize the full potential of diffusion processes through stochastic optimization techniques.