Stochastic Localization via Iterative Posterior Sampling

TL;DR

This work tackles sampling from unnormalized densities by leveraging Stochastic Localization (SL), which uses an observation process Y_t that reveals X∼π through a time-dependent drift. It introduces SLIPS, a learning-free sampler that discretizes the SL dynamics and estimates the denoiser at each step via Monte Carlo posterior sampling, anchored by an SNR-driven, adaptive time discretization. A key theoretical contribution is the duality of log-concavity, which identifies a sweet spot t_0 balancing tractable initialization with reliable posterior sampling even for non-log-concave targets; this enables practical initialization and robust MC estimation. Empirically, SLIPS demonstrates strong performance on high-dimensional Gaussian mixtures, Bayesian logistic regression, and a φ^4 lattice field model, often outperforming state-of-the-art density-based samplers while requiring minimal tuning and no parametric score models.

Abstract

Building upon score-based learning, new interest in stochastic localization techniques has recently emerged. In these models, one seeks to noise a sample from the data distribution through a stochastic process, called observation process, and progressively learns a denoiser associated to this dynamics. Apart from specific applications, the use of stochastic localization for the problem of sampling from an unnormalized target density has not been explored extensively. This work contributes to fill this gap. We consider a general stochastic localization framework and introduce an explicit class of observation processes, associated with flexible denoising schedules. We provide a complete methodology, $\textit{Stochastic Localization via Iterative Posterior Sampling}$ (SLIPS), to obtain approximate samples of this dynamics, and as a by-product, samples from the target distribution. Our scheme is based on a Markov chain Monte Carlo estimation of the denoiser and comes with detailed practical guidelines. We illustrate the benefits and applicability of SLIPS on several benchmarks of multi-modal distributions, including Gaussian mixtures in increasing dimensions, Bayesian logistic regression and a high-dimensional field system from statistical-mechanics.

Figures (18)

• Figure 1: Display of $\log\operatorname{SNR}$ for the localization schemes Geom-$\infty$ (left) and Geom (right).
• Figure 2: SNR-adapted vs uniform time discretization for the schemes Standard and Geom$(1,1)$. The uniform discretization leads to larger log-SNR differences between timesteps where the SNR increases rapidly.
• Figure 3: Duality of log-concavity: distribution of $Y^{\alpha}_t / \alpha(t)$ (up) and $q^{\alpha}_t(\cdot | y^{\alpha}_t)$ (down) for $t\in(0, T_{\text{gen}})$, where $y^{\alpha}_t$ is a realisation of the observation process (red line), for the standard scheme (left) and the Geom$(1,1)$ scheme (right). The target distribution $\pi$ is a mixture of two 1D-Gaussian distributions $\mathsf{N}(-2/3, (0.05)^2)$ and $\mathsf{N}(4/3, (0.05)^2)$ with respective weights $2/3$ and $1/3$, which density is given by the blue line. The heat map represents the likelihood of the distributions and the green line on the right edge stands for the distributions taken at the time given by the dotted green line.
• Figure 4: Metrics when sampling a bimodal Gaussian mixture with $d$ growing. Top: Relative weight estimation error. Bottom: Sliced Wasserstein distance.
• Figure 5: Estimation of the mode weight ratio of $\phi^4$ with increasing $h$ - Only $\texttt{SLIPS}$ produced correct samples.

Theorems & Definitions (34)

• Theorem 1
• Theorem 2
• Theorem 3
• Lemma 1: Tweedie's formula and extension
• proof
• Lemma 2
• proof
• Proposition 1
• proof
• Proposition 2