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

Mar 17, 2026

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

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.

Christian Belardi, Justin Lovelace, Kilian Q. Weinberger +1

Cornell University

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

Ask PDF

Apr 30, 2026

Information-geometric adaptive sampling for graph diffusion

Standard diffusion models for graph generation typically rely on uniform time-stepping, an approach that overlooks the non-homogeneous dynamics of distributional evolution on complex manifolds. In this paper, we present an information-geometric framework that reinterprets the diffusion sampling trajectory as a parametric curve on a Riemannian manifold. Our key observation is that the Fisher-Rao metric provides a principled measure of the intrinsic distance. By analyzing this metric, we derive the Drift Variation Score (DVS), a geometry-aware indicator that quantifies the instantaneous rate of distributional change. Unlike prior heuristic-based adaptive samplers, our DVS solver enforces a constant informational speed on the statistical manifold, automatically maintaining a uniform rate of distributional change along the sampling trajectory. This equal arc-length strategy ensures that each discretization step contributes equally to the information speed. Theoretical analysis verifies that DVS characterizes the local stiffness of the sampling dynamics in the Fisher-Rao sense. Experimental results on molecule and social network generation show that DVS significantly improves structural fidelity and sampling efficiency. Code is at https://github.com/kunzhan/DVS

Yuhui Lu, Wenjing Liu, Kun Zhan

School of Information Science and Engineering, Lanzhou University, Gansu, China.

Information Geometry Graph Generation Diffusion Sampling Adaptive Sampling Diffusion Timestep Scheduling Graph Diffusion

Ask PDF

Jul 6, 2026

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

Score-based Generative Models (SGMs) have achieved impressive performance in data generation across a wide range of applications. While the statistical properties of their sampling procedures are increasingly well understood, the optimization dynamics underlying their training remain less explored. SGMs are typically trained by minimizing a weighted denoising score-matching objective, yet optimization guarantees with stochastic gradients remain limited. In this work, we study Stochastic Gradient Descent (SGD) for SGMs, contributing results in two complementary regimes. For general score parameterizations, we derive a non-convex analysis of SGD for the weighted denoising score-matching objective, making explicit how the resulting optimization bound depends on the loss weighting and time-sampling distribution. We then consider overparameterized two-layer ReLU networks and develop a Neural Tangent Kernel analysis tailored to diffusion training with stochastic gradients, yielding score-approximation error bounds along the SGD trajectory. Our analysis quantifies the role of the reweighting factor in these bounds, providing a theoretical characterization of weighting choices used in practice.

Stanislas Strasman, Sobihan Surendran, Sylvain Le Corff

SU, LPSM

Nonconvex Stochastic Optimization Score-Based Generative Modeling Stochastic Gradient Descent

Ask PDF

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