Boosting Diffusion Models with an Adaptive Momentum Sampler

TL;DR

This paper introduces an adaptive momentum-based reverse sampler for diffusion probabilistic models, improving stability and output quality during image generation by smoothing the sampling process.

Contribution

The paper proposes a novel reverse sampler inspired by Adam optimizer, enhancing stability and quality in diffusion model sampling, compatible with pre-trained models.

Findings

Significant quality improvements over baselines.

Enhanced stability during sampling process.

Applicable to various pre-trained diffusion models.

Abstract

Diffusion probabilistic models (DPMs) have been shown to generate high-quality images without the need for delicate adversarial training. However, the current sampling process in DPMs is prone to violent shaking. In this paper, we present a novel reverse sampler for DPMs inspired by the widely-used Adam optimizer. Our proposed sampler can be readily applied to a pre-trained diffusion model, utilizing momentum mechanisms and adaptive updating to smooth the reverse sampling process and ensure stable generation, resulting in outputs of enhanced quality. By implicitly reusing update directions from early steps, our proposed sampler achieves a better balance between high-level semantics and low-level details. Additionally, this sampler is flexible and can be easily integrated into pre-trained DPMs regardless of the sampler used during training. Our experimental results on multiple benchmarks demonstrate that our proposed reverse sampler yields remarkable improvements over different baselines. We will make the source code available.