Iterated Diffusion Bridge Mixture

• Iterated Diffusion Bridge Mixture (IDBM) is an iterative algorithm that constructs stochastic process transports between prescribed probability measures by minimizing the KL divergence to a reference diffusion.
• It leverages mixtures of endpoint-conditioned diffusion bridges for both forward and backward transports, enabling efficient generative modeling in Euclidean and Riemannian spaces.
• The method achieves weak convergence to the dynamic Schrödinger bridge with strong theoretical guarantees, reducing simulation costs and discretization errors in practice.

The iterated diffusion bridge mixture (IDBM) is an iterative, sampling-based algorithm for constructing stochastic processes that realize transports between two prescribed probability measures, optimally approximating a reference diffusion in the sense of Kullback-Leibler divergence. IDBM provides a framework for solving the dynamic Schrödinger bridge problem in both Euclidean spaces and Riemannian manifolds, yielding at each iteration a valid transport between the marginals. The approach leverages mixtures of endpoint-conditioned diffusion bridges and enables both forward and backward transport SDE constructions, forming the foundation for practical, high-fidelity generative modeling with efficient, scalable training objectives and strong theoretical convergence guarantees (Peluchetti, 2023, Jo et al., 2023).

1. Dynamic Schrödinger Bridge and Problem Formalization

The dynamic Schrödinger bridge problem seeks a stochastic process $Q^*$ on $C([0,\tau],\mathbb{R}^d)$ that interpolates between specified initial ($\mu$) and terminal ($\nu$) probability laws, while remaining optimally close (in the sense of KL divergence) to a given reference diffusion process $P$: $$Q^* = \operatorname{arg\,min}_{Q: Q_0=\mu,\, Q_\tau=\nu}\, D_{KL}(Q\,\|\,P)$$ Given a reference SDE

$$dX_t = \mu_R(X_t, t)\,dt + \sigma_R(X_t, t)\,dW_t, \quad X_0 \sim \mu,$$

the optimal bridge law $Q^*$ is realized via a diffusion $h$-transform. For any candidate diffusion law $Q$, its time-reversal also admits an SDE characterization, and the marginal–conditional decomposition reveals that $C([0,\tau],\mathbb{R}^d)$0 is characterized by reweighting endpoint couplings of $C([0,\tau],\mathbb{R}^d)$1's bridge laws. In particular, the SDE for the bridge incorporates a drift adjustment involving $C([0,\tau],\mathbb{R}^d)$2, with $C([0,\tau],\mathbb{R}^d)$3 defined through integration against the transition kernel and one of the Schrödinger potentials (Peluchetti, 2023).

2. Algorithmic Construction: Iterated Diffusion Bridge Mixture

2.1. One-step Bridge Mixture

Given a coupling $C([0,\tau],\mathbb{R}^d)$4 of $C([0,\tau],\mathbb{R}^d)$5 on $C([0,\tau],\mathbb{R}^d)$6, form the mixture of reference bridges $C([0,\tau],\mathbb{R}^d)$7 by sampling endpoints from $C([0,\tau],\mathbb{R}^d)$8 and conditioning the reference process on their values. It is shown that there exists a single diffusion measure $C([0,\tau],\mathbb{R}^d)$9 whose time-marginals correspond to those of the mixture: $\mu$0 where $\mu$1 is the transition density of the reference process.

2.2. Forward and Backward Transports

The forward diffusion bridge mixture (DBM) SDE starts at $\mu$2, while the backward DBM begins at $\mu$3, with both drifts derived via conditional expectations over the coupling $\mu$4. Iteration alternates between updating forward and backward transports, with the coupling updated from the endpoint distributions of the most recent DBM.

2.3. IPF and Schrödinger Potentials

The iteration can be recast with Schrödinger potentials $\mu$5, which satisfy integral constraints under the reference transition kernel: $\mu$6 IPF (Iterative Proportional Fitting)–type updates successively refine these potentials, recalibrating the mixture law used in subsequent bridge mixtures (Peluchetti, 2023).

3. Convergence Theory

