Boosting Diffusion Models with an Adaptive Momentum Sampler

1 Introduction

Deep-generative modeling has emerged as a popular area of research due to its significance in comprehending and managing data. In image generation, GANs have dominated the field [Odena et al., 2017; Gong et al., 2019; Brock et al., 2018; Karras et al., 2020; Dung and Binh, 2022] since its birth in Goodfellow et al.. Compared to other loglikelihood generative models [Kingma and Welling, 2013; Hinton and Salakhutdinov, 2006; Shao et al., 2021; Guo et al., 2020; Wang et al., 2023], GANs show superiority over them in terms of both quality and diversity. However, the employment of adversarial training in GANs causes instability during its training and mode collapse, requiring specific optimization techniques and architectures.

As an alternative to adversarial training, iterative models such as Diffusion Probabilistic Models (DPMs) [Ho et al., 2020; Song et al., 2020b] have emerged as promising option. Denoising Diffusion Probabilistic Model for Diffusion (DDPM) [Ho et al., 2020] and its variants [Nichol and Dhariwal, 2021; Song et al., 2020a] are one of the streams based on the iterative process whose process includes two phases. The first phase is to diffuse the image into a predefined Gaussian noise, and the second phase will try to reverse the trajectory to recover the image. As a recent development, Denoising Diffusion Implicit Model (DDIM) [Song et al., 2020a] devises an accelerated version by turning DDPM into an implicit process and facilitates the reduction of the number of time steps during image generation. Meanwhile, the score-based scheme [Song and Ermon, 2019; Song and Ermon, 2020; Song et al., 2020b] is a theoretical version that works on the same mechanism. The main difference of the score-based model lies in its utilization of analytic tools, Stochastic Differential Equations (SDEs), to diffuse and denoise the image. Several higher-order solvers for diffusion models [Yu et al., 2022; Karras et al., 2022; Bruce et al., 2024] enable faster sampling and the use of momentum in diffusion model training [Wu et al., 2023]. The two iterative schemes of diffusion models and scorebased models are unified by the crucial connection observed in DDIM between its optimization technique and Ordinary Differential Equations (ODEs). Therefore, we focus on the diffusion models since the two categories belong to the same broader generalization.

(b) Our (DDPM + AM)

Figure 1: We introduce a new sampler featuring adaptive momentum (AM). Compared to existing reverse sampling processes in DPMs (left), our method (right) is more efficient in recovering high-level components while effectively minimizing noise in the generated images. This enhancement leads to clearer and more accurate image generation, setting our approach apart from current methodologies.

The reverse process in DPMs can be likened to the annealed Langevin dynamics method, prevalent in diffusion sampling. The sampling process of DPMs shifts the sample in the gradient direction of the log probability density and incorporates a noise element. Similarly, Stochastic Gradient Descent (SGD) updates neural networks in the loss function gradient's direction. This similarity has been utilized to improve the diffusion models sampling in many works especially in diffusion guidance sampling [Dhariwal and Nichol, 2021; Dinh et al., 2024; Dinh et al., 2023; Liu et al., 2023] Drawing inspiration from how the Adam [Kingma and Ba, 2014] optimizer improved SGD, our goal is to enhance the sample generation process of DPMs. To achieve this, we introduce a new training-free reverse sampler for diffusion models, which accumulates a velocity vector over steps to refine the update direction. Additionally, it maintains a moving average of second-order moments, allowing for an adaptive adjustment of the update pace, which effectively functions as a natural form of step-size annealing.

The first benefit of this new sampler is the foremost being a better trade-off between generating high-level semantics, such as shapes and outlines, and restoring low-level details, such as fine texture. To illustrate this point, we refer to Fig. 2, which depicts samples generated using different scales of momentum. As observed, employing a larger momentum that heavily reuses the update direction of early steps leads to augmented details while still preserving high-level information. This phenomenon aligns with the two stages of learning for likelihood-based models [Rombach et al., 2022], where the first stage focuses on learning high-frequency details, while the second stage focuses on learning conceptual compositions. With a suitable momentum value, the adaptive momentum in our sampler coordinates these two learning stages, thus effectively balancing between high-level and low-level information.

The second benefit of the proposed sampler is its flexibility to be easily integrated with existing pre-trained diffusion models. This non-parameterized sampler requires no additional training, and its momentum is dynamically adjusted on the fly during sampling. This adaptability is particularly advantageous in mitigating the train-test gap, as the sampler can be easily plugged into diffusion models pre-trained with different samplers or customized settings.

