Mirror Bridges Between Probability Measures

TL;DR

The paper proposes mirror Schrödinger bridges, a self-mapping variant of the Schrödinger bridge problem, to generate in-distribution variations proximal to a given input. By exploiting time symmetry, it develops a drift-averaging-based AMP that reduces computation (single drift model, halved iterations) while maintaining proximity control through a parameter $\sigma$. The authors provide theoretical insights, including a first-order equivalence between the reverse KL projection and the averaged drift, and demonstrate empirical gains on Gaussian, 2D, and image datasets with controllable sample proximity and improved efficiency. This approach offers a principled, scalable tool for conditional resampling and data augmentation within the data distribution.

Abstract

Resampling from a target measure whose density is unknown is a fundamental problem in mathematical statistics and machine learning. A setting that dominates the machine learning literature consists of learning a map from an easy-to-sample prior, such as the Gaussian distribution, to a target measure. Under this model, samples from the prior are pushed forward to generate a new sample on the target measure, which is often difficult to sample from directly. Of particular interest is the problem of generating a new sample that is proximate to or otherwise conditioned on a given input sample. In this paper, we propose a new model called mirror bridges to solve this problem of conditional resampling. Our key observation is that solving the Schrödinger bridge problem between a distribution and itself provides a natural way to produce new samples from conditional distributions, giving in-distribution variations of an input data point. We demonstrate how to efficiently estimate the solution to this largely overlooked version of the Schrödinger bridge problem, and we prove that under mild conditions, the difference between our estimate and the true Schrödinger bridge can be controlled explicitly. We show that our proposed method leads to significant algorithmic simplifications over existing alternatives, in addition to providing control over in-distribution variation. Empirically, we demonstrate how these benefits can be leveraged to produce proximal samples in a number of application domains.

Figures (13)

• Figure 1: We visualize the resampling of a 2D distribution obtained with our method for different values of noise $\sigma$. Higher noise value results in greater in-distribution sample variance.
• Figure 2: For each method, we plot the mean (left) and variance (middle) obtained for the terminal samples, i.e. samples obtained at time $t=T$, as well as the covariance (right) of the joint distribution, versus the number of outer iterations, averaged over 5 trials.
• Figure 3: Samples (color based on initial position) obtained using our method with various $\sigma$ values.
• Figure 4: Samples produced by our mirror bridges for MNIST, using $\sigma=1,2$.
• Figure 5: The control over in-distribution variance effect of $\sigma$ for a variety of initial samples (first row) from the empirical distribution of images in the CelebA dataset.

Theorems & Definitions (8)

• Lemma 4.1
• Proposition 4.2
• Proposition 4.3
• Proposition 4.4
• Proposition B.1
• proof
• Proposition D.1
• proof