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

1. Problem Statement

Plug-and-play (PnP) conditional generation aims to sample from a posterior distribution $p(x|y)$ using a diffusion model trained only on the marginal distribution $p(x)$, without task-specific retraining. This is achieved by guiding the sampling process using Bayes' rule:
$$\nabla_{x_t} \log p(x_t|y) = \nabla_{x_t} \log p(x_t) + \nabla_{x_t} \log p(y|x_t)$$
While the prior score ($\nabla_{x_t} \log p(x_t)$) is provided by the pre-trained diffusion model, the likelihood score ($\nabla_{x_t} \log p(y|x_t)$) is often intractable.

Existing methods approximate this likelihood score (e.g., Diffusion Posterior Sampling (DPS) or Classifier Guidance (CG)). However, these approximations introduce significant noise into the sampling dynamics, especially when:

• Conditioning information is sparse (e.g., extreme super-resolution or heavy deblurring).
• Gradients are computed through large neural networks or complex analytical operators.
• The task difficulty is high, causing guidance signals to become unstable or contradictory across timesteps.

Current state-of-the-art methods often attempt to solve this by developing more sophisticated approximation algorithms (e.g., iterative refinement, Monte Carlo smoothing, or composite frameworks), which increase computational cost and complexity.

2. Methodology: Adaptive Moment Estimation (Adam)

The authors propose a simple yet effective modification: applying Adaptive Moment Estimation (the core mechanism behind the Adam optimizer) to stabilize the noisy likelihood score gradients during the sampling process.

Instead of using the raw gradient estimate $g_t = -\nabla_{x_t} L(f_\phi(x_{0|t}), y)$ at each timestep, the method maintains exponential moving averages of the first and second moments of the gradients across the sampling trajectory:

1. First Moment (Momentum): $m_k = \beta_1 m_{k-1} + (1 - \beta_1) g_t$
2. Second Moment (Variance): $v_k = \beta_2 v_{k-1} + (1 - \beta_2) g_t^2$
3. Bias Correction: $\hat{m}_k = m_k / (1 - \beta_1^k)$, $\hat{v}_k = v_k / (1 - \beta_2^k)$
4. Stabilized Gradient: $\hat{g}_t = \hat{m}_k / (\sqrt{\hat{v}_k} + \delta)$

This stabilized gradient $\hat{g}_t$ is then used to update the sampling step. The authors implement this for two primary baselines:

• AdamDPS: Applies adaptive moments to Diffusion Posterior Sampling.
• AdamCG: Applies adaptive moments to Classifier Guidance.

Key Insight: By accumulating gradient information across steps, the method smooths the optimization trajectory and adaptively scales updates based on historical variance. This prevents the sampler from "fighting itself" (oscillating due to conflicting gradient directions) and ensures steady progress toward the condition.

3. Key Contributions

1. Simplicity and Effectiveness: Demonstrates that a simple modification (adding Adam-style moment estimation) to existing PnP methods yields state-of-the-art (SOTA) results, outperforming complex, computationally expensive frameworks.
2. Robustness to Task Difficulty: Reveals that while complex methods (like Training-Free Guidance, TFG) perform well on easy tasks, they degrade rapidly as task difficulty increases (e.g., 16x super-resolution). Adaptive moments consistently improve performance across all difficulty levels.
3. Empirical Analysis: Provides deep analysis showing that adaptive moments mitigate gradient noise, leading to higher cosine similarity between sequential guidance steps and more direct trajectories toward the ground truth.
4. Open Source: Releases code for reproducibility.

4. Experimental Results

The methods were evaluated on ImageNet, CIFAR-10, and the Cats vs. Dogs dataset across reconstruction and class-conditional generation tasks.

Reconstruction Tasks (Super-Resolution, Deblurring, Inpainting):

• Metrics: LPIPS (lower is better) and FID (lower is better).
• Performance: AdamDPS outperformed all baselines (DPS, TFG, UGD, MPGD, LGD, RED-diff) on all tasks.
• Difficulty Scaling: In challenging 16x super-resolution and high-intensity deblurring, TFG and other complex methods often performed worse than the baseline DPS. AdamDPS maintained robust improvements over DPS even in these extreme conditions.
• Qualitative: AdamDPS produced sharper details and fewer artifacts compared to the blurred or artifacted outputs of TFG and standard DPS.

Class-Conditional Sampling:

• Metrics: Classification Accuracy (harmonic mean of 3 classifiers) and FID.
• Performance:
  • On CIFAR-10, AdamDPS improved accuracy by 9.86% over the next best method (DPS).
  • On ImageNet (a much harder setting), most methods achieved near-random accuracy (~1%). AdamDPS achieved 10.49%.
  • For time-aware classifiers (AdamCG), accuracy improved by over 19 percentage points over standard CG.

Efficiency and Ablation:

• Wall Clock Time: AdamDPS adds negligible overhead compared to DPS. It is significantly faster than TFG, which requires multiple gradient computations per step ($N_{recur}(1+N_{iter})$).
• Hyperparameters: Ablation studies confirmed that both the first moment (momentum) and second moment (adaptive scaling) are essential; removing either degrades performance.
• Sampling Trajectories: Visualizations showed that DPS trajectories often stray from the target due to noisy gradients, while AdamDPS trajectories converge more directly to the ground truth.

5. Significance

This paper challenges the prevailing trend in diffusion guidance of building increasingly complex algorithmic frameworks to approximate likelihood scores. Instead, it argues that the primary bottleneck is gradient noise, which can be effectively mitigated using standard stochastic optimization techniques (Adaptive Moments).

Implications:

• Practicality: Researchers can achieve SOTA results by adding a few lines of code to existing guidance implementations without retraining models or increasing computational budgets significantly.
• Scientific Inverse Problems: The method is particularly valuable for scientific domains (e.g., medical imaging, material science) where conditioning information is sparse, noisy, and underdetermined.
• Evaluation Standards: The work highlights the need for evaluation standards that include "hard" tasks, as methods that excel on easy benchmarks often fail in realistic, difficult scenarios.

In summary, the paper establishes that stabilizing noisy gradients via adaptive moments is a more effective strategy for plug-and-play diffusion sampling than developing more complex likelihood approximations.