Adaptive Momentum Benefits Score-based Generative Models

Ziqing Wen; Ping Luo; Jiahuan Wang; Xiaoge Deng; Tao Sun; Dongsheng Li

cs.LGMay 22, 2024

Adaptive Momentum Benefits Score-based Generative Models

Authors: Ziqing Wen, Ping Luo, Jiahuan Wang, Xiaoge Deng, Tao Sun, Dongsheng Li

Organizations: College of Computer, National University of Defense Technology

Abstract

Score-based Generative Models (SGMs) have achieved remarkable success in generation tasks by establishing a diffusion process that gradually perturbs real data into Gaussian noise and then learning a reverse process to reconstruct the data from noise. However, conventional denoising methods, such as Langevin dynamics and numerical stochastic differential equation solvers, rely on randomness but suffer from slow generation speeds due to the high number of score function evaluations. To address this, we propose a plug-and-play Markov adaptive momentum sampling method, inspired by the strong connection between the SGMs' sampling process and stochastic gradient descent optimization. Our approach accelerates the transformation process without introducing additional hyperparameters. Theoretically, we establish a rigorous foundation by analyzing the relationship between the distribution of samples generated by our method and the original data distribution. This provides a novel perspective on the analysis of sampling distributions. Empirically, our method demonstrates superior efficiency, generating high-quality images/graphs in significantly fewer sampling steps, achieving a speedup of 2x to 5x while maintaining or surpassing baseline performance on various image and graph generation benchmarks. These results underscore the effectiveness of adaptive momentum in enhancing the efficiency of SGMs.

Figures & tables

Panel of Figure 1 — Fig. 1: Experiments on the strong convexity of SGM sampling trajectories. Left: The local convexity parameter $\mu$ as a function of the sampling steps, where $\mu$ remains greater than 0 throughout the entire sampling trajectory. Right: Analysis of the condition number $L/\mu$ along the sampling trajectory, where the overall condition number remains below 100, indicating a well-behaved sampling process.

Panel of Figure 2 — Fig. 2: Partial estimate of FID $\downarrow$ (2k samples) with different initial $\epsilon$ . We observe that MC (ours) can improve image quality with an appropriate initial Signal-to-Noise Ratio (SNR) hyperparameter $\alpha$ .

Sampler	NFE	50	150	250
ALS	$n_{\sigma}=1$	$57.71$	$17.16$	$12.29$
ALS	$n_{\sigma}=5$	$42.74$	$13.87$	$10.98$
Sampling time		18.64s	48.47s	79.03s
AMS	$n_{\sigma}=1$	$30.52$	$17.16$	$12.29$
AMS	$n_{\sigma}=5$	$26.37$	$11.26$	$10.08$
Sampling time		18.69s	48.50s	79.33s

TABLE I: FID

\downarrow

scores and sampling times (500 images) of different samplers on NCSN (VE) across different

n_{\sigma}

Model	Sampler	FID $\downarrow$
VE	AMS (ours)	$9.82$
VE	ALS [ 28 ]	$8.74$
VP	AMS (ours)	$3.84$
VP	ALS [ 28 ]	$4.24$
EBM	AMS (ours)	$34.1$
EBM	ALS [ 28 ]	$38.7$

TABLE II: CIFAR-10 generative performance on VE and VP.

Model	Sampler	FID $\downarrow$
VE	AMS (ours)	$9.30$
VE	ALS [ 28 ]	$10.63$
VP	AMS (ours)	$6.31$
VP	ALS [ 28 ]	$8.77$

TABLE III: CelebA generative performance on VE and VP models.

Sampler / NFE	150	195	210	270	300
RD-MC (ours)	$6.75$	3.05	2.86	2.60	$2.61$
RD-LDFM	$38.6$	$5.04$	$4.37$	$4.80$	$4.43$
Sampler / NFE	150	270	500	1000	2000
MC (ours)	$235$	$27.6$	$9.16$	$17.7$	$9.29$
RD [ 2 ]	$40.3$	$7.51$	$3.72$	$2.98$	$2.89$
LC [ 2 ]	$286$	$202$	$121$	$57.8$	$25.9$

TABLE IV: FIDs

\downarrow

for different NFEs on CIFAR-10, VE-SDE. Results marked with ∗ indicate values reported in prior works.

Sampler/NFE	150		195		210		270		300
RD-MC (ours)	9.07	3.13	9.51	1.06	9.55	1.04	9.69	0.91	9.70	0.86
RD-LDFM	7.69	25.3	9.25	2.99	9.15	2.27	9.32	3.22	9.25	2.13
Sampler/NFE	150		270		500		1000		2000
MC (ours)	2.77	246	8.65	19.7	9.97	7.16	10.1	13.2	9.74	5.71
RD	7.71	29.8	9.17	4.11	9.39	1.55	9.40	1.22	9.46	1.25
LC	1.85	308	3.31	198	5.37	99.6	7.55	41.1	9.20	16.9