Under mild regularity—nondegeneracy of the reference, finite initial KL, and Girsanov applicability—the sequence of couplings $\mu$7 and associated mixture diffusions $\mu$8 converge weakly to the Schrödinger bridge $\mu$9 as $\nu$0. The iterative KL divergences satisfy strict Pythagorean identities: $\nu$1 with monotonic decrease in $\nu$2. Equality in both KLs is attained only at the fixed point where $\nu$3 (Peluchetti, 2023).

4. Implementation, Complexity, and Manifold Extensions

4.1. Drift-Matching and Training Loss

Algorithmically, one fits a neural network $\nu$4 to approximate the drift component

$\nu$5

by minimizing a drift-matching loss aggregated over randomly sampled time points, terminal couplings, and conditional transitions. An analogous loss is minimized for the backward process, employing $\nu$6.

4.2. Simulation Cost

When the reference bridge admits a closed-form (e.g., Gaussian or linear cases), intermediate state sampling is computationally inexpensive, and per-iteration cost is dominated by minibatch size times the number of time samples. IDBM improves upon the diffusion IPF (DIPF) paradigm by reusing endpoint samples, allowing for efficient repeated conditional sampling of intermediate states and avoiding simulation–inference mismatches prevalent in low-noise regimes (Peluchetti, 2023).

4.3. Riemannian Manifold Generalization

IDBM is directly adapted to Riemannian manifolds by utilizing mixtures of diffusion bridges constructed via either logarithmic maps ("log-bridge") or spectral-geometric distances ("spectral-bridge"): $\nu$7 for a fixed endpoint $\nu$8, with the mixture drift obtained as a weighted sum over tangent vectors to sampled endpoints, where weights are determined by short-time heat kernel approximations. Efficient simulation proceeds by geodesic random walks on the manifold, and the approach readily scales to high dimensions by sub-sampling endpoints and mini-batching (Jo et al., 2023).

Reference Process	Drift Formula	Setting
Linear Gaussian	Closed-form drift, conditional means	Euclidean, closed-form
Log-bridge	Weighted sum of log-maps to endpoints	Manifold, log map required
Spectral-bridge	Spectral gradient of squared truncated distance	Manifold, spectral modes

5. First-Iteration Generative Modeling and Empirical Performance

The first backward DBM step can be employed directly as a generative model, matching both data and reference marginals exactly at training. For Euclidean reference $\nu$9, the backward generative SDE is

$P$0

or as an ODE in the zero-noise limit. In this regime, training loss modifies the marginal score-matching objective by substituting the marginal score with the conditional score. Sampling benefits from reduced discretization error, allowing coarser time steps and accelerated wall-clock performance. Empirical results on CIFAR-10 demonstrate that, at a time step $P$1, IDBM matches or surpasses classical score-based generative models (SGM) in FID, with training time reductions by a factor of three and no requirement for architectural changes (Peluchetti, 2023).

6. Connections to Riemannian Generative Modeling

Riemannian diffusion mixtures apply IDBM principles to manifold-valued data by constructing geodesic bridge mixtures, leveraging either logarithmic maps or spectral geometry. Endpoints are sampled i.i.d. from the data manifold, and mixture drifts are explicitly computed as weighted averages in the tangent bundle. Training is achieved via a two-way bridge matching objective, which avoids explicit computations of metric divergences and enables scalable, batch-oriented optimization. IDBM on manifolds allows for mini-batched endpoint updates ("iteration") or persistence across training batches, matching the statistical efficiency required for high-dimensional and non-Euclidean generative modeling (Jo et al., 2023).

7. Practical Considerations and Scalability

Metric geometry enters in all norm computations, with distances and exponentiated maps tailored to the manifold topology. The number of bridges, time discretization steps, and choice of noise schedule are tunable hyperparameters. For differentiable manifolds with closed-form geometric operations (spheres, tori), bridge simulation is straightforward; for more complex geometries, geodesic solvers or spectral truncations are employed. The IDBM approach admits both SDE and ODE sampling schemes following the learning of forward and backward drifts, offering flexibility in practical generation pipelines.

In summary, the IDBM framework unifies an iterative, weakly convergent sequence of bridge-mixture transports for the dynamic Schrödinger bridge problem, supports both Euclidean and Riemannian settings, and provides an efficient, theoretically grounded basis for high-fidelity generative modeling (Peluchetti, 2023, Jo et al., 2023).