Total Variation Guarantees for Sampling with Stochastic Localization

Abstract: Motivated by the success of score-based generative models, a number of diffusion-based algorithms have recently been proposed for the problem of sampling from a probability measure whose unnormalized density can be accessed. Among them, Grenioux et al. introduced SLIPS, a sampling algorithm based on Stochastic Localization. While SLIPS exhibits strong empirical performance, no rigorous convergence analysis has previously been provided. In this work, we close this gap by establishing the first guarantee for SLIPS in total variation distance. Under minimal assumptions on the target, our bound implies that the number of steps required to achieve an $\varepsilon$-guarantee scales linearly with the dimension, up to logarithmic factors. The analysis leverages techniques from the theory of score-based generative models and further provides theoretical insights into the empirically observed optimal choice of discretization points.

• The paper establishes non-asymptotic total variation guarantees for the SLIPS sampler with error scaling as O(d/ε² log²(d/ε)) under weak moment conditions.
• It demonstrates the optimality of a log-SNR-adapted discretization policy that minimizes discretization error compared to traditional diffusion-based approaches.
• The study highlights practical implications for high-dimensional sampling using Bayesian inference by leveraging MALA-based posterior estimation and martingale denoising techniques.

Total Variation Guarantees for Sampling with Stochastic Localization

Problem Context and Motivation

The paper addresses the fundamental problem of sampling from a high-dimensional probability measure $\pi$ specified by an unnormalized density, i.e., $\pi(x) \propto \exp(-f(x))$ for some computable function $f : \mathbb{R}^d \to \mathbb{R}$. This setting arises in Bayesian statistics, statistical physics, computational chemistry, and generative modeling. Conventional methods—including MCMC, annealing, and variational inference (VI)—face pronounced degradation in high-dimensional, multimodal scenarios, either due to slow mixing, mode collapse, or optimization landscape difficulties.

Recent advances in generative modeling, particularly diffusion-based algorithms and score-based generative models (SGMs), have demonstrated state-of-the-art empirical performance. The Stochastic Localization via Iterative Posterior Sampling (SLIPS) algorithm, proposed by Grenioux et al., exemplifies this trend by leveraging stochastic localization in combination with diffusion-based sampling. However, prior to this work, SLIPS lacked rigorous, non-asymptotic guarantees on convergence in total variation (TV) distance under weak regularity conditions—an essential step to justify its application in theory and practice.

SLIPS Algorithm and Stochastic Localization

Stochastic localization constructs a process $(Y_t)$ that interpolates between a tractable measure (Gaussian noise) and the target $\pi$, such that as $t \to \infty$, the process becomes "localized" on a sample from $\pi$ and the corresponding posterior martingale property holds. For SLIPS, the process is $Y_t = t X + \sigma B_t$, where $X \sim \pi$ and $B_t$ is Brownian motion. The key property is that $\pi(x) \propto \exp(-f(x))$0 almost surely as $\pi(x) \propto \exp(-f(x))$1.

SLIPS decomposes its operation into the following:

• The SDE characterization (with standard denoising schedule $\pi(x) \propto \exp(-f(x))$2):

$\pi(x) \propto \exp(-f(x))$3

where $\pi(x) \propto \exp(-f(x))$4 is the posterior mean of $\pi(x) \propto \exp(-f(x))$5 given $\pi(x) \propto \exp(-f(x))$6.

• Time discretization is performed using the Euler-Maruyama scheme over a non-uniform grid determined by log-SNR-adapted discretization points, reflecting local variations in the signal-to-noise ratio to minimize accumulation of discretization error.
• At each step, the posterior expectation is estimated using MCMC (specifically, MALA).

This recursion parallels diffusion SDEs in SGMs but replaces neural network score function estimation with posterior sampling.

Main Theoretical Contributions

1. Non-Asymptotic Total Variation Guarantees

The central theorem establishes that under weak moment assumptions and accurate (in $\pi(x) \propto \exp(-f(x))$7) MALA-based posterior expectation estimates, the SLIPS sampler outputs a random variable whose law is within $\pi(x) \propto \exp(-f(x))$8 TV distance of the target $\pi(x) \propto \exp(-f(x))$9 using $f : \mathbb{R}^d \to \mathbb{R}$0 steps (up to logarithmic factors):

$f : \mathbb{R}^d \to \mathbb{R}$1

