Noise-Adaptive Diffusion Sampling for Inverse Problems Without Task-Specific Tuning
For practitioners solving inverse problems with diffusion models, this method provides robust, high-quality reconstruction without task-specific tuning.
The paper proposes Noise-space Hamiltonian Monte Carlo (N-HMC) for solving inverse problems with diffusion models, avoiding local minima and manifold infeasibility. The noise-adaptive variant (NA-NHMC) handles unknown noise and achieves superior reconstruction quality across four linear and three nonlinear inverse problems, significantly outperforming recent state-of-the-art methods.
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.