UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Authors: Kaizhen Zhu
Mokai Pan
Yuexin Ma
Yanwei Fu
Jingyi Yu
Jingya Wang
Ye Shi
Venue: ICML 2025 spotlightposter (2025)
Keywords: Diffusion bridge
Doob's h-transform
Stochastic optimal control
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We present UniDB, a unified diffusion bridge framework using stochastic optimal control, significantly improving detail preservation and image quality in generative tasks with minimal code modifications.

Abstract: Recent advances in diffusion bridge models leverage Doob’s $h$-transform to establish fixed endpoints between distributions, demonstrating promising results in image translation and restoration tasks. However, these approaches frequently produce blurred or excessively smoothed image details and lack a comprehensive theoretical foundation to explain these shortcomings. To address these limitations, we propose UniDB, a unified framework for diffusion bridges based on Stochastic Optimal Control (SOC). UniDB formulates the problem through an SOC-based optimization and derives a closed-form solution for the optimal controller, thereby unifying and generalizing existing diffusion bridge models. We demonstrate that existing diffusion bridges employing Doob’s $h$-transform constitute a special case of our framework, emerging when the terminal penalty coefficient in the SOC cost function tends to infinity. By incorporating a tunable terminal penalty coefficient, UniDB achieves an optimal balance between control costs and terminal penalties, substantially improving detail preservation and output quality. Notably, UniDB seamlessly integrates with existing diffusion bridge models, requiring only minimal code modifications. Extensive experiments across diverse image restoration tasks validate the superiority and adaptability of the proposed framework. Our code is available at https://github.com/UniDB-SOC/UniDB/.

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Kaizhen Zhu
1
Mokai Pan 12*
Yuexin Ma 12 Yanwei Fu
Jingyi Yu 12 Jingya Wang" Ye Shi 1
Abstract
ing (Ho et al., 2020; Kawar al., 2022; Song al., 2020; Xia et al., 2023; Li et al., 2023), imitation learning (Wu et al., 2024; Chi et al., 2023; Ze et al., 2024) and reinforcement learning (Yang et al., 2023; Ding et al., 2024a), etc. Despite its versatility, the standard diffusion model faces limitations in transitioning between arbitrary distributions due to its inherent assumption of Gaussian noise prior. To over- come this problem, diffusion models (Dhariwal & Nichol, 2021; Ho & Salimans, 2022; Murata al., 2023; Ding et al., 2024b; Chung et al., 2022; Tang et al., 2024) often rely on meticulously designed conditioning mechanisms and clas- sifier/loss guidance to facilitate conditional sampling and ensure output alignment with a target distribution. However, these methods can cumbersome and may introduce man- ifold deviations during the sampling process. Mcanwhile, Diffusion Schrödinger Bridge (Shi et al., 2024; De Bortoli et al., 2021; Somnath et al., 2023) involves constraints that hinder direct optimization of the KI. divergence, resulting in slow convergence and limited model fitting capability.
Recent advances in diffusion bridge models leverage Doob's f--transform to establish fixed endpoints between distributions, demonstrating promising results in image translation and restora- tion tasks. However, these approaches frequently produce blurred or excessively smoothed image details and lack a comprehensive theoretical foun- dation to explain these shortcomings. To ad- dress these limitations, we propose UniDB, uni- fied framework for diffusion bridges based on Stochastic Optimal Control (SOC). UniDB for- mulates the problem through an SOC-based opti- mization and derives a closed-form solution for the optimal controller, thereby unifying and gen- eralizing existing diffusion bridge models. We demonstrate that existing diffusion bridges em- ploying Doob's h-transform constitute a special case of our framework, emerging when the termi- nal penalty coefficient in the SOC cost function tends to infinity. By incorporating a tunable ter- minal penalty coefficient, UniDB achieves an op- timal balance between control costs and terminal penalties, substantially improving detail preserva- tion and output quality. Notably, UniDB seam- lessly integrates with existing diffusion bridge models, requiring only minimal code modifica- tions. Extensive experiments across diverse im- age restoration tasks validate the superiority and adaptability of the proposed framework. Our code is available at https: / //tp://ggtttttt//// com/ UniDB-SOCa SOCUUUDU/UU////
To address this challenge, DDBMs (Zheng et al., 2024) pro- posed a diffusion bridge model using Doob's 's h-transform. This framework is specifically designed establish fixed endpoints between two distinct distributions by learning the score function of the diffusion bridge from data, and then solving the stochastic differential equation (SDE) based on these learned scores to transition from one endpoint distribu- tion to another. However, the forward SDE in DDBMs lacks the mean information of the terminal distribution, which restricts the quality of the generated images, par- ticularly in image restoration tasks. Subsequently, GOUB (Yue et al., 2023) extends this framework by integrating Doob's h-transform with mean-reverting SDE, achieving better results compared to DDBMs. Despite the promising results in diffusion bridge with Doob's 's h-transform, two fun- damental challenges persist: 1) the theoretical mechanisms by which Doob's 's h-transform governs the bridging process remain poorly understood, lacking rigorous framework to unify its empirical success; and 2) while effective for global distribution alignment, existing methods frequently degrade high-frequency details-such as sharp edges and fine tex- tures cures-resulting in outputs with blurred or oversmoothed artifacts that compromise perceptual fidelity. These limi- tations underscore the need for both theoretical grounding and enhanced detail preservation in diffusion bridges.
1. Introduction
The diffusion model has been extensively utilized across range of applications, including image generation and edit-
'Equal contribution 'ShanghaiTech University "MoE Key Laboratory of Intelligent Perception and Mlahan-Machine Col- laboration 'Fudan University. Correspondence to: Ye Shi < shanghaitcch..du..n..
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Diffusion Bridge
Low- Low-Quality (LQ)
(HQ)
æ00
2T
Diffusion Bridge via Doob's h-transform is special case when y - co in UniDB!
Diffusion Bridge via Doob's ' h-transform
UniDB
min ["] Us,y Cur
dx; [fix + ham + 9 (x,,,,, [) dt Ii dw;
ue, [et + 2 |xx - xr
2
2
bxx,,t,x,,11 Vx. logp(xr x2)
s.Ldx; (JiXe + ham gels, )dl gedwe, x xo
Ground Truth
Doob's h- trtransform
UniDB
Ground Truth Doob's h-transform
UniDB
Figure 1. Recent advances in diffusion bridge models leverage Doob's f h-traasformm to establish fixed endpoints between distributions, which introduces function into the forward process of standard stochastic differential equation (SDE) to forcibly match the endpoints. However, as shown in the figure, this method can lead to local blurring and distortion in the generated images. UniDB formulates the forward process trajectories and target endpoint matching, to produce images with more realistic details. We also find that Doob's h-transform is special case in our framework when y + oc. Therefore, our framework can scamlessly integrate with existing diffusion bridge models (with Doob's 's /s-transform).
stochastic optimal control problem and employs the penalty coefficient 7, balancing realistic SDE
In this paper, we revisit the diffusion bridges through the lens of stochastic optimal control (SOC) by introducing a novel framework called UniDB, which formulates an opti- mization problem based on SOC principles to implement dif- fusion bridges. It enables the derivation lution for the optimal controller, along with the correspond- ing training objective for the diffusion bridge. UniDB identi- fies Doob's 's h-trranffrmm as special case when the terminal penalty coefficient in the SOC cost function approaches in- finity. This explains why Doob's I--transform may result in suboptimal solutions with blurred or distorted details. To ad- dress this limitation, UniDB utilizes the penalty coefficient SOC to adjust the expressiveness the image details and enhance the authenticity of the generated outputs. Our main contributions are follows:
We derive closed-form solutions for the SOC problem, demonstrating that Doob's f--transform is merely a special case within UniDB when the terminal penalty coefficient in the SOC cost function approaches infinity. This insight reveals inherent limitations the existing diffusion bridge approaches, which UniDB overcomes. Notably, the improvement of UniDB requires minimal code modification, ensuring easy implementation.
closed-form so-
UniDB achieves state-of-the-art results in various image restoration tasks, including super-resolution (DIV2K), inpainting (CelebA-HQ), and deraining (Rain which highlights the framework's 's superior image quality and adaptability across diverse scenarios.
2. Related Work
We introduce UniDB, novel unified diffusion bridge framework based on stochastic optimal control. This framework generalizes existing diffusion bridge mod- els like DDBMs and GOUB, offering a comprehensive understanding and extension of Doob's h transform by incorporating general forward SDE forms.
Diffusion with Guidance. This technique tackles condi- tional generative tasks by leveraging a differentiable loss function for guidance without the need for additional train- ing (Chung et al., 2022; Shenoy et al., 2024; Bradley & Nakkiran, 2024). However, it often yields suboptimal im- age quality and a prolonged sampling process due to the

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
necessity of small step sizes. Most importantly, the sam- pling process is prone to manifold deviations and detail losses (Yang et al., 2024). Furthermore, enhancing the guid- ance of the diffusion typically requires the introduction of additional modules, thereby increasing the model's compu- tational complexity.
costs and terminal constraints, enhancing detail preservation and image quality.
3. Preliminaries
3.1. Denoising Diffusion Bridge Models
Diffusion Schrödinger Bridge. This approach aims de- termine stochastic process, x² that facilitates probabilistic transport between a given initial distribution P Priurr and terminal distribution Paata (Liu et al., 2023; Shi et al., 2024; De Bortoli et al., 2021) while minimizing the Kullback- Leibler (KL) divergence. However, its training process is usually intricate, involving constraints that hinder direct optimization of the KI. divergence, resulting slow con- vergence and limited model fitting capability. For instance, DSB (Somnath al., 2023) requires two independent for- ward passes during training to obtain the target distribution, thereby increasing both the complexity and time cost of training.
Starting with an initial d-dimensional data distribution Xo ~ Gdara (x), diffusion models (Song et al., 2020; Ho et al., 2020; Sohl-Dickstenn et al., 2015; Song & Ermon, 2019) construct diffusion process, which can be achieved by defining forward stochastic process evolving from X through stochastic differential equation (SDE):
dx,- f(x,, t)dt gt. dw.,
(1)
where / ranges over the interval [0, 7). f : Rd x T] - Rd is the vector-valued drift function, g : [0. 1] + R signi- fies the ccalar-valued diffusion coefficient and w, R is the Wiener process, also known as Brownian motion. To promise the transition probability p(x; x,) remains Gaus- sian, almost all the diffusion SDEs take the following linear form (Zheng et al., 2024) in (1):
Diffusion Bridge with Doob's h-transform. Recent ad- vances in diffusion bridging have demonstrated the efficacy of Doob's h-transform in enhancing transition quality be- tween arbitrary distributions. Notably, DDBMs (Zhou et al., 2023) pioneered this approach by employing a linear SDE combined with Doob's h- d-transform to construct direct dif- fusion bridges. Subsequently, GOUB (Yue et al., 2023) extends this framework by integrating Doob's fi-transform with a mean- reverting SDE, achieving sfate-of-the-art per- formance in image restoration tasks. Despite these empirical successes, the theoretical foundations of Doob's h-transform in this context remain insufficiently explored. In addition, these methods often result in images with blurred or over- smoothed features, particularly affecting the capture of high- frequency details crucial for perceptual fidelity.
f(x,, f) - f(1)x,,
(2)
where f(t) is some scalar-valued function. To realize tran- sition between arbitrary distributions, DDBMs introduces Doob's h-transformm (Särkkä & Solin, 2019), a mathematical technique applied to stochastic processes, which rectifies the drift term of the forward diffusion process to pass through a preset terminal point xT E Rd. Precisely, the forward pro- cess of diffusion bridges after Doob's h- h--rransform becomes:
dx, = [[((x,)) () + gf (xxxx 1, XT. 7)] de + gt. dw,, (3)
where h(x,, f, , T) = Xx, log p(xp(x xx) is the h func- tion. The diffusion bridge can connect the initial xo to any given terminal xy and thus is promising for various image restoration tasks. Meanwhile, its backward reverse SDE (Anderson, 1982) is given by
Diffusion with Stochastic Optimal Control. The integra- tion of SOC principles into diffusion models has emerged as promising paradigm for guiding distribution transi- tions. DIS (Berner et al., 2022) established a foundational theoretical linkage between diffusion processes and SOC, while RB-Modulation (Rout et al., 2024) operationalized SOC via a simplified SDE structure for training-free free style transfer using pre-trained diffusion models. Close to our work, DBFS (Park et al., 2024) leveraged SOC construct diffusion bridges in and also established equivalence between SOC and Doob's s h--transfomm. However, DBFS primarily extends Doob's ss h-transform to infinite Hilbert spaces via SOC, without ad- dressing its intrinsic limitations. Our analysis reveals a critical insight: Doob's h-transform corresponds a sub- optimal solution that can inherently lead to artifacts such as blurred or distorted details. To resolve this, we introduce unified SOC framework that jointly optimizes trajectory
dx. = =(((x,))) 1) + + Vx, log p(XT | xx)
(4)
gt Vx, log p(x: xT) dt + ge dw..
where w, is the reverse-time Wiener process and the un- known term Vx, log p(x; xT) can be estimated by score prediction neural network 8o (Song et al., 2020).
function spaces
3.2. Generalized Ornstein-Uhlenber Bridge
Generalized Ornstein- Uhleen--hhlebeck (GOU) process describes mean-reverting stochastic process commonly used in fi- nance, physics, and other fields in the following SDE form (Ahmad, 1988; Pavliotis & Pavliotis, 2014; Wang et al., 2018):
dx, = Or (n xx) dt + ge dw.,
(5)

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
where js is given state vector, O, denotes scalar drift co- efficient and Jt. represents the diffusion coefficient with 0z. It satisfying the specified relationship o² 2x²0, where X is given constant scalar. Based on this, Generalized Ornteen--Uhlenbeck Bridge (GOUB) is a diffusion bridge model (Yue et al., 2023), which can address image restora- tion tasks without the need for specific prior knowledge if we consider the initial state X6 to represent high-quality image and the corresponding low-quality image XT = 1a as the final condition. With the introduction of μ,, X tends to μ as time / progresses. Through Doob's s h--rransfomm, denote g:::= = - - 0_d2, 8e- 0o¹ 0_dz for simplification when 0 and o:. = A² e -20tt). the forward process of GOUB
4. Methods
4.1. Diffusion Bridges Constructed by SOC Problem
The forward SDE of the Diffusion Bridge with Doob's h- transform is enforced to pass from the predetermined origin Ze to the terminal æy.. With a similar purpose, UniDB con- arbitrary
structs SOC problem where the constraints linear SDE of the forward diffusion with a given initial state, while the objective incorporates a penalty term steering the forward diffusion trajectory towards the predetermined ter- minal ZT. Meanwhile, compared with the linear drift term (2), we combined given state vector term m with the same dimension as x, and its related coefficient hy which is a simple reformulation and generalization of the parameters in GOU process (5):
is formed as:
dx,= = -00 + - 201r
(XT - (xx) df + Ir dw,.
(6)
f(x,, () = Six, + h,m.
(10)
And the forward transition p(xx| xx.xx] is given by
Accordingly, our SOC problem with unified linear SDE (10) is formed as:
p(x: xo, xT) = x=)=))=NN =N( 1121,
1 |lu 112 dt 2 ||x xT 2212
2
8,
TT
Xo + (1 e 8,
LT
XT, o²²
min IE ue,-
of
(7)
s.L.xx, = ( + ham + gour, )dl dll + godwa, x5 = 20-
Also, GOUB presents a new reverse ODE called ODE, which directly neglects the Brownian term of (4):
(11)
According to the certainty equivalence principle (Chen et al., 2023; Rout et al., 2024), the addition of noise or perturba- tions to a linear system with quadratic costs does not change the optimal control. Therefore, we can modify the SOC problem with the deterministic ODE condition to obtain the optimal controller ui,,y follows,
dx,- = f(xx,)) t) + of V xs p(xr xx)
- - logp(x: xT) dl.
3.3. Stochastic Optimal Control
Stochastic Optimal Control (SOC) a mathematical dis- cipline that focuses on determining optimal control strate- gies for dynamic systems under uncertainty. By integrating stochastic processes with optimization theory, SOC seeks to identify the best control strategies in scenarios involving randomness, as commonly encountered in fields like finance (Geering et al., 2010) and style transfer (Rout et al., 2024). Considering the dynamics described in (I), let us exam- ine the following Linear Quadratic SOC problem (Bryson, 2018; O' O'Connell, 2003; Kappen, 2008; Chen al., 2023):
Eel 2 | 12 dt + 2 || xT12 2
min 1'
(12)
s.t. dx: = (fix: + h.m + gr dt, x = Z0.
We can derive the closed-form solution to the problem (12), which leads to the following Theorem 4.1:
Theorem 4.1. Consider the SOC problem (12), denote Se fadz, hat =
de 7-1 + e²IT e foa Id ---dz and g2. = - 2e-25ng2= denote Fe, hy and oe for simplification when = O, then the closed-form optimal controller ui,y" is
[ | dt
2 ||x xrl2 2
min E
M1, El EM
(9)
s.t. dx, (f(xi,t) t) + g.u.,)) dt ae dw,, xi To.
- = gvr DT e met" herr
(13)
di,1
where x² is the diffusion process under control, zy and Fr represent for the initial state and the preset terminal respec-
and the transition of x; from To and xrr is
tively, ||ue, 1| is the instantancous cost, 3 ||* 2T 11 is the terminal cost with its penalty coefficient 7. The SOC problem aims to design the controller U,, to drive the dy- namic system from Zo to Fr with minimum cost.
X c/
den xo +
gr xT + (he
24r borgo
m
do.7
du
du,7
(14)

---

UniDB: A Unified Diffusson Bridge Framework via Stochastic Optimal Control
The proof of Theorem 4.1 is provided in Appendix A.I. With Theorem 4.1, we can obtain an optimally controlled forward SDE connected from To to the neighborhood of the terminal FT and the transition of X, for the forward process. As for the backward process, similar to (4) and (8), the backward reverse SDE and Mean-ODE are respectively formulated as:
7 x)), which may disrupt the inherent continuity and smoothness of images. Prioritizing pixel-perfect endpoints over smooth trajectories leads to "mathematically correct but visually unrealistic" outputs. As shown in Figure 1, Doob's h-transform can lead to artifacts along edges and unnatural patterns in smooth regions. Therefore, maintain- ing the penalty coefficient y as a hyperparameter is a more effective approach.
dx,- = [ix: + ham + gell,
4.3. Training objective of UniDB
Vx. log ΦT) dt + gedwr,
(15)
In this section, we focus on constructing the training ob- jective of UniDB. According to maximum lig-likelihood (Ho et al., 2020) and conditional score matching (Song et al., 2020), the training objective is based on the forward transition p(x, XT). Thus, we begin by deriving this probability. The closed-form expression in (14) represents the mean value of the forward transition after applying repa- rameterization techniques. However, this expression lacks a noise component after the transformation based on the certainty equivalence principle. To address this issue, we employ stochastic interpolant theory (Albergo et al., 2023) to introduce noise term a(e with oo- = o O. We define similar to (7), leading to the following forward transition:
dx; = [ + hym + gont - ge Vx, logp(x (T))
dt.
(16)
4.2. Connections between SOC and Doob's h-transform
We can intuitively see from the SOC problem that when 7 + oc in Theorem 4.1, it means that the target of SDE process is precisely the predetermined endpoint (Chen et al., 2023), which is also the purpose of Doob's fi-transform and facilitates the following theorem:
Theorem 4.2. For the SOC problem (12), when - oo, the optimal controller becomes U, = 9t Vx, log p(XT | xx), and the corresponding forward and backward SDE with the linear SDE form (10) are the same as Doob's 's h-transform as in (3) and (4).
of2
p(xe
2o. ZT) N(
P = oe
den xo do.7
+
etr
+ (h.
h )m)
2 hrgh
da.7
dor7
The proof of Theorem 4.2 is presented in Appendix A.2. This theorem shows that existing diffusion bridge models using Doob's h-transform are merely special instances of our UniDB framework, which offers unified approach to diffusion bridges through the lens of SOC.
070². o
of
(18)
The derailed derivation is provided in Appendix A.4. Sim- ilar to (Yue et al., 2023) using the 1 loss form to bring improved visual quality and details the pixel level (Boyd, 2004; Hastie et al., 2017), we can derive the training ob- jective. Denote Oe,; efi de,y. assuming 14-1,0- of-1. and 1 of-1, are respectively the mean values and variances of po( xT) and p(x,-1 | Xo, Xx, (T), suppose the score Xx logp(x; IT) is parameterized as Co (X,, IT, t)/6', the final training objective is as follows,
Furthermore, using Doob's l--transform in diffusion bridge models is not necessarily optimal, as letting the terminal penalty coefficient y + oo eliminates the consideration of control costs in SOC. To support this argument, we present Proposition 4.3, which asserts that the diffusion bridge with Doob's s h-transformm is not the most effective choice.
Proposition 4.3. Consider the SOC problem (12), denote J Juu 7) & S 2 || 112 dt + ||x zrll2 as the over- all cost of the system, ui as the optimal controller (13), then
1
La = X1,x0,XxxxT
||
200-1,0
Ju 7) S Jn 20).
(17)
g M4-1,0 = =x,- - fix, - ---m---- gi
(x,, 27,t), o,
Detailed proof of Proposition 4.3 is provided in Appendix A.3. Proposition 4.3 shows that finite y achieves a lower to- tal cost not by sacrificing performance, but by optimally trad- ing minor terminal mismatches for significantly smoother and more natural diffusion paths. Doob's /i-transform re- quires larger controller ||u"",0113 112 2 |12 in SDE trajec- tory to force exact endpoint matching (the controlled target is precisely the preset endpoint ||x - 2r122 = 0 when
11-1, = Pi-1a + of (x: - Pend, O1-1,0 = gt-
(19)
Please refer to Appendix A.5 for detailed derivations. There- fore, we can recover or generate the origin image d.g through Euler sampling iterations. So far, we've 'v built the UniDB framework, which establishes and expands the forward and

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
backward process of the diffusion bridge model through SOC and comprises Doob's h-transform special case.
Algorithm 1 UniDB Training
repeat
Take a pair of images Xo z0 and xT DT t to Uniform (T)
4.4. UniDB unifies diffusion bridge models
1-1
Our UniDB is a unified framework for existing diffusion bridge models: DDBMs (VE) (Zhou et al., 2023), DDBMs (VP) (Zhou et al., 2023) and GOUB (Yue et al., 2023).
dt,7 e 6. -ip B7
f GOUB
Q1,7 = c 6: 7-' 1 +
Proposition 4.4. UniDB encompasses existing diffusion bridge models by employing different hyper-parameter spaces H as follows:
9- -1
UniDB-GOU
X, 0t + (1 -,,,7) 2T JIE
Pl, = 01,730 + (1 (t.7) IT
14- = X 0r + g c -2
DDBMs (VE) corresponds to UniDB with hyper-
(ET - xt)
नx.
parameter HyE(se- 0, h4 0, + 0o)
+ ot 872
Co (x,, IT.)) t) t GOUB
DDBMs (VP) corresponds to UniDB with hyper- parameter Hyr(ft= = - -20 h1 = 0, + 20)
14- X, 0; + 9 g --
(IT x,)
GOUB corresponds to UniDB with hyper-parameter Hoor(st = 0r, h4, = -,, Details of the proposition 4.4 are provided in Appendix A.6.
00, m - 7 + x))
- ED (x,, ST, 1) UniDB-GOU
ol
a (xe: (1e,5)
14- 1,7 3." a1-1 Take gradient descent step on VeLo until converged
4.5. An Example: UniDB-GOU
It is evident that these diffusion bridge models like DDBMs (VE), DDBMs (VP) and GOUB all based on Doob's h- transform are all special cases of UniDB with y + cc. However, according to Proposition 4.3, these models are not the effective choices. Therefore, we introduce UniDB based on the GOU process (5), hereafter referred to as UniDB-GOU which retains the penalty coefficient y as the hyper-panne parameter. Considering the SOC problem with GOU process (5), the optimally controlled forward SDE is:
Hence, only a few lines of code need to be adjusted to gen- erate more realistic images using the same training method. We provide pseudo-code Algorithm 1 and Algorithm 2 for the training and sampling process of UniDB-GOU, respec- tively. The two algorithms encapsulate the core methodolo- gies employed by our model to learn and explain how to restore HQ images from LQ images. Also, the red and the green parts highlight the main difference between UniDB and GOUB. Beyond the GOUB model, our UniDB frame- work can be similarly extended to other diffusion bridge models, such as DDBMs (VE) and DDBMs (VP). For de- tailed information UniDB- VEE and UniDB-VP, please refer to Appendix A.8.
6 20Gr
dx,=
01 + dh 7-1 + o.r
(øT xx) (20)
and the mean value of forward transition p(xx Ioo ØT) is
Pin = e 0. 1 1o
+ (1 (1-e 0 10 i:T
TT-
1 1+
1 7og
(21)
Algorithm 2 UniDB Sampling
Please refer to Appendix A.7 for detailed proof.
Input: Low- Low-Quality images Xy FT. for T to 1 do
Remark 1. It's worth noting that our UniDB model can be plugin module to the existing diffusion bridge with Doob's h-transform. Taking UniDB-GOU as an example, we highlight the key difference between UniDB-GOU and GOUB (the coefficient of zo in the mean value of forward transition and h-ffuction term) as follows:
z ~ (0. 1) if > 1, clse = 0
X,-1 = X;
26,TT (6: + g o
(IT - xx)
+ a/²² ca (x, TT, t) - ge2 + GOUB
of
0, + gt e 26,
7
X4-1 = x, -
(xT - xt)
7- + o7.
2 (TT-
x²)
gie 20.T (ET x1)
o. (x,, xr. 1) - giz +- UniDB-GOU
9th
U
7-1 + +
end for Return High-Qualit images Xo
GOUB
UniDB-GOU
(22)

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Partial enlarged
Partial enlarged
Partial enlarged
Partial enlarged
GT
GOUB
UniDB
Figure Qualitative comparison of visual results between GOUB (SDE) and UniDB (SDE) on DIV2K with zoomed-in image local regions (UniDB hased on GOU process).
Building upon equations (20) and (21), we further present proposition characterize how the penalty coefficient y affects the controlled terminal distribution as follows:
2018) and Fréchet Inception Distance (FID, lower better) (Heusel et al., 2017). For simple expressions in the follow- ing sections, UniDB (SDE) and UniDB (ODE) are applied to represent the UniDB-GOU with reverse SDE and reverse Mean-ODE, respectively. Please refer to Appendix B and C for all related implementation details and more experiment results, respectively.
Proposition 4.5. Denote the initial state distribution the terminal distribution x by the controller and the pre- defined terminal distribution FT, then
20,
|| x 11 (1+ - e-20z))))))
IT xo 1 (23)
Distance between target and control distribuutio
Impact of Penalty Weight Distribution Distance
10a
10
The detailed derivations of proposition 4.5 are provided in Appendix A.9. Notably, as approaches infinity, the control terminal converges to the predefined endpoint. However, as analyzed in Proposition 4.3, this can result in suboptimal outcomes with blurry or overly smoothed image details. To address this, it is crucial to balance the control cost and terminal term by selecting the value of y. In the follow-
10-7
10-
Inpainting
10-9
Deraining
Super Resolution
10-12
ing section, we will present comprehensive experiments to evaluate the impact of different values on the results.
10a
10a
10a
10'
10"
Penalty confficient
5. Experiments
Figure 4. The distances between target and controlled terminal distributions for different datasets (CelebA-HQ. Rain100H, and DIV2K) with diffferent penalty coefficients 7. The red shaded area and blue dotted line highlight choice of of
In this section, we evaluate our models in image restora- tion tasks including Image 4x Supee--resolltionn Image De- raining, and Image Inpainting. We take four evaluation metrics: Peak Signal-to-Noise Ratio (PSNR, higher is bet- ter) (Fardo et al., 2016), Structural Similarity Index (SSIM, higher is better) (Wang et al., 2004), Learned Perceptual Im- age Patch Similarity (LPIPS, lower is better) (Zhang al.,
5.1. Experiments Setup
According to Proposition 4.5, we first quantitatively analyze the Io-m distances between the two terminal distributions depicted in Figure 4. We computed the average distances

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Table Qualitative comparison with the relevant baselines on DIV2K, Rain100H, and CelebA-HQ 256 datasets.
Image Super-Resolution
Image Deraining
Image Inpainting
METHOD
METHOD
METHOD
PSNR*
SSIM
LPIPS,
FID.
PSNR
SSIM
LPIPS
FID.
PSNR
SSIM
LPIPS,
FID.
Bicubic
26.70
0.774
0.425
36.18
MAXIM
30.81
0,902
0.133
58.72
PromptI
30.22
0.918
0.068
32.69
DDRM
24.35
0.592
0.364
78.71
MHHNet
31.08
0.899
0.126
57.93
DDRM
27.16
0.899
0.089
37.02
IR-SDE
25.90
0.657
0.231
45.36
IR-SDE
31.65
0,904
0.047
18.64
IR-SDE
28.37
0.916
0.046
25.13
GOUB (SDE)
26.89
0.7478
0.220
20.85
GOUB (SDE)
31.96
0.9028
0.046
18.14
GOUB (SDE)
28.98
0.9067
0.037
4.30
GOUB (ODE)
28.50
0.8070
0.328
22. 14
GOUB (ODE)
34.56
0.9414
0.077
32.83
GOUB (ODE)
31.39
0.9392
0.052
12.24
UniDB (SDE)
25.46
0.6856
0.179
16.21
UniDB (SDE)
32.05
0,9036
0.045
17.65
UniDB (SDE)
29.20
0.9077
0.036
4.08
UniDB (ODE)
28.64
0.8072
0.323
22.32
UniDB (ODE)
34.68
0.9426
0.074
31.16
UNiDB (ODE)
31.67
0.9395
0.052
11.98
Deraining
Inpainting
GT
GT
GOUB
GOUB
UniDB-GOU
UniDB-GOU
Figure 3. Qualitative comparison of visual results between GOUB (SDE) and UniDB (SDE) the Rain 100H dataset on Image Deraining (Left) and CelebA-HQ dataset on Image Inpainting (Right) with zoomed-in image local regions (UniDB based on GOU process).
between high-quality and low-quality images in the three datasets (CelebA-HQ. Rain 100, and DIV2K) related to the subsequent experimental section as the distances - l in (23). As can be seen, for all three datasets, these distances remain relatively small, ranging from 10-4 to 10-10 when y is within the range of 1 x 10 subsequent experiments will focus on the 7 of this range further investigate the performance of UniDB-GOW
and Figure 2. Visually, our proposed model demonstrates significant improvement over the baseline across various metrics. It also excels by delivering superior performance in both visual quality and detail compared to other results.
Image Deraining Tasks. For image deraining tasks, we conducted the experiments based on Rain 10H datasets (Yang al., 2017). Particularly, to be consistent with other deraining models (Ren et al., 2019; Zamir et al., 2021; Luo et al., 2023; Yue et al., 2023), PSNR and SSIM scores on the Y channel (YCbCr space) are selected instead of the origin PSNR and SSIM. MAXIM (Tu et al., 2022), MHNet (Gao et al., 2025), IR-SDE (Luo et al., 2023), GOUB (SDE) (Yue et al., 2023) and GOUB (ODE) (Yue al., 2023) are chosen as the baselines. The relevant experimental results are shown in the Table 1 and Figure 3. Similarly, achieved state-of-the-art results in the deraining task. Vi- sually, it can also be observed that our model excels in capturing details such as the eyebrows, eye bags, and lips.
x 10 Therefore, our
5.2. Experimental Details
Image 4x SuSuep-Resouutiee Tasks. In super-resolution, we evaluated our models based on DIV2K dataset (Agustsson & Timofte, 2017), which contains 2K-resolution high-quality images. During the experiment, all low-resolution images were 4x bicubic upscaling to the same image size as the paired high-resolution images. For comparison, we choose Bicubic interpolantion (Kawar et al., 2022), DDRM (Kawar al., 2022). IR-SDE (Luo et al., 2023), GOUB (SDE) (Yue al., 2023) and GOUB (Mean-ODE) (Yue al., 2023) fol- lowing abbreviated as GOUB (ODE) the baselines. The qualitative and quantitative results are illustrated in Table
model
Image Inpainting Tasks. In image inpainting tasks, we evaluated our methods on CelebA-HQ 256x256 datasets

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Table Quantitative evaluation results for DIV2K, CelehA-HQ and Rain of UniDB-GOU with different penalty coefficients ?.
Different y
TASKS
METRICS
5 10
x 106
107
10
00
PSNR
24.94
24.72
25.46
25.06
26.89
SSIM
0.6419
0.6587
0.6856
0.6393
0.7478
Image 4xSuper-Resolution
LPIPS
0.234
0.199
0.179
0.289
0.220
FID
20.33
18.37
16.21
23.76
20.85
PSNR
28.73
29.15
29.20
28.65
28.98
SSIM
0,9065
0.9068
0.9077
0.9062
0,9067
Image Inpainting
LPIPS
0.038
0.036
0.036
0.039
0.037
FID
4.49
4.12
4.08
4.64
4.30
PSNR
29.44
31.96
32.00
32.05
31.96
SSIM
0.8715
0.9018
0,9029
0,9036
0,9028
Image Deraining
LPIPS
0.058
0.045
0.046
0.045
0.046
FID
24.96
18.37
17.87
17.65
18.14
(Karras, 2017). For comparison, we choose DDRM (Kawar et al., 2022), PrompuR (Potlapalli et al., 2023), IR-SDE (Luo et al., 2023), GOUB (SDE) (Yue et al., 2023) and GOUB (ODE) (Yue et al., 2023) as the baselines. As for mask type, we take 100 thin masks consistent with the baselines. The relevant experimental results are shown in Table and Figure 3. It is observed that model achieved totte-of-thc-rt results in all indicators and also delivered highly competitive outcomes on other metrics. From visual perspective, our model excels in capturing details such as faces, eyes, chins, and noses.
sion bridge models with Doob's h-transform like DDBMs and GOUB. Moreover, we demonstrate that the diffusion bridge with Doob's h--transfomm can be viewed as a specific case within UniDB when the terminal penalty coefficient approaches infinity. This insight helps elucidate why Doob's h-transform may lead to suboptimal image restoration, often resulting in blurred distorted details. By simply adjusting this terminal penalty coefficient, UniDB achieves marked improvement in image quality with minimal code modi- fications. Our experimental results underscore UniDB's superiority and versatility across various image processing tasks, particularly enhancing image details for more real- istic outputs. Despite these advantages, UniDB, like other standard diffusion bridge models, faces the challenge of computationally intensive sampling processes, especially with high-resolution images or complex restoration tasks. Future work will focus on developing strategies to acceler- ate the sampling process, enhancing UniDB's practicality particularly for real-time applications.
5.3. Ablation Study
Penalty Coefficient 7. To evaluate the specific impact of different penalty coefficients on model performance, we conducted the experiments above with several different 7- The final results are shown in Table 2. The results across all tasks show that the choice of y significantly influences the model's performance on all tasks, different optimal y for different tasks, and our UniDB achieves the best perfor- mance in almost all metrics. Particularly in super-resolution tasks, we focus more on the significantly better perceptual scores (LPIPS and FID) (Luo et al., 2023), demonstrating that UniDB ensures to capture and preserve more intricate image details and features shown in Figure 2. These findings underscore the importance of carefully tuning y to achieve the best performance for specific tasks.
A Acknowledgement
This work was supported by National Natural Science Foun- dation of China (62303319, 62406195), Shanghai Local College Capacity Building Program (23010503100), Shang- haiTech A14S Initiative SHTAI4S202404, HPC Platform of ShanghaiTech University, Core Facility Platform of Com- puter Science and Communication of ShanghaïTech Univer- sity, and MoE Key Laboratory of Intelligent Perception and Machine Collaboration ShaangaiiTech University). and Shanghai Engineering Research Center of Intelligent Vision and Imaging.
6. Conclusion
In this paper, we presented UniDB, a unified diffusion bridge framework based on stochastic optimal control prin- ciples, offering novel perspective diffusion bridges. Through this framework, we unify and extend existing diffu-

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Impact Statement
Ding, S., Hu, K., Zhang, Z., Ren, K., Zhang, W., Yu, J., Wang, J., and Shi, Y. Diffusion-based reinforcement learning via q-weighted variational policy optimization. In The Thirty-eighti Annual Conference on Neural Infor- mation Processing Systems, 2024a.
There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Ding, X., Wang, Y., Zhang, K., and Wang, Z. J. Ccdm: Continuous conditional diffusion models for image gen- eration. arXiv preprint 2XXi:24405003544 2024b.
Agustsson, and Timofte, R. Ntire 2017 challenge on single image super-resolution: Dataset and study. In the IEEE conference on computer vision and pallern recognition workshops, pp. 126-135, 2017.
Proceedings the
Fardo, A.. Conforto, V. H., de Oliveira, F. C.,, and Ro- drigues, P. S. A formal evaluation of psnr as quality mea- surement parameter for image segmentation algorithms. arXiv preprint arXiv: 1exxiv:6005.00 06006005500111 16, 2016.
Ahmad, R. Introduction to stochastic differential equations, 1988.
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint urXiv:2303008797 2023.
Gao, H., Zhang, Y., Yang, J.. and Dang, D. Mixed hierarchy network for image restoration. Pattern Recognition, 161: 111313, 2025.
Anderson, B. D. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313- 326, 1982.
Geering, H. P., Herzog, F., and Dondi, G. Stochastic optimal control with applications financial engineering. Opti- mization and Optimal Control: Theory and Applications, pp. 375 175--08,440 2010.
Berner, J., Richter, L., and Ullrich, K. An optimal con- trol perspective on diffusion-based generative modeling. arXiv preprint (.xXix:2211.01366 2022.
Hastie, T., Tibshirani, R., and Friedman, J. The elements of statistical learning: data mining, inference, and predic- tion, 2017.
Boyd, S. Convex optimization. Cambridge UP, 2004.
Bradley, A. and Nakkiran, P. Classifier-free guidance is predictor-corrector. arXiv preprint wwXXi:2408.090000 2024.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Bryson, A. E. Applied optimal control: optimization, esti- mation and control. Routledge, 2018.
Ho, and Salimans, T. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12599, 2022.
Chen, T., Gu, J., Dinh, L., Theodorou, E. A., Susskind, J., and Zhai, S. Generative modeling with phase stochastic bridges. arXiv preprint arXiv:2310.07805, 2023.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 1116840-6851,, 2020.
Chi, C., Xu, Z. Feng, S., Cousineau, E., Du, Y., Burchfiel, B., Tedrake, R., and Song. S. Diffusion policy: Visuomo- tor policy learning via action diffusion. The International Journal of Robotics Research, pp. 02783649241273668, 2023.
Kappen, H. Stochastic optimal control theory. ICML, Helsinki, Radbound University, Nijmegen, Netherlands, 2008
Karras, T. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.1011 2017.
Chung, H., Kim, J., Mccann, M. T., Klasky, M. I.,, and Ye, J. C. Diffusion posterior sampling for general noisy in- verse problems. arXiv preprint arXiv:2209. 14687, 2022.
Kawar, B., Elad, M., Ermon, S., and Song, J. Denoising diffusion restoration models. Advances in Neurul Infor- mation Processing Systems, 35:23593-23606,, 2022.
De Bortoli, V., Thortoon, J., Heng, J., and Doucet, A. Diffu- sion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34: 14:17695-1770, 2021.
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6998, 2014.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794 2021.
Kirk, D. E. Optimal control theory: an introduction. Courier Corporation, 2004.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Levine, W. Optimal control theory: An introduction. IEEE Transactions on Automatic Control, 17(3):423-423 1972.
Särkkä, S. and Solin, A. Applied stochastic disferential equations, volume 10. Cambridge University Press, 2019.
doi: 10.1109/TAC.997. 1000008.
Shenoy, R., Pan, Z., Balakrishnan, K., Cheng, Q.. Jeon, Y., Yang. H., and Kim, J. Gradient-free classifier guid- ance for diffusion model sampling. arXiv preprint arXiv:2411.15393, 2024.
Li, X., Ren, Y., Jin, X., Lan, C., Wang, X., Zeng, W., Wang, X., and Chen, Z. Diffusion models for image restora- tion and enhancement-a comprehensive survey, arXiv preprint arXiv:2308.093.8 2023.
Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. Diffu- sion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E. A., Nie, W., and Anandkumar, A. 12sb: image-to-image schrödinger bridge. In Proceedings of the 40th Inter- national Conference on Machine Learning, pp. 22042- 22062, 2023.
Dib--Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
Luo, Z., Gustafsson, F. K., Zhao, Z., Sjölund, J., and Schön, T. B. Image restoration with mean-reverting stochastic differential equations. arXiv preprint arXiv:2301.11699 2023.
Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R.,
Krause, A., and Bunne, C. Aligned diffusion schrödinger bridges. In Uncertainty in Artificial Intelligence, pp. 1985-1995. PMLR, 2023.
Murata, N., Saito, K., Lai, C.-H., Takida, Y., Uesaka, T., Mitsufuji, Y., and Ermon, S. Gibbsddrm: A partially col- lapsed gibbs sampler for solving blind inverse problems with denoising diffusion restoration. International con- ference on machine learning, pp. 25501-25522. PMLR, 2023.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXi::201.133556 2020.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffu- sion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
O' O'Connell, Conditioned random walks and the rsk cor- respondence. Journal of Physics A: Mathematical and General, 36)12))3049 2003.
Tang, J., Wang. J., Ji, K., Xu, L., Yu, J., and Shi, Y. A uni- fied diffusion framework for scene-aware human motion estimation from sparse signals. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 21251-21262, 2024.
Park, B., Choi, J., Lim, S., and Lee, J. Stochastic optimal control for diffusion bridges in function spaces. arXiv preprint arXiv:2405.20630 2024.
Tu, Z., Talebi, H., Zhang, H., Yang, F.. Milanfar, P., Bovik, A., and Li, Y. Maxim: Multi-axis mlp for image pro- cessing. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 5769-5780, 2022.
Pavliotis, G. A. and Pavliotis, G. A. Introduction to stochas- tic differential equations. Stochastic Processes and Ap- plications: Diffusion Processes, the Fokker- Fokker-Planck and Langevin Equations, pp. 55 85, 2014.
Wang. W., Cai, Y., Ding. Z., and Gui, Z. A stochastic differential equation sis epidemic model incorporating ornstein-uhlenbeck process. Physica A: Statistical Me- chanics and its Applications, 509:921 93669 2018.
Potlapalli, V., Zamir, S. W., Khan, S., and Khan, F. S. Promptir: Prompting for all-in-one blind image restora- tion. arXiv preprint arXiv:2306. /3090, 2023.
Ren, D., Zuo, W., Hu, Q., Zhu, P., and Meng, D. Progressive image deraining networks: A better and simpler baseline. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 3937-3946 2019.
Wang. Z., Bovik, A. C., Sheikh, H. R., and Simoncelli, P. Image quality assessment: from error visibility struc- tural similarity. IEEE transactions image processing, 13(4):600-613 2004.
Rout, L., Chen, Y., Ruiz, N., Kumar, A., Caramanis, C., Shakkottai, S., and Chu, W.-S. Rb-modulation: Training- free personalization of diffusion models using stochastic optimal control. arXiv preprint arXiv:2405. /740/, 2024.
Wu, S., Zhu, Y., Huang, Y., Zhu, K., Gu, J., Yu, J., Shi, Y., and Wang. J. Afforddp: Generalizable diffu- sion policy with transferable affordance. arXiv preprint 1rx1v:2412.03142, 2024.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Xia, B., Zhang, Y., Wang, S., Wang, Y., Wu, X., Tian, Y., Yang, W., and Van Gool, L. Diffir: Efficient diffu- sion model for image restoration. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 13095-13105, 2023.
Yang, L., Huang, 7.,, Lei, F., Zhong, Y., Yang. Y., Fang, C., Wen, S., Zhou, B., and Lin, Z. Policy representation via diffusion probability model for reinforcement learning. arXiv preprini 2315:2305.1322 2023.
Yang, L.. Ding, S., Cai, Y., Yu, J., Wang, J., and Shi, Y. Guidance with spherical gaussian constraint for condi- tional diffusion. In International Conference on Machine Learning, 2024.
Yang, W., Tan, R. T., Feng, J., Liu, J., Guo, Z., and Yan, S. Deep joint rain detection and removal from single im- age. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 1357-1366, 2017.
Yue, C., Peng, Z., Ma. J., Du, S., Wei, P., and Zhang, D. Image restoration through generalized cunneei bridge. arXiv preprint urXiv:2312 2. 10299, 2023.
Zamir, S. W., Arora, A., Khan, S., Hayat, M., Khan, F. S., Yang, M.-H., and Shao. L. Multi-stage progressive image restoration. In Proceedings of the IEEE/CVF conference computer vision and paltern recognition, pp. 14821- 11831,00211
Ze. Zhang, G., Zhang, K., Hu, C., Wang, M., and Xu, H. 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations. In ICRA 2024 Workshop on 3D Visual Representations for Robot Ma- nipulation, 2024.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. The unreasonable effecctiveness of deep features perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
Zheng, K., He, G., Chen, J., Bao, F., and Zhu, J. Diffusion bridge implicit models. arXiv preprint arXiv:2405.1 15885, 2024.
Zhou, L., Lou, A., Khanna, S., and Ermon, S. Denoising dif- fusion bridge models. arXiv preprint 1699922309.16648 2023.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Appendix Contents
Appendix A: Proof
Page 13
A.1 Proof of Theorem 4.1
Page 13
A.2 Proof of Theorem 4.2
Page 15
A.3 Proof of Proposition 4.3
Page 15
A.4 Derivation of the transition probability (18)
Page 17
A.5 Derivation of the training objective (19)
Page 17
A.6 Proof of Proposition 4.4
Page 19
A.7 Derivation of UniDB-GOU
Page 20
A.8 Examples of UniDB-VE VE and UniDB-VP
Page 22
A.9 Proof of Proposition 4.5
Page 22
Appendix B: Implementation Details
Page 23
Appendix C: Additional Experimental Results
Page 24
A. Proof
A.1. Proof of Theorem 4.1
Theorem 4.1. Consider the SOC problem (12), denote de, - Gfrrg... f=== Sr fadz, hws= = e h and
-h -Ia -2f dz, denote St, hn and of for simplification when = 0, then the closed-form optimal controller u is
, geeerr 2T e Xt
me!r Trar
(13)
der
and the transition of xx from zo and vT
din
et ot do.7
2fx To
X, elt
xo +
ET
+
hh,
(14)
do,7
do,7
Proof. According to Pontryagin Maximum Principle (Levine, 1972; Kirk, 2004) recipe, one can construct the Hamiltonian:
(t, Xx, Me Pi)
|u, 112 pf (fex: + h.m + gru:).
(24)
By setting:
0H
=0
n -9
(25)
ou,y
Then the value function becomes
V min (t, X1, Pt: un,)) = H(t,, xx, Pr, u.,)) = 92 |Pell + fip!'x + h.pp m.
(26)
u
2
Now, according to minimum principle theorem to obtain the following set of differential equations:
dx;
Vp.H H (x,,
Pe, Pr. .) -gÊP. fixe ham.
(27)
dp: - H (xi, Pt, ui,0) 1)= =--- -p.fe,
(28)
Xo = To.
(29)
pr = 7 (XT 7xx
TT)
(30)

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Solving the Equation (28), have:
P: = Poe ft
(31)
Pr Poe fr.
Solve the Equation (27):
dx: fixt + h,m 9.Pr
dt
d(c ((x)))x)) = -e- - -c-
g.Pr. Pt.
dt
c 1x4 Xy = .mh+ - Pook:
1 "x1 20 =mh; Poği.
e
Hence, we can get:
Xy = ctr + me hr preetogh
(32)
and
= x0 + me -preheeeeee
(33)
Take the Equation (32) into the Equation (30) and solve pr.
Pr = e xo + me! preefr ZT)
(34)
7 (efrro+ metr - IT)
pr= = 1 7e2lrl o7.
(35)
Also, take the Equation (34) into the equation (33),
X; = ed' + me F. he - ef gi oz elr + me xT
x-1 + 21 of
(36)
de. (d
etrg? xo + (he
25r hr hrr
S.
+
m
do,7
do.y
do.7
Preserve ?
- DtPr
-90 Fr eJr x0 + mefr hr 2T
7-1 e2 97
(37)
= grefrr 2T - eluix - mel" hoT
der
with the fact (36)
den xo +
fx gt xT + (he
25x hrge h do,7
X; = Li
mm)
(38)
do.7
dan
which concludes the proof of the Proposition 4.1.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
A.2. Proof of Theorem 4.2
Theorem 4.2. For the SOC problem (12), when 7 + oo, the optimal controller becomes oi 9t xx, log bgp((x] xt). and the corresponding forward and backward SDE with the linear SDE form (10) are the same Doob's 's h-transform in (3) and (4).
Proof. Since in Proposition 4.1 we have solved the control problem and the optimal controller 17 is:
ui,00 lim - gief eo x, me hET 7-000
(13)
e Per
Now we calculate the transition probability p(xr) Xx) and related function h(x,, t, xT, T). Consider Fix., 1) = x,e x,-,,, /, according to the Ito differential formula, we get:
dF - Sixe -f
fidl + - fidx,
(39)
dF = fixe - te- (fixx hmm) df + gidwe),
(40)
dF h - mdd f
(41)
xre fr - x,e = mdrr + S: e gdwgg,
(42)
+ XT ~ N (eliTx, + fr huT e gi: pirr!
(43)
Vx, logg(xT(xx)) -Vx, (XT gl melr heT)² eivr xT - e me fr
(44)
20 gir
c2lr gir
_ gee et X1 me hTTT T
90 V x. logp xi) 9: t, x,, T).
(45)
22f Per
The forward SDEs obtained through SOC and Doob's h-transform are both formed as
dx; = fixx + him + g ehe XT elvr - mee's horr) dt + gidwe.
(46)
e2fr gh.
and the both backward SDEs are
(fix
pir (xT etexx, mets her)
dx; = fixt + him + ot
of V (xxp(X<(TTT) + gedwer
(47)
22l gh.
which concludes the proof of the Theorem 4.2.
A.3. Proof of Proposition 4.3
Proposition 4.3 Consider the SOC problem (12), denote Ju 2) $ cost of the system, Wi the optimal controller (13), then
So"
|| dt + 3 ||x0 xT as the overall
J(u S J(u o0),
(48)

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Proof. According (13) and (14), denote =efr xo FT + mefrir.
" = gechr JT e me heT
dep
"i, ghe -iee-h ft etr + metr hr DT
dey
+ || 11 = g3 -22 22.12222 22r leetroo + me - xT 11
di
||,,,, 12 gee -2f 2f 2fr lall
Similarly.
|1113
01 = = gece 25 2fr
(49)
(e2trgz))²
Furthermore, take t T in (14),
y er
22r gf x-=
(----0 ivr
11 + t
x =
7-1 + c²fr g7
Do +
x- +
e gf
e2tr of
m,
(50)
which implies
7 lef
2 o7 7- + e²5x gj
7 7
||x xT =1 =1
30 +
JT xT c.r
m
m-x7i xT
7- + e2 g/
7- + 225x oj
? eds
7-1
Thr
IT + efr
m 1.
7 + e2fr g
10
7- + e2frg
x- + e²fr o7
(51)
Fr
Fo
1 + fr of
30
74 03
IT + el
11222
m 1
| xn- - IT + me 11
||1|2
(1 + ge2tr g2
(1+ + 2r fr gz
and
lim _ ||ز| rTI22 = lim
?
|a 1 = 0.
(52)
7000
2(1 + 7221 fr fr ,
Hence,
2 f.'"
1 T
(1 10-111 113) dt
1
2fr
||a 129r
1
1
(62lr 12)²
(7-1 + e2 9212
1
2
||
1207
1 2 2 fr oす
(e g )2
+ 2r
1
1 2y 2fr y7
|a
2 (c2fr
97 (1 + ye g7 )
1
ye gt
(53)
1a12
2 (62fr (f) + 26 g2 2
1
7
||11
2 2(1 + 7e2 /2
2 [1x IT l12
?
2 |1x - xr 112
lim 7 ||0
xx
00 2

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Therefore,
2
2 | - rTI12
lim
zT
11 - 2 ./
T
(lui, 11
dt
(54)
3-000
2
2 0 || 11attt
1 /'
R | xr112
2 /
1
(T
I l1att
lim
? ||2 x2112.
(55)
7-000 2
7) S Ju n
(56)
which concludes the proof of Proposition 4.3.
A.4. Derivation of the transition probability (18) Suppose Pr,y and o denote the mean value and variance of the transition probability p(x: zo, Fr), then
p(x:
zo, xr) -N
Pin ct den xo +
etr
2xx hr
oT + (hg
m
du,g
du
da,7
(18)
0.00.11
0% = 22l:ga..
off== =
57
Proof. Since Pt,y remains the same as the closed-form relationship (14), we would focus on how to obtain of. and of2
In Equation (39) of Theorem 4.2, we've obtained:
p(x; xx) ~ N (efuxx + met har, e saatt) N (eleitx + mesh oeet)
(57)
Take 52. = 22fig2 as the coefficient of the noise term, then, through Bayes' formula,
p(oT X,, Xo)p(((( 2o) p(TT x, xx)p(xx 2o)
p(x; Io, IT)
p(IT Fo)
p(uT 2o)
02
o.
o7
which concludes the derivation the the transition probability (18).
A.5. Derivation of the training objective (19)
Denote ae, = assuming 11-1,00 -2 and P1-1,7 o. are respectively the mean values and variances of po(Xt-1 Xt; .T) and p(x--1 Xo, Xt, IT), suppose the score Vx, log oogp(xx xT) is parameterized as (xo(xi, ØT, 1)/o/, the final training objective as follows.
| - th m1-smll
Lg = Exxxoxx,XT
20
h,m g,u,, + of (xlx IT,
Mt-1,0 = x,,- - fix;
(19)
of
a2, (x: P1y), 0:-100 =
M4-17 -4-17 of 01-1,

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Proof. Firstly, as for the training objective (19), according GOUB (Yue al., 2023):
Ep(x)) log po (x xr)] 2 E Bp(xo)
Ep(x, beo) [logpo (xo x1,
T
(58)
E oxx (KL ( (x Xo: Xx, (T)
IIpa (x,-1 x,, xT,
ELBO.
Accordingly,
KL (p (xt-1
Xo, Xc, 2T) Iipo (x:-1 Xt, IT))
V250, c (xc 1 My- 2/203-1
E 1xxx
log
(xt (x--1 1 1)2/203,1
V2T0y,1-1
(59)
llog 00,0-1
llog 01-1 10g0
(xt-1 H4- 1,,, 1,7
)
12/20/_1
(x1-1 122=-))/200...
0 (14-1,7 - 1,
= log 00,11 log x--
+
2
20 1
20 P
Hence, we ignore some constants and minimizing the negative ELBO, leading to the training objective:
L = Xo, Xr
20 14-10 M1-1,al1
(60)
Then, as for solving the closed form of 14-1,0- o;-1, and p through Bayes' formula.
(xt-1 10, X4, xT) p(xx Jo, x,-1, Fr (orr)r(((() (x
Xo, xT)
p(xx To: 2T))
(61)
p((x Xt-1, DT))p((11111
To. XT)
p(x, Fo, IT)
According to Appendix A.4, applying the reparameterizztion tricks:
25r ger 7- 22 o7
etr o
225 e hr
X1 c 7
Ho
+
7-1 2t g7
IT IT+
+
h
7-1 + e gt
m
+ o
A 01, 2o + II + C1, M m + J/a,
(62)
+ 25 1.T
etr 9
2r hT g e2fr gg
X;-1 cti-1
7- + 2fn of
Tu +
7-1 e o
xT +
h
m
+ 0
14-1
c2-1
01-1,-00 + 1I + 4-1 + 11-14-11
Therefore, eliminating zo obtain the relationships between Xt, Xt-1, Br, m and noise c,
X1 01-1,7
01.7 x (eer bi-1,7 01,7 oT Can C--11
01,7
of2
2 -f
0/2
(63)
a1-17
a 1/-1

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
The mean value 11-1,7 of p(x 20, Xx, XT) can be calculated as:
0
(ben - 11- ac-1.) X;
a6x FT
(c. - 11-
R1.7
m + (0, o,=
11-1-1
01-1,7
af 1n
11-1,
07²
01,
ben a..bb-111
at,CC-112
P1-1,7
11,1,
2
11-1,
11-1,
12
Xt
01-1,
2T
Ct,1
0t-1,8
m
(64)
11,7
01,70
12
-1-
X
1,71,001 14,7
0t 1,
az-
of
Pr P1
(xt - P(,))
o. a1-1
with the fact that
Ot.y 1,-
(be
0t.70r-1.
2T +
(
0t,,Ct-1,1
m,
(65)
at-1,7
a1-1,7
at-1,7
which can be easily proved by expanding and comparing the both sides of the equation.
As for p-g.-1 and 1,0- parameterized from the SDE (15):
IT e Xt me hTT hoa
x X1-
fix, + h,m
gi c e- + c2fx pi.))
vx logp(x; ixT)
ge6t
(66)
FT e me f
ot
Xt
fix: ham
gi e e (y + e2fr Th.))
co(xe, or, t) gecr.
o
where E dt,,ttt).
Hence,
fixe + h.m + oh 2T etvxx he
g co (x,, IT:
Mg.--1 = X,
e e-for (x-1 1 + e2tr gi.))
of
(67)
0- Ic:
which concludes the derivation the training objective (19).
A.6. Proof of Proposition 4.4
Proposition 4.4, UniDB encompasses existing diffusion bridge models by employing different hyper-parameter spaces H. as follows:
DDBMs (VE) corresponds to UniDB with hyper-parameter HVE = 0, h4 = 0, xc)
DDBMs (VP) corresponds to UniDB with hyper-parameter Hvp (fr --6 h4 = 0,7 - x0)
GOUB corresponds UniDB with hyper-parameter Haou(s; Oz, 1u, -0,, -4,7 14,7 0c)
Proof. DDBMs (VE).
HvE(f. = h, = 00, + 20)
SOC problem with SDE: dx = gedw; as 7 0
XT Xt dt + g.dw;
dx;
S dz
d X dt + 9,dw, with of
dx,
2 of
o o o
DDBMs (VE) with Doob's h-transform

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
DDBMs (VP).
Hvr(A = 2 hr = 00,7 - 00)
1
+ SOC problem with SDE: dx, = 2 x, + g,dw, as 7 -> oc
2 + gebl 4d= c 2 lá g. d=dxx; XT c 10 of
dx; =
dt + godw,
St sf - 2 Se g?ds
+ dx.
2 Di X + gt
xxXT xTX,
dt gedw, where at = 56 56 o2 and oZ dt
d
oi + o
ai o7OT) 07
dt XT
X1
az of
+ dxg-
2 X + 9t
dt + gidw, where SNR,
07
SNR.
- 1)
SNRT
+ DDBMs (VP) with Doob's h-trrassform
GOUB.
Hoourfe = 8,, = -r fm.m== p.? - 00)
+ SOC problem with SDE: dx, = 02 (u xx) dt + gidw, as y - 0c
26
dx:- = 02 + gt
- xx)dtt + g.d
GOUB with Doob's h-trranfform
which concludes the proof of the Proposition 4.4.
A.7. Derivation of UniDB-GOU (forward SDE (20) and mean value of forward transition (21))
Consider the SOC problem with GOU process (5), the cotimally-controlleed forward SDE is
-
dx;
0 + o6 7-1 + 66.
(or- - gidw
(20)
and the mean value of the probability p(x+|
Fo, (T) is
0. 1 1 10
1 e -8, 1 + vo% 17700
P1, = = -0
vn +
IT.
(21)
1+ 10
1 70g
Proof. Consider the SOC problem with GOU process (5) in the deterministic form:
min ue J. 2 u 11.1111d 211 2T 11
(68)
s.t. dx, = (2T xx) +e (.,)) dt, Xo 20
where the definition of and gr is the same as GOUB: μ = ェт 9 2x-0..
Similarly to the proof of Proposition A.1, according to minimum principle theorem obtain the following set of differential equations:
dx; Vp.H H (xe. Pe, ti = =0,2 0,x; 9 Pc.
(69)
dt

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
dp: dl
Xx, H. (Xe: Pr, u 1) O.Pr.
(70)
X0 -0+
(71)
pr = 7 7XT - 2T)
(72)
Solving the equation (70), we have:
P: PIe
(73)
pr Pne BT
Then we solve the equation (69):
dx: = B1T 01x: - S.P:
dt
d(c d(c - ch g.
dt
e" ! X1 I/" c²=0,dz -Pu L/ g2 dz,
e Xt 20 xT(E - 1) 2²po(E20. -1).
Hence, we can get:
XT e e x + (1 - =Tr) -2²pr -e--e0 ---2TT),
(74)
and
Xx c - Pr + (1 -1-66 e B=b))
X2c fr pr -6
(75)
Take the equation (74) into the equation (72) and solve pr.
$1 + (1 - -r))rr - 22pr(1 -2ōr) IT
pr =r=
7 8r (1r xT)
pr
1 1 712(1 - 262
Hence,
7e Fr (10 IT)
50 + (1 e )oT -2 -e (e-1-
X,
1 1 712(1
c -02r)
e -0: - 20 (6 - 6-))
722e -285 (c c E-B1)
1 -6 +
20 +
1 712(1 202))
1 + 7X2 - e-20rr
(76)
+ 7:2 6 -20iiT
1 -.. 1+ 711 c -25eTT)
e
10 +
2T
1+ 71211 e-20x)
1 1+7 e-28.)
1 1 10
1 11-e-0 + 70
0.
Tn + 1 1
FT.
1 702
which implies
P1, = -
1 + 2 Not
1 - -
1 + 7oo
30 +
IT.
(77)
1 10
1 + 707

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Then,
4i, = 9tP:
= - e OT y Pr (ID FT)
1t 7 E 26r)
br (20 IT)
70 e- or
++ 100
(78)
-ie Or 76 br 8 (x: FT)
70.0
e 20 (T - xx1
9t
+ +0
And the optimally dynamics can be:
e- 20tr
dx,-
0 Di
(DT xx)dt gedwt,
(79)
7-1 + Ti.
which concludes the derivation of UniDB-GOU (forward SDE (20) and mean value of forward transition (21)).
A.8. Examples of UniDB-VE and UniDB-VP
Similar to section 4.5, provide the other examples of UniDB-VE and UniDB-VP, highlighting the key difference of the coefficient of zo in the mean value of forward transition and h-function term between UniDB and them respectively.
UniDB-VE
7
0
7
T
2 0
TT
X
4
3T
(80)
h
9t
7-
+ 07 dz
UniDB-VE
UniDB-VP
LaaT
SNRT SNR,
00
7 SNRT
at
at
1
+ SNR,
0f07
(81)
Gt ET
X2
u
FT X,
0T
07
0 SNRT
9t
7
+ 02(
SNR, SNRT
VP
UniDR-VV
A.9. Proof of Proposition 4.5
Proposition 4.5. Denote the initial state distribution Fo, the terminal distribution x²i by the controller and the pre-defined terminal distribution IT, then
e 20
|
2T
||- 20ll3
(23)
(1 + 71 e
Proof. According to Appendix A.7, we' 'e learned that
x = e 0, + 20t
1 1- 70.
2
+
IT +
(82)
1 + 707
1 707
07
Take f = T, then
e b
c 0.
x7
In +
1
TT.
(83)
1 707
70

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Therefore, since o7 (2 e-26r)
0
e
2
| 2T112
20 +
(1
1
IT
2T
11+ 707
1 + 707
(84)
c -0r
1| 20l12
(1 712(1 - -e-265)
which concludes the proof of the Proposition 4.5.
B. Implementation Details
In Image Restoration Tasks (Image 4 Super-resolution, Image Deraining and Image Inpainting), we follow the experiment setting of GOUB (Yue et al., 2023): the same noise network which is similar U-Net structure (Chung et al., 2022), steady variance level X² 30 /255", coefficient exr 0.005 instead of zero, sampling step number T 100, 128 patch size with 8 batch size when training, Adam optimizer with 81 0.9 and $2 0.99 (Kingma, 2014), 1.2 million total training steps with 10-4 initial learning rate and decaying by half at 300, 500, 600, and 700 thousand iterations. With respect to the schedule of 0r, we choose a flipped version of cosine noise schedule (Nichol & Dhariwal, 2021; Luo et al., 2023),
COs( 1/T++ 1++ 2
0, = cos Ita 2)
(85)
where s 0.008 to achieve smooth noise schedule, St is determined through 9e = 22²0,. As for the datasets of three main experiments, we take 800 images for training and 100 for testing for the DIV2K dataset, 1800 images for training and 100 for testing for the Rain 10OH dataset, 27000 images for training and 3000 for testing for the CelebA-HQ 256x256 dataset. Our models are trained on single NVIDIA H20 GPU with 96GB memory for about 2 days.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
C. Additional Experimental Results
Here we will illustrate more experimental results.
Table Qualitative comparison between different bridge models (DDBMs (VE) and DDBMs (VP)) and ours (UniDB- VE and UniDB-VPP on DIV2K and Rain100H1 datasets.
Image 4x Super-Resolution
Image Deraining
METHOD
PSNR
SSIM
LPIPS
FID
PSNR
SSIM
LPIPS.
FID
DDBMs (VE)
23.34
0,4295
0.372
32.28
29.34
0.7654
0.185
43.22
UniDB VE
23.84
0.4454
0.357
31.29
29.46
0.7671
0.185
42.57
DDBMs (VP)
22.11
0.4059
0.491
48.09
29.58
0.828
0.113
35.46
UniDB-VP
22.42
0.4097
0.486
44.52
30.11
0.8414
0.102
33.17
Table 4. Image Ressper--esolttion Qualitative evaluation of the CelebA-Ho 256 datasets with baselines.
METHOD
Penalty Coefficient
PSNR1
SSIM1
LPIPS
FIDJ
DDBM (VE)
0
25.84
0.5099
0.381
67.98
UniDB-Ve VE
1 x 10"
25.37
0.504
0.265
53.98
DDBM (VP)
00
27.13
0.6497
0.194
42.54
UniDB-VP
1 x 10"
27.44
0.6631
0.174
42.06
GOUB
28.63
0.7776
0.104
19.02
UniDB-GOU
x 107
28.70
0.7894
0.090
17.59
Table Image
Qualitative evaluation the FFHQ 256x256 dataset with baselines.
METHOD
LPIPS
FID
DDRM
0.339
59.57
DPS
0.214
39.35
DDBM (VE)
0.239
42.85
DDBM (VP)
0.177
39.63
GOUB
0.072
21.77
UniDB-GOU
0.069
20.24
Table 6. Image Inpainting, Qualitative comparison with the relevant baselines on CelebA-HQ with thick mask.
METHOD
PSNR
SSIM
LPIPS
FID
DDRM
19.48
0.8154
0.1487
26.24
IRSDE
21.12
0.8499
0.1046
11.12
GOUB
22.59
0.8573
0.0917
8.49
UniDB-GOU
23.02
0.8571
0.0884
7.46

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
IQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
Figure Additional visual results on deraining with Rain100H datasets.
Polic
Poli
Exch
Exch
Lange
sange
Poli
Poli
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
Figure 6. Additional visual results 4x seperrtttooutooo with FFHQ datasets.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
Figure 7. Additional visual results on 4xsuper-resolution with CelebA-HQ datasets.
GT
GOUB
UniDB
GT
GOUB
UniDB
GT
GOUB
UniDB
DCA
HON
HON
HON
RS
RS
RS
H
Figure 8. Additional visual results on thin mask inpainting with CelebA-HQ datasets to show our excellence.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
LQ
UniDB
LQ
UniDB
LQ
UniDB
LQ
UniDB
Figure Additional visual results on thin mask inpainting with CelehA-HQ datasets.

---

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
GT
GOUB
UniDB
Figure 10. Additional visual results 4x super-rreuuutooo with DIV2K datasets.


Reviews:
[{"review_id": "iZ1Hx5b5yv", "reviewer": "Reviewer_33jh", "summary": "This paper proposed a diffusion-based method for image restoration problems, e.g., super-resolution, deraining, and inpainting. Given a dataset of corrupted and clean image pairs, the goal is to construct a diffusion model that at inference generates clean images given corrupted images. The proposed method is based on stochastic optimal control (SOC), which re-frames the problem as an optimization over the drift of diffusion models. Such an SOC reformulation introduces a tunable hyper-parameter ($\\gamma$) and recovers prior Doob-h-transform's based diffusion models as $\\gamma$ approaches infinity. It is shown empirically that a rather large but finite $\\gamma$ (~e7) improves performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your comments and questions. We will provide a detailed response to these concerns.\n\n> Claim: The reviewer is not convinced by the implication of Prop 4.3: $\\mathcal{J}$ is smaller for finite $\\gamma$ does not imply suboptimality in empirical performance. Intuitively, it does make sense to set to infinity since we'd like $x_T^u$ to converge exactly to the given $x_T$. \n\nThank you for raising this important point. We agree that the intuition behind Proposition 4.3 deserves clarification. Below, we address **why strict terminal matching ($\\gamma\\to\\infty$) is not ideal** for practical performance, even if it seems mathematically appealing.\n\n**1. Exact terminal matching ($\\gamma\\to\\infty$) harms image quality.**\n\nWhile Doob’s h-transform achieves $||x_T^{u_{\\infty}}-x_T||^2_2 = 0$, it requires large control ($||u_{\\infty}||^2_2\\geq||u_{\\gamma}||^2_2$) to force exact endpoint matching. The **large $u$ in SDE trajectory（$dx_t=(f_tx_t+h_t\\mathbf{m}+g_tu_{\\gamma}) dt+g_t dw_t$）may disrupt the inherent continuity and smoothness of images**. Prioritizing pixel-perfect endpoints over smooth trajectories leads to \"mathematically correct but visually unrealistic\" outputs. Our experiments (Figure 1) confirm that Doob's h-transform can lead to artifacts along edges and unnatural patterns in smooth regions.\n\n**2. Why finite $\\gamma$ works better?**\n\nBy keeping $\\gamma$ finite, UniDB explicitly balances two goals: **target matching** $||x_T^{u_{\\infty}}-x_T||^2_2$ (terminal penalty) and **trajectory smoothness** $||u||^2_2$ (controller). Proposition 4.3 shows that **finite $\\gamma$ achieves a lower total cost $\\mathcal{J}$ not by sacrificing performance**, but by optimally **trading minor terminal mismatches for significantly smoother, more natural diffusion paths**.\n\nThus, **Proposition 4.3 reflects a key insight**: real-world image generation benefits more from **stable trajectories than rigid mathematical constraints**. We will clarify this intuition in the revised manuscript.\n\n> Experiment: Comparison to baselines is insufficient: the authors should compare their method to I2SB and its follow-up works (e.g. CDDB).\n\nThank you for your feedback. Below, we clarify why our current comparisons are sufficient and address your concerns about I2SB and CDDB:\n\n1. **Direct Comparison to State-of-the-Art (DDBM > I2SB):**\n   - **DDBM** (a SOTA baseline) is **strictly superior to I2SB** (see Table 2 in [DDBM paper], where DDBM achieves FID **4.43** vs. I2SB’s **9.34** on DIODE-256×256, and **1.83** vs. I2SB’s **7.43** on Edges→Handbags-64×64).\n   - **UniDB outperforms DDBM** across multiple tasks (Tables 3–5 in Appendix E): Super-resolution (DIV2K, CelebA-HQ), Deraining (Rain100H). \n   - Critically, **DDBM is a special case of UniDB** (UniDB with $\\gamma=\\infty$), meaning our method inherently subsumes and improves upon this stronger baseline.\n\n2. **I2SB's Computational Burden.**\n   - **Training I2SB is prohibitively expensive**. It requires **16×V100 GPUs** with **more than 1 week** (*per communication with I2SB authors*).\n   - The inference time of I2SB is **90 seconds per 256×256 image** (vs. our method’s **less than 5 seconds**).\n   - Given hardware/time constraints, reproducing I2SB (and its follow-ups) for fair comparison is practically infeasible during rebuttal. Since **DDBM already outperforms I2SB**, and UniDB outperforms DDBM, we argue that **direct I2SB comparisons are redundant**.\n\n3. **CDDB is Orthogonal to Our Contribution.**\n   - **CDDB is a training-free, plug-and-play refinement module** for I2SB, not a standalone method.\n   - Such post-hoc techniques could also **integrate with UniDB** (e.g., applied during inference), but this is **beyond our paper’s scope**, which focuses on developing **a unified training framework for diffusion bridge**.\n\nWe appreciate your suggestion and will explicitly discuss the relationship between UniDB, DDBM, I2SB, CDDB in the revised manuscript.\n\n> Comment: What are x0 and xT in practice? x0 would be the clean image and xT would be the corrupted image, is this correct?\n\nYes, you're right. It's standard in diffusion-based image restoration research using x0 as the clean image and xT as the corrupted image. We'll make it clearer in the revised version.\n\n> Q1: Eq 8 and 16 need more explanation. Why is it okay to drop the Brownian motion? This seems more like an empirical trick to get better PSNR and SSIM.\n\nYes, **Dropping Brownian motion** is an empirical trick that **we followed from GOUB as stated in line 165 of our paper**. It can obtain better results on image restoration tasks, capturing more pixel details and structural perceptions of images (which contributes to better PSNR and SSIM). This phenomenon has been verified by GOUB's three ablation experiments.\n\n> Q2: There seems some typos in Eq 15,16 \\nabla p -> \\log \\nabla p.\n\nYes, you're right. We will correct it in the revised version.", "claims_and_evidence": "I'm not convinced by the implication of Prop 4.3 (the paragraph below Prop 4.3, before Sec 4.3). The SOC objective $J$ is (artificially) constructed as a surrogate for searching the drift of SDE. That $J$ is smaller for finite $\\gamma$ does *not* imply suboptimality in empirical performance. Intuitively, it does make sense to set $\\gamma$ to infinity since we'd like xT^u to converge exactly to the given xT. Any finite $\\gamma$ would fail to achieve that, as shown in Prop 4.5. What's presented in this paper is somewhat counter-intuitive but also interesting, yet it's rather empirically observed than theoretically justified.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Thm 4.1 is expected but still quite nice to see the actual analytic form! Though, I did not check proof carefully.", "experimental_designs_or_analyses": "Comparison to baselines is insufficient. GOUB is a special case of UniDB (these values re-appear in Table 2 last column) and DDRM is not quite comparable as a non-learning-based method and requiring knowing additionally corruption type. The authors should compare their method to I2SB and its follow-up works (e.g., CDDB https://arxiv.org/abs/2305.19809). These works are also SOC-inspired method for solving image restoration.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "- What are x0 and xT in practice? To my best guess, x0 would be the clean image and xT would be the corrupted image. So in Eq 12 we learn a forward process that bring x0 close to xT for every (x0, xT) pair, and then reverse it with Eq 15. Is this correct? It'd be better to have clarification in Sec 3.", "questions_for_authors": "- Eq 8 and 16 need more explanation. Why is it okay to drop the Brownian motion? This seems more like an empirical trick to get better PSNR and SSIM.\n\n- There seems to be some typos in Eq 15,16 \\nabla p --> \\log \\nabla p", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3hkFqdAyH8", "reviewer": "Reviewer_AzN7", "summary": "This paper proposes a framework that unifies and extends various diffusion bridge methods by way of stochastic optimal control. In the case of linear dynamics, they derive a computationally tractable method, which can be thought of as a regularization of previous methods by the introduction of a new hyperparameter.  Implementing this change requires only a minimal modification to existing code. They show that, by tuning this new hyperparameter, improved performance can be achieved in a number of benchmark examples.\n\n## update after rebuttal\nMy assessment has not changed substantially. I have maintained my score at 4: Accept\n\nIn addition, earlier I posted a followup discussion regarding the proof of Prop. 4.3 under an incorrect heading, so the authors were unable to see it. I believe it would further improve the paper if the following was addressed, though it would not change my score either way:\n\n I do still believe that  Proposition 4.3 is much more elementary than it is made out to be. Below is a more detailed version of the comment in my original review,  on which the authors could base a revised elementary proof.\n\nStarting with the $\\gamma=\\infty$ case, the optimal control $u_{t,\\infty}^*$ must lead to $x_T^{u_\\infty^*}=x_T$, else the cost is infinite, and therefore\n\n$J(u_{t,\\infty}^*,\\infty)=\\int_0^T\\frac{1}{2}\\|u_{t,\\infty}^*\\|^2_2dt$\n\nNow looking at the finite $\\gamma$ case, we can bound the minimum by the value at the control $u_{t,\\infty}^*$:\n\n$J(u_{t,\\gamma}^*,\\gamma)=\\min_u\\{\\int_0^T \\frac{1}{2}\\|u_t\\|^2_2dt+\\frac{\\gamma}{2}\\|x_T^{u}-x_T\\|\\} \\leq \\int_0^T \\frac{1}{2}\\|u_{t,\\infty}^*\\|^2_2dt+\\frac{\\gamma}{2}\\|x_T^{u_\\infty^*}-x_T\\|$\n\nThe SDE determining $x_t$ doesn't depend on $\\gamma$, hence in the above expression we still have $x_T^{u_\\infty^*}=x_T$.  Therefore we arrive at\n\n$J(u_{t,\\gamma}^*,\\gamma)\\leq  \\int_0^T \\frac{1}{2}\\|u_{t,\\infty}^*\\|^2_2dt=J(u_{t,\\infty}^*,\\infty)$", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your comment.\n\n> Claim & Question 2: The reviewer finds that result to be mathematically trivial and also to be disconnected from saying anything about how well the method performs in practice. However, The reviewer doesn’t think this proposition is integral to their work, and it could be removed (or de-emphasized) without any negative effects. The authors need to provide a better intuitive discussion about why (17) should be expected to correspond to better performance in the experiments in section 5 or else remove this proposition and the surrounding paragraph from the main text.\n\nWe appreciate this feedback and can move Proposition 4.3 to Appendix. \n\nHere we add more analysis to help **understand the practical implications of our UniDB**. When $\\gamma$ approaches $\\infty$, our UniDB reduces to Doob's h-transform, and **the control term $u$ in Eq (12) becomes ineffective**, leading to $||u_{{\\gamma}}||^2_2 \\leq ||u_{{\\infty}}||^2_2$. Although Doob's h-transform ensures $x_T^u$ reaches the target endpoint $x_T$ exactly, i.e., $||x_T^{u_{{\\infty}}}-x^T||^2_2 = 0$, it may force the model to preserve even harmful noise/artifacts in the target. This is becuase Doob's h-transform will apply disproportionately large control inputs $||u_{{\\infty}}||^2_2$ to achieve such exact matching. The **large $u$ in SDE trajectory may disrupt the inherent continuity and smoothness of images**. The over-control in Doob's h-transform violates the natural statistical properties of images, **prioritizing mathematical precision (pixel-perfect endpoints) over visual authenticity (realistic SDE trajectories)**. As shown in our **Figure 1**, Doob's h-transform can lead to artifacts along edges and unnatural patterns in smooth regions. \n\nMoreover, we want to emphasize that the discovery of **Proposition 4.3 is non-trivial**. Proposition 4.3 and related mathematical derivations in Appendix A.3 show that $\\mathcal{J}(u_{{\\gamma}},\\gamma) \\leq \\mathcal{J}(u_{{\\infty}},\\infty)$. **Generally, one cannot determine the magnitude between $\\mathcal{J}(u_{{\\gamma}},\\gamma)$ and $\\mathcal{J}(u_{{\\infty}},\\infty)$** because although $||u_{{\\gamma}}||^2_2 \\leq ||u_{{\\infty}}||^2_2$ but $||x_T^{u_{{\\gamma}}}-x^T||^2_2 \\geq ||x_T^{u_{{\\infty}}}-x^T||^2_2 = 0$. We use strict mathematical derivations in Appendix A.3 to show that $\\mathcal{J}(u_{{\\gamma}},\\gamma) \\leq \\mathcal{J}(u_{{\\infty}},\\infty)$ is true. **Though Proposition 4.3 does not directly mean better image quality by UniDB, it can show that UniDB leads to better overall performance considering both realistic SDE trajectories and target endpoint matching**. \n\n> Question 1: Is there a reason why the UniDB (SDE) performs well for LPIPS and FID while the UniDB (ODE) performs well for PSNR and SSIM?\n\nThis is a common phenomenon in various Diffusion models [R1-R3]. As analyzed in [R1], \"although solvers for the probability flow ODE allow fast sampling, their samples typically have higher (worse) FID scores than those from SDE solvers if no corrector is used\". In the paper SDE-Drag [R2] and GOUB [R3], they compare the performance of the SDE and ODE models and get the similar phenomenon: SDE model performs well for LPIPS and FID. Particularly in GOUB [R3], the experimental results also demonstrate the better performance of ODE model for PSNR and SSIM.\n\n[R1]: Song et al. \"Score-Based Generative Modeling through Stochastic Differential Equations.\", ICLR 2021.\n\n[R2]: Nie et al. \"The Blessing of Randomness: SDE Beats ODE in General Diffusion-based Image Editing\", ICLR 2024.\n\n[R3]: Yue et al. \"Image restoration through generalized ornstein-uhlenbeck bridge\", ICML 2024.\n\n> Comments Or Suggestions: Equations 9 and 11 require expected values, yes? They are still stochastic control problems; Above equation 19, the formula for the score should have a log; Computations in Eq 36 and 37 have extraneous commas at the end of lines; Line 705 appears to be repeated.\n\nYes, you're right. Thanks for pointing out these typos and we will correct them in the revised version.", "claims_and_evidence": "The majority of their claims are supported by clear and convincing evidence.  Generalization through the methods of of stochastic control is well-grounded theoretically, and they show improved performance empirically on a number of convincing benchmarks\n\nThe one claim that I don’t find substantiated revolves around their Proposition 4.3.  They claim that this proposition lends theoretical support to the observation that introducing their new hyperparameter leads to improvements in practice. I find that result to be mathematically trivial and also to be disconnected from saying anything about how well the method performs in practice. However, I don’t think this proposition is integral to their work, and it could be removed (or de-emphasized) without any negative effects.", "methods_and_evaluation_criteria": "I find the methods and evaluation criterial to make sense for the problem.", "theoretical_claims": "I checked the proof of theorem 4.1 and didn’t find any substantial issues.", "experimental_designs_or_analyses": "I reviewed the experimental designs in Section 5 and did not observe any issues.", "relation_to_broader_scientific_literature": "The key contribution is a reformulation of diffusion bridge methods in the language of stochastic control, which motivates a natural one-parameter family of methods extending previous work, especially GOUB (Yue et al., ICML, 2024).  I find this to be a natural and interesting extension.", "essential_references_not_discussed": "I am not aware of any missing essential references.", "other_strengths_and_weaknesses": "I found the contribution to be original and well-motivated theoretically, the writing to be relatively clear, and the performance improvements to be nontrivial.", "other_comments_or_suggestions": "1) Equations 9 and 11 require expected values, yes? They are still stochastic control problems.\n2) Above equation 19, the formula for the score should have a log.\n3) Computations in Eq 36 and 37   have extraneous commas at the end of lines.\n4) Line 705 appears to be repeated.", "questions_for_authors": "1) Regarding Proposition 4.3, perhaps I misunderstand, but I believe this proposition is trivial from the definitions.  Under the hard constraint ($\\gamma=\\infty$) the terminal cost is enforced to be exactly zero, hence so both costs agree on $u^*_{t,\\infty}$.  Therefore the minimal cost when $\\gamma<\\infty$ can't be larger than the cost of the control $u^*_{t,\\infty}$.  I'm not sure this fact is deserving of its own proposition.\nI am also not convinced that comparing the minimal cost values in this way says anything about why one method performs better than another in any operational sense.  I think the answer probably lies more in the direction of the soft constraint being a numerically better behaved regularization than the hard constraint.   I think the authors need to provide a better intuitive discussion about why (17) should be expected to correspond to better performance in the experiments in section 5 or else remove this proposition and the surrounding paragraph from the main text, as well as alter the discussion in Section 4.5.\n\n2) Could you comment further on the performance of UniDB (SDE) vs UniDB (ODE) in table 1. Specifically, is there a reason why the former performs well for LPIPS and FID while the latter performs well for PSNR and SSIM?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MWga6WquW0", "reviewer": "Reviewer_9hKM", "summary": "This paper introduces UniDB, a diffusion bridge model framework that utilizes\nStochastic Optimal Control (SOC) for process optimization, providing an ana\u0002lytical solution for the optimal controller. UniDB generalizes existing diffusion\nbridge models by showing that Doob’s h-transform is a special case where the\npenalty coefficient γ tends to infinity. By adjusting the trade-off between control\ncosts and terminal penalties, UniDB improves image detail and quality while\nmaintaining compatibility with existing models. Experimental results demon\u0002strate its effectiveness in image restoration tasks.\n\n## update after rebuttal\nI have carefully read the rebuttal and would like to maintain my original score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your feedback and inquiries.\n\n> Claims And Evidence 1: Additional explanation is needed regarding whether the optimal controller in LQ SOC directly contributes to producing sharper and more detailed images.\n\nThe over-control in Doob's h-transform violates the natural statistical properties of images, **prioritizing mathematical precision (pixel-perfect endpoints) over visual authenticity (realistic SDE trajectories)**. UniDB leads to better overall performance considering **both realistic SDE trajectories and target endpoint matching**. A comprehensive analysis can be found in response for **Reviewer 33jh**'s \nClaim and response for **Reviewer AzN7**'s Claim and Question 2. \n\n> Claims And Evidence 2: It remains unclear whether there is a systematic analysis regarding the choice of $\\gamma$. Based on the results in Table 2, it appears that using an optimal controller with finite $\\gamma$ does not always lead to improvements in actual evaluation metrics.\n\nIn **Figure 2** of our paper, we illustrate that a range of $10^5$ to $10^9$ for $\\gamma$ is pretty good. What we want to emphasize is not that there is a single best value of $\\gamma$ for each dataset, but rather that most value of $\\gamma$ chosen **within this interval can yield a better result than Doob's h-transform ($\\gamma = \\infty$)**. While suboptimal $\\gamma$ values may occasionally degrade performance, our approach provides **a simple, efficient, and nearly cost-free way** to improve model performance, which works well in most cases. \n\nIt is worth noting that when $\\gamma$ is set to $1e7$, our model **outperforms all baselines across multiple tasks**, including super-resolution (on DIV2K, CelebA, and FFHQ), deraining on (Rain100H), and inpainting on (CelebA).\n\n> Question: Could the authors provide more details on the derivation of this part: $\\frac{\\gamma}{2}\\left\\||\\mathbf{x}_T^u-x_T\\right\\||_2^2=\\frac{\\gamma\\||a\\||_2^2}{2\\left(1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2\\right)^2}$? \n\nYes, below is the detailed proof: At line 770 of our paper, we defined $a=e^{\\bar{f}_T}x_0-x_T+\\mathbf{m}e^{\\bar{f}_T}\\bar{h}_T$ for simple denotation. The first equation of Eq. 51 is a simple simplification by substituting the expression $\\mathbf{x}_T^u$ from Eq. 50 into Eq. 51. A more detailed proof is as follows:\n\n$\\left\\||\\mathbf{x}_T^u-x_T\\right\\||_2^2=\\left\\|\\left\\|\\left(\\frac{\\gamma^{-1} e^{f_T}}{\\gamma^{-1}+e^{2 f_T} \\bar{g}_T^2}\\right) x_0+\\left(\\frac{e^{2 \\bar{f}_T} \\bar{g}_T^2}{\\gamma^{-1}+e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) x_T+e^{\\bar{f}_T}\\left(\\frac{\\gamma^{-1} \\bar{h}_T}{\\gamma^{-1}+e^{2 f_T} \\bar{g}_T^2}\\right) \\mathbf{m}-x_T\\right\\|\\right\\|_2^2$\n\n$=\\left\\|\\left\\|\\left(\\frac{\\gamma^{-1} e^{\\bar{f}_T}}{\\gamma^{-1}+e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) x_0-\\left(\\frac{\\gamma^{-1}}{\\gamma^{-1}+e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) x_T+e^{\\bar{f}_T}\\left(\\frac{\\gamma^{-1} \\bar{h}_T}{\\gamma^{-1}+e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) \\mathbf{m}\\right\\|\\right\\|_2^2$\n\n$=\\left\\|\\left\\|\\left(\\frac{e^{\\bar{f}_T}}{1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) x_0-\\left(\\frac{1}{1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) x_T+e^{\\bar{f}_T}\\left(\\frac{\\bar{h}_T}{1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2}\\right) \\mathbf{m}\\right\\|\\right\\|_2^2$\n\n$\\begin{aligned} & =\\frac{\\left\\|\\left\\|e^{\\bar{f}_T} x_0-x_T+\\mathbf{m} e^{\\bar{f}_T} \\bar{h}_T\\right\\|\\right\\|_2^2}{\\left(1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2\\right)^2} \\\\\\\\ & =\\frac{\\||a\\||_2^2}{\\left(1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2\\right)^2}\\end{aligned}$\n\nThen, just mutiplying $\\frac{\\gamma}{2}$ on both sides of the equation, we can learn that $\\frac{\\gamma}{2}\\left\\||\\mathbf{x}_T^u-x_T\\right\\||_2^2=\\frac{\\gamma\\||a\\||_2^2}{2\\left(1+\\gamma e^{2 \\bar{f}_T} \\bar{g}_T^2\\right)^2}$. We will add the detailed proof in the revised version. \n\n> Suggestion: There are minor typos regarding the connection between the statements and the proofs.\n\nYes, you are right. Thanks for pointing out them and we will correct them in the revised version.", "claims_and_evidence": "UniDB claims to generate higher-quality images than Doob’s h-transform and\nsupports this claim with experimental evidence on tasks such as super-resolution,\ninpainting, and deraining. However, additional explanation is needed regarding\nwhether the optimal controller in LQ SOC directly contributes to producing\nsharper and more detailed images.\nFurthermore, through Proposition 4.3, it is shown from an LQ SOC perspec\u0002tive that the optimal controller obtained with a finite γ is preferable to that of\nthe infinite case. However, it remains unclear whether there is a systematic anal\u0002ysis regarding the choice of γ. In Proposition 4.5 and Figure 2, a sufficiently large\nγ is selected to minimize differences in terminal point positions. Nevertheless,\nbased on the results in Table 2, it appears that using an optimal controller with\nfinite γ does not always lead to improvements in actual evaluation metrics. I\nwould appreciate further clarification on this point.", "methods_and_evaluation_criteria": "The authors demonstrate their approach’s strength through evaluation criteria\nwidely used in prior works", "theoretical_claims": "The detailed proofs of the Theorems and Propositions in the main paper are clearly described in the Appendix. However, I noticed some minor issues when cross-referencing the statements with their proofs.\n\n1. There are minor typos regarding the connection between the statements and the proofs. The authors are encouraged to check the appendix number references carefully.\n\n2. In equation (53) of Appendix A.3, could the authors provide more details on the derivation of this part? \n   $  \n    \\frac{1}{2} \\displaystyle\\frac{\\gamma}{(1+\\gamma e ^ {2 \\bar{f}_T} \\bar{g}_T^2)^2} ||a||_2^2 \n    = \\frac{\\gamma}{2} || \\mathbf{x}_T^u - x_T ||_2^2\n    $\n    The intermediate steps would help improve clarity.", "experimental_designs_or_analyses": "This paper follows a standard experimental design and evaluation process, so I\nbelieve there are no issues in this regard.", "relation_to_broader_scientific_literature": "No.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Please refer to the questions and comments provided above.", "other_comments_or_suggestions": "Please refer to the questions and comments provided above.", "questions_for_authors": "Please refer to the questions and comments provided above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NzxLEbaEsI", "reviewer": "Reviewer_G2ra", "summary": "This paper proposes a unified framework, UniDB, of diffusion bridge models based on Stochastic Optimal Control. This framework enhances the quality and detail of generated images by balancing control cost and terminal penalty.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedbacks and comments.\n\n> Claim 1: Choice of $L_1$ norm in the training objective (Equation 19): The paper does not provide a clear justification for using the $L_1$ norm instead of other alternatives like $L_2$ norm. The authors should explain whether this choice is based on empirical performance, theoretical considerations, or robustness to outliers.\n\nThanks for your feedback. **Using $L_1$ norm is a common trick** in the field of image restoration with **many prior works [G1-G4]**. We just follow to use the same trick as we mentioned in **line 270 of our paper**. The reason that choosing $L_1$ norm over other alternatives like $L_2$ norm in image restoration tasks have been analyzed in [G2-G4]. In [G2], \"when the network is trained using $L_1$ as a cost function, instead of the traditional $L_2$, the average quality of the output images is superior for all the quality metrics considered.\" [G3, G4] establish the $L_1$ loss's superiority over $L_2$ through improved convergence and outlier robustness.\n\n[G1] Yue et al. \"Image restoration through generalized ornstein-uhlenbeck bridge\", ICML 2024.\n\n[G2] Zhao et al. \"Loss Functions for Image Restoration with Neural Networks\", IEEE Transactions on Computational Imaging 2017. \n\n[G3] Lim et al. \"Enhanced Deep Residual Networks for Single Image Super-Resolution\", CVPR 2017 workshop.\n\n[G4] Mu et al. \"Riemannian Loss for Image Restoration\", CVPR 2019 workshop.\n\n> Claim 2: Introduction of the state vector term $\\mathbf{m}$ in linear SDE: One of the paper's key novelties is introducing the $\\mathbf{m}$ term in the linear SDE form, but it is not explicitly explained how it is computed or designed in the main context. Further elaboration on the motivation, computation, and impact of $\\mathbf{m}$ would improve the clarity of the contribution.\n\nWe want to clarify that we've **never** claimed the introduction of the state vector term $\\mathbf{m}$ in linear SDE is the **key novelty of our framework**. Instead, our main novelty lies in constructing the diffusion bridge in the form of stochastic optimal control and revealing that Doob's h-transform is a special case of ours. The introduction is just a **simple reformulation** referenced by prior works IR-SDE [G5] and GOUB [G1]: we reformulated the drift term $\\theta_t(\\mu - x_t)$ into $f_t x_t + h_t \\mathbf{m}$. Particularly in our experiments of UniDB-GOU, for a fair comparison, we set $\\mathbf{m} = \\mu$ which is identical to GOUB [G1], ensuring consistency with baselines. \n\n[G5] Luo et al. \"Image Restoration with Mean-Reverting Stochastic Differential Equations\", ICML 2023.\n\n> Experimental Designs Or Analyses: The reviewer suggests that the authors should include DDBMs as a benchmark for comparison. \n\nActually, we **have compared DDBMs** as a benchmark in **Appendix E \"Additional Experimental Results\"** of our paper, theoretically analyzing the application of UniDB to DDBMs (Appendix: A.8. Examples of UniDB-VE and UniDB-VP) and conducting extensive experiments (Tables 3, 4, and 5). Specifically, it outperforms DDBMs on both LPIPS and FID metrics, with gains reaching up to ~20\\% in some cases.\n\n> Essential References Not Discussed: a similar approach was explored last year in: Zhang, Shaorong, et al. \"Exploring the Design Space of Diffusion Bridge Models via Stochasticity Control\". To differentiate from prior work, the paper should explicitly highlight the key distinctions between this work and Zhang et al. (2024), particularly in how Stochastic Optimal Control is formulated and applied.\n\nOur work takes a completely different approach from the prior work Zhang et al. (2024) (following simply denoted as SDB), specifically: \n   - **Different purposes.** SDB mainly focuses on solving **singularities and accelerations** for training and sampling caused by the neglection of the impact of noise in sampling SDEs. In contrast, our UniDB aims to **address the issues resulting from Doob's h-transform (e.g. artifacts along edges and unnatural patterns)** that occur in the existing diffusion bridge models. \n   - **Different methods.** Although both that paper and our work mention \"Stochastic Control\" in titles, **the two \"Stochastic Control\" are totally different**. SDB is more inclined towards **stochastic process**, adding noise into the base distribution and stochasticity to the reverse process which **is still based on Doob's h-transform**. While our UniDB focuses on stochastic **optimal** control, modeling the forward process as an **optimization problem** to analyze the drawbacks of Doob's h-transform. \n\nTherefore, the two articles are **fundamentally different** in nature.", "claims_and_evidence": "Claim 1: UniDB helps to understand and generalize Doob’s $h$-transform.\nThe theoretical derivations in the paper demonstrate that Doob’s $h$-transform is a special case of UniDB when the penalty coefficient approaches infinity. This provides a solid mathematical foundation for the claim. \n\nClaim 2: UniDB improves image quality by allowing the design of different controllers $u_{t, \\gamma}$.\nThe claim is supported by a lot of qualitative and quantified experimental results, demonstrating that controller design enhances image quality.", "methods_and_evaluation_criteria": "The use of PSNR, SSIM, LPIPS, and FID as evaluation metrics in experiments on three high-resolution datasets (CelebA-HQ, Rain100H, and DIV2K) is reasonable. These metrics comprehensively assess image quality, covering aspects of pixel-wise similarity (PSNR, SSIM), perceptual quality (LPIPS), and generative diversity (FID). The chosen datasets are also well-suited for high-resolution image generation and restoration tasks, ensuring a robust evaluation.", "theoretical_claims": "The mathematical derivations in the paper are complete and generally reliable. However, two aspects require further clarification from the authors:\n\n1. Choice of $L_1$ norm in the training objective (Equation 19):\nThe paper does not provide a clear justification for using the $L_1$ norm instead of other alternatives like $L_2$ norm. The authors should explain whether this choice is based on empirical performance, theoretical considerations, or robustness to outliers.\n\n2. Introduction of the state vector term $m$ in linear SDE:\nOne of the paper's key novelties is introducing the $m$ term in the linear SDE form, but it is not explicitly explained how it is computed or designed in the main context. Further elaboration on the motivation, computation, and impact of $m$ would improve the clarity of the contribution.", "experimental_designs_or_analyses": "The experiments are comprehensive and well-designed, covering three different tasks to ensure robustness and generalizability. Additionally, the paper conducts an ablation study on the key penalty coefficient $\\gamma$, which helps evaluate its impact on model performance.\n\nHowever, I suggest that the authors include DDBM as a benchmark for comparison. Since DDBM is also a Doob’s $h$-transform-based model and is mentioned in the **Preliminaries**, it would be beneficial to compare the proposed method against it. This would provide a clearer assessment of the advantages and potential improvements offered by the proposed framework.", "relation_to_broader_scientific_literature": "1. UniDB as a Generalization of Doob’s h-Transform:\nDoob’s h-transform has been widely studied in stochastic processes and has been applied in bridge modeling [1][2]. The paper demonstrates that Doob’s $h$-transform is a special case of the proposed UniDB framework when the penalty coefficient tends to infinity, providing a broader theoretical foundation for diffusion bridge models.\n\n2. Controller Design for Improved Image Generation:\nStochastic Optimal Control has been explored for diffusion bridge models [3], but existing approaches often lead to artifacts such as\nblurred or distorted details. By allowing the design of different controllers $u_{t, \\gamma}$, UniDB provides greater control over generation quality, leading to improved image fidelity and diversity across multiple datasets.\n\n[1] Zhou, Linqi, et al. \"Denoising diffusion bridge models.\", 2023.\n\n[2] Yue, C., et al. \"Image restoration through generalized ornstein-uhlenbeck bridge\", 2024.\n\n[3] Park, B., et al. Stochastic optimal control for diffusion bridges in function spaces, 2024.", "essential_references_not_discussed": "The key contribution of this paper is the introduction of Stochastic Optimal Control into the DDBM theoretical framework. However, a similar approach was explored last year in:\n\n[1] Zhang, Shaorong, et al. \"Exploring the Design Space of Diffusion Bridge Models via Stochasticity Control.\" arXiv preprint arXiv:2410.21553 (2024).\n\nTo differentiate from prior work, the paper should explicitly highlight the key distinctions between this work and Zhang et al. (2024), particularly in how Stochastic Optimal Control is formulated and applied.", "other_strengths_and_weaknesses": "The paper is well-structured and clearly organized, making it easy for readers to follow. However, the novelty is sceptical, as a similar approach was explored last year.", "other_comments_or_suggestions": "See the above comments.", "questions_for_authors": "See the above comments.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors reformulated a diffusion bridge matching problem as a stochastic control problem and developed a UniDB framework. 

The proposed UniDB approach focuses on stochastic optimal control, modeling the forward process as an optimization problem to analyze the drawbacks of Doob's h-transform.

As a result they arrived to a one-dimensional parametric family of solutions. Thanks to an additional parameter it is possible to improve quality of the distribution matching results.

Proofs are correct. Experimental validation is solid. The approach demonstrated the best accuracy on a number of relatively complex benchmarks.

The reviewers did come to a consensus here, so the decision is relatively easy.
Decision: Accept (spotlight poster)