Control Consistency Losses for Diffusion Bridges

1. Problem Statement

The paper addresses the challenge of simulating diffusion bridges: conditioned stochastic processes $(X_t)_{t \in [0,T]}$ that start at a known state $x_0$ and are forced to terminate at a specific state $x_T$.

• Context: This is crucial in fields like computational chemistry (rare event sampling, transition paths), finance, and shape evolution.
• The Challenge: Standard simulation methods (like MCMC or SMC) struggle when the conditioning event is rare (i.e., the unconditioned dynamics rarely reach $x_T$).
• Limitations of Existing Neural Methods:
  • Methods learning from uncontrolled dynamics (e.g., Score Matching, Malliavin calculus) fail for rare events because training data rarely covers the region near $x_T$.
  • Methods learning from controlled dynamics via KL-divergence minimization (e.g., Neural Guided Diffusion Bridges, NGDB) require backpropagation through simulated trajectories. This is computationally expensive and hinders scalability.

2. Methodology: Control Consistency Diffusion Bridge (CCDB)

The authors propose a novel framework based on a self-consistency property of the optimal control drift $u(t, x)$. The core idea is to learn the control without differentiating through the simulation path.

A. Theoretical Foundation: Self-Consistency

The paper derives a property linking the control at an earlier time $s$ to the control at a later time $t$ along the solution trajectory.

• Jacobian Process: They utilize the Jacobian process $J_{t|s} = \nabla_{X_s} X_t$, which measures the sensitivity of the state $X_t$ to perturbations in $X_s$.
• Property (SC1): For the optimal control $u^*$, the following holds:
  $$u^*(s, x) = \mathbb{E}[J_{t|s}^\top u^*(t, X_t) \mid X_s = x]$$
  This implies that the control value at time $s$ is the expected value of the future control value, transformed by the Jacobian of the flow.
• Generalization (SC2): To handle cases where the Jacobian grows large (causing numerical instability), they derive a generalized property using an auxiliary process with an added drift $\tilde{b}$. A specific case where $\tilde{b} = -b$ yields:
  $$u^*(s, x) = \mathbb{E}\left[ \int_s^t \nabla b(r, X_r)^\top u^*(r, X_r) dr + u^*(t, X_t) \mid X_s = x \right]$$
  This avoids the explicit Jacobian matrix $J_{t|s}$, replacing it with an integral of the drift gradient, which is more stable in high-energy landscapes.

B. Algorithm Design

The proposed algorithm, Control Consistency Diffusion Bridge (CCDB), combines the self-consistency property with terminal constraints:

1. Terminal Condition: The process is constructed to jump directly to $x_T$ at the final time step, enforcing the bridge constraint by construction.
2. Gradient Form: The control is parameterized to ensure it is a gradient field (satisfying necessary conditions for the diffusion bridge).
3. Fixed-Point Loss: Instead of minimizing KL-divergence via backpropagation, the model minimizes a self-consistency loss:
   $$L(\theta) = \mathbb{E}_{s \sim U[0,T]} \left[ \| u_\theta(s, X_s) - U_s \|^2 \right]$$
   Where $U_s$ is a target computed by integrating future control values weighted by the Jacobian (or the generalized form) backwards along the simulated trajectory.
   • Crucial Innovation: The target $U_s$ is computed using a stop-gradient on the trajectory. The regression is performed against a fixed target derived from the current estimate of the control. This eliminates the need to backpropagate through the stochastic differential equation (SDE) solver.

C. Training Targets and Schedules

The authors introduce a weighting schedule $\alpha_t$ to control how far into the future the consistency is enforced:

• Next-step: Enforces consistency only between adjacent time steps (low variance, slower propagation).
• Average: Enforces consistency over the remaining time horizon (balanced).
• Final-step: Enforces consistency only with the terminal step (high variance, fast propagation).
  The Average schedule is generally found to be the most robust.

3. Key Contributions

1. Novel Self-Consistency Framework: Derivation of a necessary and sufficient condition (Theorem 3.1) for a controlled diffusion to be a diffusion bridge, based on the self-consistency of the control drift and a terminal condition.
2. Scalable Training: Introduction of a fixed-point style loss that allows training on controlled dynamics without backpropagation through trajectories. This significantly reduces computational cost compared to KL-based methods.
3. Stability Mechanism: The derivation of the generalized self-consistency property (SC2) to mitigate numerical instability caused by exploding Jacobian terms in high-energy or rare-event settings.
4. Theoretical Unification: Connecting diffusion bridge learning to broader Stochastic Optimal Control (SOC) literature, showing that the method recovers recent "Adjoint Matching" objectives as special cases.

4. Experimental Results

The method (CCDB) was evaluated against state-of-the-art baselines, including NGDB (Yang et al., 2025), Score-Diffusion Bridge (SDB), Forward-Bridge (FB), and BEL.

• Performance: CCDB achieves performance competitive with NGDB (the strongest baseline) across various settings, including:
  • Ornstein-Uhlenbeck (OU) Bridges: Standard benchmark; CCDB converges faster and achieves lower KL divergence near the terminal time.
  • Cox-Ingersoll-Ross (CIR) Model: Successfully handles spatially dependent diffusion coefficients.
  • Double-Well Potential: A rare-event scenario with high energy barriers. Here, SC2 significantly outperforms SC1 and other methods, demonstrating the importance of the generalized property.
  • Cell Differentiation: Handles non-gradient drifts and multimodal distributions.
  • Transition Path Sampling (TPS): Applied to Müller-Brown potentials and Alanine Dipeptide (molecular dynamics). CCDB remains stable in 100-dimensional settings and produces physically plausible transition paths.
• Efficiency:
  • Speed: CCDB is approximately 3x faster than NGDB in training time (e.g., 230s vs 800s for OU bridges).
  • Reason: The elimination of trajectory backpropagation allows for larger batch sizes and more efficient gradient updates.
• Failure Modes of Baselines: Methods relying on uncontrolled dynamics (SDB, FB, BEL) failed to learn bridges for rare events (e.g., Double-Well, Multimodal Cell Diffusion) because the training data did not cover the transition region.

5. Significance and Impact

• Rare Event Sampling: The ability to learn diffusion bridges for rare events without the computational bottleneck of trajectory backpropagation makes this method highly valuable for computational chemistry and physics, where simulating transition paths between metastable states is critical.
• Scalability: By decoupling the gradient computation from the trajectory simulation, the method scales better to high-dimensional problems and complex potentials.
• Theoretical Insight: The work provides a unified view of diffusion bridges and stochastic optimal control, suggesting that self-consistency losses are a powerful, generalizable tool for learning controlled stochastic processes.
• Future Directions: The authors suggest extending the framework to hypoelliptic diffusions (underdamped dynamics) and incorporating domain-specific initialization strategies for even more complex molecular systems.

In summary, this paper presents a mathematically rigorous and computationally efficient approach to learning diffusion bridges, overcoming the scalability limitations of previous neural methods while maintaining high accuracy in challenging rare-event scenarios.