Experimental results on five common benchmarks, CIFAR-10 [Krizhevsky et al., 2009], CelebA [Liu et al., 2018], ImageNet [Deng et al., 2009], LSUN [Yu et al., 2015] and CelebA-HQ [Karras et al., 2017], show that the proposed adaptive momentum sampler improves both DDPM and DDIM in terms of generating quality. In summary, the contribution of this paper is to newly propose a balanced, flexible, and highly-performing sampler that can be easily applied to pre-trained diffusion models.

2 Preliminary

2.1 Denoising Diffusion Probabilistic Model

Generally, Gaussian diffusion models are utilized to approximate the data distribution x0 ∼ q(x0) by pθ(x0). pθ(x0) is modelled to be the form of latent variables models:

where x1, x2, ..., xT are latent variables with the same dimensions with x0.

Thus, the diffusion process from a data distribution to a Gaussian distribution for time step t can be expressed as:

where the αt := ∏ti=0 (1 - βi) and ϵ ∼ N(0, I). The Ho et al.[Ho et al., 2020] trains a U-net [Ronneberger et al., 2015] model θ to fit the data distribution x0 by maximizing the following variational lower-bound:

where the q (x1:T | x0) is a certain inference distribution, which could be calculated by the Bayes theorem with x0 and xT , over the latent variable.

2.2 Denoising Diffusion Implicit Model

DDPM's dependence on the Markovian process has two main problems. First, due to the Markovian property, the generative process is forced to have a similar number of time steps to the diffusion process. Secondly, the stochastic process causes uncertainty in the synthetic images. This challenges the image interpolation in the latent space.

DDIM generalizes the DDPM as a Non-Markovian process. Different from the Eq. 1, xt-1 is conditioned by both xt and x0:

Where qσ(xt-1|xt, x0) is chosen so that qσ(xT |x0) = N (√αT x0, (1 - αT )I):

From Eq. 2, xt can be obtained by sampling x0 ∼ q(x0) and ϵt ∼ N(0, I). By training the model ϵ(t) θ to predict ϵt at each time step, x0 predicting function f (t) θ is defined as:

This results in the whole generative process as:

θ is optimized as Eq. 3 with qσ(x1:T |x0) is defined in Eq. 4. Based on the defined generative process, given a sample xt, we could sample the xt-1 as:

where σt = η√(1-αt-1)/(1 -αt)√1-αt/αt-1 and η = 0 [Song et al., 2020a] or η = 1 [Song et al., 2020a] or η = √(1-αt)/(1 -αt-1) [Ho et al., 2020]. Yang Song et al.[Song and Ermon, 2019] have proved that when η = √(1 - αt)/(1 - αt-1) the Eq. 5 is essentially a different discretization to the same reverse-time SDEs and Jiaming Song [Song et al., 2020a] shows that when η = 0 the Eq. 5 similarity to Euler integration for solving ODEs. We provide when the η = 1 is the one type of DDPM, the reverse process is approximate solution of the same reverse-time SDEs in Appendix B.

3 Methodology

In this paper, we propose a training-free Adaptive Momentum Sampler to generate high-quality images, which can build long-term communications between the previously explored path and the current denoising direction. The first part of this section presents how to integrate the basic momentum method into the origin inference process of DDPM/DDIM. After that, we illustrate how to effectively adjust the pacing rate for denoising with the moving average of the squared prior increments. In each part, we also show the design of how to trade-off between high- and low-level information.

3.1 Momentum Sampler

We employ the momentum method to enhance the reverse process in diffusion models. This method is commonly utilized to accelerate gradient flow optimization in neural network training. By incorporating momentum, we aim to significantly improve the efficiency and effectiveness of the reverse sampling process in diffusion models, leading to more optimized outcomes. The memorization of previously searched directions contributes to the optimizer's overcoming narrow valleys, small humps. We find that the momentum method can also benefit the reverse process of diffusion SDEs and ODEs.

Inspired by SGD with momentum in optimizing neural networks, we replace the Euler-Maruyama method with the momentum method in the reverse process of diffusion models to avoid the plight mentioned. We find that the prediction of x0 is relatively stable and tractable during the sampling phase. Therefore, the sampling process contains some trend information, enabling our momentum sampler to utilize such information and improve the sampling procedure.

We first repeat the Eq. 5 as an incremental form of x. For brevity, we let the ¯xt = xt √αt and the µ = (√ 1-αt-1-σ2 t αt-1 - √ 1-αt αt ). And then, we can formulate the increment of ¯xt and rewrite the DDPM (or DDIM) iteration

Eq. 5, as:

As demonstrated in the equation referenced as Eq. 7, the updating process in dispersed sampling can be viewed as an increment involving the current xt and the previous xt-1. Similarly, the SGD optimizer updates the network's parameters

