Tutorial notebook. Source
.ipynb on GitHub.
Sequential Monte Carlo¶
SMC moves a population of particles through a sequence of distributions that ends at the target. Every step reweights the particles towards the next distribution, resamples them and moves them. The population copes with several modes, and the weights give an estimate of the normalising constant (the evidence) at no extra cost.
probjax has two SMC kernels, both on the same contract as MCMC, so run,
adaptors and jax.vmap work unchanged:
smctempers from the prior to the posterior and moves the particles with any MCMC kernel. The schedule can be fixed or adaptive, and the move adapts to the population as it goes.diffusion_smcfollows the reverse SDE of a diffusion model, either a pretrained network (neural.diffusion_guide) or Monte Carlo score estimates of the target (score_guide).
The notebook runs on CPU in a few minutes. Particle counts are modest; raise them for harder problems.
import time
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from flax import nnx
from jax.scipy.stats import multivariate_normal, norm
jax.config.update("jax_platform_name", "cpu")
from probjax.inference import (
diffusion_smc, hmc, population_diagnostics, run, score_guide, smc,
summarize_replicates,
)
from probjax.inference.neural import diffusion_guide
from probjax.nn import VP, TimeMLP
from probjax.nn.generative.diffusion.config import VPNoiseSchedule
key = jax.random.key(42)
def next_key():
global key
key, result = jax.random.split(key)
return result
1. A posterior with a known evidence¶
A Gaussian prior and a correlated Gaussian likelihood give a Gaussian posterior and a closed-form evidence \(p(y) = \mathcal N(y; 0, I + \Sigma)\), so every estimate below can be checked.
smc(logprior, loglikelihood, move=...) targets
\(\pi_t(x) \propto p(x)\, p(y \mid x)^t\) and raises \(t\) from 0 to 1. Each step
reweights by the likelihood gained, resamples, and runs a few steps of the
move kernel, here HMC. The SMC kernel's params are the move's params,
shared by all particles.
observation = jnp.array([1.5, -0.75])
noise_cov = jnp.array([[0.4, 0.15], [0.15, 0.7]])
logprior = lambda x: norm.logpdf(x).sum()
loglikelihood = lambda x: multivariate_normal.logpdf(observation, x, noise_cov)
posterior_cov = jnp.linalg.inv(jnp.eye(2) + jnp.linalg.inv(noise_cov))
posterior_mean = posterior_cov @ jnp.linalg.solve(noise_cov, observation)
exact_log_evidence = multivariate_normal.logpdf(
observation, jnp.zeros(2), jnp.eye(2) + noise_cov
)
num_particles = 2048
move = hmc(num_integration_steps=4)
fixed = smc(logprior, loglikelihood, move=move, num_mcmc_steps=3)
state = fixed.init(next_key(), jax.random.normal(next_key(), (num_particles, 2)))
params = fixed.init_params(state, step_size=0.3)
schedule = jnp.linspace(0.0, 1.0, 21)[1:]
start = time.perf_counter()
result = run(next_key(), fixed, state, params, 20, args=(schedule,),
collect=("ess", "acceptance_rate"))
result.state.log_evidence.block_until_ready()
print(f"{time.perf_counter() - start:.1f}s including compilation")
print("log evidence:", float(result.state.log_evidence), " exact:", float(exact_log_evidence))
print("mean:", result.state.particles.mean(0), " exact:", posterior_mean)
0.4s including compilation
log evidence: -3.3175902366638184 exact: -3.316518545150757
mean: [ 1.1304522 -0.52671814] exact: [ 1.1293744 -0.5408271]
Fixed, adaptive and waste-free schedules¶
A fixed schedule spends the same effort on every step, however much the
likelihood changes there. With target_ess each step instead picks the
largest temperature that keeps that fraction of effective particles. Once the
temperature reaches 1, further steps change nothing, so the step count passed
to run is a cap.
waste_free=p keeps every state of the move chains, not only the last (Dau
and Chopin, 2022): it resamples \(N/p\) particles and runs \(p - 1\) moves from
each. The kept states are correlated, so \(N\) overstates the information, but
nothing the moves computed is thrown away.
kernels = {
"fixed, 20 steps": (fixed, dict(args=(schedule,)), 20),
"adaptive": (smc(logprior, loglikelihood, move=move, num_mcmc_steps=3,
target_ess=0.7), {}, 30),
"adaptive, waste-free": (smc(logprior, loglikelihood, move=move,
target_ess=0.7, waste_free=4), {}, 30),
}
fig, axes = plt.subplots(1, 2, figsize=(11, 3.5))
for name, (kernel, options, num_steps) in kernels.items():
result = run(next_key(), kernel, state, params, num_steps,
collect=("temperature", "ess"), **options)
temperatures = np.asarray(result.trace["temperature"])
steps = int(np.argmax(temperatures >= 1.0)) + 1
print(f"{name:22s} steps {steps:2d} log evidence {float(result.state.log_evidence):.3f}")
axes[0].plot(temperatures[:steps], "o-", label=name)
axes[1].plot(result.trace["ess"][:steps], "o-", label=name)
print(f"{'exact':22s} log evidence {float(exact_log_evidence):.3f}")
axes[0].set(xlabel="SMC step", ylabel="temperature", title="Schedules")
axes[1].set(xlabel="SMC step", ylabel="ESS / N after reweighting", title="Weight degeneracy")
axes[0].legend()
plt.tight_layout()
plt.show()
fixed, 20 steps steps 20 log evidence -3.320
adaptive steps 3 log evidence -3.297
adaptive, waste-free steps 3 log evidence -3.323
exact log evidence -3.317
The adaptive schedule takes small steps where the likelihood first bites and large ones at the end, and holds the ESS at the target throughout. All three evidence estimates match the exact value to a few hundredths.
The estimate is random, so the honest error bar comes from independent
repetitions. The kernel is a pure function, so jax.vmap runs them in one
compiled call; summarize_replicates gives the Monte Carlo standard error
across runs. Evidence is averaged on its own scale, not in logs.
adaptive = kernels["adaptive"][0]
def one_run(key):
init_key, run_key = jax.random.split(key)
state = adaptive.init(init_key, jax.random.normal(init_key, (512, 2)))
result = run(run_key, adaptive, state, params, 30)
return jnp.exp(result.state.log_evidence), result.state.particles[:, 0].mean()
evidence, means = jax.jit(jax.vmap(one_run))(jax.random.split(next_key(), 16))
for label, values, exact in (("evidence", evidence, jnp.exp(exact_log_evidence)),
("posterior mean x[0]", means, posterior_mean[0])):
stats = summarize_replicates(values)
print(f"{label:20s} {float(stats.mean):.4f} ± {float(stats.standard_error):.4f}"
f" exact {float(exact):.4f}")
evidence 0.0363 ± 0.0005 exact 0.0363
posterior mean x[0] 1.1235 ± 0.0058 exact 1.1294
2. Adapting the move to the population¶
A move that suits the prior can be hopeless for the posterior. Below, the likelihood pins one coordinate to within 0.02 and leaves the other loose. HMC with the prior's unit mass matrix and a fixed step size starts rejecting almost every proposal as the posterior narrows, and the particles stop moving.
The population after each step is a sample from the current target, so it
tells us the next target's shape. kernel.default_adaptation() uses it after
every step:
particle_adaptorsets HMC's inverse mass matrix to the particle variances (the full covariance if you start from a full matrix);acceptance_rate_adaptorscales the step size by \(\exp(\text{acceptance} - 0.65)\).
Both are ordinary adaptors, passed to run like any other.
sharp_noise = jnp.array([0.02, 1.0])
sharp_observation = jnp.array([0.8, -0.4])
sharp_likelihood = lambda x: norm.logpdf(sharp_observation, x, sharp_noise).sum()
sharp_exact = norm.logpdf(sharp_observation, 0.0, jnp.sqrt(1 + sharp_noise**2)).sum()
sharp = smc(logprior, sharp_likelihood, move=move, num_mcmc_steps=3, target_ess=0.7)
sharp_state = sharp.init(next_key(), jax.random.normal(next_key(), (1024, 2)))
sharp_params = sharp.init_params(sharp_state, step_size=0.3)
def sharp_run(key, adapt):
init_key, run_key = jax.random.split(key)
state = sharp.init(init_key, jax.random.normal(init_key, (1024, 2)))
adaptor = sharp.default_adaptation() if adapt else None
result = run(run_key, sharp, state, sharp_params, 40, adaptor=adaptor,
collect=("temperature", "acceptance_rate"))
return result.state.log_evidence, result.trace, result.params
fig, ax = plt.subplots(figsize=(7, 3.5))
for adapt in (False, True):
log_evidence, trace, final = jax.jit(jax.vmap(lambda k: sharp_run(k, adapt)))(
jax.random.split(next_key(), 8)
)
label = "adapted" if adapt else "fixed move"
steps = int(np.argmax(np.asarray(trace["temperature"][0]) >= 1.0)) + 1
ax.plot(trace["acceptance_rate"][0, :steps], "o-", label=label)
print(f"{label:11s} log evidence {float(log_evidence.mean()):7.3f} ± {float(log_evidence.std()):.3f}"
f" (exact {float(sharp_exact):.3f})")
print("adapted inverse mass matrix:", final.inverse_mass_matrix[0],
" posterior variances:", 1 / (1 + 1 / sharp_noise**2))
ax.set(xlabel="SMC step", ylabel="mean acceptance rate", title="HMC moves along the path")
ax.legend()
plt.tight_layout()
plt.show()
fixed move log evidence -2.473 ± 0.117 (exact -2.545)
adapted log evidence -2.527 ± 0.062 (exact -2.545)
adapted inverse mass matrix: [4.236471e-04 5.405653e-01] posterior variances: [3.9984006e-04 5.0000000e-01]
Without adaptation the acceptance rate collapses to zero, the particles stop moving, and resampling only copies them: the evidence estimates scatter twice as much and sit off the exact value. The adapted move keeps accepting, and its mass matrix ends at the posterior variances.
3. Data tempering: one observation at a time¶
Instead of raising one temperature, observations can enter one by one. Let the log-likelihood return one term per observation and give the kernel a vector temperature: the target becomes \(p(x) \exp(\sum_i t_i \log p(y_i \mid x))\). A schedule of 0/1 masks then adds the data in order, and every intermediate population is the posterior given the data so far.
num_observations = 12
data = jnp.array([1.0, -0.5]) + 0.7 * jax.random.normal(next_key(), (num_observations, 2))
per_observation = lambda x: norm.logpdf(data, x, 0.7).sum(axis=1)
streaming = smc(logprior, per_observation, move=move, num_mcmc_steps=3)
state = streaming.init(next_key(), jax.random.normal(next_key(), (1024, 2)),
temperature=jnp.zeros(num_observations))
masks = jnp.tril(jnp.ones((num_observations, num_observations)))
result = run(next_key(), streaming, state, params, num_observations, args=(masks,),
collect=("particles",))
seen = jnp.arange(1, num_observations + 1)
exact_var = 1 / (1 + seen / 0.7**2)
exact_mean = exact_var * jnp.cumsum(data[:, 0]) / 0.7**2
particles = result.trace["particles"][..., 0]
fig, ax = plt.subplots(figsize=(7, 3.5))
ax.errorbar(seen, particles.mean(1), 2 * particles.std(1), fmt="o", label="SMC mean ± 2 sd")
ax.plot(seen, exact_mean, "k-", label="exact mean")
ax.fill_between(seen, exact_mean - 2 * jnp.sqrt(exact_var), exact_mean + 2 * jnp.sqrt(exact_var),
color="k", alpha=0.1, label="exact ± 2 sd")
ax.set(xlabel="observations seen", ylabel="x[0]", title="Posterior as the data arrive")
ax.legend()
plt.tight_layout()
plt.show()
4. Multimodal and curved targets¶
Tempering from a broad reference makes SMC robust to several modes: the population spreads over all of them while the target is still flat, and the reweighting keeps the right proportions as they separate. Here the "likelihood" is the target divided by a broad Gaussian reference, so the path runs from that reference to the target.
def funnel(z):
return norm.logpdf(z[1], 0.0, 3.0) + norm.logpdf(z[0], 0.0, jnp.exp(z[1] / 2))
mixture_means = jnp.array([[6.0, 0.0], [0.0, 6.0], [-6.0, 0.0], [0.0, -6.0], [0.0, 0.0]])
def mixture(z):
components = jax.vmap(lambda m: norm.logpdf(z, m, 0.7).sum())(mixture_means)
return jax.scipy.special.logsumexp(components) - jnp.log(5.0)
turns = jnp.linspace(0.0, 6 * jnp.pi, 150)
spiral_means = 0.15 * jnp.stack([turns * jnp.cos(turns), turns * jnp.sin(turns)], -1)
def spiral(z):
components = jax.vmap(lambda m: norm.logpdf(z, m, 0.15).sum())(spiral_means)
return jax.scipy.special.logsumexp(components) - jnp.log(150.0)
reference = lambda x: norm.logpdf(x, 0.0, 5.0).sum()
gallery = [("Funnel", funnel, ((-12, 12), (-6, 6))),
("Five modes", mixture, ((-9, 9), (-9, 9))),
("Spiral", spiral, ((-3.5, 3.5), (-3.5, 3.5)))]
fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))
for ax, (name, target, limits) in zip(axes, gallery):
kernel = smc(reference, lambda x, target=target: target(x) - reference(x),
move=move, num_mcmc_steps=5, target_ess=0.8)
state = kernel.init(next_key(), 5.0 * jax.random.normal(next_key(), (2048, 2)))
result = run(next_key(), kernel, state, kernel.init_params(state, step_size=0.1),
100, adaptor=kernel.default_adaptation(), collect=("temperature",))
steps = int(np.argmax(np.asarray(result.trace["temperature"]) >= 1.0)) + 1
x = result.state.particles
ax.hist2d(np.asarray(x[:, 0]), np.asarray(x[:, 1]), bins=60, range=limits, cmap="Blues")
gx, gy = np.meshgrid(np.linspace(*limits[0], 120), np.linspace(*limits[1], 120))
log_density = np.asarray(jax.vmap(target)(jnp.stack([gx.ravel(), gy.ravel()], -1)))
levels = log_density.max() - np.array([8.0, 4.0, 1.0])
ax.contour(gx, gy, log_density.reshape(gx.shape), levels=levels, colors="black",
linewidths=0.6, linestyles="solid")
ax.set(title=f"{name}: {steps} steps, log Z = {float(result.state.log_evidence):.2f}")
plt.tight_layout()
plt.show()
print("All three targets are normalised, so log Z should be close to 0.")
All three targets are normalised, so log Z should be close to 0.
5. Reverse diffusion SMC: no training needed¶
diffusion_smc takes a different route to the target. Particles start as
Gaussian noise and follow the reverse SDE of a diffusion process down to the
data. A twist \(\psi_t(x)\) looks ahead to the target and weights the
particles; at the last time it is the target itself, so the sampler is
consistent whatever the twist, and a good twist only needs fewer particles.
score_guide builds the reverse SDE from the target alone (reverse diffusion
SMC, Wu et al. 2025). At every particle it estimates the score of the noised
target by importance sampling, and the same samples give the twist: an
estimate of the noised target's density. That costs num_inner target
evaluations per particle and step, but needs neither a tempering path nor
MCMC moves. Too few inner samples bias the score at small noise levels, which
shows up as a low evidence estimate; 256 are enough here.
noise = VPNoiseSchedule(t_min=1e-3, t_max=1.0)
guide = score_guide(mixture, noise, (2,), num_inner=256)
times = jnp.linspace(1.0, 1e-3, 101)
reverse = diffusion_smc(guide, times)
start = time.perf_counter()
state = reverse.init(next_key(), 1024)
result = run(next_key(), reverse, state, None, len(times) - 1, collect=("ess", "resampled"))
weights = np.asarray(jnp.exp(result.state.log_weights))
print(f"{time.perf_counter() - start:.1f}s; log Z = {float(result.state.log_evidence):.3f} (exact 0)")
nearest = jnp.argmin(jnp.linalg.norm(result.state.particles[:, None] - mixture_means, axis=-1), 1)
mass = np.bincount(np.asarray(nearest), weights=weights, minlength=5)
print("mass per mode:", np.round(mass, 3), " (exact 0.2 each)")
print("resampled at", int(result.trace["resampled"].sum()), "of", len(times) - 1, "steps")
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
x = np.asarray(result.state.particles)
axes[0].scatter(x[:, 0], x[:, 1], s=4 + 400 * weights, alpha=0.5)
axes[0].set(title="Reverse diffusion SMC on the five modes", xlim=(-9, 9), ylim=(-9, 9))
axes[1].plot(times[1:], result.trace["ess"], "o-", ms=3)
axes[1].invert_xaxis()
axes[1].set(xlabel="diffusion time t", ylabel="ESS / N", title="Weights along the reverse path")
plt.tight_layout()
plt.show()
2.0s; log Z = -0.157 (exact 0)
mass per mode: [0.204 0.178 0.233 0.173 0.212] (exact 0.2 each)
resampled at 4 of 100 steps
6. Steering a pretrained diffusion model¶
A diffusion model trained on data is a prior you can sample but not evaluate.
neural.diffusion_guide(model, V) turns it into a guide for the target
\(p_\phi(x)\, e^{V(x)}\): an observation model, a reward or a constraint, with no
retraining. The twist is \(V\) at the model's denoised estimate of the clean
point, faded in as the noise goes, and its gradient steers each reverse step
(twisted diffusion sampler, Wu et al. 2023). The guide's params are the
network weights, split out of the model the same way neural.density does
for flows.
We train a small VP model on a ring of eight Gaussians, so that the exact posterior is known, and condition it on a noisy observation \(y\) of \(x\).
angles = jnp.linspace(0, 2 * jnp.pi, 8, endpoint=False)
ring_means = 3.0 * jnp.stack([jnp.cos(angles), jnp.sin(angles)], -1)
ring_std = 0.3
def ring_samples(key, n):
component_key, noise_key = jax.random.split(key)
component = jax.random.randint(component_key, (n,), 0, 8)
return ring_means[component] + ring_std * jax.random.normal(noise_key, (n, 2))
start = time.perf_counter()
model = VP(TimeMLP(2, rngs=nnx.Rngs(0)), event_spec=2, num_steps=64)
losses = model.fit(next_key(), ring_samples(next_key(), 8000),
num_steps=15_000, batch_size=256,).losses
print(f"trained in {time.perf_counter() - start:.0f}s; loss {float(losses[:200].mean()):.2f} -> {float(losses[-200:].mean()):.2f}")
times = model.solver_cfg.solve_schedule(model.train_cfg.t_min, model.train_cfg.t_max, 64)
trained in 31s; loss 3.07 -> 2.03
The observation \(y = (1.5, 0)\) with noise 0.3 sits between two modes of the
ring and far from where the prior puts its mass. We compare the twisted
sampler with twist=False, which simulates the prior and weights by the
likelihood only at the end, over 16 independent runs of 1000 particles each.
y, obs_noise = jnp.array([1.5, 0.0]), 0.3
potential = lambda x: norm.logpdf(y, x, obs_noise).sum()
# Exact posterior of the true ring prior: a reweighted Gaussian mixture.
component_log_weights = norm.logpdf(y, ring_means, jnp.sqrt(ring_std**2 + obs_noise**2)).sum(-1)
exact_log_evidence = jax.scipy.special.logsumexp(component_log_weights) - jnp.log(8.0)
shrink = ring_std**2 / (ring_std**2 + obs_noise**2)
component_means = ring_means + shrink * (y - ring_means)
component_std = jnp.sqrt(shrink) * obs_noise
def exact_posterior(key, n):
component_key, noise_key = jax.random.split(key)
component = jax.random.categorical(component_key, component_log_weights, shape=(n,))
return component_means[component] + component_std * jax.random.normal(noise_key, (n, 2))
samplers = {}
for twist in (False, True):
guide, model_params = diffusion_guide(model, potential, twist=twist)
kernel = diffusion_smc(guide, times)
def one_run(key, kernel=kernel):
init_key, run_key = jax.random.split(key)
state = kernel.init(init_key, 1000, model_params)
result = run(run_key, kernel, state, model_params, len(times) - 1, collect=("ess",))
return result.state, result.trace["ess"]
states, ess = jax.jit(jax.vmap(one_run))(jax.random.split(next_key(), 16))
samplers["twisted" if twist else "no twist"] = (states, ess)
print(f"{'twisted' if twist else 'no twist':9s} log Z {float(states.log_evidence.mean()):6.2f}"
f" ± {float(states.log_evidence.std()):.2f} lowest ESS {float(ess.min(1).mean()):.3f}")
print(f"exact log Z {float(exact_log_evidence):6.2f} (for the true ring; the model is close)")
no twist log Z -8.49 ± 0.50 lowest ESS 0.005
twisted log Z -8.59 ± 0.14 lowest ESS 0.115
exact log Z -8.45 (for the true ring; the model is close)
fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))
truth = np.asarray(exact_posterior(next_key(), 2000))
axes[0].scatter(truth[:, 0], truth[:, 1], s=3, alpha=0.4, color="k")
axes[0].scatter(*np.asarray(ring_means).T, marker="x", color="C3", label="prior modes")
axes[0].set(title="Exact posterior", xlim=(-4, 4), ylim=(-4, 4))
axes[0].legend()
states, ess = samplers["twisted"]
x = np.asarray(states.particles[0])
w = np.asarray(jnp.exp(states.log_weights[0]))
axes[1].scatter(x[:, 0], x[:, 1], s=3 + 300 * w, alpha=0.4)
axes[1].set(title="Twisted diffusion SMC (one run)", xlim=(-4, 4), ylim=(-4, 4))
for name, (_, ess) in samplers.items():
axes[2].plot(times[1:], ess.mean(0), label=name)
axes[2].invert_xaxis()
axes[2].set(xlabel="diffusion time t", ylabel="mean ESS / N", title="Weight degeneracy")
axes[2].legend()
plt.tight_layout()
plt.show()
Without a twist the particles follow the prior and only meet the likelihood at the end. The final reweighting then rests on the few particles that happen to land near \(y\). The twisted sampler steers the particles towards the observation while there is still noise to spend, keeps more of them effective, and its evidence estimate varies less between runs.
The twist is a heuristic. For an observation sitting on a mode of the prior, plain prior simulation already does well, and the twist can do slightly worse. Either way the answer is consistent; only the particle budget changes.
7. Diagnostics beyond the weights¶
The weight ESS says nothing about duplicates: after resampling, many particles
can be copies of one parent. population_diagnostics counts distinct
particles and distinct ancestors, which exposes that collapse.
state = fixed.init(next_key(), jax.random.normal(next_key(), (num_particles, 2)))
result = run(next_key(), fixed, state, params, 20, args=(schedule,), collect=("ancestors",))
d = population_diagnostics(result.state.particles, jnp.full(num_particles, -jnp.log(num_particles)),
result.trace["ancestors"][-1])
print(f"weight ESS {float(d.ess):.0f}, distinct particles {int(d.unique_particles)}, "
f"parents in the last step {int(d.unique_ancestors)} of {num_particles}")
weight ESS 2048, distinct particles 2048, parents in the last step 2022 of 2048
The MCMC moves make every particle distinct again, even though the last resampling drew from fewer distinct parents.
When to use what¶
| Situation | Kernel |
|---|---|
| Posterior with a prior and a likelihood; you want the evidence | smc, target_ess=... |
| Badly scaled or changing geometry along the path | add kernel.default_adaptation() |
| Data arriving over time, fixed parameters | smc with a per-observation likelihood and mask schedule |
| Expensive moves; reuse every MCMC state | waste_free=p |
| Large populations that do not fit in memory at once | batch_size=... |
| Unnormalised target, no tempering path, cheap evaluations | diffusion_smc with score_guide |
| Condition a trained diffusion model on an observation or reward | diffusion_smc with neural.diffusion_guide |
For filtering, SMC² and trajectories over physical time see the temporal SMC notebook.