TABLE V: IS

\uparrow

, KID (

\times 10^{-3}

)

\downarrow

for different NFEs on CIFAR-10, VE-SDE.

Dataset	NFE	RD-MC	RD-LC	DDIM	Analytic-DPM
FFHQ	150	$11.2$	$50.0$	10.9	11.6
	300	9.39	$13.4$	10.1	9.88
	450	7.38	$7.50$	8.75	7.47
ImageNet	150	5.46	$59.7$	6.44	10.2
	300	3.52	$11.3$	5.08	7.21
	450	2.90	$5.17$	4.11	3.66

TABLE VI: FIDs for different NFEs on FFHQ, ImageNet, and Bedroom.

Dataset		Ego-small				Community-small				Grid
Dataset		Real, $4\leq\|V\|\leq 18$				Synthetic, $12\leq\|V\|\leq 20$				Synthetic, $100\leq\|V\|\leq 400$
	Methods	Deg. $\downarrow$	Clus. $\downarrow$	Orbit $\downarrow$	Avg.	Deg. $\downarrow$	Clus. $\downarrow$	Orbit $\downarrow$	Avg.	Deg. $\downarrow$	Clus. $\downarrow$	Orbit $\downarrow$	Avg.
Autoreg.	GraphRNN	$0.060$	$0.220$	$0.003$	$0.104$	$0.080$	$0.120$	$0.040$	$0.080$	0.064	$0.043$	0.021	0.043
	DeepGMG	$0.040$	$0.100$	$0.020$	$0.053$	$0.220$	$0.950$	$0.400$	$0.523$	$-$	$-$	$-$	$-$
	GRAN	$0.020$	$0.126$	$\underline{0.010}$	$0.052$	$0.070$	$0.045$	$0.021$	$0.045$	$-$	$-$	$-$	$-$
	GraphDF	$0.040$	$0.130$	$\underline{0.010}$	$0.060$	$0.060$	$0.120$	$0.030$	$0.070$	$-$	$-$	$-$	$-$

TABLE VII: Quantitative results

\downarrow

on generic graph datasets.

NFE	Methods	Deg.	Clus.	Orbit	Time/s
300	EM	$0.074$	$0.111$	$\underline{0.007}$	$2.99$
	EM-LD	$\underline{0.061}$	$0.122$	$0.031$	$2.86$
	EM-LDFM	$0.065$	$0.122$	$0.031$	$2.94$
	EM-MC (ours)	$0.067$	$\underline{0.109}$	$0.031$	$2.83$
450	EM	$0.049$	$0.121$	$0.013$	$4.28$
	EM-LD	$0.042$	$0.071$	$0.005$	$4.09$

TABLE VIII: Quantitative results on Community-small.

Dataset	NFE	$10^{-1}$	$10^{-2}$	$10^{-3}$	$10^{-4}$	$10^{-5}$
CIFAR-10	150	$7.21$	$6.76$	$6.75$	$6.75$	$6.75$
CIFAR-10	300	$3.14$	$2.66$	$2.61$	$2.61$	$2.61$
ImageNet	150	$5.51$	$5.46$	$5.46$	$5.46$	$5.46$
ImageNet	300	$4.18$	$3.66$	$3.52$	$3.52$	$3.52$

TABLE IX: Ablation study on hyperparameter

\delta

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Sampler	NFE $=150$	NFE $=195$	NFE $=210$	NFE $=270$	NFE $=300$
RD-MC	$\alpha=0.20,n_{\sigma}=2$	$\alpha=0.19,n_{\sigma}=2$	$\alpha=0.21,n_{\sigma}=2$	$\alpha=0.29,n_{\sigma}=2$	$\alpha=0.265,n_{\sigma}=2$
RD-CDL	$\alpha=0.07,n_{\sigma}=2$	$\alpha=0.06,n_{\sigma}=2$	$\alpha=0.05,n_{\sigma}=2$	$\alpha=0.04,n_{\sigma}=2$	$\alpha=0.03,n_{\sigma}=2$
	NFE $=150$	NFE $=270$	NFE $=500$	NFE $=1000$	NFE $=2000$
MC	$\alpha=0.30,n_{\sigma}=5$	$\alpha=0.20,n_{\sigma}=5$	$\alpha=0.085,n_{\sigma}=5$	$\alpha=0.05,n_{\sigma}=5$	$\alpha=0.0315,n_{\sigma}=5$
LC	$\alpha=0.44,n_{\sigma}=5$	$\alpha=0.38,n_{\sigma}=5$	$\alpha=0.32,n_{\sigma}=5$	$\alpha=0.26,n_{\sigma}=5$	$\alpha=0.22,n_{\sigma}=5$
CDL	$\alpha=0.12,n_{\sigma}=5$	$\alpha=0.08,n_{\sigma}=5$	$\alpha=0.056,n_{\sigma}=5$	$\alpha=0.039,n_{\sigma}=5$	-