(ω) in the direction of the loss function (Q), following the formula ω := ω -▽Q(ω). This reveals a strong similarity between the formulas used in SGD and the reverse process of diffusion models. Therefore, our objective is to adjust the reverse process in a way akin to how the SGD + momentum optimizer refines SGD, which involves enhancing the update mechanism in the reverse process to optimize the generation of samples in diffusion models. Instead of only relying on the single last step result to generate samples, we incorporate the effects of prior directions as follows:

where the a is the damping coefficient, and the a is a historydependent velocity coefficient. Since the superposition property of the Gaussian distribution, we use the spherical difference, a2 + b2 = 1, to balance the momentum and current value. The mt is the current momentum of t-th iteration, and the mT = 0. The strength of the momentum method is controlled by the a and b coefficients.

3.2 Adaptive Momentum Sampler

In this section, we study how to adapt the learning rate based on the running average of recent magnitudes of the increment, d¯xt. Besides the momentum-based sampling process only depending on the advantage of momentum by using a moving average, the idea of Root Mean Square Propagation [Mukkamala and Hein, 2017] (RMSProp) can also be applied to the probabilistic diffusion models generation process. To be more specific, we imitate RMSProp to automatically tailor the step size with a decaying average on each pixel. In this way, the process solver will focus on the most recently observed partial gradients and discard history from the extreme past.

We propose an Adaptive Momentum sampling process that adapts the learning rates of each pixel per time step The adaptive adjusting of denoising step size is based not only on the average first moment but also on the average of the second moments of the gradients. Based on the momentum sampling process, we use a vt to store an exponentially decaying average of prior squared increment:

where the c is the decay rate to control the moving averages of squared increment, and the vt is the current averages of prior squared increment, and the vT = 1, and the ζ is a smoothing term (a small number) to avoid any division by 0. We use this moving average of second-order moments to make the reverse process concentrate on the current high-frequency information to avoid the image being over-smooth. Compared with the low-frequency information, the high-frequency information, such as the pattern and fine texture of images, is more sensitive to the second-order increment d¯x. In this assumption, the pacing rate of the high-frequency information will be adaptively magnified, while the pacing rate of the lowfrequency information is relatively stable. The adaptive momentum sampler can future balance between high-frequency and low-frequency information.

Algorithm 1 illustrates the adaptive momentum sampler process. Specifically, we use the momentum increment, mt-1, to replace the current increment of the t-th iteration of noised image. The a and b build a trade-off between the former and current time-dependent score function, demonstrating a strong numerical solver of the reverse process. As shown in Fig. 2, this adjustment can force the reverse sampling process to focus only on the early steps to generate more high-level semantics, such as shapes and outlines, and smooth the low-level information, such as the detailed pattern and texture. The c controls the exponential decay proportion for the second-moment estimates. In order to avoid excessive changes in the magnitude of the momentum increment for each pixel, we amass the square of the d¯xt's Frobenius norm and set the vT = 1 during the implement. The basic momentum sampler could be considered as a particular case of the adaptive momentum solver when c = 1 all the time. If a = 0, b = 1, and c = 1 for all the t-th iterations, the sampling process would degenerate to the original DDPM or DDIM generation process.

3.3 Synthetic Data and Analysis

In this section, we will compare the generation process of the vanilla sampler with our proposed adaptive momentum sampler, both in practice and in theory.

Visualization. To perform a quantitative analysis, we trained a diffusion model to generate 1-dimension toy data with two points (x = -2 and x = 4) over 200 timesteps. In Fig. 3, we provide a visual representation of the differences between the original DDPM sampler and our proposed adaptive momentum sampler during the sampling process. The heatmaps in the background display the distribution of values for 10,000 samples generated by the original DDPM sampler and our adaptive momentum sampler, respectively. The lighter the color in the background, the greater the number of samples. Specifically, the heat map generated by the adaptive momentum sampler is lighter, especially around the points x = -2 and x = 4, indicating that it can more precisely produce the value of point. Furthermore, the generation

Algorithm 1 Sampling with the Adaptive Momentum

trace of the adaptive momentum sampler is more stable, exhibiting less trembling, while the generation trace of the original DDPM sampler (gold line) features significant ineffective vibrations. Notably, for the gold line, the final value generated by the vanilla DDPM sampler is inaccurate. In contrast, the adaptive momentum sampler can more easily escape from narrow valleys, ensuring continuity in the generation process and producing a more accurate final value at last step.

