PaperlayerView exact revision on arXiv

Source-linked AI summary

Noise, Denoise, Correct: MCMC Posterior Sampling with Diffusion Priors in Three Steps

So Takao, Gregory David Bellchambers, Luke Ye, Sanmitra Ghosh, Michalis Michaelides

arXiv:2610.09407v1cs.LG

TL;DR

Posterior sampling with pretrained diffusion priors remains difficult when the prior density is unavailable and forward models are nonlinear or non-differentiable. Diffusion waltz uses noising–denoising proposals with Metropolis–Hastings correction and adds a conditional-noising extension; on Navier–Stokes initial condition recovery, it matches or outperforms strong derivative-free baselines in accuracy and calibration.

  • Problem

    Posterior sampling remains challenging when pretrained generative priors have no tractable density, and many guidance methods rely on likelihood approximations or differentiable forward models.

    Source: §1 · p. 1

  • Method

    Diffusion waltz uses a noising–denoising MCMC proposal with a Metropolis–Hastings correction, and its conditional-noising extension injects observations while preserving prior-free correction.

    Source: §1 · p. 1, §2.2 · p. 2

  • Results

    On Navier–Stokes initial condition recovery, diffusion waltz matches or outperforms strong derivative-free baselines in both accuracy and calibration.

    Source: §5 · p. 4

  • Takeaways & Limitations

    Diffusion waltz provides an asymptotically exact posterior-sampling framework, assuming a perfect denoiser, that handles non-differentiable forward models.

    Source: §5 · p. 4

    Exact posterior sampling assumes a perfect denoiser; in practice, the denoising step is approximate, and SDE discretization error is usually left without further correction.

    Source: §2.1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

Inverse problems seek to recover unknown states from noisy, incomplete observations, and pretrained generative priors leave posterior sampling difficult when their density is unavailable or the forward model is non-differentiable. Diffusion waltz uses a noising–denoising proposal with a Metropolis–Hastings correction, and can inject observations into that proposal while retaining exact posterior targeting.

1 Introduction

  • Inverse problems recover an unknown state from noisy, incomplete observations, which may not contain enough information to determine that state exactly.The Bayesian formulation places a prior over the state and targets its posterior.Source: §1 · p. 1
  • Pretrained generative priors make posterior sampling difficult when their density is unavailable, while many inference-time guidance methods use biased likelihood-score approximations or require differentiable forward models.Scientific applications often demand accurate posterior estimation, and realistic forward operators may not be differentiable.Source: §1 · p. 1
  • Diffusion waltz uses a three-step MCMC cycle to target the posterior exactly, assuming a perfect denoiser.The proposal noises the current state, denoises to propose a new state, then applies a Metropolis–Hastings accept/reject step.Source: §1 · p. 1
  • Its Metropolis–Hastings correction does not require evaluating the prior density, and observation information can be injected into the proposal without requiring a differentiable forward model.The modified proposal retains an exact Metropolis–Hastings correction that does not involve the prior density.Source: §1 · p. 1

2 Method

Diffusion waltz turns noising and denoising into an MCMC proposal corrected to target the posterior without evaluating the prior density. Conditional noising incorporates observations into the proposal, and EKCN supplies a gradient-free ensemble Kalman update for that step.

2.1 Diffusion waltz for posterior sampling

  • RWMH waltz uses noising and denoising to propose a move, then applies a Metropolis–Hastings correction whose likelihood-only acceptance ratio gives posterior convergence without evaluating the prior density.The method builds on a stochastic interpolant between data and reference distributions; in practice, the denoising step is approximate.Source: §2.1 · p. 2
  • The denoising step is approximate in practice, and discretization error from solving the denoising SDE can in principle be corrected with further Metropolis–Hastings steps.The authors report that this additional correction is usually unnecessary for satisfactory results.Source: §2.1 · p. 2

2.2 Conditional noising

  • Conditional noising replaces the usual noising step with a kernel conditioned on observations, while retaining a tractable Metropolis–Hastings acceptance probability that does not require the prior density.The acceptance probability is tractable provided the conditional noising kernel can be evaluated.Source: §2.2 · p. 2

2.3 The ensemble Kalman conditional noising (EKCN) kernel

  • EKCN estimates its Kalman gain from an ensemble without gradients of the forward operator and cycles through a Kalman update, SDEdit projection, and Metropolis–Hastings correction.The ensemble Kalman Levenberg–Marquardt update nudges the forward-model output toward the observations.Source: §2.3 · pp. 2–3

3 Related works

Related work includes concurrent RWMH-waltz methods, other ways to inject observation information, and several MCMC posterior samplers using pretrained diffusion priors.

3 Related works

  • Concurrent work considers RWMH waltz, while related approaches inject observations during denoising, often requiring gradients, and other diffusion-prior MCMC families include sequential Monte Carlo, split-Gibbs, and elliptical slice sampling.The paper compares against several of these methods in its experiments.Source: §3 · p. 3

