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.
