Table of Contents
Fetching ...

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.

Score-based Generative Models with Adaptive Momentum

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.
Paper Structure (28 sections, 5 theorems, 54 equations, 18 figures, 13 tables, 4 algorithms)

This paper contains 28 sections, 5 theorems, 54 equations, 18 figures, 13 tables, 4 algorithms.

Key Result

Proposition 1

For any $\alpha \in(0,\frac{2}{L}]$, the Markov chain $(\mathbf{z}^t)_{t\ge 0}$ defined by the recursion Eq. markovz admits a unique stationary distribution $\pi_{\alpha}^*$ such that $\pi_{\alpha}^* R_{\alpha} = \pi_{\alpha}^*$. Additionally, for all $\mathbf{\hat{z}}\in\mathbb R^{2d},t\in \mathbb

Figures (18)

  • Figure 1: Partial estimate of FID (2k samples) with different initial $\epsilon$. We can see that if an appropriate $\epsilon$ is chosen, MC can bring a considerable improvement in image quality.
  • Figure 2: AMS generated samples after denoising on synthetic 2D experiments. For each subplot, the left is the data point without denoise, the right is the expected denoised samples. The orange point is real data, and the green is generated points.
  • Figure 3: FIDs obtained from NCSN2 with different NFEs. AMS can lead to better FIDs with those small NFEs, e.g. 50, 100, 150, 200, 250. Besides, AMS is only slightly behind ALS in large NFEs. For experimental details, please refer to Appendix \ref{['app:experimental details']}.
  • Figure 4: CIFAR10 samples for different NFEs, result obtained with our sampler from NCSN2.
  • Figure 5: Partial estimate of FID (2k samples) with different initial snrs , NFE $=270$.
  • ...and 13 more figures

Theorems & Definitions (8)

  • Proposition 1
  • Proposition 2
  • Proposition 3
  • Lemma 1
  • Lemma 2: Strong convexity case
  • proof
  • proof
  • proof