Localized Schrödinger Bridge Sampler
This addresses the bottleneck of high sample complexity in sampling methods for researchers in machine learning and statistics, though it is incremental as it builds on existing Schrödinger bridge and Langevin sampler approaches.
The paper tackles the problem of sampling from unknown distributions using training samples by proposing a localization strategy that reduces the exponential sample complexity dependence on dimension, replacing a high-dimensional Schrödinger bridge with low-dimensional ones, and demonstrates performance on high-dimensional Gaussian and stochastic problems.
We consider the problem of sampling from an unknown distribution for which only a sufficiently large number of training samples are available. In this paper, we build on previous work combining Schrödinger bridges and plug & play Langevin samplers. A key bottleneck of these approaches is the exponential dependence of the required training samples on the dimension, $d$, of the ambient state space. We propose a localization strategy which exploits conditional independence of conditional expectation values. Localization thus replaces a single high-dimensional Schrödinger bridge problem by $d$ low-dimensional Schrödinger bridge problems over the available training samples. In this context, a connection to multi-head self attention transformer architectures is established. As for the original Schrödinger bridge sampling approach, the localized sampler is stable and geometric ergodic. The sampler also naturally extends to conditional sampling and to Bayesian inference. We demonstrate the performance of our proposed scheme through experiments on a high-dimensional Gaussian problem, on a temporal stochastic process, and on a stochastic subgrid-scale parametrization conditional sampling problem. We also extend the idea of localization to plug & play Langevin samplers using kernel-based denoising in combination with Tweedie's formula.