Score-based Generative Models with Adaptive Momentum

TL;DR

This paper introduces an adaptive momentum sampling method for score-based generative models, significantly speeding up the sampling process while maintaining high-quality data generation.

Contribution

It proposes a novel adaptive momentum sampling technique inspired by SGD, with theoretical convergence guarantees and empirical improvements in speed and quality.

Findings

2 to 5 times faster sampling speed

Produces more faithful images and graphs with fewer steps

Achieves competitive scores on generation tasks

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.