Control Consistency Losses for Diffusion Bridges

Samuel Howard; Nikolas Nüsken; Jakiw Pidstrigach

stat.MLDec 4, 2025

Control Consistency Losses for Diffusion Bridges

Authors: Samuel Howard, Nikolas Nüsken, Jakiw Pidstrigach

Organizations: Department of Statistics, University of Oxford · Department of Mathematics, King’s College London

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 propose a novel approach for learning diffusion bridges based on a self-consistency property of the optimal control. The resulting algorithm learns the conditioned dynamics in an iterative online manner, and exhibits strong performance in a range of empirical settings without requiring differentiation through simulated trajectories. Beyond the diffusion bridge setting, we draw connections between our self-consistency framework and recent advances in the wider stochastic optimal control literature.

Figures & tables

Panel of Figure 3 — Figure 3: Comparison of our proposed approach CCDB (using self-consistency properties ( SC1 ) and ( SC2 )), and Neural Guided Diffusion Bridge (NGDB) ( Yang et al., 2025 ) , in the Ornstein-Uhlenbeck bridge experiment (mean $\pm$ std, over 5 runs).

	CIR model, ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Double Well, $v=3.0$ , ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Cell Diffusion, ${\mathrm{KL}}(\mathbb{P}^{\theta}\\|\mathbb{P})~\downarrow$
			Normal	Rare	Multimodal
CCDB (Ours) , using ( SC1 )	0.023 $\pm$ 0.010 •	0.148 $\pm$ 0.073 • / $\blacktriangle$	6.53 $\pm$ 0.14 •	66.1 $\pm$ 0.9 •	792.8 $\pm$ 0.9 • / $\blacktriangle$
CCDB (Ours) , using ( SC2 )	0.055 $\pm$ 0.009 • / $\blacktriangle$	0.041 $\pm$ 0.013 •	6.53 $\pm$ 0.17 •	66.0 $\pm$ 0.5 •	792.4 $\pm$ 1.0 • / $\blacktriangle$
NGDB ( Yang et al., 2025 )	0.051 $\pm$ 0.012 • / $\blacktriangle$	0.098 $\pm$ 0.047 • / $\blacktriangle$	6.48 $\pm$ 0.15 •	65.5 $\pm$ 0.6 •	783.4 $\pm$ 0.8 •
SDB ( Heng et al., 2025 )	0.104 $\pm$ 0.011 $\blacktriangle$	0.229 $\pm$ 0.024 $\blacktriangle$	5.56 $\pm$ 0.39*	75.0 $\pm$ 5.1 $\blacktriangle$	1845 $\pm$ 230 $\blacksquare$
FB ( Baker et al., 2025 )	0.074 $\pm$ 0.011 $\blacktriangle$	-	6.51 $\pm$ 0.14 •	237 $\pm$ 98 $\blacksquare$	4836 $\pm$ 954 $\blacksquare$

Table 1: Quantitative results for experiments (mean

\pm

std over 5 runs). Methods that perform strongly are indicated by • , those with reasonable performance by

\blacktriangle

, and methods that fail to converge to the correct solution by

\blacksquare

. We note the result marked with an asterisk* is likely unreliable as it gave an extremely large control towards the endpoint, causing the integration to underestimate the correct value.

Table 4

Panel of Figure 4 — Figure 4: (Left) Learned trajectories for the Müller–Brown potential, for our method using ( SC2 ). (Middle) Ramachandran plot showing obtained paths for alanine dipeptide. (Right) Energy of an individual transition path for alanine dipeptide.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

{\mathrm{d}}X_{t}^{\theta}=\big(b+(\sigma\sigma^{\top})u_{\theta}\big)(t,X_{t}^{\theta}){\mathrm{d}}t+\sigma(t,X_{t}^{\theta}){\mathrm{d}}B_{t},

(43)

Appendix

Algorithm 2 Control Consistency Diffusion Bridge (Generalised version)

