LightSBB-M: Bridging Schrödinger and Bass for Generative Diffusion Modeling

AI summary

Overview

• Research area: Generative modeling with stochastic processes, specifically Schrödinger Bridge (SB) methods, optimal transport, and stochastic control, applied to diffusion-based generative models.
• Technical level: Advanced. The paper builds on entropic optimal transport, Schrödinger systems, Girsanov's theorem, Fenchel–Legendre duality, and convex (inf/sup) convolution operators, and it assumes familiarity with diffusion Schrödinger Bridge solvers such as Sinkhorn, Iterative Markovian Fitting, and LightSB-M.
• Scope (one sentence): The paper introduces LightSBB-M, an algorithm that solves the Schrödinger–Bass Bridge (a bridge formulation in which drift and volatility are jointly optimized) in a few iterations and applies it to synthetic transport benchmarks and unpaired image-to-image translation.

What This Paper Is About

Classical Schrödinger Bridge generative models transport a source distribution to a target distribution using a fixed volatility, and they require the path measure to have finite entropy relative to a Brownian prior, which rules out targets with heavy tails or singular behavior. Bass martingale transport provides an alternative that optimizes over volatility instead, but it imposes zero drift and does not fit the diffusion generative framework. The Schrödinger–Bass Bridge (SBB) unifies both by jointly optimizing drift and volatility through a parameter β, but it had received little computational attention, particularly in high dimensions. This paper's goal is to make the SBB problem practically solvable and scalable for generative modeling.

Key Contributions

1. A practical, scalable algorithm for the SBB transport plan. The authors propose LightSBB-M, which computes the optimal Schrödinger–Bass Bridge problem in high-dimensional generative modeling settings, where extending existing SB solvers is difficult because of stochastic volatility and implicit transport maps.
2. Explicit, tractable expressions for the optimal SBB controls. Using a dual representation of the SBB objective, the method yields analytic expressions for the optimal drift and volatility, enabling simulation-free sampling and avoiding costly SDE discretization.
3. A tunable interpolation parameter β. The formulation incorporates a parameter β > 0 that interpolates between pure drift behavior (the classical Schrödinger Bridge, β → ∞) and pure volatility behavior (Bass martingale transport, β → 0).
4. Numerical validation on synthetic and real-world tasks. The authors report extensive experiments on synthetic and high-dimensional datasets showing improved transport accuracy and generative performance compared with state-of-the-art Schrödinger Bridge and diffusion-based baselines.

Main Findings

• Lowest 2-Wasserstein distance on synthetic data. LightSBB-M achieves the lowest 2-Wasserstein distance on synthetic datasets against state-of-the-art SB and diffusion baselines, with an average improvement of 19%.
• Generative capability on unpaired image translation. The framework is illustrated on an adult → child faces unpaired image-to-image translation task in FFHQ. (Quantitative image metrics such as FID are not reported in the available content.)
• Fast convergence in few iterations. LightSBB-M computes the optimal SBB transport plan in only a few iterations; empirically, the alternating procedure converged in a small number of outer iterations (five in most experiments) and consistently produced stable solutions.
• Volatility is recovered implicitly, not learned directly. LightSBB-M does not parameterize the SBB volatility σ*_t directly. The full matrix-valued volatility in X-space is recovered through the learned transport map, while the transformed process Y has constant volatility √ε I_d. This is described as a multivariate, time-dependent Lamperti-type reparameterization, which avoids learning or evaluating a d × d volatility field.
• Unconstrained MLP parameterization suffices. In additional experiments, structure-preserving neural architectures enforcing the monotonicity structure of the admissible map were tested; they did not yield a meaningful improvement in transport accuracy while increasing training time, so the simpler MLP parameterization was retained.
• Motivation from heavy-tailed targets. When μ₀ = δ₀ and μ_T = 𝒯(2), the classical SB entropy term diverges because x² p_{𝒯(2)}(x) ~ 1/|x| for large |x|, making classical SB infeasible, whereas the SBB objective remains well posed.
• Stated limitations. The β ↓ 0 limit is understood only as a formal interpolation in this paper; a genuine zero-drift limit requires (μ₀, μ_T) to satisfy convex order, and the authors do not prove that SBB optimizers converge to the specific Bass optimizer among all martingale couplings. The condition βT > 1 is required for dual attainment. The absence of a formal convergence proof concerns the outer alternating loop only; the inner SB subproblem is solved via LightSB-M, which carries theoretical guarantees.

Methodology in Plain English

The authors start from the SBB objective, which penalizes both the size of the drift (with weight 1) and the deviation of the volatility from the constant √ε I_d (with weight β). By working with a dual representation of this objective, they obtain closed-form expressions for the optimal drift and the optimal volatility in terms of derivatives of a dual potential function v.

The construction has a useful structural consequence: there is a change of variables (a map 𝒴) under which the process becomes an ordinary Schrödinger Bridge with constant volatility √ε in the transformed "Y-space." The bridge in Y-space is then solved with LightSB-M, which characterizes the SB as a single optimal projection and parametrizes the potential φ as a Gaussian mixture, giving a closed-form analytic drift that requires no neural network. In practice, the drift s_θ is computed from equation (14), where the Gaussian-mixture parameters are {α_j, r_j, Σ_j}, with A^t_j = t/(ε(T−t)) I_d + (1/ε)Σ_j^{-1} and h_j(t,y) = y/(ε(T−t)) + (1/ε)Σ_j^{-1} r_j.