Second-Order Approximation. We conducted a preliminary analysis to compare our method with the original method. As per Song et al. [Song et al., 2020b], the sampling process of DDPM/DDIM uses a first-order approximation of numerical SDEs/ODEs solvers. Building on this conclusion, we demonstrate in Appendix B that our method represents a second-order approximation of SDEs/ODEs. Instead of utilizing the Euler-Maruyama method for numerical solving SDEs/ODEs, we predict the next result by utilizing the midpoint method with previously generated results and the current state. This approach allows us to demonstrate that our method converges to the actual sample of the pre-trained model of DDPM.

In general, the adaptive momentum sampler typically employs historical data to facilitate a smoother generation process, as demonstrated by the gold line in Fig 3 for DDPM+AM. By contrast, the gold line in Fig 3 for DDPM illustrates early time steps characterized by wild swings, which can lead to misguided outcomes and large variance values for the final point.

4 Experiments

In this section, we quantitatively demonstrate the effectiveness of our proposed sampling method. First, we provide the experimental setup. Next, the method is quantitatively evaluated against the baseline samplers. The improvement compared to the baseline will be highlighted on different datasets. Lastly, an ablation study will be offered to inspect the impact of each hyper-parameter. The code is publicly available at github.com/ShinyGua/DPMs-with-Adam

4.1 Experimental Setup

Datasets. Following most of the setup in DDPM/DDIM and Latent Diffusion Models (LDM) [Rombach et al., 2022], we utilize CIFAR10 [Krizhevsky et al., 2009] (32 × 32), CelebA [Liu et al., 2018] (64 × 64), ImageNet [Deng et al.,

2009] (64 × 64), LSUN [Yu et al., 2015] (256 × 256) and CelebA-HQ [Karras et al., 2017] (256×256) in experiments.

Configurations. We mainly employ the original DDIM, DDPM, Analytic-DPM [Bao et al., 2022] and LDM in our comparisons as these schemes are solid and compact baselines. Nonetheless, our proposed schemes can be applied to other improved variants as well. For DDPM/DDIM and Analytic-DPM, the same pre-trained diffusion models are used for the generation. The pre-trained models for CIFAR10, LSUN and CelebA-HQ are collected from the DDPM [Ho et al., 2020], and the pre-trained models for CelebA and ImageNet are collected from DDIM [Song et al., 2020a], and IDDPM [Song and Ermon, 2020] respectively. Besides, we also use the pre-trained models from LDM [Rombach et al., 2022] for high-resolution image generation. The sampling steps are set to 4000 on ImageNet and 1000 on other datasets. For the Eq. 5, we select three η values, the η = 0 will be equivalent to DDIM, η = 1 will be a DDPM case in [Song et al., 2020a], and η = ˆ η is another DDPM case in its original paper [Ho et al., 2020]. For settings where the sampling step does not equal the original pretrained model, we follow the same strategy proposed in the DDIM, which replaces the αt with the corresponding scaled ατ . All the hyperparameters for the sampling process are presented in Appendix. Our experiments run on one node with 8 NVIDIA A100 GPUs. We use 'AM' to denote sampling with our proposed adaptive momentum.

Measurements. Similar to many other generative models [Goodfellow et al., 2014; Bao et al., 2022; Nichol and Dhariwal, 2021; Ho et al., 2020], we mainly adopt two evaluation metrics which are Frechet Inception Score (FID) and Inception Score (IS) [Lucic et al., 2018; Borji, 2019]. IS is highly correlated with human-annotators [Salimans et al., 2016]. Nevertheless, this measure is often referred to as a method to measure inter-class diversity and is less sensitive to the diversity of the images inside one label [Lucic et al., 2018; Borji, 2019]. In contrast, FID can detect intra-class mode collapsing as well [Lucic et al., 2018]. Thus, FID is considered a better measure for image generation tasks. A lower FID value indicates better performance, while a larger IS is better. For each experiment, we draw 50K samples for evaluation.

4.2 Overall Performance

Our proposed sampling scheme is first compared with DDPM and DDIM in Table 1 and 2. The same pre-trained diffusion DDPM/DDIM is used on each dataset to generate the data with vanilla samplers and our adaptive momentum samplers (the '+AM' rows). Analytic-DPM [Bao et al., 2022], and LDM [Rombach et al., 2022] sampler is also used as a baseline to justify our 'AM' on different diffusion schemes.

The results in Table 1 show the improvement of the adaptive momentum sampling over baselines on most lowresolution datasets (lower than 64x64) regarding FID. For the CelebA dataset, the proposed method performs significantly better than DDPM with η = 1. It not only outperforms the baselines by around 50% but also achieves the state-of-theart for this dataset which reaches around 2.29 FID. For CIFAR10, although the FID has been very low for the baselines, we still have some rooms to improve from 3.16 to 3.03.