\mathcal{U}_{T-\delta t}=u(T-\delta t,X_{T-\delta t})=(\sigma\sigma^{\top})(T-\delta t,X_{T-\delta t})^{-1}\Big(\frac{X_{T}-X_{T-\delta t}}{\delta t}-b(T-\delta t,X_{T-\delta t})\Big)

(45)

Appendix

Algorithm 3 Computing training targets along a single trajectory (generalised version)

Figure 5: Comparison of performance on the double-well experiment, for different barrier height parameters

v

Figure 6: Visualisation of the learned trajectories for each method in the normal , rare , and multimodal cell diffusion examples. The experimental setting is taken from Yang et al. (2025) , using the cell differentiation model from Wang et al. (2011) . Results in the leftmost column are obtained using the MCMC-based guided proposal approach from Schauer et al. (2017) .

Panel of Figure 7 — Figure 7: Comparison of the trajectories obtained using our method, in the Müller-Brown examples.

	${\mathrm{KL}}(\mathbb{P}^{\theta}\\|\mathbb{P})\downarrow$	Max Energy
Ours (using ( SC2 ))	27.0 $\pm$ 1.2	-35.2 $\pm$ 5.13
NGDB ( Yang et al., 2025 )	29.1 $\pm$ 0.8	-33.7 $\pm$ 5.8
DL ( Du et al., 2024 )	39.2 $\pm$ 0.4	-32.0 $\pm$ 6.2

Appendix

Table 4: Comparison of results for the 100

d

Müller–Brown experiment. Metrics are evaluated using the leading two dimensions corresponding to the Müller–Brown potential. The values for

{\mathrm{KL}}(\mathbb{P}^{\theta}\|\mathbb{P})

show mean

\pm

std over 5 runs, while the Max Energy values report the average mean and average std over the 5 runs.

Panel of Figure 8 — Figure 8: Visualisations of the obtained trajectories transitioning from the $C5$ state (upper left) to the $C7ax$ state (lower right), displayed on Ramachandran plots showing the dihedral angles. (Left) Our method, using ( SC2 ). (Right) NGDB.

	CIR model, ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Double Well, $v=3.0$ , ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Cell Diffusion, ${\mathrm{KL}}(\mathbb{P}^{\theta}\\|\mathbb{P})~\downarrow$
			Normal	Rare	Multimodal
Using ( SC1 ), next-step	0.021 $\pm$ 0.007	0.126 $\pm$ 0.069	7.38 $\pm$ 0.49	74.6 $\pm$ 14.1	959.9 $\pm$ 55.3
Using ( SC1 ), average	0.023 $\pm$ 0.010	0.148 $\pm$ 0.073	6.53 $\pm$ 0.14	66.1 $\pm$ 0.9	792.8 $\pm$ 0.9
Using ( SC1 ), endpoint	0.026 $\pm$ 0.005	3.49 $\pm$ 2.40	6.79 $\pm$ 0.10	69.1 $\pm$ 1.1	793.0 $\pm$ 1.0
Using ( SC2 ), next-step	0.042 $\pm$ 0.022	0.067 $\pm$ 0.044	6.95 $\pm$ 0.29	68.6 $\pm$ 1.9	1136.2 $\pm$ 288.9
Using ( SC2 ), average	0.055 $\pm$ 0.009	0.041 $\pm$ 0.013	6.53 $\pm$ 0.17	66.0 $\pm$ 0.5	792.4 $\pm$ 1.0

Appendix

Table 6: Quantitative results for experiments, comparing

\alpha

-schedules (mean

\pm

std, over 5 runs).

	CIR model, ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Double Well, $v=3.0$ , ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Cell Diffusion, ${\mathrm{KL}}(\mathbb{P}^{\theta}\\|\mathbb{P})~\downarrow$
			Normal	Rare	Multimodal
