Joint stochastic localization and applications

TL;DR

This work advances the algorithmic understanding of stochastic localization (SL) by unifying existing SL schemes under Eldan's $\alpha$-scheme and introducing a joint SL framework to construct couplings between distributions. It shows how joint SL induces new transport-like distances $\mathsf{d}_{\mathrm{SL}_\alpha}$ and, in particular, that $\mathsf{d}_{\mathrm{SL}_0}$ is topologically equivalent to the $2$-Wasserstein distance on compact supports, while connections to Gaussian KL and score-matching provide links to established divergences and diffusion-model objectives. The paper then demonstrates practical applications: extending optimal normal couplings to log-concave distributions, bounding $W_2$ via joint SL couplings, and proposing a distribution-estimation method based on minimizing $\mathsf{d}_{\mathrm{SL}_\alpha}$. Numerical simulations corroborate localization-rate behavior and the efficacy of the induced couplings, suggesting SL-based distances as computationally appealing tools for transport, divergence approximation, and distribution learning. The results open avenues for scalable, Monte Carlo estimable transport-like distances and for integrating SL into diffusion-model training and distribution-fitting tasks.

Abstract

Stochastic localization is a pathwise analysis technique originating from convex geometry. This paper explores certain algorithmic aspects of stochastic localization as a computational tool. First, we unify various existing stochastic localization schemes and discuss their localization rates and regularization. We then introduce a joint stochastic localization framework for constructing couplings between probability distributions. As an initial application, we extend the optimal couplings between normal distributions under the 2-Wasserstein distance to log-concave distributions and derive a normal approximation result. As a further application, we introduce a family of distributional distances based on the couplings induced by joint stochastic localization. Under a specific choice of the localization process, the induced distance is topologically equivalent to the 2-Wasserstein distance for probability measures supported on a common compact set. Moreover, weighted versions of this distance are related to several statistical divergences commonly used in practice. The proposed distances also motivate new methods for distribution estimation that are of independent interest.

Figures (6)

• Figure 1: A trajectory of the observation process ${\theta}_t = tX + W_t$ with $X = (1, 1)^\top$ (left). As $t$ grows, the signal-to-noise ratio increases and thus ${\theta}_t$ becomes increasingly informative of the unobserved signal $X$. This can be measured by computing the $\ell_2$ distance (error) between the normalized signal $\frac{{\theta}_t}{t}$ and $X$ (right).
• Figure 2: Joint SL of point cloud on the Yin and Yang Taichi logo using the (regularized) extrapolation scheme in \ref{['cd']}. The initial distributions are uniform on the blue and red points, and time increases in the diagram as one goes from left to right. The size of the points is proportional to their weights.
• Figure 3: Localization rates of Eldan's $\alpha$-scheme for two different distributions in ${\mathbb R}^2$: a uniform distribution over $[-1 ,1]^2$ (left) and a Gaussian mixture with three components (right).
• Figure 4: Simulated data from $\mu$ and $\nu$ in Case 1 (top) and Case 2 (bottom). A line connecting two points represents a joint sample obtained in the MC simulation of the corresponding joint SL schemes.
• Figure 5: Estimated bounds on the $W_2$-distance using different joint SL schemes in Case 1 (left) and Case 2 (right). The $95\%$ CI is also given based on the estimated variances.

Theorems & Definitions (34)

• Remark 2.1
• Theorem 3.1: Localization rate of Eldan's $\alpha$-scheme
• Remark 3.2
• proof : Proof of Theorem \ref{['thm:01']}
• Example 1: Finite-time localization of Eldan's $\frac{1}{2}$-scheme
• Definition 1: Log-concave distributions on ${\mathbb R}^d$
• Theorem 3.3
• proof
• Theorem 4.1
• proof