Steering Diffusion Models to Rare Events with Sequential Monte Carlo

AI summary

Overview

Research area: Statistical machine learning / generative modelling — specifically sequential Monte Carlo (SMC) inference on score-based diffusion models, applied to rare-event probability estimation in scientific surrogate models (stat.ML).

Technical level: Advanced. The paper relies on diffusion model theory (reverse-time SDEs, score functions, Tweedie's formula), sequential Monte Carlo (importance weights, resampling, rejuvenation kernels, normalisation-constant estimators), and stochastic integration schemes.

Scope: The paper introduces DireSMC (Diffusion Importance Sampling of Rare Events), an SMC scheme that steers a population of weighted particles through the denoising trajectory of a pretrained diffusion model toward a rare event, returning both representative samples of that event and an unbiased estimate of its probability. It is validated on a Gaussian mixture with analytical ground truth and on a score-based climate emulator for compound hot-and-dry extremes over the region of the 2021 Pacific Northwest heatwave.

What This Paper Is About

Diffusion models are increasingly used as cheap surrogates for expensive simulators in weather prediction, molecular dynamics and materials design, but estimating the probability of a rare event under such a model by plain Monte Carlo is intractable, because stabilising the estimate requires a number of samples that grows as 1/p₀[E]. The paper's goal is to produce, within a single denoising run, both representative samples of the rare event and a calibrated, unbiased probability estimate, with a variance governed by the quality of the steering rather than by how rare the event is.

Key Contributions

1. DireSMC algorithm. The authors construct an unbiased sequential Monte Carlo estimator of p₀[E] that runs inside the denoising process, using a sequence of reward-tilted marginals π_t to guide particles toward the rare event while accumulating importance weights. Both samples and probability estimates come out of the same run.

2. Reward design for rare events. They show how to design the reward for threshold-exceedance events of the form E = {x : φ(x) ≥ τ}, building a smooth surrogate whose look-ahead guidance has a closed form under a Tweedie posterior approximation. This design is presented as easily extensible to a wide range of user-defined rare events because it rests on an analytical relaxation of the event set.

3. Theoretical guarantees. Theorem 1 establishes that the estimator is unbiased for any reward (however poorly it steers) under the stated domination and invariance conditions, that the conditional sample approximation converges almost surely as N → ∞, and that with an exact score and p_data[∂E] = 0 the expectation tends to p_data[E] as K → ∞. Lemma 1 gives a closed-form log incremental weight in which the numerical integrator enters only through a scalar ζ_k.

4. Empirical validation on two test beds. A Gaussian mixture with analytical ground truth, and a score-based climate emulator for compound hot-and-dry extremes over the region of the 2021 Pacific Northwest heatwave.

Main Findings

• Toy problem is accurate across rarities. On a Gaussian mixture where ground truth is available, the method obtains probability estimates that do not degrade across 10⁻³ to 10⁻⁵ orders of rarity.

• Climate emulator results. On the score-based climate emulator, the method obtains accurate rare-event probabilities on a range of rarities from 10⁻³ to 10⁻⁵, for compound hot-and-dry extremes over the region of the 2021 Pacific Northwest heatwave.

• Speed-ups over Monte Carlo. Net speed-ups of 9× to 1413× over Monte Carlo are reported across that range of rarities (10⁻³ to 10⁻⁵).

• Variance is governed by steering quality, not rarity. The paper states that the variance of the estimator is controlled by how good the steering is, rather than by how rare the event is — the opposite of the Monte Carlo regime, where the relative variance is (1 − p₀[E])/(N p₀[E]) and N ∝ 1/p₀[E] samples are needed for fixed relative error.

• Unbiasedness holds regardless of reward quality. Because Theorem 1 gives unbiasedness for any reward, a poorly chosen reward degrades variance (potentially to an impractical number of samples, per Chatterjee and Diaconis, 2018) but does not bias the estimate.

• Weights are cheap to track. Tracking the importance weights requires no additional network calls, because both ∇_x r_k and the noise term ϵ_k are already available when a particle is steered.

• Integrator dependence is a single scalar. For one-step explicit integrators with Gaussian transitions and means affine in the score, guidance shifts the transition mean but leaves the covariance invariant; the integrator enters the weights only through ζ_k = c_k/√v_k. For Euler–Maruyama with a VP-SDE schedule, c_k = v_k = σ²_{t_k} Δt_k, giving ζ_k = σ_{t_k}√(Δt_k), which recovers the weights of He et al. (2026). The behaviour does not extend to higher-order schemes that evaluate the score at a noisy intermediate point, such as the stochastic Heun integrator of Karras et al. (2022).

• Not reported in the available content. Specific particle counts N, the ESS threshold values η used, per-rarity speed-up breakdowns, wall-clock times, and the full derivation of the reward path r_{K−1}, …, r₁ (the reward selection section is truncated) are not reported in the content summarised here.

Methodology in Plain English

Diffusion models generate samples by starting from Gaussian noise at t = 1 and denoising down to data at t = 0, following the learned score ∇_x log p_t. To focus on a rare event E, the authors "tilt" the intermediate distributions by a reward: at each step they define γ_k(x) = p_k(x) e^{r_k(x)} and its normalised form π_k. The reward is zero at the noise end (r_K = 0, so γ_K = p_ref) and equal to the rare-event reward at the data end.

Rather than pushing a single sample, they run a population of N weighted particles (sequential Monte Carlo). Each particle is propagated using a guided reverse step — the ordinary reverse drift plus the gradient of the reward, ∇x r_t, added to the score, so the guidance is exactly a tilt of the score. Each particle's weight is then updated by an incremental factor. Because the naive weight involves the intractable ratio p{k−1}/p_k, they choose the auxiliary forward kernel to be the time reversal of the base chain, which cancels that ratio and leaves a computable expression in which the only remaining ingredient is the ratio of the guided and base transition kernels; for Gaussian one-step integrators this has the closed form given in Lemma 1.

To prevent weight degeneracy, particles are resampled when the effective sample size falls below a threshold ηN, using systematic resampling, with weights reset to the average of the ancestors' weights. Because resampling duplicates particles and reduces diversity, each resampling is followed by a rejuvenation step using a π_k-invariant MCMC kernel; the authors use the Metropolis-adjusted diffusion model (MADM) kernel of Lam et al. (2026), which needs only the score and r_k. At the end, summing the final weights against the reward and the event indicator gives the probability estimate, and the weighted particle cloud gives samples from p₀(·|E).

The reward itself is designed for threshold events E = {x : φ(x) ≥ τ} through a smooth surrogate whose gradient provides "look-ahead" guidance in closed form under a Tweedie posterior approximation, allowing the method to extend to user-defined rare events via an analytical relaxation of the event set.

Why This Matters

Impact on research. The method inverts the usual scaling story: instead of requiring a number of samples proportional to 1/p₀[E], the estimator's variance depends on steering quality, while unbiasedness is preserved for any reward. That combination — unbiasedness plus steering-controlled variance — makes it possible to attach error bars to rare-event probabilities from generative surrogates rather than only producing plausible-looking samples. It also connects diffusion steering, SMC and classical rare-event simulation (importance sampling, splitting-style sequences of events), and builds on and compares against existing approaches such as Xie et al. (2026) and Manshausen et al. (2026), whose compute costs the authors describe as especially problematic in rare-event settings.

Real-world applications (as highlighted or exemplified by the paper):

• Climate and environmental risk: compound hot-and-dry extremes, illustrated on the region of the 2021 Pacific Northwest heatwave, with the same framing extended to other environmental-science rare events.
• Molecular simulation and materials design, where diffusion models act as surrogates for expensive simulators and metastable-state occupancies are the target quantity.
• Numerical weather prediction, where rare extremes drive planning and resource allocation for failure modes.
• Finance, engineering and autonomous driving, domains the paper lists as relying on rare-event probabilities — for example, for risk management and safety-critical system design (Glasserman and Li, 2005; Au and Wang, 2014; O'Kelly et al., 2018).

Industry relevance. Practitioners who use or deploy generative surrogates need calibrated probabilities to inform planning and resource allocation (Kaplan and Garrick, 1981; Rubino et al., 2009). Reported net speed-ups of 9× to 1413× over Monte Carlo, obtained within the denoising process rather than as a separate post-hoc procedure, matter directly for compute budgets in operational forecasting, risk analytics and scientific simulation pipelines.

Future Directions

• Extending beyond threshold-exceedance events. The paper's reward construction is presented for events E = {x : φ(x) ≥ τ}, and posits easy extension to a wide range of user-defined rare events via an analytical relaxation of the event set; demonstrating this breadth empirically is a natural next step.

• Higher-order integrators. Lemma 1 covers one-step explicit integrators whose transitions are Gaussian with means affine in the score, but explicitly does not cover higher-order schemes that evaluate the score at a noisy intermediate point (e.g. the stochastic Heun integrator of Karras et al., 2022). Deriving usable weights for such samplers is left open.

• Reward path selection. The paper defers the choice of the intermediate rewards r_{K−1}, …, r₁ to its reward-selection section; since the variance of the estimator is governed by steering quality, understanding how to choose or adapt that path to minimise variance is a central open question.

• Scaling the empirical evidence. The reported validation covers a Gaussian mixture with analytical ground truth and one score-based climate emulator across rarities from 10⁻³ to 10⁻⁵. Whether the unbiasedness-plus-low-variance behaviour persists in higher-dimensional settings such as molecular simulation and materials design, and at greater levels of rarity, is not established in the content available here.

Target Audience

This paper is aimed at researchers and practitioners with a background in probabilistic machine learning and Monte Carlo methods: statisticians and ML researchers working on diffusion models, sequential Monte Carlo or rare-event simulation; climate and weather scientists using diffusion-based emulators for extreme-event risk; and computational scientists in molecular dynamics and materials design who need calibrated probabilities of rare configurations rather than just samples. Readers seeking a beginner-level tutorial on diffusion models would need supplementary background, as the paper assumes fluency with reverse-time SDEs, score tilting and SMC weight mechanics.

Authors’ abstract

Diffusion models are increasingly used as surrogates for expensive simulators in weather prediction, molecular dynamics, and materials design. In these models, computing the probability $p_0[E]$ of an event $E$ is difficult, especially when the event of interest is rare. A stable estimate using Monte Carlo becomes computationally intractable, requiring a growing sample size $\propto\!1/p_0[E]$ to compensate for an increasing rarity. In this paper, we present Diffusion Importance Sampling of Rare Events or DireSMC, a sequential Monte Carlo scheme that guides a population of weighted samples towards the rare event, giving access not only to samples but also to a calibrated estimate of its probability. We set up our guidance using an analytical relaxation of the event set, allowing the method to easily extend to a wide range of user-defined rare events. We validate our method on a toy problem with analytical solutions and on a score-based climate emulator, where we obtain accurate rare-event probabilities on a range of rarities from $10^{-3}$ to $10^{-5}$, achieving net speed-ups of $9\times$ to $1413\times$ over Monte Carlo.