Contact Wasserstein Geodesics for Non-Conservative Schr\"odinger Bridges

TL;DR

This paper introduces the non-conservative generalized Schr"odinger bridge (NCGSB) and contact Wasserstein geodesics (CWG), enabling modeling of energy-varying stochastic processes with efficient, non-iterative neural network implementations.

Contribution

It proposes a novel energy-varying reformulation of Schr"odinger bridges using contact Hamiltonian mechanics and develops a practical, neural network-based solver for Wasserstein geodesics.

Findings

CWG achieves near-linear complexity in computation.

The framework effectively models diverse real-world stochastic processes.

Validated on manifold navigation, molecular dynamics, and image generation.

Abstract

The Schr\"odinger 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\"odinger 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.