Geometry-Correct Diffusion Posterior Sampling with Denoiser-Pullback Curvature Guidance and Manifold-Aligned Damping

Abstract: Diffusion posterior sampling conditions diffusion priors on measurements, but data-consistency updates are typically scaled by hand-tuned guidance weights and can destabilize sampling under stiff, operator-dependent curvature. We replace scalar guidance with a per-noise-level damped Gauss--Newton correction computed in diffusion-state coordinates. The correction pulls likelihood gradients back through the denoiser, uses a one-sided curvature model that avoids forward denoiser Jacobians, and applies diffusion-calibrated rank-one damping aligned with the denoiser residual. Each correction is solved with matrix-free GMRES using automatic differentiation, and sampling proceeds with a variance-preserving Langevin transition with a closed-form drift/noise split. On FFHQ and ImageNet across inverse problems, it achieves competitive PSNR/SSIM/LPIPS while running markedly faster than most of the compared baselines; on accelerated MRI reconstruction, it achieves the best PSNR/SSIM among the compared baselines.

• The paper introduces CLAMP, a novel posterior sampler that employs a denoiser-pullback Gauss–Newton method for geometry-correct likelihood guidance.
• It integrates a variance-preserving Langevin framework with rank-one damping to address stiffness and nonlinearities in inverse imaging.
• Experimental results across various imaging tasks and MRI data demonstrate competitive reconstruction quality and improved runtime efficiency.

Geometry-Correct Diffusion Posterior Sampling: CLAMP

Introduction

The paper "Geometry-Correct Diffusion Posterior Sampling with Denoiser-Pullback Curvature Guidance and Manifold-Aligned Damping" (2605.27990) addresses a longstanding problem in diffusion-based approaches to inverse imaging: effectively and efficiently sampling the posterior conditioned on indirect measurements. Standard diffusion posterior samplers typically inject hand-tuned scalar guidance to enforce measurement consistency, but their updates are brittle when facing stiff or anisotropic likelihood geometry, particularly with nonlinear forward models. This work introduces CLAMP (Curvature-aware Langevin with Aligned Manifold Pullback), a posterior sampler that computes matrix-scaled, geometry-correct likelihood corrections using a denoiser-pullback Gauss–Newton quadratic model with rank-one manifold-aligned damping, all integrated into a variance-preserving Langevin sampling framework.

Methodology

CLAMP’s approach consists of several algorithmic innovations that specifically handle the geometric and numerical challenges that arise during posterior sampling in diffusion-based inverse problems.

Denoiser-Pullback Curvature Guidance

The likelihood term in the posterior is defined on the clean signal domain, whereas the sample trajectory traverses the noisy diffusion variable space. CLAMP addresses the coordinate mismatch by differentiating through the denoiser $F_\theta$: likelihood gradients are “pulled back” to the state space via the denoiser Jacobian, ensuring coordinate correctness under local anisotropy. The measurement sensitivity is mapped to diffusion-state coordinates by the chain $\mathbf{J}_F^\top \mathbf{J}_A^\top \mathbf{r}$, where $\mathbf{J}_A$ is the Jacobian of the forward model $A$ and $\mathbf{J}_F$ of the denoiser.

Curvature-Scaled Correction via One-Sided Gauss–Newton Approximation

Rather than relying on scalar step sizes, CLAMP solves a damped Gauss–Newton system for the likelihood correction at each noise level. The key approximation in CLAMP’s curvature operator drops the expensive forward denoiser-Jacobian product inside the curvature while retaining the denoiser pullback in the adjoint, yielding favorable computational properties (resolving the system with only VJP/JVP primitives via autodiff). The correction is stabilized using a prior-aligned, rank-one damping metric that penalizes motion along the denoiser-residual direction, adapting strength to the diffusion noise schedule.

Variance-Preserving Langevin Sampling

Sampling proceeds by a variance-preserving Langevin transition, with the split between drift and noise injection determined in closed form via fluctuation–dissipation coupling. This avoids explicit stochasticity scheduling and maintains exploration consistent with prescribed diffusion chain statistics.

Matrix-Free Implementation

