Score-based Generative Models with Adaptive Momentum
Ziqing Wen, Xiaoge Deng, Ping Luo, Tao Sun, Dongsheng Li
TL;DR
This work addresses the slow sampling of score-based generative models by introducing Adaptive Momentum Sampling (AMS), which injects momentum into stochastic Langevin-type samplers without adding hyperparameters. Through a Markov-chain analysis, the authors establish convergence properties and show how an adaptive momentum update accelerates the reverse diffusion, enabling 2–5x speedups at small NFEs while preserving sample quality. Empirically, AMS improves image fidelity on CIFAR-10 and CelebA-HQ and enhances graph generation on several datasets, often outperforming baselines in low-NFE regimes. The approach provides a practical, theoretically grounded accelerator for diffusion-based generative models with broad applicability to images and graphs, and is accompanied by public code for reproducibility.
Abstract
Score-based generative models have demonstrated significant practical success in data-generating tasks. The models establish a diffusion process that perturbs the ground truth data to Gaussian noise and then learn the reverse process to transform noise into data. However, existing denoising methods such as Langevin dynamic and numerical stochastic differential equation solvers enjoy randomness but generate data slowly with a large number of score function evaluations, and the ordinary differential equation solvers enjoy faster sampling speed but no randomness may influence the sample quality. To this end, motivated by the Stochastic Gradient Descent (SGD) optimization methods and the high connection between the model sampling process with the SGD, we propose adaptive momentum sampling to accelerate the transforming process without introducing additional hyperparameters. Theoretically, we proved our method promises convergence under given conditions. In addition, we empirically show that our sampler can produce more faithful images/graphs in small sampling steps with 2 to 5 times speed up and obtain competitive scores compared to the baselines on image and graph generation tasks.
