Enhancing Diffusion Model Stability for Image Restoration via Gradient Management

TL;DR

This paper identifies gradient instabilities in diffusion-based image restoration and introduces SPGD, a novel method that manages gradients to improve stability and achieve state-of-the-art results.

Contribution

The paper presents SPGD, a new gradient management technique with warm-up and smoothing strategies, to enhance diffusion model stability for image restoration.

Findings

SPGD improves stability across diverse restoration tasks.

State-of-the-art quantitative performance achieved.

Visually superior restoration results with SPGD.

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 https://github.com/74587887/SPGD.