Table of Contents
Fetching ...

Control Consistency Losses for Diffusion Bridges

Samuel Howard, Nikolas Nüsken, Jakiw Pidstrigach

TL;DR

This work tackles diffusion-bridge sampling by exploiting a self-consistency property of conditioned diffusion dynamics. It introduces a gradient-form control framework and a fixed-point, online self-consistency loss that trains neural drift terms without backpropagating through simulated trajectories, enabling scalable learning even for rare conditioning events. Theoretical results characterize the diffusion bridge as the unique gradient-form process satisfying both the terminal-bridge constraint and the self-consistency relation, with a generalised version incorporating auxiliary drifts to improve numerical stability. Empirically, the method achieves competitive KL divergences to ground-truth diffusion bridges across Ornstein-Uhlenbeck, double-well, and Müller–Brown potentials, often at substantially lower computational cost than state-of-the-art baselines, and demonstrates robustness to ablations like different α-schedules and STL adjustments. This approach offers a scalable, principled alternative for diffusion-bridge problems with potential impact in computational chemistry and beyond.

Abstract

Simulating the conditioned dynamics of diffusion processes, given their initial and terminal states, is an important but challenging problem in the sciences. The difficulty is particularly pronounced for rare events, for which the unconditioned dynamics rarely reach the terminal state. In this work, we leverage a self-consistency property of the conditioned dynamics to learn the diffusion bridge in an iterative online manner, and demonstrate promising empirical results in a range of settings.

Control Consistency Losses for Diffusion Bridges

TL;DR

This work tackles diffusion-bridge sampling by exploiting a self-consistency property of conditioned diffusion dynamics. It introduces a gradient-form control framework and a fixed-point, online self-consistency loss that trains neural drift terms without backpropagating through simulated trajectories, enabling scalable learning even for rare conditioning events. Theoretical results characterize the diffusion bridge as the unique gradient-form process satisfying both the terminal-bridge constraint and the self-consistency relation, with a generalised version incorporating auxiliary drifts to improve numerical stability. Empirically, the method achieves competitive KL divergences to ground-truth diffusion bridges across Ornstein-Uhlenbeck, double-well, and Müller–Brown potentials, often at substantially lower computational cost than state-of-the-art baselines, and demonstrates robustness to ablations like different α-schedules and STL adjustments. This approach offers a scalable, principled alternative for diffusion-bridge problems with potential impact in computational chemistry and beyond.

Abstract

Simulating the conditioned dynamics of diffusion processes, given their initial and terminal states, is an important but challenging problem in the sciences. The difficulty is particularly pronounced for rare events, for which the unconditioned dynamics rarely reach the terminal state. In this work, we leverage a self-consistency property of the conditioned dynamics to learn the diffusion bridge in an iterative online manner, and demonstrate promising empirical results in a range of settings.

Paper Structure

This paper contains 43 sections, 10 theorems, 54 equations, 9 figures, 4 tables, 3 algorithms.

Key Result

Theorem 2.1

The control drift $u^*$ of the solution satisfies Moreover, the control drift satisfies the following self-consistency property,

Figures (9)

  • Figure 1: The self-consistency property $u(s, X_s) = \mathbb{E}_{\mathbb{P}^u} [u (t, X_t) J_{t|s} |X_s]$ relates the control value at an earlier time $s$ to its value at a later time $t$, along the trajectories.
  • Figure 2: The diffusion bridge is the unique process of form \ref{['eq:controlled_process']} bridging from $x_0$ to $x_T$ with $u$ of gradient form, that satisfies the self-consistency property \ref{['eq:self consistent']}.
  • Figure 3: Comparison of our proposed approach with Neural Guided Diffusion Bridge (NGDB) yang2025_guidedbridges in the Ornstein-Uhlenbeck bridge experiment (mean$\pm$ std, over 5 runs).
  • Figure 4: Visualisation of obtained bridge trajectories for the 1d double-well potential.
  • Figure 5: Visualisation of obtained bridge trajectories for the 2$d$ Müller-Brown potential.
  • ...and 4 more figures

Theorems & Definitions (17)

  • Theorem 2.1
  • Theorem 3.1
  • Theorem B.1: Generalised self-consistency property
  • Theorem B.2
  • Lemma D.1
  • proof
  • Theorem D.1
  • proof
  • Definition D.1
  • Theorem D.2
  • ...and 7 more