[Paper Review] Sampling, Diffusions, and Stochastic Localization

Andrea Montanari|arXiv (Cornell University)|May 18, 2023

Markov Chains and Monte Carlo Methods4 citations

TL;DR

This paper establishes a theoretical connection between denoising diffusion models and stochastic localization, showing that standard denoising diffusions are a special case of stochastic localization processes. It generalizes stochastic localization to broader settings, demonstrates how this viewpoint simplifies model learning, and proves convergence to the target distribution under linear estimation approximations in Gaussian models.

ABSTRACT

Diffusions are a successful technique to sample from high-dimensional distributions. The target distribution can be either explicitly given or learnt from a collection of samples. They implement a diffusion process whose endpoint is a sample from the target distribution. The drift of the diffusion process is typically represented as a neural network. Stochastic localization is a successful technique to prove mixing of Markov Chains and other functional inequalities in high dimension. An algorithmic version of stochastic localization was recently proposed in order to sample from certain statistical mechanics models. This expository article has three objectives: $(i)$~Generalize the algorithmic construction to other stochastic localization processes. This construction is both simple and broadly applicable; $(ii)$~Clarify the connection between diffusions and stochastic localization. This allows to derive several known sampling schemes in a unified fashion; $(iii)$~Describe the insights that follow from this unified viewpoint.

Motivation & Objective

• To generalize the algorithmic stochastic localization framework from [EAMS22] to broader classes of stochastic processes.
• To clarify the theoretical relationship between diffusion-based sampling and stochastic localization, showing that standard denoising diffusions are a specific instance of the latter.
• To demonstrate how the stochastic localization viewpoint simplifies the learning of drift functions in diffusion models.
• To provide new insights into the design and analysis of sampling algorithms via the lens of stochastic localization.
• To establish convergence guarantees for a linear-estimator-based instantiation of the method in high-dimensional Gaussian models.

Proposed method

• Proposes a generalization of stochastic localization processes beyond the original framework, allowing for time-inhomogeneous and non-reversible dynamics.
• Derives the time-reversed SDE for a diffusion process to construct a sampling path from a known initial distribution to the target measure.
• Introduces a continuous-time stochastic process where the drift is the conditional expectation of a linear transformation of the target variable given the current state.
• Uses a linear, shift-invariant estimator structure to approximate the optimal drift, leveraging circulant matrix structure for computational efficiency.
• Applies a finite-window convolution approximation to the optimal estimator, reducing complexity while preserving asymptotic consistency.
• Employs Kullback-Leibler divergence as a loss function to guide the learning of drift components in the stochastic localization framework.

Experimental results

Research questions

• RQ1 Can the stochastic localization framework be generalized beyond the original construction to include a wider class of diffusion processes?
• RQ2 How are standard denoising diffusion models related to stochastic localization processes?
• RQ3 Can the stochastic localization viewpoint simplify the training of score-based generative models by reducing the complexity of the learned drift function?
• RQ4 What is the convergence behavior of a linear-estimator-based stochastic localization process in high-dimensional Gaussian models?
• RQ5 Does approximating the optimal conditional expectation with finite-window convolutions preserve asymptotic convergence to the target distribution?

Key findings

• The paper proves that standard denoising diffusions are a special case of stochastic localization processes, establishing a formal theoretical link between two major sampling paradigms.
• A linear, shift-invariant estimator with finite-window convolution approximation achieves asymptotic convergence to the target distribution in high-dimensional Gaussian models.
• The limiting Wasserstein and total variation distances between the target and generated distributions both converge to zero as time t approaches infinity.
• The proposed estimator for the drift component is Bayes-optimal up to O(1/n) terms in the Gaussian setting, justifying its use as a principled approximation.
• The method simplifies the learning task by reducing the need for complex neural network architectures, as the optimal drift can be approximated via simple linear functions with circulant structure.
• The framework allows for the design of more interpretable and analyzable sampling algorithms by unifying diffusion and localization perspectives.