for an error $f : \mathbb{R}^d \to \mathbb{R}$2 and dimension $f : \mathbb{R}^d \to \mathbb{R}$3, provided certain initialization and step-wise estimation errors are controlled. This is in stark contrast to classical methods, whose guarantees scale polynomially or exponentially worse in $f : \mathbb{R}^d \to \mathbb{R}$4 for multimodal targets.

Notably, the result is quantitative and explicitly separates sources of error into (i) information error (finite $f : \mathbb{R}^d \to \mathbb{R}$5 approximation of ideal localization), (ii) discretization error (via a novel pathwise SDE analysis reminiscent of recent SGM bounds), and (iii) posterior expectation estimation error using MALA.

2. Optimality and Discretization Policy

The bound on TV distance depends critically on a complexity term $f : \mathbb{R}^d \to \mathbb{R}$6 reflecting the choice of discretization grid. The work proves that the log-SNR-adapted grid, setting $f : \mathbb{R}^d \to \mathbb{R}$7, is, among a relevant class, optimal in minimizing $f : \mathbb{R}^d \to \mathbb{R}$8. This result provides theoretical foundation for a practice commonly used in state-of-the-art diffusion models and generative samplers.

3. Comparison with Related Diffusion-based Samplers

The guarantees are contrasted with those for Reverse Diffusion Monte Carlo (RDMC) [see "Reverse Diffusion Monte Carlo" (2404.00000)] and SGM-based protocols. A key contradictory claim is that, unlike SGM analyses requiring Lipschitz-continuity of the score ($f : \mathbb{R}^d \to \mathbb{R}$9-smoothness), the SLIPS guarantee does not require Lipschitz regularity of the score, nor strong convexity, only moment and $(Y_t)$0-type conditions. In cases where the SGM Lipschitz constant grows with $(Y_t)$1, SLIPS can outperform in dimensional scaling.

4. Martingale Structure and Exactness

A technical point elucidated is that, for linear denoising schedules (i.e., $(Y_t)$2), the optimal denoiser is actually a martingale with respect to the observation process; this is leveraged to prove sharp pathwise and discretization error bounds.

Implications and Limitations

From a practical standpoint, the result gives the first rigorous justification for deploying SLIPS and similar SDE-based samplers in high-dimensional, multi-modal settings where both classical MCMC and standard VI approaches falter.

Since the analysis conditions on accurate posterior expectation estimates (via MALA), it is conditional rather than fully non-asymptotic in the total computational cost. The remaining open problem is to obtain non-exponential-in-dimension mixing time guarantees for the inner MCMC scheme so as to close the overall computational complexity guarantee for SLIPS.

From a theoretical perspective, the connection to stochastic localization makes explicit the deep link between generative modeling, diffusion-based sampling, and martingale representations in high dimensional inference. The separation into information and discretization error, and especially the pathwise SDE analysis via Girsanov's theorem, reflects the transfer of advanced SGM analysis techniques into the sampling domain.

The flexibility of non-uniform discretization schedules (validated here for log-SNR) is expected to influence both the design of future samplers and the theoretical analysis of diffusive sampling protocols beyond Gaussian/OU cases.

Speculation on Future Developments

The coupling of rigorous TV guarantees with minimal smoothness assumptions hints that stochastic localization-based samplers may form the basis for the next generation of provable samplers in high-dimensional inference and generative modeling. Key directions include:

• Developing dimension-free or polynomial-in-dimension mixing guarantees for the inner MCMC (e.g., MALA or HMC) used for posterior expectation estimation. This could enable unconditional, fully non-asymptotic total complexity results.
• Extending and unifying SDE-based analysis to cover a broader class of non-linear denoising schedules (e.g., those used in most practical score-based models).
• Exploring connections to new classes of samplers based on optimal transport, normalizing flows, and reinforcement-learning-inspired control of diffusions.
• Leveraging the stochastic localization framework for fine-grained control of tradeoffs between computational tractability and sampling bias in wild multi-modal targets.

Conclusion

This work rigorously demonstrates that the SLIPS algorithm, when equipped with a log-SNR-adapted discretization and accurate posterior estimators, can achieve total variation error guarantees scaling linearly in dimension under minimal regularity conditions. The analysis synthesizes techniques from SGM theory, stochastic localization, and diffusion-based analysis, offering both theoretical insight and practical guidance for high-dimensional sampling. The results clarify necessary conditions for efficient sampling, the potential for further improvement, and point to lasting connections between generative modeling and theoretical sampling methodology (2603.29555).