Table 2 shows that our adaptive momentum scheme leads to significant improvements in most baselines when applied to high-resolution datasets such as CelebA-HQ, LSUN Church, and LSUN Bedroom. We also demonstrate the compatibility of our approach with a popular high-resolution model, Latent Diffusion Models, which indicates that our method can be easily incorporated into other DPM-based methods. These results demonstrate the flexibility of our adaptive momentum scheme in improving the performance of various generative models.

Fig. 4 presents a qualitative comparison of synthetic samples generated using different methods. In comparison to the original DDPM sampler, our proposed adaptive momentum sampler produces images with more realistic details and pronounced object outlines, confirming its ability to balance high-level and low-level information more effectively. For instance, our method enhances image details while preserving original patterns, such as the recovered person structure in row 2 column 4 and the enhanced facial details (forehead hair) in the last picture of the 2nd row.

4.3 Ablation Study

This subsection investigates different aspects of the settings, including different time steps and hyperparameters' effects.

Different Sampling Steps. Table 3 illustrates the performance of our method when fewer sample steps are used for acceleration on CIFAR10. The proposed method could perform better or be comparable when using different sampling steps. On most of the settings, we can provide up to 50% performance enhancements over the original DDPM sampler. It is observed that the improvement is significant on the larger number of sample steps while less evident on the smaller ones. This makes sense, possibly because using a small number of sampling steps is not always sufficient for accurate momentum estimation and noise suppression.

Hyper-Parameters. Considering about the effects of b and c in the Eq. 8 and 9, we setup with different values as in Table 4 on CIFAR10. The first five lines illustrate how the FID of the generated images changes as the b increases when c = 0. In general, the image quality will first increase and then decrease. When the b = 0.15, the FID achieves the lowest value for image generation in different sampling steps, which is the best trade-off point. The last four lines show how the b influence the sampling results when b = 0.15. Similarly, the FID first decreases and then increases when the b decrease. The last four rows reflect the impact of c with b = 0.15, where the most of best performance is obtained at c = 0.01. Empirically, different denoising steps can share the same hyper-parameters as shown in Table 3 and Appendix. Thus, we only need to find bmax like 10-3 , 10-2 or 10-1 in the small denoising steps, e.g., 100 or 200. Most of hyperparameter values are shared across different datasets.

Adaptivity of the Momentum. Table 4 also presents the results of an ablation study on our proposed method, the Adaptive Momentum Sampler. The first five rows display the outcomes when the sampling process only utilizes the moving average of the moment increment. Moreover, the comparison between the first and fourth rows reveals that our proposed method outperforms the original reverse process sampler, underscoring the efficacy of the momentum idea. Furthermore, comparing the fourth and eighth rows, we observe that introducing adaptivity through the moving average of second-order moments further enhances the performance of our adaptive momentum sampler, which demonstrates its excellent generalization ability. Thus, both non-adaptive and adaptive momentum samplers prove advantageous for DPMs.

4.4 Rate-Distortion Trade-Off

The rate-distortion curves of a trained model using both adaptive momentum inference and baseline inference schemes are plotted in Fig. 5. It was criticized in the literature that diffusion models tend to spend much more rates, i.i., model capacity, on restoring imperceptible distortions than high-level semantics [Ho et al., 2020; Rombach et al., 2022], which is also observed in Fig. 5 for the baseline scheme. In contrast, the newly adaptive proposed momentum sampling is able to straighten the curve and thus strike a better balance between generating high-level semantics and low-level details. This could provide a possible explanation for the improved generation performance observed in our experiments. Note that it is reasonable for the red curve to have larger rates, which are computed by the cumulative sum of the variational bound terms [Ho et al., 2020], since the momentum inference is being applied to a model pre-trained without momentum. We expect that incorporating an adaptive momentum strategy into the training process could further enhance the rate-distortion trade-off as well as the generation quality in future works.

5 Conclusion

In this work, we propose a training-free sampler to improve the generated images of Diffusion Generative Models, especially for discrete discrete time steps settings. Through extensive experiments, we found out that the adaptive momentum sampler use the history the denoising trajectory to improve the quality of the generated images. In future works, we will expand the scheme to continuous settings as well as with solid theoretical versions.

Acknowledgements

This work was supported in part by the Australian Research Council under Projects DP210101859 and FT230100549. The authors acknowledge the use of the National Computational Infrastructure (NCI) which is supported by the Australian Government and accessed through the NCI AI Flagship Scheme and Sydney Informatics Hub HPC Allocation Scheme.