cs.LGOct 7, 2026
SaveNoise, Denoise, Correct: MCMC Posterior Sampling with Diffusion Priors in Three Steps
Organizations: PhysicsX London, UK
Abstract
Pretrained diffusion models are powerful priors for inverse problems, but posterior sampling under nonlinear, non-differentiable forward models remain hard. We introduce diffusion waltz, an MCMC method using SDEdit-style noising-denoising as a proposal, corrected via Metropolis-Hastings for exact posterior sampling without prior evaluation. We further propose injecting observations into the proposal while preserving exactness, using a gradient-free ensemble Kalman update. On a non-differentiable Navier-Stokes initial condition recovery task, diffusion waltz outperforms existing baselines across different noise and nonlinearity regimes.
Figures & tables
| Low nonlinearity ( ) | High nonlinearity ( ) | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Low noise | High noise | Low noise | High noise | |||||||||
| Methods | Rel | CRPS | SSR | Rel | CRPS | SSR | Rel | CRPS | SSR | Rel | CRPS | SSR |
| EKI | 0.517 | 0.370 | 0.024 | 0.832 | 0.560 | 0.126 | 1.059 | 0.754 | 0.031 | 0.865 | 0.522 | 0.267 |
| EnKG | 0.129 | 0.086 | 0.076 | 0.617 | 0.414 | 0.071 | 0.418 | 0.279 | 0.098 | 0.812 | 0.523 | 0.084 |
| FK Steering | 0.787 | 0.538 | 0.012 | 0.782 | 0.500 | 0.101 | 1.383 | 1.000 | 0.016 | 0.992 | 0.511 | 1.084 |
| ESS-Flow | 0.178 | 0.080 | 1.405 | 0.248 | 0.118 | 1.179 | 0.506 | 0.242 | 0.987 | 0.491 | 0.233 | 0.987 |
Table 1 : Mean Rel- , CRPS and SSR across five unseen test examples (for an expanded version with standard deviations and gradient-based baselines, see Tables 5 – 6 ). Best results in bold , second best underlined .
Figure 1 : Results on the Navier-Stokes (NS) initial condition recovery problem in the high nonlinearity / high noise regime. Despite the heavily corrupted observation, diffusion waltz with ensemble Kalman conditional noising (DW-EKCN) recovers the ground truth accurately with well-calibrated uncertainties.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 2 : Illustration of the ensemble Kalman conditional noising diffusion waltz (DW-EKCN). This applies a single step of ensemble Kalman Levenberg-Marquardt (EK-LM) update to minimise data misfit, followed by applying SDEdit to project it back onto the data manifold. Earlier in the iterations, acceptance rate is high and convergence is fast as most proposed samples have lower data misfit.
| solver | |||||||
|---|---|---|---|---|---|---|---|
| 0.3 | 0.3 | 0.3 | 0.3 | 0.7 | 0.01 | 1.0 | LM |
Table 2 : Choice of adaptive MCMC hyperparameters
| Low | High | |
|---|---|---|
| Nonlinearity | ||
| Noise |
Table 3 : Experimental regimes
| Grid size | Smoothness | Lengthscale | Variance |
|---|---|---|---|
| 64 | 0.5 | 0.1 | 1.0 |
Table 4 : Choice of Matérn GP reference hyperparameters
| Observation noise | Low noise | High noise | ||||
|---|---|---|---|---|---|---|
| Methods | Rel | CRPS | SSR | Rel | CRPS | SSR |
| DPS | 0.075 (0.011) | 0.039 (0.008) | 1.393 (0.084) | 0.275 (0.078) | 0.127 (0.023) | 0.855 (0.162) |
| MMPS | 0.064 (0.009) | 0.032 (0.006) | 0.983 (0.095) | 0.254 (0.063) | 0.116 (0.020) | 0.927 (0.189) |
| DAPS | 0.401 (0.013) | 0.182 (0.037) | 1.188 (0.105) | 0.710 (0.062) | 0.338 (0.028) | 1.342 (0.068) |
| EKI | 0.517 (0.077) | 0.370 (0.086) | 0.024 (0.006) | 0.832 (0.064) | 0.560 (0.042) | 0.126 (0.009) |
| EnKG | 0.129 (0.014) | 0.086 (0.009) | 0.076 (0.019) | 0.617 (0.120) | 0.414 (0.049) | 0.071 (0.009) |
Table 5 : Results in the low nonlinearity regime ( ). Best results in bold , second best underlined . Gradient based methods are displayed in gray.
| Observation noise | Low noise | High noise | ||||
|---|---|---|---|---|---|---|
| Methods | Rel | CRPS | SSR | Rel | CRPS | SSR |
| DPS | 0.472 (0.084) | 0.201 (0.069) | 1.834 (0.285) | 0.506 (0.084) | 0.233 (0.079) | 1.244 (0.250) |
| MMPS | 0.560 (0.090) | 0.276 (0.077) | 1.572 (0.382) | 0.523 (0.070) | 0.241 (0.069) | 1.461 (0.292) |
| DAPS | 0.639 (0.071) | 0.327 (0.071) | 1.510 (0.281) | 0.697 (0.057) | 0.351 (0.075) | 1.358 (0.258) |
| EKI | 1.059 (0.069) | 0.754 (0.104) | 0.031 (0.006) | 0.865 (0.086) | 0.522 (0.069) | 0.267 (0.041) |
| EnKG | 0.418 (0.153) | 0.279 (0.140) | 0.098 (0.035) | 0.812 (0.157) | 0.523 (0.070) | 0.084 (0.006) |
Table 6 : Results in the high nonlinearity regime ( ). Best results in bold , second best underlined . Gradient based methods are displayed in gray.
Figure 3 : Comparison of diffusion waltz and baseline methods on the Navier-Stokes (NS) initial condition recovery problem in the low nonlinearity ( ) / low noise ( ) regime.
Figure 4 : Comparison of diffusion waltz and baseline methods on the Navier-Stokes (NS) initial condition recovery problem in the low nonlinearity ( ) / high noise ( ) regime.
Figure 5 : Comparison of diffusion waltz and baseline methods on the Navier-Stokes (NS) initial condition recovery problem in the high nonlinearity ( ) / low noise ( ) regime.
Figure 6 : MCMC traces for the potential . We compare DW-RWMH (left) and DW-EKCN (right), across all four regimes. Many chains in DW-RWMH get stuck, unable to find a path to high likelihood/low potential regions, leading to slow convergence. Conditional noising in DW-EKCN alleviates this problem, leading to rapid convergence within the first 50 steps.
Figure 7 : Instantaneous acceptance rates for DW-RWMH (left) and DW-EKCN (right), across all four regimes. DW-RWMH shows steady decay with MCMC steps. In the low noise regime, DW-EKCN has high acceptance rates initially, fluctuates around while the chains converge, then hits . In the high noise regime, acceptance rate is high initially, then immediately relaxes to .
Explore similar work
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.
A Gradient Flow Approach to Solving Inverse Problems with Latent Diffusion Models
Solving ill-posed inverse problems requires powerful and flexible priors. We propose leveraging pretrained latent diffusion models for this task through a new training-free approach, termed Diffusion-regularized Wasserstein Gradient Flow (DWGF). Specifically, we formulate the posterior sampling problem as a Wasserstein gradient flow in the latent space of an expected negative log posterior objective, regularized by a Kullback-Leibler divergence to the diffusion prior. We demonstrate the performance of our method on standard benchmarks using StableDiffusion (Rombach et al., 2022) as the prior.
Provable diffusion-based posterior sampling for linear inverse problems via DDIM
Diffusion-based methods have achieved remarkable empirical success in solving inverse problems. However, many existing posterior samplers either lack rigorous theoretical guarantees or incur substantial computational overhead. We propose a simple and efficient algorithm, called \pddim, for solving linear inverse problems with diffusion priors via a DDIM-type sampler. Our method requires only lightweight, coordinate-wise modifications to the standard DDIM update, while explicitly incorporating the measurement model. The key idea is to perform posterior sampling separately along each singular direction of the measurement operator: for each direction, the sampler follows the learned diffusion prior when the observation signal-to-noise ratio (SNR) is below the corresponding diffusion SNR, and switches to a calibrated measurement-based predictor otherwise. We prove that the proposed sampler converges to the Bayesian posterior conditioned on the measurements. Empirical results show that the proposed sampler performs favorably against existing diffusion-based posterior samplers across a range of image restoration tasks, achieving the best performance on the majority of evaluation metrics considered. Overall, our results convert posterior sampling for noisy linear inverse problems to simple coordinate-wise DDIM updates, yielding an efficient, easy-to-implement algorithm with provable posterior consistency.