Diffusion models guided to find rare events much faster

Steering Diffusion Models to Rare Events with Sequential Monte Carlo

Machine Learning

Summary

Some computer models help predict unusual events, but spotting very rare ones can take a lot of time because they happen so infrequently. The authors came up with a smarter way to focus on these rare cases by guiding many guesses closer to the rare event, making it easier to estimate how likely they are. Their method works with different kinds of rare events and was tested on simple examples and climate models, showing it can be hundreds to over a thousand times faster than usual methods.

What this means in practice

Authors

Aavash Subedi, Tim Reichelt, Christopher Williams, Philip Stier, Yee Whye Teh, Saifuddin Syed

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.