Noise-Adaptive Diffusion Sampling for Inverse Problems Without Task-Specific Tuning

TL;DR

This paper introduces Noise-space Hamiltonian Monte Carlo (N-HMC), a novel diffusion sampling method for inverse problems that explores the solution space effectively without task-specific tuning, demonstrating superior results.

Contribution

The paper proposes N-HMC and its noise-adaptive extension NA-NHMC, enabling robust inverse problem solutions without task-specific tuning and with comprehensive solution space exploration.

Findings

NA-NHMC outperforms state-of-the-art methods in reconstruction quality.

The approach is robust across different hyperparameters and initializations.

Theoretical analysis supports the effectiveness of the proposed method.

Abstract

Diffusion models (DMs) have recently shown remarkable performance on inverse problems (IPs). Optimization-based methods can fast solve IPs using DMs as powerful regularizers, but they are susceptible to local minima and noise overfitting. Although DMs can provide strong priors for Bayesian approaches, enforcing measurement consistency during the denoising process leads to manifold infeasibility issues. We propose Noise-space Hamiltonian Monte Carlo (N-HMC), a posterior sampling method that treats reverse diffusion as a deterministic mapping from initial noise to clean images. N-HMC enables comprehensive exploration of the solution space, avoiding local optima. By moving inference entirely into the initial-noise space, N-HMC keeps proposals on the learned data manifold. We provide a comprehensive theoretical analysis of our approach and extend the framework to a noise-adaptive variant (NA-NHMC) that effectively handles IPs with unknown noise type and level. Extensive experiments across four linear and three nonlinear inverse problems demonstrate that NA-NHMC achieves superior reconstruction quality with robust performance across different hyperparameters and initializations, significantly outperforming recent state-of-the-art methods. The code is available at https://github.com/NA-HMC/NA-HMC.