4 Experiments

On Navier–Stokes initial-condition recovery, DW-EKCN performs strongly across noise and nonlinearity regimes, while conditional noising accelerates convergence. The experiments also compare derivative-free baselines and show a high-noise, high-nonlinearity reconstruction.

4 Experiments

  • DW-EKCN performs strongly on all scores across all regimes, consistently ranking first or second.Table 1 reports mean Rel-ℓ2, CRPS, and SSR across five unseen test examples.Source: §4 · p. 3, Table 1 · p. 4
  • In the high-nonlinearity, high-noise regime, DW-EKCN recovers the ground truth accurately with well-calibrated uncertainties despite the heavily corrupted observation.Source: Figure 1 · p. 4
  • Blade and ESS-Flow have slow inference, while FK Steering and MFM-GF show weight collapse; MCMC approaches mostly outperform gradient-based guidance methods.Source: §4 · p. 4
  • Conditional noising guides ensemble members toward high-likelihood regions, with convergence within about 50 MCMC steps, especially in the highly nonlinear setting.Acceptance rates later decay, settling at 10–20% in the high-noise setting and below 5% in the low-noise regime.Source: §4 · p. 4

5 Conclusion

Diffusion waltz is an asymptotically exact MCMC framework, assuming a perfect denoiser, that handles non-differentiable forward models; its conditional-noising extension injects observation information while preserving exactness. The accompanying stochastic-interpolant discussion describes the reverse-time marginal law and relationships among drift, score, and clean-sample estimates in a Gaussian-reference case.

5 Conclusion

  • Diffusion waltz provides a plug-and-play framework for posterior sampling via MCMC, handles non-differentiable forward models, and is asymptotically exact assuming a perfect denoiser.Its conditional-noising extension injects observation information into the proposal while preserving exactness and substantially improving mixing; on Navier–Stokes recovery it matches or outperforms strong derivative-free baselines in accuracy and calibration.Source: §5 · p. 4

A Stochastic interpolants

  • A stochastic interpolant connects data and reference measures through a process whose reverse-time SDE shares its marginal law when solved backward from the reference distribution.The interpolant uses monotonic C2 functions with specified endpoint conditions; its marginal velocity and score can be approximated by neural networks.Source: A Stochastic interpolants · pp. 7–8

A.1 Three-sides of the same coin

  • With a Gaussian reference, the drift, score, and conditional clean-sample estimate are related linearly, including through Tweedie’s formula.Source: §A.1 · p. 8

A.2 Choice of SDE to sample from p(x0 | xt)

The selected reverse-time SDE is designed to recover the conditional law needed for diffusion waltz denoising. The appendix also derives MH acceptance probabilities and notes a covariance choice used in the Navier–Stokes experiments.

A.2 Choice of SDE to sample from p(x0 | xt)

  • The reverse-time SDE’s terminal law matches the interpolant conditional p(x0 | xt), providing the distribution needed for the denoising step.This result follows for the specified choice of C_t.Source: §A.2 · p. 8
  • A general covariance C is used to accommodate the functional flow matching architecture in the Navier–Stokes experiments.That architecture uses a Matérn GP covariance rather than white noise.Source: §A.2 · p. 8

B Acceptance probability computations

B.1 RWMH diffusion waltz

The appendix establishes how diffusion-waltz proposals yield prior-free MH acceptance calculations and how conditional noising preserves posterior reversibility. It also develops ensemble-based, derivative-free updates for nonlinear data-misfit problems.

B.1 RWMH diffusion waltz

  • The RWMH proposal is reversible with respect to the diffusion prior, reducing its MH acceptance probability to a likelihood ratio without evaluating prior densities.Source: §B.1 · p. 9

B.2 Conditional noising waltz

  • The conditional noising waltz kernel is reversible with respect to the posterior; ordinary RWMH acceptance follows when the noising kernel ignores y.Source: §B.2 · p. 11

C Statistical linearisation and ensemble Kalman update rules

  • Statistical linearisation fits a least-squares linear approximation to A(x), yielding an ensemble-based Jacobian approximation that supports derivative-free optimisation.The linearisation is described in [Zheng et al., 2025a].Source: C Statistical linearisation and… · p. 11

C.1 Ensemble Kalman update rule and the Levenberg-Marquardt solver

  • The ensemble Kalman update gives closed-form estimates for the regularised least-squares solution and its curvature, enabling a gradient-free Levenberg–Marquardt update.Source: §C.1 · pp. 11–12

C.2 Details on the EKCN kernel

The EKCN proposal combines an ensemble Kalman data-misfit update with noising and denoising to return proposals to the data manifold. Its implementation reuses the gain for the MH calculation, while the all-member particle update has an exactness caveat.