The remaining unknown is the transport map 𝒴. Since 𝒴_t is the inverse of 𝒳_t(y) = y + (1/β)s*_t(y), the authors learn this inverse with a neural network 𝒵_θ̃, which avoids solving a fixed-point equation at each step.

Training alternates between two regression steps, initialized with 𝒴⁰ = 𝒵⁰_θ̃ = I_d:

1. Endpoint sampling: draw pairs (y₀, y_T) by pushing source and target samples through the current map 𝒵^k_θ̃ at times 0 and T.
2. Bridge sampling: draw an intermediate point y_t from the Brownian bridge with mean interpolation (T−t)/T · y₀ + t/T · y_T and variance σ_t² = t(T−t)/T.
3. Regression on θ: minimize a bridge-matching (denoising score matching style) loss, which is the KL projection onto the set of Schrödinger Bridges.
4. Regression on θ̃: update the inverse-map network using a loss comparing 𝒵^k_θ̃ evaluated at 𝒳₀(x₀) and 𝒳_T(x_T) against x₀ and x_T.

For inference, a new sample X₀ ~ μ₀ is mapped to Y₀ = 𝒵^K_θ̃(X₀), Y_T is sampled from the learned coupling π_{v_θ}(Y_T | Y₀), and the output is recovered as X_T = Y_T + (1/β)s^K_θ(T, Y_T). Since the drift is not well defined at t = T, the authors approximate it at T̃ = T − δ for a small δ > 0. Simulating the SDE with a numerical solver such as Euler–Maruyama is offered as an alternative, but is described as generally more time-consuming and as introducing additional discretization errors.

Practical guidance on hyperparameters: the loss rules out arbitrarily small T because the target term would explode, while excessively large T makes the noisy marginal μ_T nearly indistinguishable from μ₀ and inflates the variance of the optimal control, degrading sample quality; the authors therefore choose T away from extremal values and avoid too small β. Appendix B is said to contain alternative algorithms, including a simplification of Algorithm 1 when β is large and a Sinkhorn-based solver for the SBB problem.

Why This Matters

• Impact on research: The work advances the interface between optimal transport theory and modern generative modeling by making stochastic-volatility transport computationally tractable at scale. It connects the SBB formulation of Henry-Labordère et al. (2026) with the LightSB-M solver of Gushchin et al. (2024), and offers a complementary alternative to heavy-tailed generative modeling approaches that replace Gaussian noise with α-stable noise or Student-t perturbation kernels (Shariatian et al., 2025; Pandey et al., 2025): rather than changing the noise family, SBB keeps a constant-volatility bridge in the transformed Y-space and gains flexibility through a learned transport-induced, state-dependent volatility in the original space.
• Real-world applications (as framed by the authors):
  • Generative modeling under distributional constraints.
  • Heavy-tailed data generation, where classical SB is infeasible.
  • Time-series modeling with heteroskedasticity.
  • Unpaired image-to-image translation, demonstrated on adult → child faces in FFHQ.
• Industry relevance: Three of the six authors are affiliated with financial institutions (BNP Paribas and Qube RT), and the authors explicitly motivate the framework by settings such as heavy tails and heteroskedastic time series, which are typical of financial data. The released code repository lowers the barrier for practitioners to adopt the solver.

Future Directions

• Establishing convergence to the Bass limit. The authors state that the β ↓ 0 limit is treated only as a formal interpolation and that a genuine zero-drift limit requires (μ₀, μ_T) to satisfy convex order; proving convergence of SBB optimizers to the specific Bass optimizer among martingale couplings is left open.
• Convergence theory for the alternating loop. The absence of a formal convergence proof concerns the outer alternating loop; only the inner SB subproblem carries theoretical guarantees through LightSB-M.
• Addressing the poorly conditioned regime. The authors identify small β as the numerically delicate case, where the inverse map becomes poorly conditioned.
• Extending beyond images. The authors state that the framework opens directions for generative modeling under distributional constraints, heavy-tailed data, and time-series with heteroskedasticity, settings that remain challenging for standard diffusion models.

Target Audience

This paper is aimed at machine learning researchers working on diffusion models, Schrödinger Bridge methods, and generative modeling; at optimal transport and stochastic control theorists interested in drift-and-volatility transport problems; and at quantitatively trained practitioners in finance and other data-driven fields who need generative models for heavy-tailed or heteroskedastic data. Readers without a background in stochastic calculus, entropic optimal transport, and duality will find the theoretical sections demanding.

Authors’ abstract

The Schrodinger Bridge and Bass (SBB) formulation, which jointly controls drift and volatility, is an established extension of the classical Schrodinger Bridge (SB). Building on this framework, we introduce LightSBB-M, an algorithm that computes the optimal SBB transport plan in only a few iterations. The method exploits a dual representation of the SBB objective to obtain analytic expressions for the optimal drift and volatility, and it incorporates a tunable parameter beta greater than zero that interpolates between pure drift (the Schrodinger Bridge) and pure volatility (Bass martingale transport). We show that LightSBB-M achieves the lowest 2-Wasserstein distance on synthetic datasets against state-of-the-art SB and diffusion baselines with up to 32 percent improvement. We also illustrate the generative capability of the framework on an unpaired image-to-image translation task (adult to child faces in FFHQ). These findings demonstrate that LightSBB-M provides a scalable, high-fidelity SBB solver that outperforms existing SB and diffusion baselines across both synthetic and real-world generative tasks. The code is available at https://github.com/alexouadi/LightSBB-M.