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

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

cs.LGMar 17, 2026

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

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

Organizations: Cornell University

Abstract

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.

Figures & tables

Figure 1: Left: The KL divergence between each method’s empirical distribution and the target distribution as a function of the guidance noise coefficient

\zeta

. Right: Visualization of the empirical and target distributions at

\zeta=0.175

Figure 2: Qualitative comparison of AdamDPS, DPS, and TFG on Cats dataset for super resolution at 12x downsampling and Gaussian deblurring at blur intensity 9.

Figure 3: Reconstruction performance measured in LPIPS and FID, where lower is better for both. Comparison on ImageNet and Cats dataset for super resolution at 16x downsampling, Gaussian deblurring at blur intensity 12, and inpainting with a 90% random mask.

Figure 4: Class-conditional sampling performance measured in classification accuracy and FID, where higher accuracy and lower FID is better. Accuracy is computed as the harmonic mean across three held-out classifiers. Left & Center: Comparison of plug-and-play methods with a standard classifier on CIFAR-10 and ImageNet, respectively. Right: Comparison of plug-and-play methods with a time-aware classifier on ImageNet.

Figure 5: Relative improvement over DPS on ImageNet as task difficulty increases for super resolution and Gaussian deblurring.

Figure 6: Relative improvement over DPS for super resolution at 16x downsampling on ImageNet, across DDPM and DDIM sampling step budgets.

Figure 7: Left: Ablation of Adam

\beta_{1}

\beta_{2}

for super resolution at 16x downsampling, Gaussian deblurring at blur intensity 12, and inpainting at 90% random mask on the Cats dataset. Right: Wall clock comparison on a single H100 GPU of 100 step class-conditional sampling with a standard classifier on ImageNet for a batch of 8 256x256 images.

Figure 8: Sampling trajectories for DPS and AdamDPS projected onto two dimensions for super resolution at 16x downsampling on ImageNet. The y-axis is defined by the difference between the initial noise and the target, and the x-axis by the difference between the AdamDPS and DPS solutions. Contours depict the MSE loss surface with respect to the target.

Figure 9: Cosine similarity between sequential guidance terms

g_{t}

and

g_{s}

throughout sampling for DPS and AdamDPS. Guidance terms collected from 16x super resolution task on ImageNet. Shading denotes the 25th to 75th percentile.

Figure 10: The guidance loss throughout sampling for class-conditional and super resolution tasks on ImageNet. Shading denotes the 25th to 75th percentile.

Figure 11: Qualitative example of sampling trajectory for super resolution at 16x downsampling on ImageNet. Displayed are the intermediate clean data predictions $x_{0|t}$ throughout sampling.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

	Inpainting			Super Resolution 4x			Super Resolution 16x			Gaussian Deblur 3			Gaussian Deblur 12
Method	LPIPS	FID	IS	LPIPS	FID	IS	LPIPS	FID	IS	LPIPS	FID	IS	LPIPS	FID	IS
$\Pi$ GDM	0.30	66.83	36.85	0.48	63.73	26.34	0.48	49.41	30.11	0.48	65.04	24.69	0.50	51.97	27.76
RED-diff	0.31	75.86	36.60	0.20	30.42	102.09	0.57	100.30	9.66	0.27	40.13	76.44	0.52	209.90	5.85
DPS	0.16	27.79	132.32	0.19	26.43	128.03	0.30	35.55	60.02	0.23	27.66	109.29	0.34	40.25	45.07
LGD	0.18	37.37	89.52	0.24	31.97	105.63	0.32	38.10	49.85	0.31	37.02	74.39	0.39	46.11	36.46
MPGD	0.45	142.47	14.47	0.16	27.75	120.52	0.35	52.66	35.93	0.23	31.18	107.18	0.38	52.11	34.34

Appendix

Table 1: Reconstruction on ImageNet

	Super Resolution 8x			Super Resolution 12x			Gaussian Deblur 6			Gaussian Deblur 9
Method	LPIPS	FID	IS	LPIPS	FID	IS	LPIPS	FID	IS	LPIPS	FID	IS
DPS	0.23	29.02	88.64	0.26	32.63	71.02	0.26	31.39	71.41	0.30	36.65	54.54
TFG ${}_{N_{\text{iter}}=1}$	0.22	31.76	79.87	0.27	41.54	52.95	0.29	35.43	63.57	0.33	42.91	46.36
TFG ${}_{N_{\text{iter}}=2}$	0.22	31.56	80.09	0.27	37.08	57.60	0.26	33.21	72.71	0.32	40.15	49.85
TFG ${}_{N_{\text{iter}}=4}$	0.22	32.26	76.58	0.27	39.84	53.69	0.26	33.58	73.29	0.31	39.58	49.63
AdamDPS	0.19	24.36	117.35	0.23	27.60	85.87	0.23	27.26	90.08	0.28	33.63	62.14