All corrections are computed as linear system solves in a matrix-free manner using GMRES, operating with only Jacobian-vector and vector-Jacobian products (autodiff or numerical). The method is fully compatible with both pixel and latent parameterizations.

Figure 1: Qualitative reconstructions across linear and nonlinear inverse imaging problems demonstrate high-fidelity, artifact-free outputs from CLAMP in both pixel and latent spaces, including accelerated MRI.

Experimental Results

Extensive empirical evaluation is performed across FFHQ and ImageNet datasets on a suite of inverse problems: super-resolution, inpainting (box/random), Gaussian and motion deblurring, phase retrieval, nonlinear deblurring, and high dynamic range reconstruction. CLAMP attains competitive or leading PSNR/SSIM/LPIPS metrics against recent baselines, such as SITCOM, DAPS, DMPlug, DCDP, PSLD, and ReSample, matching or outperforming them while requiring substantively less runtime per sample.

Figure 2: CLAMP achieves Pareto-optimal trade-offs between image quality (PSNR/SSIM/LPIPS) and runtime over motion and nonlinear deblurring as the computational budget is varied.

CLAMP is particularly robust under nonlinear forward models, where instability is common. It remains stable and effective due to the curvature-aware scaling and prior-aligned damping, whereas scalar-guided samplers are often forced to take small steps to avoid instability, resulting in slow convergence or poor reconstruction. Ablation studies confirm the necessity of each design component: variance-preserving Langevin propagation, prior-aligned damping, denoiser-pullback curvature, and GMRES for nonsymmetric systems.

Figure 3: Sample reconstructions at $512\times512$ resolution, showing that CLAMP maintains stability and visual fidelity for high-resolution, high-dynamic-range and random inpainting tasks.

Figure 4: MRI reconstructions from $\times8$ undersampled measurements; CLAMP produces sharper anatomical details and edges compared to other diffusion-based baselines.

On multi-coil accelerated MRI (SKM-TEA dataset), CLAMP achieves the best PSNR/SSIM among diffusion-based baselines for both $\times4$ and $\times8$ acceleration, demonstrating practical relevance beyond natural images.

Numerical and Calibration Analysis

CLAMP yields normalized measurement residuals tightly concentrated around the ideal value, indicating statistically consistent coverage of the posterior distribution under the assumed likelihood. Large-update diagnostics show near-optimality relative to the locally linearized problem, and hyperparameter sweeps confirm the efficacy of prior-aligned damping over isotropic or orthogonal alternatives.

Figure 5: Empirical histogram of normalized measurement residual $r_s$ for box inpainting; CLAMP samples are tightly calibrated to the measurement noise level.

Figure 6: CLAMP enables diverse posterior sampling under severely ill-posed scenarios (e.g., $\mathbf{J}_F^\top \mathbf{J}_A^\top \mathbf{r}$0 super-resolution or extreme inpainting), consistently yielding visually plausible and measurement-consistent reconstructions.

Theoretical and Practical Implications

CLAMP provides a principled resolution to the coordinate and curvature mismatch endemic in heuristic diffusion posterior sampling, significantly reducing the need for per-operator guidance tuning. It achieves direction-wise, data-consistency-driven correction while maintaining diffusion-aligned step control. The combination of curvature scaling, coordinate correctness, and schedule-aligned damping should generalize smoothly to other domains such as 3D and video, subject to practical considerations regarding Jacobian action computations.

On the theoretical side, CLAMP bridges classical Gauss–Newton optimization and stochastic posterior sampling in diffusion geometry, motivating future exploration of non-differentiable forward operators (where Jacobian actions are not tractable) and higher-dimensional modalities. The matrix-free structure lends itself well to scaling, and further research may consider global curvature models or learning damping metrics from data.

Conclusion

CLAMP introduces a robust, efficient posterior sampling framework for inverse problems using diffusion models. Its denoiser-pullback curvature guidance and manifold-aligned damping yield geometry-correct, matrix-scaled corrections, addressing both stiffness and coordinate mismatch. Across a range of domains and tasks—including accelerated MRI—CLAMP delivers superior or competitive reconstruction quality while markedly improving inference speed and reliability, making it an attractive framework for future work in diffusion-based inference and uncertainty quantification in imaging and beyond.