Boosting Diffusion Models with Moving Average Sampling in Frequency Domain

TL;DR

This paper reinterpret the iterative denoising process as model optimization and leverage a moving average mechanism to ensemble all the prior samples and lead to superior performances compared to the baselines, with almost negligible additional complexity cost.

Abstract

Diffusion models have recently brought a powerful revolution in image generation. Despite showing impressive generative capabilities, most of these models rely on the current sample to denoise the next one, possibly resulting in denoising instability. In this paper, we reinterpret the iterative denoising process as model optimization and leverage a moving average mechanism to ensemble all the prior samples. Instead of simply applying moving average to the denoised samples at different timesteps, we first map the denoised samples to data space and then perform moving average to avoid distribution shift across timesteps. In view that diffusion models evolve the recovery from low-frequency components to high-frequency details, we further decompose the samples into different frequency components and execute moving average separately on each component. We name the complete approach "Moving Average Sampling in Frequency domain (MASF)". MASF could be seamlessly integrated into mainstream pre-trained diffusion models and sampling schedules. Extensive experiments on both unconditional and conditional diffusion models demonstrate that our MASF leads to superior performances compared to the baselines, with almost negligible additional complexity cost.

Figures (6)

• Figure 1: We utilize a diffusion model from ADM adm pre-trained on ImageNet-64 to sample from white noise and capture the intermediate output as denoised sample $\hbox{\boldmath $x$}_t$. Subsequently, we plot the pixel value of $\hbox{\boldmath $x$}_t$ with respect to generative timesteps. The starting point and ending point of the trajectory of $\hbox{\boldmath $x$}_t$ are regarded as the ground truth for the noisy image and clean image, respectively. With that we calculate the ground truth of $\hbox{\boldmath $x$}_t$ at any timestep $t$ using the formulation defined in forward process.
• Figure 2: The evolution of (a) denoised sample $\hbox{\boldmath $x$}_t$, (b) estimated sample in data space $\hbox{\boldmath $x$}_0^{t}$ and (c) four subbands in frequency domain of $\hbox{\boldmath $x$}_0^{t}$ along denoising process. Here each group of subbands is achieved via wavelet decomposition, yielding four different frequency components: $ll$ ($\nwarrow$), $lh$ ($\nearrow$), $hl$ ($\swarrow$), and $hh$ ($\searrow$).
• Figure 3: The overall framework of our Moving Average Sampling in Frequency domain (MASF) for denoising stabilization. At each denoising timestep $t$, MASF first maps the denoised sample $\hbox{\boldmath $x$}_t$ into data space, leading to the estimated sample $\hbox{\boldmath $x$}_0^{t}$. We then perform frequency decomposition of $\hbox{\boldmath $x$}_0^{t}$ via Discrete Wavelet Transformation (DWT) and achieve four subbands ($\hbox{\boldmath $x$}^t_{\{ll, lh, hl, hh\}}$). After that, MASF updates each frequency component (e.g., the low-frequency component $\hbox{\boldmath $x$}^t_{ll}$) through moving average over prior samples, pursuing harmonized stabilization along with frequency evolution. The refined subbands $\bar{\hbox{\boldmath $x$}}^t_{\{ll, lh, hl, hh\}}$ are finally converted back to image domain via Inverse DWT (IDWT) to trigger the subsequent denoising process.
• Figure 4: Evolution ($l_2$ norm) of different frequency subbands during denoising process. (a) The low-frequency subband oscillates sharply at the beginning and stabilizes after a certain timestep. (b) In contrast, the $l_2$ norm of high-frequency subbands drops rapidly into a small value and then increases steadily along with the denoising process. (We use the model from ADM adm pre-trained on ImageNet-64)
• Figure 5: The generated images on MS-COCO using DPM-Solver++ (top) and DPM-Solver++ plus MASF (bottom).