Technical Summary: LightSBB-M

Problem Statement

Generative modeling via stochastic processes has largely relied on the classical Schrödinger Bridge (SB) and diffusion models. The SB formulation constructs a controlled stochastic dynamics transporting a source distribution $\mu_0$ to a target $\mu_T$ by minimizing the relative entropy with respect to a Brownian prior. However, the classical SB framework imposes a strict requirement: the target distribution must have finite relative entropy with respect to the Wiener measure ($KL(P \| W_\epsilon) < \infty$). This constraint renders the classical SB infeasible for heavy-tailed or singular target distributions, as the entropy term diverges in such cases.

Conversely, Bass martingale transport offers flexibility by optimizing over volatility while enforcing a zero-drift constraint, but it lacks explicit drift control and departs from standard diffusion mechanisms. The Schrödinger-Bass Bridge (SBB) problem, introduced by Henry-Labordere et al. (2026), unifies these paradigms by jointly optimizing over both drift and volatility. It introduces a tunable parameter $\beta > 0$ that interpolates between the classical SB limit ($\beta \to \infty$) and the Bass martingale transport regime ($\beta \to 0$). Despite its theoretical advantages—particularly its ability to handle heavy-tailed targets without the finite-entropy restriction—the SBB problem has remained computationally intractable in high dimensions due to the complexity of stochastic volatility and implicit transport maps.

Methodology: LightSBB-M

The authors propose LightSBB-M, a simulation-efficient algorithm designed to compute the optimal SBB transport plan in a few iterations. The method leverages a dual representation of the SBB objective to derive analytic expressions for the optimal drift and volatility, avoiding costly Stochastic Differential Equation (SDE) discretization.

Core Theoretical Framework

The SBB problem minimizes a quadratic cost $J(P)$ over the drift $\alpha_t$ and volatility $\sigma_t$:
$$J(P) = \mathbb{E}_P \left[ \int_0^T \left( \|\alpha_t\|^2 + \beta \|\sigma_t - \sqrt{\epsilon}I_d\|^2 \right) dt \right]$$
where $\beta$ controls the trade-off between drift and volatility.

The authors establish a dual representation of this problem. By maximizing a Lagrangian functional over potentials $(v, \psi)$, they derive explicit feedback policies for the optimal controls:

1. Optimal Drift: $\alpha^*_t = \nabla_x v^*(t, x)$
2. Optimal Volatility: $\sigma^*_t = \sqrt{\epsilon} \left( I_d - \frac{D^2_x v^*(t, x)}{\beta} \right)^{-1}$

Crucially, the SBB system can be viewed as a "stretched" Schrödinger Bridge. By defining a transformed process $Y_t = Y_t(X_t)$, the authors show that $Y_t$ follows a standard Brownian motion with constant volatility $\sqrt{\epsilon}$ in a transformed space, while the original process $X_t$ is recovered via a learned transport map. This transformation converts a state-dependent diffusion in $X$-space into a constant-volatility process in $Y$-space.

Algorithmic Implementation

The LightSBB-M algorithm (Algorithm 1) operates via an alternating optimization procedure:

1. Transport Map Learning: A neural network $Z_{\tilde{\theta}}$ is trained to approximate the inverse transport map $Y_t = Z_{\tilde{\theta}}(t, X_t)$, effectively learning the mapping from the original space to the constant-volatility $Y$-space.
2. Drift Learning: Using a Gaussian-mixture parametrization for the potential (inspired by LightSB-M), the algorithm learns the score-drift $s_\theta$ in the $Y$-space by minimizing a denoising score matching loss (bridge matching).
3. Alternating Updates: The algorithm alternates between updating the drift parameters $\theta$ (to match the SB dynamics in $Y$-space) and the transport map parameters $\tilde{\theta}$ (to ensure the endpoints in $X$-space match $\mu_0$ and $\mu_T$).

This approach allows for simulation-free sampling at inference time. Once trained, one can sample $Y_T$ from the learned coupling and recover $X_T$ via the inverse map, bypassing the need to solve the complex SDE for $X_t$ directly.

Key Contributions

1. Scalable SBB Solver: The proposal of LightSBB-M, a practical algorithm that computes the optimal SBB transport plan efficiently, enabling generative modeling with jointly learned drift and volatility.
2. Analytic Control Expressions: The derivation of explicit, tractable expressions for optimal SBB controls, which facilitate simulation-free sampling and avoid the computational burden of SDE discretization.
3. Empirical Superiority: Extensive numerical experiments demonstrating that LightSBB-M achieves state-of-the-art performance on both synthetic and high-dimensional real-world datasets, outperforming existing SB and diffusion baselines.

Experimental Results

Synthetic Datasets

On low-dimensional synthetic datasets (8gaussians, moons), LightSBB-M was evaluated against state-of-the-art SB solvers (e.g., LightSB-M, DSBM) and flow-based models.

• Performance: LightSBB-M achieved the lowest 2-Wasserstein distance ($W_2$) across all benchmark tasks.
• Improvement: The method demonstrated an average 19% improvement in transport accuracy compared to the best SB and diffusion baselines, while also reducing result variance.
• Heavy-Tailed Handling: In experiments transporting a Gaussian to a Student-t distribution (a case where classical SB is ill-posed due to infinite entropy), LightSBB-M accurately recovered the target distribution, whereas classical SB produced significant artifacts and spurious modes.

Real-World Application: Image-to-Image Translation

The framework was applied to unpaired image-to-image translation (adult to child faces) on the FFHQ dataset (1024x1024 resolution) within the latent space of an Adversarial Latent AutoEncoder (ALAE).

• Age Translation: LightSBB-M significantly improved the quality of age translation. At low noise ($\epsilon=0.1$), the percentage of generated faces estimated as under 18 years old doubled compared to the LightSB-M baseline (41.2% vs. 20.7%). At higher noise ($\epsilon=1$), this metric nearly quadrupled (76.0% vs. 19.3%).
• Identity Preservation: While LightSBB-M showed a modest decrease in identity similarity (cosine similarity) compared to LightSB-M, this is attributed to the generation of genuinely more child-like faces rather than a failure of the model.
• Fidelity: The Fréchet Inception Distance (FID) remained comparable to the baseline (20.9 vs. 21.6), confirming that the improved diversity and age translation did not compromise image quality.
• Efficiency: Unlike conventional diffusion models requiring backward-forward sampling, LightSBB-M generates samples in a single forward pass, avoiding error accumulation from reverse diffusion steps.

Significance and Claims

The paper positions LightSBB-M as a significant advancement at the intersection of optimal transport theory and modern generative modeling. Its primary significance lies in making stochastic-volatility transport computationally tractable at scale.

The authors claim that by jointly controlling drift and volatility, the framework:

• Expands Applicability: It can model target distributions with heavy tails or singularities that are inaccessible to classical SB methods due to finite-entropy constraints.
• Offers Flexibility: The parameter $\beta$ allows practitioners to interpolate between pure drift control (SB) and pure volatility control (Bass), tuning the trade-off between sample diversity and fidelity.
• Improves Performance: It provides a scalable, high-fidelity solver that outperforms existing baselines in both transport accuracy and generative quality.

The authors note that while the algorithm converges empirically in few iterations (typically 5), a formal convergence proof for the outer alternating loop remains an open theoretical question. They suggest future work could focus on tightening iteration-complexity bounds and extending the framework to time-series data with heteroskedasticity.