Diffusion Bridges: UniDB Framework

• Diffusion bridges are stochastic processes, defined via time-inhomogeneous SDEs or deterministic ODEs, that interpolate between two endpoint probability distributions.
• The UniDB framework employs a unified path-space least-squares objective to fit parameterized drift fields, effectively recovering various bridge algorithms like flow matching and Schrödinger bridges.
• Algorithmic specializations, including OT-FM and DSBM, offer practical trade-offs between computational efficiency and marginal accuracy, benefiting generative modeling and optimal transport applications.

A diffusion bridge is a stochastic process—typically defined as a time-inhomogeneous stochastic differential equation (SDE) or, in certain cases, an ordinary differential equation (ODE)—whose law interpolates between two given probability distributions at its endpoints. The "Unified Diffusion Bridge" (UniDB) framework provides a comprehensive formalism that subsumes flow matching, Schrödinger bridge algorithms, and their modern deep variants. UniDB constructs a single path-space least-squares objective whose minimization, for a suitable choice of reference process and endpoint coupling, recovers all popular bridge algorithms and delivers a principled, flexible approach for generative modeling, data translation, and optimal transport between distributions (Kim, 27 Mar 2025).

1. Mathematical Formalism of Diffusion Bridges

The general diffusion bridge problem defines two endpoint distributions π₀ and π_T over ℝ^d, and seeks a process {X_t} for t ∈ [0,T] such that X₀ ∼ π₀ and X_T ∼ π_T. The process is governed by a time-inhomogeneous SDE or ODE:

• Forward SDE:

$$dX_t = f(t, X_t)\,dt + \sigma(t)\,dW_t, \quad X_0 \sim \pi_0, \quad X_T \sim \pi_T$$

• Reverse SDE:

$$dX_t = \bar{f}(t, X_t)\,dt + \sigma(T-t)\,d\overleftarrow{W}_t,$$

Setting σ(t) ≡ 0 yields a deterministic ODE, recovering classical flow matching (FM). The Schrödinger bridge problem, a classical instance, imposes a further constraint: it seeks a controlled path measure P^u for the SDE that minimizes KL(P^u || P^ref), subject to matching the desired endpoint marginals. This approach encapsulates the entropy-regularized optimal control interpretation of the bridge.

2. Unified Path-Space Optimization Objective

UniDB introduces a unified path-space matching objective by parametrizing a "pinned" transitional law P_t(x | x₀, x_T) and a coupling Q(x₀, x_T), such that:

$$P_t(x) = \int P_t(x\,|\,x₀, x_T)\, Q(x₀, x_T)\, dx₀\,dx_T$$

The drift field $u_t(x\,|\,x₀,x_T)$ ensuring P_t(* | x₀, x_T) as the marginal is obtained analytically for the chosen bridge process (ODE or SDE; e.g., linear Brownian bridge or deterministic linear flow). UniDB fits a parameterized vector field v_θ(t, x) by minimizing the squared-error loss:

$$L(\theta) = \mathbb{E}_{t, (x₀,x_T), X_t}\left\|u_t(X_t\,|\,x₀,x_T) - v_θ(t,X_t)\right\|^2$$

For Schrödinger bridges, this is equivalent to a Girsanov-based likelihood loss:

$$L_{SB}(\theta) = \mathbb{E}_{P^{v_θ}} \left[ \frac{1}{2\sigma_{ref}^2}\int_0^T \|v_θ(t,X_t)\|^2 dt \right ]$$

along with a terminal density constraint.

3. Algorithmic Structure and Specializations

The canonical UniDB algorithm proceeds as follows:

Algorithm UniDB(π₀, π_T) Inputs: samples from π₀, π_T Repeat until convergence: 1. Choose path family P_t(x|x₀,x_T) and coupling Q(x₀,x_T) with marginals enforcing π₀, π_T. 2. Set diffusion σ ≥ 0; compute drift u_t(x|x₀,x_T) matching the conditional law. 3. Fit v_θ(t,x) via regression: minimize E_{t,(x₀,x_T),X_t} [∥u_t(X_t|x₀,x_T) − v_θ(t,X_t)∥²]. Return: SDE dX_t = v_θ(t,X_t)dt + σdW_t.

Specific choices yield:

• Flow Matching (FM): σ = 0, P_t = δ(x - [(1-t)x₀ + t x_T]), Q = π₀(x₀)π_T(x_T).
• OT-FM: σ = 0, P_t Gaussian on line segment, Q = OT coupling.
• SB-FM: σ = 0 or small, P_t is a Gaussian Brownian bridge, Q = mini-batch entropic OT.
• DSBM: σ = σ_ref, P_t is a full Brownian bridge, Q iteratively updated to match marginals of previous fit.

4. Theoretical Guarantees and Convergence

The central theoretical result (Theorem 3.1) establishes: under regularity and universal-approximation assumptions, each iteration of UniDB constructs an SDE whose time-marginals match the prescribed P_t(·), attaining π₀ at t=0 and πT at t=T. Fokker-Planck analysis confirms the legality of the marginal flows; regression ensures the drift vθ(t,x) converges to its conditional expectation given the bridge trajectory. Existence and uniqueness are inherited from Schrödinger bridge theory; iterative updates in DSBM are underpinned by monotonic KL-descent in path space.

5. Practical Implementation Aspects

• Mini-batch Matching: For OT-based and entropy-based bridges, optimal transport (Sinkhorn, etc.) is performed on small batches, mitigating computational complexity while introducing manageable bias.
• ODE vs SDE Sampling: σ=0 requires deterministic solvers (Runge–Kutta), σ>0 (stochastic) is handled via Euler–Maruyama or higher-order SDE solvers. Sampling noise in SDE improves robustness and exactness of marginals.
• Computational Trade-Offs: SDE bridges (DSBM) necessitate double training (forward and reverse) but recover exact endpoint and intermediate marginals. ODE flows are faster but only guarantee marginal alignment in L₂ sense.

6. Extensions and Broader Applications

UniDB's path-space least-squares structure permits multiple generalizations:

• Cost-aware bridges: Inserting a state-dependent cost ℓ(x,t) enables tailored generative flows (generalized Schrödinger bridge).
• Conditional Generation: By selecting π₀ as the encoding distribution and π_T as a target class, UniDB trains conditional flows for structured generation, including alternatives to conditional GANs.
• Manifold-constrained bridges: Employing manifold geometry (Riemann-Brownian motion) yields bridges that adapt to non-Euclidean data domains, enabling shape and geometry-oriented generative modeling.

UniDB thus unifies ODE-based flow matching, SDE-based Schrödinger bridges, mini-batch optimal transport methods, and deep stochastic bridge algorithms into a single framework with shared theoretical foundation and algorithmic machinery. This facilitates principled model development in generative modeling, high-dimensional optimal transport, and stochastic optimal control, with broad utility in image generation, translation, restoration, and manifold-valued data synthesis (Kim, 27 Mar 2025).