Using ( SC1 ), without STL	0.023 $\pm$ 0.010	0.148 $\pm$ 0.073	6.53 $\pm$ 0.14	66.1 $\pm$ 0.9	792.8 $\pm$ 0.9
Using ( SC1 ), with STL	0.016 $\pm$ 0.001	0.042 $\pm$ 0.024	6.47 $\pm$ 0.16	65.7 $\pm$ 0.6	792.4 $\pm$ 1.0
Using ( SC2 ), without STL	0.055 $\pm$ 0.009	0.041 $\pm$ 0.013	6.53 $\pm$ 0.17	66.0 $\pm$ 0.5	792.4 $\pm$ 1.0
Using ( SC2 ), with STL	0.051 $\pm$ 0.007	0.022 $\pm$ 0.015	6.48 $\pm$ 0.14	66.1 $\pm$ 0.5	792.6 $\pm$ 0.8

Appendix

Table 7: Quantitative results for experiments, comparing performance with and without STL adjustments (mean

\pm

std, over 5 runs).

	CIR model, ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Double Well, $v=3.0$ , ${\mathrm{KL}}(\mathbb{P}^{*}\\|\mathbb{P}^{\theta})~\downarrow$	Cell Diffusion, ${\mathrm{KL}}(\mathbb{P}^{\theta}\\|\mathbb{P})~\downarrow$
			Normal	Rare	Multimodal
Using ( SC1 ), with $\eta_{\theta}=\nabla\psi_{\theta}$	0.023 $\pm$ 0.010	0.148 $\pm$ 0.073	6.53 $\pm$ 0.14	66.1 $\pm$ 0.9	792.8 $\pm$ 0.9
Using ( SC1 ), with general $\eta_{\theta}$	0.025 $\pm$ 0.009	0.080 $\pm$ 0.033	6.49 $\pm$ 0.12	66.2 $\pm$ 0.6	979.5 $\pm$ 2.0
Using ( SC2 ), with $\eta_{\theta}=\nabla\psi_{\theta}$	0.055 $\pm$ 0.009	0.041 $\pm$ 0.013	6.53 $\pm$ 0.17	66.0 $\pm$ 0.5	792.4 $\pm$ 1.0
Using ( SC2 ), with general $\eta_{\theta}$	0.057 $\pm$ 0.008	0.028 $\pm$ 0.013	6.51 $\pm$ 0.15	66.1 $\pm$ 0.5	792.8 $\pm$ 1.1

Appendix

Table 8: Quantitative results for experiments, comparing parameterisations using a general neural adjustment

\eta_{\theta}

with a gradient-form neural adjustment

\eta_{\theta}=\nabla\psi_{\theta}

(mean

\pm

std, over 5 runs).

	$v=2$	$v=3$	$v=4$
$\tilde{b}=0$ ( SC1 )	0.039 $\pm$ 0.011	0.204 $\pm$ 0.100	0.584 $\pm$ 0.346
$\tilde{b}=-\tfrac{1}{3}b$	0.036 $\pm$ 0.012	0.091 $\pm$ 0.032	0.192 $\pm$ 0.045
$\tilde{b}=-\tfrac{2}{3}b$	0.034 $\pm$ 0.006	0.072 $\pm$ 0.026	0.146 $\pm$ 0.033
$\tilde{b}=-b$ ( SC2 )	0.035 $\pm$ 0.006	0.081 $\pm$ 0.004	0.153 $\pm$ 0.016

Appendix

Table 9: Results for different choices of

\tilde{b}

interpolating from ( SC1 ) to ( SC2 ), for different barrier heights

v

in the double-well experiment (mean

\pm

std over 5 runs, using average

\alpha

-schedule, 1e-4 learning rate, and 2000 training steps).

Training step	200	400	600	800	1000
Using ( SC1 )	0.12	0.078	0.067	0.052	0.046
Using ( SC2 )	0.15	0.088	0.070	0.061	0.055

Appendix

Table 10: Convergence speed of training objectives ( SC1 ) and ( SC2 ) in the Ornstein-Uhlenbeck example (KL divergence to ground-truth, mean over 5 runs).

