Stabilizing Diffusion Guidance with Adaptive Moments: When Optimization Meets Generative Sampling

Diffusion models have established themselves as the dominant paradigm in generative modeling, powering applications from text-to-image synthesis to molecular design. Yet conditioning these models on external information without retraining remains a persistent challenge. The framework of plug-and-play conditional generation offers an elegant solution: it enables sampling from conditional distributions using only unconditional models by guiding the reverse diffusion process toward desired conditions. However, this approach introduces a fundamental tension. The guidance signal, derived from an approximation of the intractable likelihood score, injects noise into the sampling dynamics that can destabilize generation. In their paper "Adaptive Moments are Surprisingly Effective for Plug-and-Play Diffusion Sampling," Belardi et al. propose a remarkably simple intervention that addresses this instability. By applying adaptive moment estimation, the optimization technique underlying Adam, to stabilize noisy guidance gradients, they achieve state-of-the-art results across image restoration tasks while revealing important insights about task difficulty and method robustness.

The Noise Problem in Plug-and-Play Sampling

To understand the contribution, we must first examine the mechanics of conditional diffusion sampling. The reverse process of a diffusion model follows annealed Langevin dynamics, gradually denoising a sample by following the score function, which points toward high probability regions of the data distribution. When conditioning on some observation y, Bayes' rule decomposes the posterior score into two components: the prior score provided by the unconditional diffusion model, and the likelihood score representing the gradient of log p(y|x_t). While the prior score is well approximated by the trained noise prediction network, the likelihood score is typically intractable.

Existing plug-and-play methods approximate this term through various strategies. Diffusion Posterior Sampling (DPS) uses a Monte Carlo approximation based on the Gaussian transition kernel, while more recent approaches like Universal Guidance for Diffusion Models (UGD) and Training-Free Guidance (TFG) incorporate increasingly complex algorithmic machinery involving latent space computations, iterative refinement, and multiple gradient estimators. Yet all these methods share a common vulnerability: they estimate the likelihood score independently at each timestep using noisy, high-variance approximations. Because diffusion sampling proceeds through hundreds or thousands of iterative steps, these independent approximations accumulate error, causing the sample trajectory to deviate from the true posterior.

The key insight of Belardi et al. is that we should view guidance not as a series of independent optimization problems but as a sequential estimation task where information from earlier steps can inform later ones. Rather than developing ever more sophisticated single-step approximations, they ask whether we can better utilize the noisy estimates we already have.

Adaptive Moments as Temporal Regularization

The proposed solution applies adaptive moment estimation to the likelihood score gradients across sampling timesteps. In standard stochastic optimization, Adam maintains exponential moving averages of first and second moments of gradients to stabilize parameter updates. Here, the authors apply the same principle to the sampling trajectory itself. For each timestep t, they maintain running estimates m_t and v_t of the first and second moments of the likelihood score gradient ∇_{x_t} log p(y|x_t), computed as:

m_t = β_1 m_{t-1} + (1 - β_1) g_t v_t = β_2 v_{t-1} + (1 - β_2) g_t^2

where g_t represents the noisy likelihood score estimate at step t. The actual guidance signal used in the sampling update is then normalized using these accumulated statistics, incorporating bias correction terms analogous to standard Adam.

This approach functions as temporal smoothing of the guidance signal. By maintaining memory of past gradient estimates, the method dampens high-frequency noise while preserving the underlying trend of the likelihood score. The second moment estimation proves particularly crucial; it adapts the step size of the guidance component based on the historical variance of the estimates, effectively performing automatic annealing of the guidance strength across different regions of the sampling trajectory and different noise levels.

What makes this technique notable is its orthogonality to existing improvements. It does not replace DPS or Classifier Guidance (CG); rather, it wraps around them, stabilizing their outputs with minimal computational overhead. The memory requirements are modest, requiring only storage of momentum buffers comparable in size to the diffusion latent itself.

Task Difficulty and the Robustness Spectrum

The empirical evaluation reveals findings that extend beyond the primary contribution. On standard image restoration benchmarks including super-resolution and inpainting, the adaptive moment approach consistently improves upon baseline DPS, often matching or exceeding the performance of UGD and TFG while avoiding their computational complexity. On class-conditional generation using pretrained classifiers, the method similarly demonstrates superior alignment between generated samples and target classes.

However, the most revealing experiments concern varying task difficulty. The authors systematically evaluate performance across super-resolution factors ranging from 4× to 16×. Under mild degradation (4×), many guidance methods perform adequately, with complex techniques like TFG showing marginal improvements over simpler baselines. As difficulty increases to 16× super-resolution, where the conditioning signal becomes extremely sparse and the likelihood score approximation correspondingly noisier, a striking pattern emerges. Sophisticated methods that compose multiple gradient estimates and iterative refinements degrade rapidly, often performing worse than vanilla DPS. In contrast, the adaptive moment approach maintains consistent improvements over DPS across all difficulty levels, with the relative gap widening as tasks become harder.

This suggests that complex guidance methods, despite their theoretical appeal, may be overfitting to the noise characteristics of easier tasks or introducing compounding errors when approximations break down. The adaptive moment approach, by explicitly modeling the noise structure through second moment estimation, demonstrates greater robustness to the extreme variance regimes encountered in difficult inverse problems.

Your Take: Optimization Meets Sampling

This work represents a fascinating collision between two traditionally distinct communities: stochastic optimization and diffusion-based sampling. The success of adaptive moments here suggests we should reconsider the theoretical framing of plug-and-play guidance. Rather than viewing each timestep as drawing an independent noisy sample from the true likelihood score, we might conceptualize the entire sampling trajectory as a stochastic optimization process over the posterior distribution, where momentum and adaptive step sizes provide the stability necessary for convergence.

The results also speak to a broader pattern in machine learning research regarding simplicity versus complexity. UGD and TFG represent the accumulation of algorithmic sophistication, layering multiple techniques in an attempt to squeeze out performance gains. Yet this paper demonstrates that a single, well-chosen bias-variance tradeoff, implemented through moment estimation, can outperform these intricate pipelines. This aligns with the observation that in high-noise regimes, variance reduction often dominates bias reduction in importance.

From a theoretical perspective, the paper opens questions about the convergence properties of momentum-based Langevin dynamics. Standard analyses of Langevin Monte Carlo assume unbiased gradient estimates with bounded variance. Plug-and-play diffusion violates these assumptions severely; the likelihood score estimates are biased by the Gaussian approximation of the noising process and exhibit heteroscedastic variance across timesteps. The empirical success of adaptive moments suggests that the method may be implicitly correcting for these distributional shifts, but formal convergence guarantees under these conditions remain undeveloped.

Conclusion

Belardi et al. demonstrate that stabilizing noisy guidance gradients through adaptive moment estimation offers a simple yet powerful improvement to plug-and-play diffusion sampling. Their method achieves state-of-the-art results while revealing that task difficulty serves as a crucial axis for evaluating guidance robustness. As diffusion models expand into scientific domains involving noisy measurements and ill-posed inverse problems, the principles of gradient stabilization and adaptive step sizing may prove essential. Future work might extend these ideas to other variance reduction techniques from the optimization literature, or develop theoretical frameworks for understanding momentum in non-convex, stochastic sampling landscapes. For practitioners, the message is clear: before adding architectural complexity, consider whether you are making optimal use of the noisy signals you already possess.