Appendix

Table 2: Reconstruction on ImageNet (Intermediate Difficulties)

	Inpainting		Super Resolution 4x		Super Resolution 16x		Gaussian Deblur 3		Gaussian Deblur 12
Method	LPIPS	FID	LPIPS	FID	LPIPS	FID	LPIPS	FID	LPIPS	FID
$\Pi$ GDM	0.32	80.77	0.40	60.26	0.44	52.76	0.42	58.80	0.46	58.48
RED-diff	0.22	40.72	0.12	22.04	0.35	103.73	0.16	32.82	0.40	120.05
DPS	0.13	17.13	0.14	17.74	0.25	30.68	0.18	23.58	0.29	34.74
LGD	0.17	35.29	0.14	19.63	0.27	35.15	0.20	25.72	0.33	40.10
MPGD	0.42	76.60	0.09	15.70	0.28	44.16	0.14	20.88	0.32	49.50

Appendix

Table 3: Reconstruction on Cats

	Super Resolution 8x		Super Resolution 12x		Gaussian Deblur 6		Gaussian Deblur 9
Method	LPIPS	FID	LPIPS	FID	LPIPS	FID	LPIPS	FID
DPS	0.21	25.36	0.22	25.22	0.22	26.99	0.26	31.44
TFG ${}_{N_{\text{iter}}=1}$	0.17	23.70	0.22	31.46	0.24	31.36	0.27	38.37
TFG ${}_{N_{\text{iter}}=2}$	0.16	22.10	0.22	29.77	0.21	27.40	0.27	36.95
TFG ${}_{N_{\text{iter}}=4}$	0.17	23.18	0.21	29.02	0.20	26.25	0.25	33.75
AdamDPS	0.16	19.45	0.20	22.64	0.20	24.32	0.23	26.48

Appendix

Table 4: Reconstruction on Cats (Intermediate Difficulties)

	DDPM								DDIM
	LPIPS				FID				LPIPS				FID
Method	100	50	25	12	100	50	25	12	100	50	25	12	100	50	25	12
DPS	0.30	0.33	0.42	0.58	35.55	44.54	62.21	154.29	0.31	0.34	0.45	0.57	41.50	43.62	51.34	58.01
TFG ${}_{N_{\text{iter}}=4}$	0.32	0.34	0.39	0.50	42.92	44.55	52.35	118.59	0.35	0.35	0.35	0.41	58.38	59.20	53.84	60.34
AdamDPS	0.27	0.30	0.35	0.46	30.16	35.63	49.05	98.71	0.29	0.30	0.35	0.44	39.60	37.18	46.35	77.33

Appendix

Table 5: Sampling Step Ablation for Super Resolution at 16x Downsampling on ImageNet

	DDPM								DDIM
	LPIPS				FID				LPIPS				FID
Method	100	50	25	12	100	50	25	12	100	50	25	12	100	50	25	12
DPS	0.25	0.29	0.35	0.52	30.68	40.15	58.25	184.52	0.29	0.29	0.37	0.51	30.31	32.35	42.73	67.90
TFG ${}_{N_{\text{iter}}=4}$	0.26	0.29	0.34	0.40	38.83	39.72	44.83	70.17	0.35	0.34	0.33	0.34	55.81	51.70	49.85	47.33
AdamDPS	0.24	0.26	0.29	0.38	27.62	30.32	34.15	72.19	0.27	0.27	0.29	0.37	27.53	27.20	32.51	61.72

Appendix

Table 6: Sampling Step Ablation for Super Resolution at 16x Downsampling on Cats

	Guidance Classifier Accuracy (%)	Held-out Classifier Accuracy (%)			FID
Method		ViT	ViT Finetuned	ConvNeXt
RED-diff	10.7	9.9	10.2	10.7	203.8
DPS	42.8	39.4	40.3	42.8	32.3
LGD	21.8	19.6	18.1	21.8	31.0
MPGD	26.7	24.9	24.3	26.7	34.1
UGD ${}_{N_{\text{iter}}=1}$	9.3	9.7	10.2	9.3	25.4