Appendix

TABLE X: Used hyperparameters for NCSN++ (VE-SDE) sampling/evaluation. We observe that the optimal

\alpha

hyperparameter is, to some extent, linearly negatively correlated with the NFE.

NFE	150	300	450
FFHQ	$\alpha=0.08,n_{\sigma}=2$	$\alpha=0.04,n_{\sigma}=2$	$\alpha=0.027,n_{\sigma}=2$
ImageNet	$\alpha=0.08,n_{\sigma}=2$	$\alpha=0.04,n_{\sigma}=2$	$\alpha=0.027,n_{\sigma}=2$
Bedroom	$\alpha=0.08,n_{\sigma}=2$	$\alpha=0.04,n_{\sigma}=2$	$\alpha=0.027,n_{\sigma}=2$

Appendix

TABLE XI: Hyperparameters for FFHQ, ImageNet (EDM, VE) sampling/evaluation.

Fig. 3: Additional samples on CIFAR10 from NCSN++ (VE-SDE) using RD-MC sampler with NFE $=210$ , we report an FID of $2.86$ with this setting.

Fig. 4: Additional samples on CIFAR10 from NCSN++ (VE-SDE) using MC sampler with NFE $=500$ , we report an FID of $9.16$ with this setting.

Fig. 5: Additional samples on CIFAR10 from NSCN using AMS sampler with NFE $=1000$ . We report an FID of $9.82$ with this setting.

Fig. 6: Additional samples on CIFAR10 from DDPM++ cont. (VP-SDE) using RD-MC sampler with NFE $=2000$ .

Fig. 7: Additional samples on $256\times 256$ CelebA-HQ from NCSN++ (VE-SDE) using RD-MC sampler with NFE $=450$ .

Fig. 8: Additional samples on $256\times 256$ CelebA-HQ from NCSN++ (VE-SDE) using RD-MC sampler with NFE $=900$ .

Fig. 9: Additional samples on $256\times 256$ bedroom from NCSN++ (VE-SDE) using RD-MC sampler with NFE $=500$ .

Fig. 10: Additional samples on $256\times 256$ church from NCSN++ (VE-SDE) using RD-MC sampler with NFE $=300$ .

Fig. 11: Additional samples on $64\times 64$ FFHQ from EDM with VE schedule using RD-MC sampler with NFE $=450$ .

Fig. 12: Additional samples on $64\times 64$ ImageNet from EDM with VE schedule using RD-MC sampler with NFE $=450$ .

Diffusion Model Acceleration

Diffusion model acceleration reduces the training and sampling cost of iterative generative models while preserving output quality.

See all 141 papers

Adaptive Moments are Surprisingly Effective for Plug-and-Play Diffusion Sampling
Mar 17, 2026Christian Belardi, Justin Lovelace, Kilian Q. Weinberger +1Diffusion Sampling Diffusion Model Guidance
Cornell University
Information-geometric adaptive sampling for graph diffusion
Apr 30, 2026Yuhui Lu, Wenjing Liu, Kun ZhanInformation Geometry Graph Generation
School of Information Science and Engineering, Lanzhou University, Gansu, China.
Non-asymptotic Convergence of Stochastic Gradient Descent in Score-based Generative Models
Jul 6, 2026Stanislas Strasman, Sobihan Surendran, Sylvain Le CorffNonconvex Stochastic Optimization Score-Based Generative Modeling
SU, LPSM

Mar 17, 2026J/K move · Enter open · S save

Adaptive Moments are Surprisingly Effective for Plug-and-Play Diffusion Sampling

Christian Belardi, Justin Lovelace, Kilian Q. Weinberger, Carla P. Gomes

Cornell University

Diffusion Sampling Diffusion Model Guidance Diffusion-Based Image Generation Diffusion-Based Denoising Diffusion Models for Image Restoration

Open paper AskPDF arXiv

Guided diffusion sampling relies on approximating often intractable likelihood scores, which introduces significant noise into the sampling dynamics. We propose using adaptive moment estimation to stabilize these noisy likelihood scores during sampling. Despite its simplicity, our approach achieves state-of-the-art results on image restoration and class-conditional generation tasks, outperforming more complicated methods, which are often computationally more expensive. We provide empirical analysis of our method on both synthetic and real data, demonstrating that mitigating gradient noise through adaptive moments offers an effective way to improve alignment.

Abstract

Diffusion Model Acceleration

Explore similar work

Adaptive Moments are Surprisingly Effective for Plug-and-Play Diffusion Sampling

Information-geometric adaptive sampling for graph diffusion

Non-asymptotic Convergence of Stochastic Gradient Descent in Score-based Generative Models