Figure 10: Connections between self-consistency properties for stochastic optimal control, showing the integrand $\Phi(X^{u})$ such that ${u(s,X_{s}^{u})={\mathbb{E}}[\Phi(X^{u})|X_{s}^{u}=x]}$ . (Left) The expressions contain the Jacobian process $J_{t|s}$ for the base dynamics $X$ , analogous to ( SC1 ) from the main paper. (Middle) Generalised versions, containing the Jacobian $\tilde{J}_{t|s}$ for an auxiliary process $\tilde{X}$ . These are analogous to the generalised self-consistency property in ( SC-gen ). (Right) The expressions take $\tilde{J}_{t|s}={\mathrm{Id}}$ , analogous to ( SC2 ) from the main paper. Top line: When restricted to the terminal time $t=T$ , these self-consistency relations recover recently introduced objectives in the SOC and generative modelling literature ( Domingo-Enrich et al., 2025 ; Domingo-Enrich, 2024 ) . Bottom Line: In the case of zero running costs, the SOC self-consistency properties coincide with those presented in the main paper, which form the basis of our method for learning diffusion bridge dynamics.

Explore similar work

Analytic Bridge Diffusions for Controlled Path Generation
May 3, 2026Michael ChertkovStochastic Optimal Control Schrödinger Bridges
Program in Applied Mathematics & Department of Mathematics University of Arizona, AZ 85721, USA
Data-to-Energy Stochastic Dynamics
Sep 30, 2025Kirill Tamogashev, Esmeralda S. WhitammerSchrödinger Bridges Diffusion Models
University of Edinburgh
Resolving Endpoint Underfitting in Diffusion Bridges via Noise Alignment
May 27, 2026Yurong Gao, Zicheng Zhang, Congying Han +2Diffusion Models for Image Restoration Image-to-Image Translation
University of Chinese Academy of Sciences

May 3, 2026J/K move · Enter open · S save

Analytic Bridge Diffusions for Controlled Path Generation

Michael Chertkov

Program in Applied Mathematics & Department of Mathematics University of Arizona, AZ 85721, USA

Stochastic Optimal Control Schrödinger Bridges Optimal Control Optimal Transport

Open paper AskPDF arXiv

Most modern bridge-diffusion methods achieve finite-time transport by specifying an interpolation, Schrödinger-bridge, or stochastic-control objective and then learning the associated score or drift field with a neural network. In contrast, we identify a restricted but sufficiently broad and analytically solvable class in which the score, intermediate marginals, and protocol gradients are available in closed form without inner stochastic simulation loops and without neural networks in the optimization loop. We recast the classical linear--quadratic--Gaussian (LQG) stochastic-control structure as a transport problem of the Path Integral Diffusion (PID) type. In classical LQG control, linear dynamics, Gaussian noise, and quadratic costs lead to Riccati equations and closed-form optimal feedback. In LQ-GM-PID, we retain the linear--quadratic stochastic-control backbone, but replace terminal state regulation by a prescribed terminal probability density and allow both the initial and terminal laws to be Gaussian Mixtures (GM). Moreover, LQ-GM-PID turns bridge diffusion from a tool for terminal target matching alone into a tool for path shaping. We demonstrate this on a 2D corridor task, a 2D multi-entrance transport task, and a high-dimensional scaling study with

d=32

and

M=16

Gaussian-mixture terminal modes, all with sub-50,ms analytic precompute on a laptop. We position LQ-GM-PID as an analytically solvable reference model for the state-of-the-art neural bridge-diffusion and generative-transport methods: a controlled setting in which neural approximations, score estimates, path-shaping objectives, and protocol-learning procedures can be tested against exact quantities.

Abstract

Explore similar work

Analytic Bridge Diffusions for Controlled Path Generation

Data-to-Energy Stochastic Dynamics

Resolving Endpoint Underfitting in Diffusion Bridges via Noise Alignment