Appendix

Table 7: Class-Conditional (Standard Classifier) on CIFAR10

	Guidance Classifier Top-10 Accuracy (%)	Held-out Classifier Top-10 Accuracy (%)			FID
Method		ViT	DINO v2	MobileViT v2
RED-diff	12.7	23.0	3.2	1.4	272.7
DPS	0.7	0.8	1.2	1.0	43.0
LGD	0.7	0.7	1.2	1.1	42.5
MPGD	0.9	0.8	1.2	1.4	43.6
UGD ${}_{N_{\text{iter}}=1}$	0.8	0.6	0.7	0.7	42.3

Appendix

Table 8: Class-Conditional (Standard Classifier) on ImageNet

	Guidance Classifier Accuracy (%)	Held-out Classifier Accuracy (%)			FID
Method		ViT	DINO v2	MobileViT v2
CG	61.0	63.4	58.3	51.0	30.8
AdamCG	82.5	84.2	77.8	68.1	29.6

Appendix

Table 9: Class-Conditional (Time-Aware Classifier) on ImageNet

Figure 12: Additional qualitative results for super resolution at 16x downsampling on ImageNet.

Figure 13: Additional qualitative results for Gaussian deblurring at blur intensity 12 on ImageNet.

Figure 14: Additional qualitative results for super resolution at 12x downsampling on Cats.

Figure 15: Additional qualitative results for Gaussian deblurring at blur intensity 9 on Cats.

Figure 16: Additional qualitative results for class-conditional sampling with a standard classifier on CIFAR10.

Figure 17: Additional qualitative results for class-conditional sampling with a standard classifier on ImageNet.

Figure 18: Additional qualitative results for class-conditional sampling with a time-aware classifier on ImageNet.

Diffusion Sampling

Diffusion sampling studies how to turn a trained diffusion model into high-quality samples efficiently and reliably.

See all 258 papers

General and Efficient Steering of Diffusion Models
Feb 11, 2026Qingsong Wang, Mikhail Belkin, Yusu WangDiffusion Model Guidance Activation Steering
Halıcıoğlu Data Science Institute, University of California San Diego, La Jolla, CA, USA
Denoising, Fast and Slow: Difficulty-Aware Adaptive Sampling for Image Generation
Apr 21, 2026Johannes Schusterbauer, Ming Gui, Yusong Li +3Diffusion Model Sampling Diffusion Sampling
Ctrl-Z Sampling: Scaling Diffusion Sampling with Controlled Random Zigzag Explorations
Jun 25, 2025Shunqi Mao, Wei Guo, Chaoyi Zhang +3Diffusion Sampling Conditional Diffusion Models
School of Computer Science, The University of Sydney, Australia

Feb 11, 2026J/K move · Enter open · S save

General and Efficient Steering of Diffusion Models

Qingsong Wang, Mikhail Belkin, Yusu Wang

Halıcıoğlu Data Science Institute, University of California San Diego, La Jolla, CA, USA

Diffusion Model Guidance Activation Steering Controllable Image Generation Diffusion Model Inference Acceleration Diffusion-Based Denoising

Open paper AskPDF arXiv

Steering diffusion models toward conditions unseen during training typically requires either retraining with conditional inputs or per-step gradient computations, both of which incur substantial computational overhead. We present Noise-Aligned RFM Steering (NA-RFM), a general recipe for efficiently steering diffusion models without gradient guidance during inference, enabling fast controllable generation. The method combines two offline-computed signals: noise alignment, a high-noise correction from PCA statistics of the target examples and the full data, and Recursive Feature Machine (RFM) activation steering, which learns a target-discriminative direction from labeled forward-process activations. During sampling, noise alignment provides coarse control at high noise, while the RFM direction is reused over intermediate/late timesteps through lightweight activation edits. Experiments on CIFAR-10, ImageNet, CelebA, and fine-grained bird species show improved target accuracy over gradient-based post-hoc guidance baselines, improved FID on the class-guidance benchmarks, and substantial inference speedups. Code: https://github.com/isotrivial/na-rfm.

Abstract

Diffusion Sampling

Explore similar work

General and Efficient Steering of Diffusion Models

Denoising, Fast and Slow: Difficulty-Aware Adaptive Sampling for Image Generation

Ctrl-Z Sampling: Scaling Diffusion Sampling with Controlled Random Zigzag Explorations