C.2 Details on the EKCN kernel

  • EKCN first applies an ensemble Kalman LM update to reduce data misfit, then uses SDEdit to project the proposal back onto the data manifold.The figure caption reports high early acceptance and fast convergence because most proposed samples have lower data misfit.Source: §C.2 · pp. 12–13, Figure 2 · p. 13
  • The pushthrough identity reduces the matrix inverse to size J × J, making the gain computation tractable when the observation dimension is high.Source: §C.2 · p. 13
  • The same Kalman gain is reused in both conditional proposal terms so that the same kernel appears in the MH acceptance calculation.Source: §C.2 · pp. 13–14
  • Using all ensemble members makes the interacting-particle update subtly break exactness, although the authors report faster convergence and better empirical performance.Splitting particles into groups is given as a way to overcome this issue.Source: §C.2 · p. 14

D Adaptive MCMC tuning

The method adaptively tunes its noising time and ensemble Kalman step size, while preventing noising time from collapsing toward zero. The experiments vary Navier–Stokes evolution time and observation noise, using a trajectory-based vorticity dataset.

D Adaptive MCMC tuning

  • As t approaches zero, the conditional-noising term can dominate the log acceptance probability and drive acceptance close to zero, so tuning clips t away from zero.The method also tunes the Kalman step size h jointly with t; setting h = 0 recovers RWMH waltz.Source: D Adaptive MCMC tuning · p. 15

E.1 Problem statement

  • The experiments vary T and σy to control nonlinearity and noise, with observations formed by evolving the initial field, downsampling, and adding Gaussian noise.The downsampling factor is fixed at k = 4.Source: §E.1 · p. 15

E.2 Data

  • The data comprise about 20,000 snapshots from 200 Navier–Stokes vorticity trajectories, split by trajectory into 180 training and 20 validation trajectories.Trajectories are thinned by sampling every five timesteps; the original fields have 64 × 64 resolution.Source: §E.2 · p. 15

E.3 Model architecture

The experiments use a functional flow-matching prior with a Fourier Neural Operator and evaluate reconstruction, ensemble quality, and uncertainty calibration. The appendix also specifies the MCMC settings and the comparison baselines.

E.3 Model architecture

  • The functional flow-matching prior parameterizes its marginal velocity with a Fourier Neural Operator using four Fourier layers and 32 retained modes per spatial axis.The network uses a hidden width of 64, with the input field concatenated with a positional grid and flow-matching time.Source: §E.3 · p. 16

E.4 Metrics

  • Rel-ℓ2 measures ensemble-mean reconstruction error, CRPS scores the full ensemble, and SSR assesses whether ensemble spread matches its actual error.SSR near 1 indicates calibrated spread; values below 1 indicate underdispersion and values above 1 indicate overdispersion.Source: §E.4 · pp. 16–17

E.5 Diffusion waltz details

  • Diffusion waltz uses 100 parallel chains for 250 steps, discards 249 burn-in steps, and computes metrics from the final ensemble members.The ensemble has 100 members, and the denoising SDE is discretized with Euler–Maruyama.Source: §E.5 · p. 17

E.6 Baseline details

  • EKI minimizes data fit without prior information and its particles can collapse, whereas EnKG guides diffusion denoising but also risks particle collapse.EKI uses a Levenberg–Marquardt solver; EnKG uses ensemble Kalman updates during the denoising corrector.Source: §E.6 · p. 17
  • FK Steering uses sequential Monte Carlo and is asymptotically exact as particle count increases, while ESS-Flow is exact up to ODE discretization error.ESS-Flow’s delayed acceptance can require many forward-model evaluations.Source: §E.6 · pp. 17–18
  • Blade combines Gibbs sampling with an unadjusted Langevin likelihood step, while MFM-GF estimates likelihood scores without differentiating through the likelihood.Blade approximates Langevin gradients by statistical linearization; MFM-GF uses importance sampling.Source: §E.6 · p. 18

E.7 Comparison with gradient-based guidance methods

DPS and MMPS perform best in the low-nonlinearity, low-noise setting, but are mostly outperformed by MCMC-based methods in other settings, especially at high nonlinearity. The guidance comparisons use differentiable surrogates because the Navier–Stokes forward operator is non-differentiable.

E.7 Comparison with gradient-based guidance methods

  • Because the forward operator is non-differentiable, the gradient-based methods are evaluated using trained FNO surrogates rather than the solver directly.The paper applies functional modifications to the methods to accommodate its functional flow-matching prior.Source: §E.7 · p. 18
  • DPS and MMPS perform best at low nonlinearity and low noise, but MCMC-based methods mostly outperform them in other settings, especially at high nonlinearity.The authors attribute their weaker high-nonlinearity performance to increasingly significant posterior bias; the results are summarized for low and high nonlinearity in Tables 5 and 6.Source: §E.7 · p. 18, Table 5 · p. 19, Table 6 · p. 19
  • DAPS performs poorly on this problem, which the authors note is consistent with its difficulty on PDE forward models.Source: §E.7 · p. 18

Related papers

Loading 2610.09407v1…