Diffusion models for inverse problems and data assimilation

Before this lesson: lesson 11 (the score split of classifier guidance; here the classifier becomes a likelihood), lesson 08 (Tweedie's formula and its second-order form), lesson 09 (the PF-ODE and DDIM), lesson 06 (Langevin corrector steps) and lesson 14 (rank histograms and CRPS). You also need the multivariate Gaussian (conditioning a joint Gaussian). The classical side of data assimilation is in ml-for-science lesson 06 (Kalman filter, EnKF, 4D-Var). Time: about 100 minutes, plus the notebook.

---

Why this matters #

In an inverse problem you recover an unknown $x$ from measurements $y = \mathcal A(x) + n$, where the forward operator $\mathcal A$ maps the unknown to what the sensors see and $n$ is measurement noise. Examples: deblurring and inpainting, MRI and CT, seismic imaging, ground-motion maps from a sparse station network, and data assimilation for weather, which estimates the state of the atmosphere from observations. These problems are ill-posed: many $x$ fit $y$ equally well. The posterior $p(x \mid y) \propto p(y \mid x)\,p(x)$ is the honest answer, and the prior $p(x)$ decides which solutions are plausible.

A diffusion model is a strong learned prior. The promise is to train the prior once and then solve many inverse problems zero-shot: only the likelihood term changes at sampling time. The problem is that the likelihood term needed at noise level $t$ is intractable. Every method approximates it, and the errors make the posterior over-confident or biased. In science you report uncertainty, so you must know what each method computes, where it is exact, and how to test it.

1. Predict first #

2. The picture #

Notation. Noising $x_t = a_t x_0 + \sigma_t\varepsilon$ with $\varepsilon \sim \mathcal N(0, I)$; in VP form $a_t = \sqrt{\bar\alpha_t}$ and $\sigma_t = \sqrt{1 - \bar\alpha_t}$. The measurement noise has standard deviation $\sigma_y$. Do not confuse the two: $\sigma_t$ is the diffusion noise that the sampler adds and removes, $\sigma_y$ is the noise of the real sensors.

Bayes at every noise level #

The noising does not involve $y$. So the noised posterior is $p_t(x_t \mid y) \propto p_t(x_t)\,p_t(y \mid x_t)$, and its score splits into two terms:

This is the same split as classifier guidance in lesson 11, with the measurement model in place of a classifier. If both terms are exact at every level, any sampler of lesson 09 (reverse SDE, DDIM, PF-ODE) run with this score produces samples of $p(x_0 \mid y)$.

The noisy likelihood is the likelihood blurred by the denoiser's uncertainty #

• At low noise, $x_t$ almost fixes $x_0$, so $p_t(y \mid x_t) \approx p(y \mid x_0 = x_t/a_t)$: sharp.
• At high noise, $x_t$ says little about $x_0$. Then $p(x_0 \mid x_t)$ is close to the prior, and $p_t(y \mid x_t)$ is almost flat in $x_t$: the guidance should be weak.
• At high noise $p(x_0 \mid x_t)$ can also have several modes. Its mean $\hat x_0(x_t)$ can then lie between them, where neither mode is.

The simplest case. A scalar prior $x_0 \sim \mathcal N(0, S)$ and $y = x_0 + n$. Everything is Gaussian, and (Theorem 2 below)

DPS keeps the mean and drops $\mathrm{Var}(x_0 \mid x_t)$. That single omission explains most of its behavior.

Interactive figure: loading…

3. Definitions and main results #

Setup. The unknown $x_0 \in \mathbb R^d$ has the prior $p_0$, learned by a diffusion model. The observation is $y = \mathcal A(x_0) + n$ with $n \sim \mathcal N(0, \sigma_y^2 I_k)$. In the linear case $\mathcal A(x) = Ax$ with $A \in \mathbb R^{k \times d}$, often with $k < d$ (fewer measurements than unknowns). The denoiser is the Tweedie mean (lesson 08): $\hat x_0(x_t) = \mathbb E[x_0 \mid x_t] = (x_t + \sigma_t^2\nabla\log p_t(x_t))/a_t$. The results below hold for any $a_t, \sigma_t > 0$.

Theorem 1 (the posterior score splits) #

Proof.

1. Given $x_0$, the forward noise $\varepsilon$ and the measurement noise $n$ are independent. So $y$ and $x_t$ are independent given $x_0$: $p(y \mid x_0, x_t) = p(y \mid x_0)$.
2. Marginalize over $x_0$: $p(y \mid x_t) = \int p(y \mid x_0, x_t)\,p(x_0 \mid x_t)\,dx_0 = \int p(y \mid x_0)\,p(x_0 \mid x_t)\,dx_0$.
3. Bayes' rule: $p_t(x_t \mid y) = p_t(x_t)\,p_t(y \mid x_t)/p(y)$. Take the log and the gradient in $x_t$; $p(y)$ does not depend on $x_t$. $\blacksquare$

Step 1 also shows that, given $y$, the forward process is the same Gaussian kernel. So the posterior is reversed by the usual reverse process with the conditional score.

Why $p_t(y \mid x_t)$ is intractable. $p(x_0 \mid x_t)$ is the full denoising posterior. The network gives only its mean (and, through derivatives, its covariance). Sampling from $p(x_0 \mid x_t)$ would mean running the whole reverse process from $t$ to 0 inside every step, and the gradient of the log of a Monte Carlo estimate is biased. For a nonlinear $\mathcal A$, even a Gaussian $p(x_0 \mid x_t)$ gives no closed form. So zero-shot methods approximate. One case is exactly solvable.

Theorem 2 (the exact noisy likelihood of a linear-Gaussian problem) #

Let $x_0 \sim \mathcal N(m, S)$, $y = Ax_0 + n$ with $n \sim \mathcal N(0, \sigma_y^2 I)$, $x_t = a_t x_0 + \sigma_t\varepsilon$, and $\Sigma_t = a_t^2 S + \sigma_t^2 I$. Then

Proof.

1. $(x_0, x_t)$ is jointly Gaussian with means $m$ and $a_t m$, $\mathrm{Cov}(x_0, x_t) = a_t S$ and $\mathrm{Cov}(x_t) = \Sigma_t$. Gaussian conditioning (the Schur complement) gives $\mu_{0|t}$ and $\Sigma_{0|t}$. Note that $\Sigma_{0|t}$ does not depend on $x_t$.
2. By Theorem 1, given $x_t$ we have $y = Ax_0 + n$ with $x_0 \sim \mathcal N(\mu_{0|t}, \Sigma_{0|t})$ and $n$ independent. So $y \mid x_t$ is Gaussian with mean $A\mu_{0|t}$ and covariance $A\Sigma_{0|t}A^\top + \sigma_y^2 I = C_t$.
3. $\mu_{0|t}$ is affine in $x_t$ with Jacobian $J = a_t S\Sigma_t^{-1}$, and $C_t$ does not depend on $x_t$. The gradient of $-\tfrac12(y - A\mu_{0|t})^\top C_t^{-1}(y - A\mu_{0|t})$ is $J^\top A^\top C_t^{-1}(y - A\mu_{0|t})$. ($J$ is symmetric, because $S$ and $\Sigma_t = a_t^2S + \sigma_t^2I$ commute.) $\blacksquare$

Two checks.

• Adding the prior score $-\Sigma_t^{-1}(x_t - a_t m)$ gives the score of $\mathcal N(a_t\mu_{\text{post}},\ a_t^2\Sigma_{\text{post}} + \sigma_t^2 I)$, where $\mathcal N(\mu_{\text{post}}, \Sigma_{\text{post}})$ is the exact posterior of $x_0$ given $y$ (the Kalman update). So the noised posterior is the exact posterior pushed through the same noising. Section 5 checks this numerically.
• Second-order Tweedie. $\Sigma_{0|t} = S\Sigma_t^{-1}(\Sigma_t - a_t^2S) = \sigma_t^2 S\Sigma_t^{-1} = (\sigma_t^2/a_t)\,J$. This is the general identity $\mathrm{Cov}(x_0 \mid x_t) = (\sigma_t^2/a_t)\,\partial\hat x_0/\partial x_t$, which holds for any prior (lesson 08): the Jacobian of the denoiser gives the posterior covariance.

The family of approximations #

Replace $p(x_0 \mid x_t)$ by a Gaussian $\mathcal N(\hat x_0(x_t), \tilde\Sigma_t)$. For a linear $A$ this gives $p_t(y \mid x_t) \approx \mathcal N(y;\ A\hat x_0,\ \sigma_y^2 I + A\tilde\Sigma_tA^\top)$ and, treating $\tilde\Sigma_t$ as constant in $x_t$,

The methods differ only in $\tilde\Sigma_t$:

Method	$\tilde\Sigma_t$	Exact when
DPS (Chung et al. 2023)	$0$	only in the limit $\bar\alpha \to 1$
ΠGDM (Song et al. 2023)	$r_t^2 I$ with $r_t^2 = \sigma_t^2$ in VP form	data $\sim \mathcal N(0, I)$
Moment matching: TMPD (Boys et al. 2023), MMPS (Rozet et al. 2024)	$(\sigma_t^2/a_t)\,\partial\hat x_0/\partial x_t$ (second-order Tweedie)	$p(x_0 \mid x_t)$ Gaussian, for example a Gaussian prior
SDA (Rozet and Louppe 2023)	$(\sigma_t^2/a_t^2)\,\Gamma$ with a fixed $\Gamma$ (for example $10^{-2} I$)	a heuristic

(ΠGDM writes $r_t^2 = \tilde\sigma^2/(1 + \tilde\sigma^2)$ with the VE noise level $\tilde\sigma = \sigma_t/a_t$; in VP form this is $\sigma_t^2$.)

DPS (diffusion posterior sampling, Chung et al. 2023) uses $\nabla_{x_t}\log\mathcal N\big(y;\ \mathcal A(\hat x_0(x_t)),\ \sigma_y^2 I\big)$. It works for any differentiable $\mathcal A$, linear or not. It needs a vector–Jacobian product through the denoiser network at every step, so each step costs a forward and a backward pass. The exact log-likelihood has the factor $1/(2\sigma_y^2)$; in practice DPS replaces it by a step size $\zeta_i = \zeta'/\|y - \mathcal A(\hat x_0)\|$ with a tuned $\zeta'$ (the paper uses values from 0.1 to 1). Its biases:

1. Too strong at high noise. In the linear-Gaussian case DPS has the right direction but a scale $(\sigma_y^2 + \Sigma_{0|t})/\sigma_y^2$ too large (scalar case). As $\bar\alpha \to 0$ this tends to $(\sigma_y^2 + S)/\sigma_y^2$, which is huge for precise measurements. Early steps are pushed to fit $y$ too soon.
2. Over-confident samples. The samples are not calibrated posterior samples: their spread is too small and their mean can be biased. In the 1D problem of section 4 (exact posterior $\mathcal N(1.6, 0.8)$), DPS with the true $\sigma_y$, no step-size tuning and an accurate PF-ODE solver ends at mean $2.56$ and variance $0.073$: it overshoots $y = 2$ and has a tenth of the right variance (notebook exercise 2). The figure below shows the same effects in 2D.
3. The mean lies between modes. For a multimodal $p(x_0 \mid x_t)$, guidance through $\hat x_0$ can point the wrong way; section 5 shows a sign error.
4. Step-size tuning hides the scale error. The result can look good while it is not a posterior.

Interactive figure: loading…

Covariance-corrected methods. ΠGDM is cheap (an isotropic variance) and exact for white data of unit variance. For data with another variance its $r_t$ is wrong, most at high noise: there the true $\mathrm{Var}(x_0 \mid x_t)$ tends to the data variance, while $r_t^2$ tends to 1. Moment matching uses the second-order Tweedie covariance, so it is exact for Gaussian priors and recovers Theorem 2. It never forms the $d \times d$ Jacobian: it solves $(\sigma_y^2 I + A\tilde\Sigma_tA^\top)v = y - A\hat x_0$ by conjugate gradients, where each product with $\partial\hat x_0/\partial x_t$ costs one Jacobian–vector product through the network. Its remaining bias: $p(x_0 \mid x_t)$ is not Gaussian at high noise.

Zero-shot guidance, amortized training and SDEdit #

• Zero-shot guidance (all methods above). One unconditional prior serves every $\mathcal A$ and $\sigma_y$ at test time. It costs gradients through the network at every step, and the posterior is approximate.
• Amortized conditional model. Train a conditional diffusion model of $p(x_0 \mid y)$ with $y$ as a network input, on simulated or paired data (super-resolution models, CorrDiff downscaling). It is fast at test time (no gradients through the network) and can be well calibrated for the operators and noise levels it was trained on. It must be retrained when the operator, the noise level or the sensor layout changes, and it has an amortization gap: one network must fit the posterior for every $y$.
• SDEdit (Meng et al. 2022). Take a guess (for example $y$ itself, or $A^\top y$), noise it to a level $t_0$, and run the unconditional reverse process. This projects the guess toward the prior and keeps its coarse structure; $t_0$ trades faithfulness against realism. It uses no likelihood and no noise model, so it is not posterior sampling. It fails when $A$ is far from the identity, for example with sparse sensors: the guess then does not look like a noisy sample of the data.

Score-based data assimilation #

Data assimilation estimates the trajectory $x_{1:L}$ of a dynamical system from sparse, noisy observations $y = \mathcal A(x_{1:L}) + n$. The classical tools are the Kalman filters (Gaussian, sequential updates) and 4D-Var (the most probable trajectory of a time window, found with an adjoint model); see ml-for-science lesson 06 and the Kalman filter in the time-series track.

Score-based data assimilation (SDA, Rozet and Louppe 2023) has four parts:

1. A prior over trajectories. A diffusion model trained on simulated or reanalysis trajectories. It learns the joint law of whole windows, not a one-step forecast.
2. Local scores for long trajectories (their section 3.1). The dynamics are Markov, so the score of state $x_i$ depends mainly on a window around $i$. A score network is trained on short windows $x_{i-k:i+k}$ and slid along the trajectory to assemble the full score, for any length $L \gg k$.
3. Observation guidance (section 3.2): $p(y \mid x_t) \approx \mathcal N\big(y;\ \mathcal A(\hat x_0),\ \Sigma_y + (\sigma_t^2/a_t^2)\,\mathbf A\Gamma\mathbf A^\top\big)$, with $\mathbf A$ the Jacobian of $\mathcal A$ (as in the extended Kalman filter). The extra covariance term keeps the guidance from being too strong at high noise, without tuning a step size.
4. Predictor–corrector sampling (section 3.3): DDIM-type steps plus a few Langevin corrector steps per level (lesson 06), so that errors of the approximate score do not accumulate.

The link to the Kalman smoother. If the prior over trajectories is Gaussian and the observations are linear, the exact posterior over the whole trajectory is Gaussian, and its mean and variance are those of the Kalman smoother (Theorem 2 with $x_0$ = the trajectory). So a diffusion sampler with the exact noisy likelihood must reproduce the Kalman smoother. This is the known-posterior test of SDA.

Interactive figure: loading…

Compared with the EnKF, SDA needs no Gaussian analysis step and no localization or inflation tuning, and it smooths a whole window at once. Compared with 4D-Var, it needs no tangent-linear or adjoint model, and it returns samples instead of one trajectory. Both classical methods also handle a new observation operator without retraining. The costs of SDA: a trained prior that is only as good as its training trajectories, many network evaluations, and an approximate posterior that must be tested. Manshausen et al. (2024) applied SDA, with a diffusion prior trained on km-scale weather analyses, to observations of 40 weather stations in a central-US test region; they report 10% lower errors at left-out stations than a baseline, and also an ensemble that is not spread out enough.

Checking posterior calibration #

A posterior method can produce sharp, realistic samples and still be wrong. The tests, from the strongest to the weakest:

1. Known-posterior tests. Linear-Gaussian problems (as in the figures), Gaussian-process priors with a matched covariance, or small problems where MCMC gives the reference. Compare means, covariances or whole distributions.
2. Simulation-based calibration (SBC; Talts et al. 2018). Repeat many times: draw $x^*$ from the prior (or from held-out simulations), simulate $y = \mathcal A(x^*) + n$, draw $L$ posterior samples, and record the rank of $x^*$ among them (per variable or for a summary statistic). For an exact sampler the ranks are uniform on $\{0, \dots, L\}$: given $y$, the truth $x^*$ and the $L$ samples are exchangeable draws from $p(x \mid y)$. Read the rank histogram as in lesson 14: a U-shape means over-confident (like DPS), a dome under-confident, a slope biased.
3. Coverage of central credible intervals at several levels (50%, 90%). With $n$ test cases, the sampling error of a coverage $p$ is $\sqrt{p(1-p)/n}$.
4. Proper scores on held-out synthetic truths (CRPS per variable, energy and variogram scores for fields, spread–skill), and on real data, scores at held-out stations.

A trap. SBC averages over the prior. A sampler that ignores $y$ and returns prior samples also gives uniform ranks. So SBC and coverage test calibration but not sharpness. Always add a proper score, which rewards both.

Answers to section 1. (1) Too strong. In the linear-Gaussian case the exact likelihood score has the covariance $\sigma_y^2 + \Sigma_{0|t}$ and DPS has $\sigma_y^2$, so DPS is larger by the factor $(\sigma_y^2 + \Sigma_{0|t})/\sigma_y^2$, up to $(\sigma_y^2 + S)/\sigma_y^2$ as $\bar\alpha \to 0$. (2) The amortized model takes $y$ as an input with a fixed layout, so it must be retrained or fine-tuned. A zero-shot method only changes $\mathcal A$ (one more row of $A$).

4. Worked example #

Scalar prior $x_0 \sim \mathcal N(0, 4)$, observation $y = x_0 + n$ with $\sigma_y^2 = 1$, observed $y = 2$.

(a) Exact posterior. Precisions add: $1/4 + 1/1 = 1.25$, so the posterior variance is $1/1.25 = 0.8$. The mean is $0.8 \cdot (0/4 + 2/1) = 1.6$.

(b) One noise level, $\bar\alpha = 0.2$: $a_t = \sqrt{0.2} = 0.4472$ and $\sigma_t^2 = 0.8$.

1. $\Sigma_t = 0.2 \cdot 4 + 0.8 = 1.6$.
2. $J = a_t S/\Sigma_t = 0.4472 \cdot 4/1.6 = 1.118$.
3. $\Sigma_{0|t} = 4 - 0.2 \cdot 16/1.6 = 4 - 2 = 2.0$. Second-order Tweedie: $(\sigma_t^2/a_t)J = (0.8/0.4472) \cdot 1.118 = 2.0$ ✓.
4. $p_t(y \mid x_t) = \mathcal N(y;\ 1.118\,x_t,\ 1 + 2) = \mathcal N(y;\ 1.118\,x_t,\ 3)$.

(c) Likelihood scores at $x_t = 0.5$. $\mu_{0|t} = 1.118 \cdot 0.5 = 0.559$, so $y - \mu_{0|t} = 1.441$.

• Exact: $1.118 \cdot 1.441/3 = 0.537$.
• DPS ($\tilde\Sigma = 0$): $1.118 \cdot 1.441/1 = 1.611$, three times too large.
• ΠGDM ($r_t^2 = \sigma_t^2 = 0.8$, which assumes unit prior variance; here it is 4): $1.118 \cdot 1.441/1.8 = 0.895$, still 1.67 times too large.
• Moment matching ($\tilde\Sigma = 2.0$): $0.537$, exact.

(d) Check with the noised posterior. The prior score is $-x_t/\Sigma_t = -0.5/1.6 = -0.3125$. The sum is $-0.3125 + 0.537 = 0.2245$. The noised posterior is $\mathcal N(a_t \cdot 1.6,\ a_t^2 \cdot 0.8 + \sigma_t^2) = \mathcal N(0.7155,\ 0.96)$, and its score at $0.5$ is $(0.7155 - 0.5)/0.96 = 0.2245$ ✓.

(e) How the DPS error depends on the noise level. The ratio DPS/exact is $1 + \Sigma_{0|t}/\sigma_y^2$: $5$ for $\bar\alpha \to 0$, $3$ at $\bar\alpha = 0.2$, $1.235$ at $\bar\alpha = 0.8$, $1.010$ at $\bar\alpha = 0.99$.

5. In code #

NumPy and SciPy check Theorem 2 and the score split in 2D, compare the exact likelihood score with DPS and moment matching for a prior with two modes, and run a calibration check.

import numpy as np
from scipy.special import logsumexp
from scipy.stats import kstest, multivariate_normal, norm

# (1) Theorem 2 in 2D: the formula vs direct Gaussian conditioning of the joint (x_t, y)
S = np.array([[1.0, 0.6], [0.6, 2.0]]); m = np.array([0.5, -1.0])
A = np.array([[1.0, 0.0]]); sy2 = 0.25                  # observe the first coordinate only
ab = 0.3; a, s2 = np.sqrt(ab), 1 - ab                   # a_t and sigma_t^2
St = a**2 * S + s2 * np.eye(2)
J = a * S @ np.linalg.inv(St)
Sig0t = S - a**2 * S @ np.linalg.inv(St) @ S
C = sy2 + A @ Sig0t @ A.T
xt = np.array([0.2, 0.7])
mu0t = m + J @ (xt - a * m)
mean_y = A @ m + a * A @ S @ np.linalg.solve(St, xt - a * m)          # Cov(y, x_t) = a A S
var_y = A @ S @ A.T + sy2 - a**2 * A @ S @ np.linalg.solve(St, S @ A.T)
print(np.allclose(A @ mu0t, mean_y), np.allclose(C, var_y))            # True True
print(np.allclose(Sig0t, s2 / a * J))                                  # True (second-order Tweedie)

# The score split: prior score + likelihood score = score of the noised exact posterior
y = np.array([1.2])
lik = J.T @ A.T @ np.linalg.solve(C, y - A @ mu0t)
prior = -np.linalg.solve(St, xt - a * m)
K = S @ A.T @ np.linalg.inv(A @ S @ A.T + sy2)                         # Kalman gain
mp, Sp = m + K @ (y - A @ m), S - K @ A @ S                            # exact posterior of x0
post = -np.linalg.solve(a**2 * Sp + s2 * np.eye(2), xt - a * mp)
print(np.allclose(prior + lik, post))                                  # True

# (2) A prior with two modes (at -2 and +2): exact likelihood score vs DPS vs moment matching
w, mk, v0, sy2, y = np.array([0.5, 0.5]), np.array([-2.0, 2.0]), 0.25, 0.25, 1.0

def log_pt(x, ab):                     # log p_t(x_t): a mixture of Gaussians
    return logsumexp(np.log(w) + norm.logpdf(x, np.sqrt(ab) * mk, np.sqrt(ab * v0 + 1 - ab)))

def log_joint(x, ab):                  # log p_t(x_t, y): a mixture of 2D Gaussians
    a = np.sqrt(ab)
    cov = np.array([[ab * v0 + 1 - ab, a * v0], [a * v0, v0 + sy2]])
    return logsumexp([np.log(wk) + multivariate_normal.logpdf([x, y], [a * mu, mu], cov) for wk, mu in zip(w, mk)])

def scores(ab, x, h=1e-5):
    a, s2 = np.sqrt(ab), 1 - ab
    d = lambda f, z: (f(z + h) - f(z - h)) / (2 * h)
    exact = d(lambda z: log_joint(z, ab) - log_pt(z, ab), x)
    x0_hat = lambda z: (z + s2 * d(lambda u: log_pt(u, ab), z)) / a   # Tweedie mean
    J = (x0_hat(x + 1e-3) - x0_hat(x - 1e-3)) / 2e-3                # d x0_hat / d x_t
    var = s2 / a * J                                                  # second-order Tweedie
    return exact, J * (y - x0_hat(x)) / sy2, J * (y - x0_hat(x)) / (sy2 + var)

print("alpha_bar   exact    DPS   moment-matched")
for ab in [0.05, 0.3, 0.7, 0.95, 0.999]:
    print(f"{ab:9.3f}" + "".join(f"{v:8.3f}" for v in scores(ab, 0.3)))
#     0.050   0.284   2.717   0.157
#     0.300   0.528   2.078   0.162
#     0.700   0.330  -1.059  -0.259
#     0.950   1.344   1.644   1.363
#     0.999   2.758   2.769   2.758

Near $\bar\alpha = 1$ all three agree. At high noise DPS is far too strong. At $\bar\alpha = 0.7$ both Gaussian approximations have the wrong sign. There $p(x_0 \mid x_t)$ puts 89% of its mass near $+1.40$ and 11% near $-1.13$. Its mean, $1.12$, is above $y = 1$, so both approximations push $x_t$ down. The exact score pushes $x_t$ up: a larger $x_t$ moves mass away from the mode near $-1.13$, which explains $y = 1$ very badly. A single Gaussian cannot represent "which mode".

Third, a calibration check on the problem of section 4. The value $u = F_{\text{post}}(x^*)$ of the truth under the claimed posterior distribution function (the PIT, probability integral transform) is the continuous version of the rank:

rng = np.random.default_rng(0)
n = 20000
x_true = rng.normal(0.0, 2.0, n)                        # x ~ N(0, 4)
y = x_true + rng.normal(0.0, 1.0, n)                    # exact posterior: N(0.8 y, 0.8)
samplers = {"exact": (0.8 * y, np.sqrt(0.8)),
            "too narrow": (0.8 * y, 0.5 * np.sqrt(0.8)),
            "biased": (0.8 * y + 0.3, np.sqrt(0.8))}
for name, (mu, sd) in samplers.items():
    u = norm.cdf(x_true, mu, sd)
    cover = [np.mean(np.abs(u - 0.5) <= p / 2) for p in (0.5, 0.9)]
    hist = np.histogram(u, bins=5, range=(0, 1))[0]
    print(f"{name:10s} cover50={cover[0]:.3f} cover90={cover[1]:.3f} KS p={kstest(u, 'uniform').pvalue:.2f}", hist)
# exact      cover50=0.504 cover90=0.898 KS p=0.84 [3952 4075 4038 3888 4047]
# too narrow cover50=0.268 cover90=0.590 KS p=0.00 [6694 2310 2074 2228 6694]
# biased     cover50=0.475 cover90=0.882 KS p=0.00 [6098 4623 3726 3117 2436]

The narrow sampler covers only 27% and 59% instead of 50% and 90% (theory: 0.264 and 0.589), and its histogram is U-shaped. The biased one has nearly correct coverage but a sloped histogram: coverage alone would miss it.

6. Where it shows up #

• Imaging. Super-resolution, inpainting and deblurring benchmarks (DPS, ΠGDM), and MRI and CT reconstruction with score-based priors. In latent diffusion (lesson 13) the guidance also passes through the decoder (latent DPS).
• Weather and climate. SDA was demonstrated on the Lorenz 63 system and a Kolmogorov flow; Manshausen et al. (2024) assimilate weather-station data at km scale. Amortized conditional diffusion is used for downscaling (CorrDiff) and forecasting (GenCast conditions on previous states).
• Seismology. Diffusion priors regularize full-waveform inversion. Conditioning generated ground-motion fields on station recordings is the same inverse problem: $A$ selects the station locations, $\sigma_y$ models recording and site noise, and a calibrated posterior gives uncertainty maps between stations.
• Inverse problems with adjoints (ml-for-science lesson 04). A diffusion prior replaces a hand-made regularizer; the forward model and its gradient stay the same.
• Practice. Start with a linear-Gaussian version of your problem, where the exact posterior is known. Then run SBC on simulations, then check held-out real observations.

7. Common mistakes #

• Reading DPS samples as posterior samples. They are over-confident, and the step size was tuned.
• ΠGDM's $r_t$ for data that is not unit-variance. It assumes $\mathrm{Var}(x_0) = 1$; standardize the data or rescale $r_t$.
• Mixing VP and VE variables. ΠGDM's noise level is $\sigma_t/a_t$ in VP terms, and SDA's covariance is $(\sigma_t^2/a_t^2)\Gamma$. A wrong scale makes the guidance too strong or too weak by that factor.
• Judging samples by PSNR to the truth. The posterior mean maximizes PSNR. A correct posterior sample has twice the expected squared error of the posterior mean ($2\,\mathrm{tr}\,\Sigma_{\text{post}}$ against $\mathrm{tr}\,\Sigma_{\text{post}}$).
• Calling SDEdit posterior sampling. It has no likelihood and no noise model.
• SBC or coverage alone. A sampler that ignores $y$ passes both; add a proper score.
• Testing on the prior's own training simulations. That hides a wrong prior; use held-out data.

8. Check yourself #

1. Same scalar model as section 4, at $\bar\alpha = 0.8$. Give $\Sigma_t$, $\Sigma_{0|t}$ and the ratio DPS/exact.
   Answer

   $\Sigma_t = 0.8 \cdot 4 + 0.2 = 3.4$; $\Sigma_{0|t} = 4 - 0.8 \cdot 16/3.4 = 4 - 3.765 = 0.235$; ratio $1 + 0.235/1 = 1.235$. The DPS error shrinks as the noise goes down.
2. Why does $\nabla\log p(y)$ not appear in the posterior score, and why does the forward process not depend on $y$?
   Answer

   $p(y)$ is a constant in $x_t$, so its gradient is 0. Given $x_0$, the forward noise and the measurement noise are independent, so $p(x_t \mid x_0, y) = p(x_t \mid x_0)$: the same Gaussian kernel.
3. For the Gaussian prior, show that $\Sigma_{0|t} = (\sigma_t^2/a_t)\,J$ with $J = a_t S\Sigma_t^{-1}$.
   Answer

   $\Sigma_{0|t} = S - a_t^2S\Sigma_t^{-1}S = S\Sigma_t^{-1}(\Sigma_t - a_t^2S) = S\Sigma_t^{-1}\sigma_t^2 = (\sigma_t^2/a_t)(a_tS\Sigma_t^{-1})$, using that $S$ and $\Sigma_t$ commute.
4. 2D prior $\mathcal N(0, S)$ with $S = \begin{pmatrix}1&\rho\\\rho&1\end{pmatrix}$; you observe only $x_1$ ($A = (1, 0)$). Does the likelihood score move $x_2$?
   Answer

   Yes if $\rho \ne 0$. $A^\top C_t^{-1}(y - A\mu_{0|t})$ points along $e_1$, but $J^\top = a_tS\Sigma_t^{-1}$ has off-diagonal entries when $\rho \ne 0$, so it spreads the update to $x_2$. For $\rho = 0$, $J$ is diagonal and $x_2$ keeps its prior: an unobserved, independent variable learns nothing.
5. Why does SDEdit fail when only 5% of the pixels are observed?
   Answer

   The starting guess (observed pixels, zeros elsewhere) is not a noisy version of a real image. Noising it to $t_0$ and denoising either keeps the holes (small $t_0$) or forgets the observations (large $t_0$). Without a likelihood term nothing enforces $y \approx Ax$ with the right noise model.
6. In an SBC run (1000 synthetic cases, 99 posterior samples each), the rank histogram of the truth is U-shaped. What does it mean? What would a sampler that ignores $y$ give?
   Answer

   U-shape: the truth often falls outside the samples, so the posterior is over-confident (as with DPS). A sampler that returns prior samples gives a flat histogram, because averaged over the prior it is calibrated. A proper score such as the CRPS exposes it: the prior is far less sharp than the posterior.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the posterior score at noise level $t$ for an inverse problem.

   Show answer

2. Define the noisy likelihood $p_t(y \mid x_t)$.

   Show answer

3. Why is $p_t(y \mid x_t)$ intractable in general?

   Show answer

4. Give $p_t(y \mid x_t)$ for $x_0 \sim \mathcal N(m, S)$ and $y = Ax_0 + n$, $n \sim \mathcal N(0, \sigma_y^2 I)$.

   Show answer

5. What approximation does DPS make?

   Show answer

6. Scalar linear-Gaussian case: by which factor is the DPS likelihood score too large?

   Show answer

7. State the second-order Tweedie identity used by moment-matching methods.

   Show answer

8. What is SDEdit, and why is it not posterior sampling?

   Show answer

9. What does a U-shaped SBC rank histogram mean?

   Show answer

10. Why can SBC and coverage not detect a sampler that ignores $y$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How do you solve an inverse problem with an unconditional diffusion model?" — Split the posterior score (Theorem 1); approximate $\nabla\log p_t(y \mid x_t)$ through the Tweedie mean (DPS) or a Gaussian with a covariance; run the usual sampler.
• "What does DPS approximate, and when does it fail?" — It sets $\mathrm{Cov}(x_0 \mid x_t) = 0$: too strong at high noise, over-confident samples, wrong direction for a multimodal $p(x_0 \mid x_t)$.
• "Derive $p_t(y \mid x_t)$ for a linear-Gaussian problem." — Theorem 2: $\mathcal N(A\mu_{0|t}, \sigma_y^2 I + A\Sigma_{0|t}A^\top)$.
• "Guidance or a conditional model?" — Guidance: one prior, any operator, approximate, slow. Conditional: fast and calibrated for its training operators, but retrained for new ones.
• "How do you know your generative data assimilation is calibrated?" — Known-posterior tests, SBC ranks and coverage, proper scores on synthetic truths, held-out stations.

10. Practice #

The notebook for this lesson has five exercises, each with a check cell and a folded solution:

1. Compute the exact noisy likelihood of a linear-Gaussian problem and check it against Monte Carlo.
2. Run a diffusion sampler with the exact and with the DPS likelihood score on the 1D Gaussian problem and measure the posterior variance bias of DPS.
3. Sample the posterior of the 8-Gaussians prior with the exact conditional score and compare the mode weights with the closed-form posterior mixture.
4. Write a coverage and rank-histogram test, and apply it to an exact, a too-narrow and a biased sampler.
5. Check that the Kalman update gives exactly the Gaussian posterior (gain form against precision form).

11. Watch and read #

Hyungjin Chung - Adapting and Regularizing Diffusion Models for Inverse ProblemsVisual Machines Group, UCLA · 51 min

The author of DPS explains approximate posterior sampling (chapter "Approximate posterior sampling") and then solvers in latent space.

Diffusion Models for Solving Inverse Problems (Jiaming Song, NVIDIA)Yingzhen Li · 64 min

Covers DDRM for linear problems and ΠGDM; watch for how ΠGDM chooses the variance $r_t^2$.

Diffusion Models for Inverse ProblemsGenerative Memory Lab · 42 min

Hyungjin Chung presents DPS and related methods; compare their approximations with the table in section 3.

• Daras et al., A Survey on Diffusion Models for Inverse Problems (2024): section 2.4 (conditional models vs pre-trained priors), section 3.1 (DPS, ΠGDM, moment matching).
• Chung et al., Diffusion Posterior Sampling for General Noisy Inverse Problems (2023): section 3.1 and Algorithm 1.
• Song, Vahdat, Mardani, Kautz, Pseudoinverse-Guided Diffusion Models for Inverse Problems (ΠGDM, ICLR 2023).
• Boys et al., Tweedie Moment Projected Diffusions For Inverse Problems (2023) and Rozet et al., Learning Diffusion Priors from Observations by Expectation Maximization (2024): moment matching with the second-order Tweedie covariance.
• Rozet and Louppe, Score-based Data Assimilation (2023): local scores, the likelihood approximation and predictor–corrector sampling.
• Manshausen et al., Generative Data Assimilation of Sparse Weather Station Observations at Kilometer Scales (2024).
• Meng et al., SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations (2022).
• Talts et al., Validating Bayesian Inference Algorithms with Simulation-Based Calibration (2018).