Contact Wasserstein Geodesics for Non-Conservative Schrödinger Bridges

TL;DR

This work generalizes the Schrödinger Bridge to non-conservative dynamics by introducing the non-conservative generalized Schrödinger bridge (NCGSB) grounded in contact Hamiltonian mechanics. It recasts the bridge computation as geodesics on the Wasserstein manifold with a Jacobi metric, enabling energy to vary along the path via a scalar state $z^t$ and damping $\gamma$, and proposes Contact Wasserstein Geodesics (CWG) as a fast, non-iterative solver implemented as a ResNet over discrete segments. CWG leverages a two-stage training regime and supports guided generation by incorporating task-specific penalties into the metric, achieving near-linear complexity $\mathcal{O}(N K (T_{\text{sh}}+D(LW+\log N)))$. Empirical results across LiDAR, molecular dynamics-like cell data, sea temperatures, robotics, and unpaired image translation demonstrate improved fidelity, guided controllability, and computational efficiency compared to prior SB solvers. The framework offers a versatile, scalable approach for modeling complex intermediate dynamics in stochastic processes where energy is not conserved.

Abstract

The Schrödinger Bridge provides a principled framework for modeling stochastic processes between distributions; however, existing methods are limited by energy-conservation assumptions, which constrains the bridge's shape preventing it from model varying-energy phenomena. To overcome this, we introduce the non-conservative generalized Schrödinger bridge (NCGSB), a novel, energy-varying reformulation based on contact Hamiltonian mechanics. By allowing energy to change over time, the NCGSB provides a broader class of real-world stochastic processes, capturing richer and more faithful intermediate dynamics. By parameterizing the Wasserstein manifold, we lift the bridge problem to a tractable geodesic computation in a finite-dimensional space. Unlike computationally expensive iterative solutions, our contact Wasserstein geodesic (CWG) is naturally implemented via a ResNet architecture and relies on a non-iterative solver with near-linear complexity. Furthermore, CWG supports guided generation by modulating a task-specific distance metric. We validate our framework on tasks including manifold navigation, molecular dynamics predictions, and image generation, demonstrating its practical benefits and versatility.

Figures (18)

• Figure 1: Probability paths obtained under energy-conserving (), energy-decreasing (), and energy-increasing conditions () (details in App. \ref{['app:teaser']}). Energy variation increases modeling flexibility in applications where distributions at intermediate time steps are of interest.
• Figure 2: Visualization of the ResNet transformation. Two successive pushforwards ${\rho}_\theta^{t_{k-1}} \to {\rho}_\theta^{t_k} \to {\rho}_\theta^{t_{k+1}}$ on $\mathcal{P}^+(\mathcal{M})$ are shown as local updates $\partial_t {\rho}^{t_k}_\theta$, $\partial_t {\rho}^{t_{k+1}}_\theta$ on tangent spaces. Each update is parameterized by $\theta^k, \theta^{k+1} \in \Theta$, defining local coordinates on $\mathcal{T}\mathcal{P}^+(\mathcal{M})$. This coordinate system is not unique.
• Figure 3: Two-Moons (top) and Checkerboard (bottom) benchmarks with guided variants (right).
• Figure 4: LiDAR Manifold Navigation: CWG before and after guidance (top), CWG vs GSBM (bottom).
• Figure 5: Predictions from CWG (ours, top), GSBM (middle), and DSBM (bottom). The red row shows marginal samples.

Theorems & Definitions (2)

• Proposition 1
• Proposition 2