Enhancing Diffusion Model Stability for Image Restoration via Gradient Management

2025-07-09

AI Summary

Diffusion models in Bayesian image restoration suffer from gradient misalignment—conflicting directions between prior and likelihood gradients—and temporal instability of likelihood gradients, leading to unstable sampling. This work is the first to systematically characterize this gradient dynamic mismatch mechanism and proposes a gradient management framework: progressive likelihood warm-up mitigates initial gradient conflict, while adaptive directional momentum smoothing suppresses gradient oscillations. The method integrates seamlessly into standard diffusion sampling without requiring auxiliary networks or additional training. Evaluated on denoising and super-resolution tasks, it achieves state-of-the-art performance, with significant PSNR and SSIM improvements over existing methods. Sampling trajectories exhibit enhanced stability, and both quantitative metrics and human perceptual assessments confirm superior visual quality compared to current Bayesian-guided approaches.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Bayesian Learning
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

Diffusion models have shown remarkable promise for image restoration by leveraging powerful priors. Prominent methods typically frame the restoration problem within a Bayesian inference framework, which iteratively combines a denoising step with a likelihood guidance step. However, the interactions between these two components in the generation process remain underexplored. In this paper, we analyze the underlying gradient dynamics of these components and identify significant instabilities. Specifically, we demonstrate conflicts between the prior and likelihood gradient directions, alongside temporal fluctuations in the likelihood gradient itself. We show that these instabilities disrupt the generative process and compromise restoration performance. To address these issues, we propose Stabilized Progressive Gradient Diffusion (SPGD), a novel gradient management technique. SPGD integrates two synergistic components: (1) a progressive likelihood warm-up strategy to mitigate gradient conflicts; and (2) adaptive directional momentum (ADM) smoothing to reduce fluctuations in the likelihood gradient. Extensive experiments across diverse restoration tasks demonstrate that SPGD significantly enhances generation stability, leading to state-of-the-art performance in quantitative metrics and visually superior results. Code is available at href{https://github.com/74587887/SPGD}{here}.

Problem

• Analyzes gradient conflicts in diffusion models for image restoration
• Identifies instability from prior-likelihood gradient direction mismatches
• Addresses temporal fluctuations in likelihood gradients during generation

Innovation

• Progressive likelihood warm-up strategy
• Adaptive directional momentum smoothing
• Stabilized Progressive Gradient Diffusion technique