Improving Diffusion Posterior Samplers with Lagged Temporal Corrections for Image Restoration

Technical Summary: Improving Diffusion Posterior Samplers with Lagged Temporal Corrections (LAMP)

Problem Statement

Diffusion-based posterior sampling (PS) has emerged as a leading framework for solving imaging inverse problems (e.g., deblurring, super-resolution, inpainting). These methods combine a learned diffusion prior with measurement constraints to sample from the posterior distribution $p(x_0|y)$, where $y = Kx_0 + e$. Standard PS formulations, such as DiffPIR and DDRM, operate by iteratively updating the current estimate $x_t$ using a denoising step from the diffusion model and a data-consistency correction $D_t$.

However, the authors identify a fundamental structural limitation in these standard approaches: they rely on instantaneous, step-wise data-consistent estimates ($D_t$) that are recomputed at every reverse iteration. This design introduces temporal variability and instability in the reverse dynamics. From a dynamical systems perspective, standard PS updates correspond to a first-order discretization of the diffusion dynamics plus a residual correction term. This residual term acts as a time-varying forcing signal that perturbs the trajectory, causing it to deviate from the smooth, well-defined continuous-time dynamics assumed by high-order solvers (like DPM-Solver) used in unconditional generation. Consequently, the lack of temporal memory in standard PS methods leads to inconsistent updates and suboptimal stability, particularly in ill-posed regimes.

Methodology: LAMP

The paper proposes LAMP (LAgged Multistep Posterior), a modular plug-in scheme that addresses these instabilities by introducing a lagged temporal correction.

1. Dynamical Reinterpretation: The authors reinterpret the standard PS update as a first-order discretization of a forced diffusion dynamics. They show that the update $x^{PS}_{t-\Delta t}$ can be decomposed into a first-order exponential step driven by a measurement-aware target $D_t$ and a residual correction term proportional to the mismatch $(D_t - \hat{x}_{0|t})$.
2. Second-Order Discretization: To improve stability, LAMP employs a second-order discretization of the underlying dynamics. Unlike standard high-order solvers that assume a smooth ODE, LAMP explicitly accounts for the residual-dependent forcing term inherent to inverse problems.
3. Lagged Correction Mechanism: The core innovation is the introduction of a temporal correction based on the variation of consecutive estimates. The LAMP update combines the second-order discretization with the residual correction:
   $$x^{LAMP}_{t-\Delta t} = x^{PS}_{t-\Delta t} + \alpha_{t-\Delta t} A_1(h) \gamma \frac{D_t - D_{t+\Delta t}}{h_{prev}}$$
   Here, the term $(D_t - D_{t+\Delta t})/h_{prev}$ acts as a discrete temporal derivative. By choosing a negative coefficient $\gamma$, the method applies a lagged correction, pulling the current estimate slightly toward previous ones rather than extrapolating. This smooths the trajectory and counteracts the variability induced by instantaneous corrections.
4. Structural Preservation: Crucially, the authors prove (Proposition 1) that LAMP preserves the standard posterior sampling structure. The update can be rewritten in the canonical form $x_{t-\Delta t} = \alpha_{t-\Delta t} \tilde{D}_t + \sigma_{t-\Delta t} \epsilon_\Theta(x_t, t)$, where $\tilde{D}_t$ is a temporally filtered estimate defined as a convex combination of the current and previous data-consistent estimates: $\tilde{D}_t = (1-\beta_t)D_t + \beta_t D_{t+\Delta t}$. This ensures LAMP is a modular extension compatible with any existing PS backbone (e.g., DiffPIR, DDRM) without modifying the denoiser or training procedure.

Theoretical Analysis

The paper provides a one-step risk analysis (Proposition 2) characterizing the bias-variance trade-off of the lagged correction. The analysis shows that LAMP improves the reverse transition when the reduction in estimation error variance (due to temporal smoothing) outweighs the bias introduced by the lag (due to the drift of the posterior target). The authors demonstrate that this condition is typically satisfied in the later stages of the diffusion process, motivating a hybrid strategy where LAMP is activated after a short warm-up phase.

Key Contributions

• Identification of Structural Mismatch: The authors highlight that standard ODE-based interpretations of diffusion solvers fail to capture the residual-dependent forcing terms present in inverse problem posterior sampling, leading to temporal instability.
• LAMP Framework: They propose LAMP, a second-order corrected scheme that integrates a lagged temporal correction into the posterior sampling update.
• Modularity and Compatibility: LAMP is proven to be a valid posterior sampler that can be implemented as a plug-in over any existing PS method (e.g., DiffPIR, DDRM) without retraining or changing the forward model.
• Theoretical Justification: A local error analysis establishes the conditions under which the lagged correction reduces risk, specifically when adjacent posterior targets evolve smoothly relative to the estimation noise.

Experimental Results

The authors evaluate LAMP on standard linear inverse problems (Gaussian deblurring, motion deblurring, and super-resolution) across three datasets: FFHQ, ImageNet, and CelebA. They compare LAMP-enhanced versions (DiffPIR-LAMP, DDRM-LAMP) against their baselines and naive higher-order variants (1M, 2M).

• Performance: LAMP consistently improves reconstruction quality over strong baselines. In the noisy setting ($\sigma_y = 0.05$), DDRM-LAMP achieves the best overall performance, showing consistent gains in PSNR, SSIM, and LPIPS across all tasks. DiffPIR-LAMP also shows improvements, particularly in deblurring and super-resolution.
• Efficiency: These improvements are achieved without increasing the number of denoising evaluations (NFE), maintaining the computational cost of the base methods (e.g., 20 NFE for DDRM, 100 NFE for DiffPIR).
• Comparison with Naive Higher-Order: Naive first- and second-order ODE discretizations (without the specific residual correction) consistently degrade performance, confirming that the proposed lagged residual correction is the key factor for improvement.
• Qualitative Results: Visual comparisons show that LAMP produces smoother textures, sharper boundaries, and better preservation of fine details (e.g., hair texture in CelebA, scales in Iguana images) compared to baselines, reducing graininess and blurring artifacts.

Significance and Claims

The paper claims that temporal filtering of measurement-aware estimates is a critical, previously underexplored mechanism for stabilizing diffusion-based inverse solvers. By reframing posterior sampling as a dynamical system with a residual-dependent forcing term, the authors demonstrate that a simple lagged correction can significantly enhance reconstruction quality and stability.

The significance lies in the modularity of the approach: LAMP does not require new priors, retraining, or complex conditioning strategies. Instead, it offers a general-purpose, low-cost enhancement to existing state-of-the-art posterior samplers. The authors modestly note that LAMP inherits the limitations of its backbone methods (e.g., it cannot compensate for fundamental model mismatches) and that its effectiveness depends on the temporal behavior of the underlying sampler's correction sequence. However, the results suggest that for stable baselines, LAMP provides a consistent and perceptually meaningful improvement in data fidelity, particularly for fine-scale structures.