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

Yingzhi Xia^1*, Setthakorn Tanomkiattikun^1,3*, Liangli Zhen^1✉, Zaiwang Gu²

¹Institute of High Performance Computing, Agency for Science, Technology and Research, Singapore ²Institute for Infocomm Research, Agency for Science, Technology and Research, Singapore ³Johns Hopkins University

* Equal contribution. | ✉ Corresponding author.

Xia_Yingzhi@a-star.edu.sg | zhen_liangli@a-star.edu.sg | Gu_Zaiwang@a-star.edu.sg | stanomk1@jhu.edu

Motivation and Method

Comparison of existing methods and NA-NHMC

Figure. Comparison of existing methods and their limitations with the N-HMC method. (a) Iterative Guidance Methods (DPS) lead to manifold infeasibility. (b) Stochastic MAP methods (ReSample) are susceptible to overfitting to noise. (c) Deterministic MAP methods (DMPlug) become trapped in a local mode. (d) Our method performs sampling in the noise space $x T$ and maps samples to images via a deterministic mapping $x 0 = D(x T)$ .

Abstract

Diffusion models have shown strong performance in inverse problems, but existing approaches often suffer from manifold infeasibility, noise overfitting, or local-mode collapse. We propose Noise-space Hamiltonian Monte Carlo (N-HMC), which treats reverse diffusion as a deterministic mapping from initial Gaussian noise to clean images and performs posterior sampling entirely in the noise space. This formulation keeps proposals on the learned manifold and enables robust posterior exploration. We further introduce Noise-Adaptive N-HMC (NA-NHMC), which marginalizes unknown measurement noise variance and removes task-specific tuning of noise-level hyperparameters. Across linear and nonlinear inverse problems on FFHQ and ImageNet, NA-NHMC delivers strong reconstruction quality and robust behavior under unknown noise.

Sampling Dynamics Across Inverse Problems

Three representative inverse problems. Each task card shows reference image, measurement, and NA-NHMC convergence trajectory.

Phase Retrieval

Setting: FFHQ 256x256, observation noise sigma_y=0.01. Detail: Fourier-magnitude measurements cause severe ambiguity; trajectory shows gradual manifold-aligned recovery.

Random Inpainting (92%)

Setting: FFHQ 256x256, random mask ratio 92%, observation noise sigma_y=0.05. Detail: Most pixels are missing; NA-NHMC samples progressively restore global structure and local facial textures.

HDR Reconstruction

Setting: FFHQ 256x256, nonlinear camera response, observation noise sigma_y=0.05. Detail: Trajectory illustrates iterative correction of exposure and contrast while preserving semantic consistency.

Experimental Results

Main Quantitative Comparison

Key benchmark table from the paper. It summarizes performance over representative linear and nonlinear inverse problems. NA-NHMC consistently improves robustness, especially in difficult nonlinear settings and unknown-noise scenarios.

Main quantitative table from paper — Main result table highlighting competitive or best performance of NA-NHMC across tasks and metrics.

Nonlinear Deblurring Visual Comparison

Under high noise (σ_y=0.2), NA-NHMC recovers sharper facial details with fewer artifacts than prior baselines.

Nonlinear deblurring qualitative results — Nonlinear deblurring results on FFHQ (256×256) with σ_y=0.2.

Unknown Noise Robustness

With the same hyperparameter setting used for Gaussian experiments, NA-NHMC remains strong under impulse and speckle noise.

Unknown noise qualitative results — Nonlinear deblurring on FFHQ under impulse noise (top) and speckle noise (bottom).

Initialization Robustness

Compared with DPS and DAPS, NA-NHMC shows lower variance across repeated runs while preserving reconstruction fidelity.

Heatmap robustness visualization — Top: MAE heatmaps. Bottom: standard deviation heatmaps over 100 runs.

Convergence in Highly Ill-Posed Task

The annealed noise schedule improves early exploration and raises successful recovery rate in phase retrieval.

Phase retrieval convergence — Phase retrieval on FFHQ (256×256), σ_y=0.01; median with 5th–95th percentile band.

BibTeX

@inproceedings{ tanomkiattikun2026noiseadaptive, title={Noise-Adaptive Diffusion Sampling for Inverse Problems Without Task-Specific Tuning}, author={Yingzhi Xia and Setthakorn Tanomkiattikun and Liangli Zhen and Zaiwang Gu}, booktitle={The Fourteenth International Conference on Learning Representations}, year={2026}, url={https://openreview.net/forum?id=Yfk4ex3Z1G} }