Inverse Bridge Matching Distillation
Authors: Nikita Gushchin
David Li
Daniil Selikhanovych
Evgeny Burnaev
Dmitry Baranchuk
Alexander Korotin
Venue: ICML 2025 poster (2025)
Keywords: Diffusion Bridges
Bridge Matching
Diffusion Bridges Distillation
Image Translation
Inverse Bridge Matching
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Learning diffusion bridge models is easy; making them fast and practical is an art. Diffusion bridge models (DBMs) are a promising extension of diffusion models for applications in image-to-image translation. However, like many modern diffusion and flow models, DBMs suffer from the problem of slow inference. To address it, we propose a novel distillation technique based on the inverse bridge matching formulation and derive the tractable objective to solve it in practice. Unlike previously developed DBM distillation techniques, the proposed method can distill both conditional and unconditional types of DBMs, distill models in a one-step generator, and use only the corrupted images for training. We evaluate our approach for both conditional and unconditional types of bridge matching on a wide set of setups, including super-resolution, JPEG restoration, sketch-to-image, and other tasks, and show that our distillation technique allows us to accelerate the inference of DBMs from 4x to 100x and even provide better generation quality than used teacher model depending on particular setup. We provide the code at https://github.com/ngushchin/IBMD

Inverse Bridge Matching Distillation
Nikita Gushchin" 12
David Li
Daniil Selikhanovych 145
Evgeny Burnaev 12 Dmitry Baranchuk
Alexander Korotin
Abstract
Input
IBMD (Ours)
Teacher
Super-resolution
Learning diffusion bridge models is easy; mak- ing them fast and practical is an art. Diffusion
bridge models (DBMs) are promising extension of diffusion models for applications in image-to- image translation. However, like many modern diffuusinn and flow models, DBMs suffer from the problem of slow inference. To address it, we propose novel distillation technique based on the inverse bridge matching formulation and derive the tractable objective to solve it in prac- tice. Unlike previously developed DBM distil- lation techniques, the proposed method can dis- till both conditional and unconditional types of DBMs, distill models in one-step generator, and use only the corrupted images for training. We evaluate our approach for both conditional and unconditional types of bridge matching on a wide set of setups, including super-resolution, JPEG restoration, sketch-to-image, and other tasks, and show that distillation technique allows accelerate the inference of DBMs from 4x to 100x and even provide better generation qual- ity than used teacher model depending on par- ticular setup. We provide the code at https: //githubs
JPEG restoration
Inpainting
Normal-te to- Image
IBMD.
1. Introduction
Diffusion Bridge Models (DBMs) represent a specialized class of diffusion models designed for data-to-data tasks, such image-to-image translation. Unlike standard diffu- sion models, which operate by mapping noise data (Ho
Sketch-to- Imagge Image
et al., 2020; Sohl-Dickstein et al., 2015), DBMs construct diffusion processes directly between two data distributions (Peluchetti, 2023a; Liu et al., 2022b; Somnath et al., 2023;
Figure 1. Outputs of DBMs models distilled Inverse Bridge Matching Distillation (IBMD) approach various image-to- image translation tasks datasets (85). Teachers use NFE> 500 steps, while IBMD distilled models use NFES 4.
"Equal contribution Skolkovo Institute of Science and Technol- ogy Artificiil Intelligence Research Institute 'Moscow Institute of Physics and Technology *Yandex Research 'HSE University. Correspondence to: Nikita Gushchin <n. gushchineee @ skoltech.nu>, Alexander Korotin <o..oorttin skoltech.ru>.
Zhou et al., 2024a; Yue et al., 2024; Shi et al., 2023; De Bor- toli et al., 2023). This approach allows DBMs to modify only the necessary components the data, starting from an input sample rather than generating entirely from Gaus-
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Inverse Bridge Matching Distillation
sian noise. As result, DBMs have demonstrated impressive performance in image-to-image translation problems.
require the target data domain to perform distillation.
4. Better quality of distilled models. Our distillation tech- nique is tested on a wide set of image-to-image problems for conditional and unconditional DBMs in both one and multi-step regimes. It demonstrates improvements com- pared to the previous acceleration approaches including DBIM (Zheng et al., 2024) and CDBM (He et al., 2024),
The rapid development of DBMs has led to two dominant ap- proaches, usually considered separately. The first branch of approaches (Peluchetti, 2023a; Liu et al., 2022b; 2023a; Shi et al., 2023; Somnath et al., 2023) considered the construc- tion of diffusion between two arbitrary data distributions performing Unconditional Bridge Matching (also called the Markovian projection) of process given by mixture of diffusion bridges. The application of this branch includes different data like images (Liu et al., 2023a; Li et al., 2023), audio (Kong et al., 2025) and biological tasks (Somnath et al., 2023; Tong et al., 2024) not only in paired but also in unpaired setups using its relation to the Schrödinger Bridge problem (Shi et al., 2023; Gushchin al., 2024). The second direction follows a framework closer to classical diffusion models, using forward diffusion to gradually map to the point of different distibution rather than mapping distribu- tion to distribution as in previous case (Zhou et al., 2024a; Yue et al., 2024). While these directions differ in theoretical formulation, their practical implementations are closely re- lated; for instance, models based on forward diffusion can be seen as performing Conditional Bridge Matching with additional drift conditions (De Bortoli et al., 2023).
2. Background
In this paper, we propose universal distillation frame- work for both conditional and unconditional DBMs. To not repeat fully analogical results for both cases, we denote by this color the additional conditioning on ØT used for the conditional models, i.e. for the unconditional case this conditioning is not used.
2.1. Bridge Matching
We start by recalling the bridge matching method (Peluchetti, 2023b;a; Liu et al., 2022b; Shi et al., 2023). Consider two probability distributions p(æo) and p(ær) on RD dimensional space, which represent target and source domains, respectively. For example, in image inverse problem, p(zo) represents the distribution of clean im- ages and p(zr) the distribution of corrupted images. Also consider a coupling p(zo. of these two distributions, which is a probability distribution on TR. x RD. Cou- pling p(zo, can be provided by paired data or con- structed synthetically, i... just using the independent distri- bution p(z0, æT) = p(zo)p((T) Bridge Matching aims to construct the diffusion that transforms source distribution p(rr) to target distribution p(ro) based on given coupling p(ro, .T) and specified diffusion bridge.
Similar to classical DMs. DBMs also exhibit multi-step sequential inference, limiting their adoption in practice. De- spite the impressive quality shown by DBMs in the practical tasks, only few approaches were developed for their accel- eration, including more advanced sampling schemes (Zheng et al., 2024; Wang et al., 2024) and consistency distillation (He et al., 2024), adapted for bridge models. While these approaches significantly improve the efficiency of DBMs, some unsolved issues remain. The first one is that the de- veloped distillation approaches are directly applicable only for DBMs based on the Conditional Bridge Matching, i.e., no universal distillation method can accelerate any DBMs. Also, due to some specific theoretical aspects of DBMs, con- sistency distillation cannot be used to obtain the single-step model (He et al., 2024, Section 3.4).
Diffusion bridges. Consider forward time diffusion Q called "Prior" on time horizon [o, T] represented by the stochastic differential equation (SDE):
Prior Q dx: = sx ()dt + g(()dwr, f(z,,t) : RD x [0, + RR, g(t): : [0, T| -R.
(1)
Contributions. To address the mentioeet issues of DBMs acceleration, we propose a new distillation technique based on the inverse bridge matching problem, which has several advantages compared to existing methods:
where f( 1) is a drift function, g(t) is the noise schedule function and dw, is the differential of the standard Wiener process. By (( we denote the transition probability density of prior process Q from time s to time t. Diffussion bridge is a conditional process Q (igg.ar. which is obtained by pinning down starting and ending points xo and xT. This diffusion bridge can be derived from prior process Q using the Doob-h transform (Doob & Doob, 1984):
1. Universal Distillation. Our distillation technique is ap- plicable to DBMs trained with both conditional and un- conditional regimes, making it the first distillation ap- proach introduced for unconditional DBMs.
2. Single-Step and Multi-step Distillation. Our distilla- tion is capable of distilling DBMs into generators with any specified number of steps, including the distillation of DBMs into one-step generators.
Diffusion Bridge Q Fo, XT are fixed, dx.- = (f + (t) log (( Jt + g(t)dur,
(2)
For this diffusion bridge we denote the distribution at time t of the diffusion bridge Qi by q xT))
3. Target data-free distillation, Our method does not
2

---

Inverse Bridge Matching Distillation
For conditional Bridge Matching model
IT p(æT)
Tt g (zy 27, z0)
zo ~ p(
T
Z
2
Tol (25 tym
Es
2
Figure Overview of (Conditional) Bridge Matching with Sy reparamcterization. The process begins by sampling a pair (2o. TT) from the data coupling p p(ro, .r). An intermediate sample ze is then drawn from the diffusion bridge [[z 2r) a random time t r U',,T), The model Ep is trained with an MSE loss to reconstruct zu from It. In the conditional setting (dashed red path), Tu is also conditioned on 27 as additional input, leveraging information about the terminal state improve reconstruction.
Mixture of bridges. Bridge Matching procedure starts with creating a mixture of bridges process IT. This process is represented as follows:
Time moment here is sampled according to the uniform distribution on the interval [o, T).
Relation Between Flow and Bridge Matching. The Flow Matching (Liu et al., 2023b; Lipman et al., 2023) can be the limiting case - 0 of the Bridge Matching for particular example see (Shi et al., 2023, Appendix A.
Mixture of Bridges IT:
11(-) / mo. (-)((z0, IT xT)
(3)
Practically speaking, the definition (3) means that to sample from a mixture of bridges IT, one first samples the pair (zo; XT) ~ p(x0, IT) from data coupling and then samples trajectory from the bridge (-).
2.2. Augmented (Conditional) Bridge Matching and Denoising Diffusion Bridge Models (DDBM)
For a given coupling p(ra, IT) = p( (r )p(T)), one can use an alternative approach to build data-to-data diffusion. Consider a set of Bridge Matching problems indexed by ør between Po = p(r) and p(ær) = 6gy (x) (delta mea- sure centered 2T)) This approach called Augmented Bridge Matching (De Bortoli et al., 2023). The key dif- ference of this version in practice is that it introduces the condition of the drift function v" (æz, t, xr)) on the starting point Ty in the reverse time diffusion (5):
Bridge Matching problem. The mixture of bridges II can- not be used for data-to-data translation since it requires first to sample pair of data and then just inserts the trajectory. In turn, we are interested in constructing a diffusion, which can start from any sample ZT ~ p(xr) and gradually trans- form it to zo ~ p(rn). This can be done by solving the Bridge Matching problem (Shi et al., 2023, Proposition 2)
Bridge Matching problem:
(4)
dr. (f(())) (t) (z,, xr) ]dt + g())
(III
def arg min (IIII MEM
The drift function or can be recovered almost in the same way just by the addition of this condition on Zy:
where M the set of Markovian processes associated with some SDE and (1 (M) the KL-divergence between constructed mixture of bridges II and diffusion M. It is known that the solution of Bridge Matching is the reversed- time SDE (Shi et al., 2023, Proposition 9):
Augmented (Conditional) Bridge Matching Problem. min Exog [llva [ (, - Vz, log (1
(zo: TT) ~ p(zo: TT), and x; ~ g( [zo zT).
The SDE of Bridge Matching solution: : dr;= = (x+) - g² (Ie, d) + g(t)dū, IT ~ (TT)),
(5)
Since the difference is the addition of conditioning on ÆT. call this approach Conditional Bridge Matching.
Relation to DDBM. As was shown in the Augmented Bridge Matching (De Bortoli et al., 2023), the conditional Bridge Matching is equivalent to the Denoising Diffusion Bridge Model (DDBM) proposed (Zhou et al., 2024a). The difference is that in DDBM, the authors learn the score function of s s(xe, ar, t)) conditioned ør of process for which 2g p( (ry) and g(x;) q( ær): Then, it is combined with the drift of forward Doob-h transform (5) to get the reverse SDE drift v(x,, t, ITT:
where is standard Wiener process when time flows backward from é = T' to O, and dt is an infinitesimal negative timestep. The drift function v* is obtained solving the following problem (Shi et al., 2023; Liu et al., 2023a):
Bridge Matching problem with a tractable objective: (6)
min minE E... [llvs (x1, t) - V., log q(
(zo, .T) ~ p(2o, t U(0, T),, ot ~ g(( xo, xT)
e( (T) = s(1,, xr, 1) - Vz, log ((r(()))
3

---

Inverse Bridge Matching Distillation
dr, (1 l)dl o² (1)o(1, wr) + (()dur
and Gaussian transitional kernels. Thanks to it, for each øT to sample, one use the probability flow ODE and ODE- solvers or hybrid solvers to accelerate sampling (Zhou et al., 2024a) or use consistency distillation of bridge models (He et al., 2024). Another beneficial property is that one can con- sider a non-Markovian forward process to develop a more efficient sampling scheme proposed in DBIM (Zheng et al., 2024) similar to Denoising Diffusion Implicit Models (Song et al., 2021). However, in the Unconditional Bridge Match- ing problem, the forward diffusion process, which maps p(zo) to p(ær) without conditioning on specific point 2T, is unknown. Hence, the abovementioned methods cannot be used to accelerate this model type.
or reverse probability flow ODE drift:
VODE (,T) = 2 - Tr, 1)- Vae log g ((r)))))
dx, (1(aa, (1)onnn( (,rr) )dl,
which is used for consistency distillation in (He et al., 2024).
2.3. Practical aspects of Bridge Matching
Priors used in practice. In practice (He et al., 2024; Zhou et al., 2023; Zheng et al., 2024), the drift of the prior pro- cess is usually set to be f(x,, t) = ff) i., it depends linearly on z4. For this process the transitional distribution q( = (xslarrto I) ) is Gaussian, where:
3. IBMD: Inverse Bridge Matching Distillation
This section describes proposed universal approach to distill the both Unconditional and (Conditional) Bridge Matching models v* (called the teacher model) into few- step generator using only the corrupted data pr(xr). The key idea of our method is to consider the inverse problem of finding the mixture of bridges IIg for which Bridge Match- ing provides the solution Uo with the same drift as the given teacher model v². We formulate this task the optimiza- tion problem (83.1). However, gradient methods cannot solve this optimization problem directly due to the absence of tractable gradient estimation. To avoid this problem, we prove a theorem that allows us to reformulate the inverse problem in the tractable objective for gradient optimiza- tion ($3.2). Then, we present the fully analogical results for the Conditional Bridge Matching case in (53.3). Next, we present the multi-step version of distillation ($3.5) and the final algorithm (83.4). We provide the proofs for all considered theorems and propositions in Appendix A.
f(t) d'log at
g²(()) = dt - dlogar ar of
dop
dt
dt
The bridge process distribution is also a Gaussian (00 = N(arlaar + bero, I) with coefficients:
04 SNRT a2 oT SNR,
SNR, ) at of 1 SNRT SNR,
b at (1- [_SRRR)
a o/ is the signal-to-nol ratio at time t.
where SNR,-
Data prediction reparameterization. The regression target of the loss function (6) for the priors with the drift v( 1) given by V.. log q( (zo)
Zu- Hence, one can
of
use the parametrization v( L, #T) and solve the equivalent problem:
3i ao- (.ar))
of
3.1. Bridge Matching Distillation as Inverse Problem
Reparametrized (Conditional) Bridge Matching problem:
In this section, we focus on the derivation of our distilla- tion method for the case of Unconditional Bridge Match- ing. Consider the fitted teacher model v² v'(xe, t), which is an SDE drift of some process M = BM(II*), where II constructed using some data coupling p" (xT,, p" (z (r (p)rp())) ( (eulerr)((2)) p(ær) and aim to find such H. build on po (To, TT), that BM(II") we parametrize po(zo (IT) by the stochastic generator GolØT, 2), ~ (O,, which generates samples based on input TT ~ p(ar) and the gaussian noise z. Now, we formulate the inverse problem as follows:
min Ee [A(t)] 20 (,,,t,,7)) t, - ==0l
(7)
(T) ~ p(z0, oT); L ~~ U(00, T)), xe ~ g( (ro, xT);
where X(1) is any positive weighting function. Note that Ør used only for the Conditional Bridge Matching model.
We parametrize pa pa 2T)
2.4. Difference Between Acceleration of Unconditional and Conditional DBMs
BM(Ilo). In practice,
Since both conditional and unconditional approaches learn drifts of SDEs, they share the same problems of long in- ference. However, these models significantly differ in the approaches that can accelerate them. The source of this difference is that Conditional Bridge Matching considers the set of problems of reversing diffusion, which gradually transforms distribution p(zy (TT) to the fixed point ør. Fur- thermore, the forward diffusion has simple analytical drift
min (L.
(8)
Note, that since the objective (8) is the KL-divergence be- tween BM(Ils) and M it is equal to 0 if and only (M())) and M* coincide, Furthermore, using the disinte-
4

---

Inverse Bridge Matching Distillation
For condilional Bridge Matching model distillation
20
2T ~ p(Tr)
Z7g ~ Po (2 (T)
Go(2,ær) Ør)
To ( 1ar0
C
q( IT, Fo)
Fo (215 (.
Zg
Figure 3. Overview of our method Inverse Bridge Matching Distillation (IBMD). The goal distill trained (Conditional) Bridge Matching model into a generator Co(z, wr). which learns to produce samples using the corrupted data p(ær). Generator Go(z, xr) defines the coupling p= .r) = p= (rr)p(Z)) and we aim learn the generator in such way that Bridge Matching with pa (a Tr) produces the same (Conditional) Bridge Matching model ig Tg- To do so, we learn bridge model 26 using coupling ps in the same way as the teacher model was learned. Then, we use our novel objective given Theorem 3.2 to update generator model Go.
gration and Girsanov theorem (Vargas et al., 2021; Pavon & Wakolbinger, 1991), we have the following result:
This is the general result that can applied with any diffusion bridge. For the priors with with drift f(x,, t) = f(t)a,, we present its reparameterized version.
Proposition 3.1 (Inverse Bridge Matching problem). The inverse problem (8) is equivalent to
Proposition 3.3 (Reparameterizeed tractable inverse prob- lem). Using the reparameterizattion (62.3) the prior with the linear drift f(x,,) = f(t)xz, the inverse problem in Theorem 3.2 is equivalent to:
min E [A(t)]l((zn.t) t) v² (IT,,))13) s.t.
(9)
v = arg min E2,, [Io' (Ie, - Vv, log ( (m0) 13].
min Ex [A(1)| =ol13]-
(zo, xT) ~ (0 æT), t ~ U(0,T)), x ~ q(<e|2o. 2T).
min E na 1.00 [A(t] 20 (2,, (z,,t) - =zol"1
where A(t) is any positive functionn
function.
(aro, TT) ~ pa (0 XT), , ~~ U10, TJ), 22 ~ g( (20, TT))
Thus, this is the constrained problem, where the drift is the result of Bridge Matching for coupling pa (zo, æT) parametrized by the generator Gg. Unfortunately, there is
Thanks to the unconstrained reformulation, this problem ad- mits explicit gradients with respect to the generator Gg. as all samples ( IT, 2z)) are obtained via reparameterizable transformations: xo Go 2) with 2 ~ (O 7), and ~ g( (o. where g(xx, In, aT) is Gaussian dis- tribution (under priors with linear drift f t) f()) This enables differentiability of the entire objective, which involves an expectation over pe (o(xo, W. and allows opti- mization using standard gradient-based methods.
no clear way to use this objective efficiently for optimizing generator Ge since would require gradient backpropaga- tion through the argmin of the Bridge Matching problem.
3.2. Tractable objective for the inverse problem
In this section, we introduce our new unconstrained refor- mulation for the inverse problem (9), which admits direct optimization using gradient methods:
Interpretation of the auxiliary model á. Note that the minimal value of the inner problem is the averaged variance of zo ~ p Ie, TT):
Theorem 3.2 (Tractable inverse problem reformulation).
The constrained inverse problem (9) w.r.t 0 is equivalent to the unconstrained optimization problem:
min Burittrr (x()) 1 (88(111))) 200
=
min E., r.o (()))l"((eet) a. log (xg(((1(20)11)
Exi,, (ACe) [| [(pootoo [20) x0l1
min t. (Z)), - Va, log log (No, 2T) po(zo. 2T), U(O,T)), T1), ot g(.e TT).
Er,, (X(t) Eps(anolan)
||| B (ao
zool
of (2olax))
Where the constraint in the original inverse problem (9) is relaxed by introducing the inner bridge matching problem.
For t = T, this is directly the variance of the generator To ~ po IT). Since this part comes with a negative
5

---

Inverse Bridge Matching Distillation
sign in the objective, its minimization enforces the generator to produce more diverse outputs and avoid collapsing.
4. Related work
Diffusion Bridge Models (DBMs) acceleration. Unlike a wide scope of acceleration methods developed for classical diffusion/flow models, only a few approaches were devel- oped for DBM acceleration. Acceleration methods include more advanced samplers (Zheng et al., 2024; Wang et al., 2024) based on a reformulated forward diffusion process as a non-markovian process inspired by Denoising Diffusion Implicit Models (Song et al., 2021). Also, there is a distilla- tion method based on the distilling probability into few steps using consistency models (He et al., 2024), which is applicable only for conditional DBMs. However, for theoretical reasons (He et al., 2024, Section 3.4), consis- tency models for Diffusion Bridges cannot be distilled into one-step generators. Unlike existing distillation methods, our method is applicable to both conditional and uncondi- tional DBMs and can distill into a one-step generator.
3.3. Distillation of conditional Bridge Matching models
Since Conditional Bridge Matching is, in essence, set of Unconditional Bridge Matching problems for each ær ($2.2), the analogical results hold just by adding the condi- tioning on 2r for v, i.e., using v(ze, f, Fr) or Zo. i.e. using (o 1, .T). Here, we provide the final reparametrized formulation, which we use in our experiments:
ODE
Theorem 3.4 (Reparameterized tractable inverse problem for conditional bridge matching).
min Be. ( 2,7) - x0|1
(10)
0
min E Exg [A(1]] (1), (2,, 1, (,) - =oll"]]
(mo, xT) ~ po(zo. t ~ ((0,T)), x ~ (0: vT).
Related diffusion and flow models distillation techniques. Among the methods developed for the distillation of classi- cal diffusion and flow models, the most related to our work are methods based on simultaneous training of few-step generators and auxiliary "fake" model, that predict score or drift function for the generator (Yin et al., 2024b;a; Zhou et al., 2024b; Huang et al., 2024). Unlike these approaches, we consider the distillation of Diffusion Bridge Models the generalization of flow and diffusion models.
where A(1) is some positive weight functiin.
To use in practice, we parameterize Fo t, .r) by a neural network with an additional condition on Ty.
3.4. Algorithm
We provide a one-step Algorithm that solves the inverse Bridge Matching problem in the reformulated version that we use in our experiments. We provide visual abstract of it in Figure 3. Note that a teacher in the velocity parame- terization 2n (xz, f) can be easily reparameterized (52.3) in co-prediction model using 3et (xt, t) = 0f0" (w..))+
Furthermore, previous distillation methods for diffusion and flow models rely on marginal- based losses such as Fisher divergence, these approaches do not account for the full structure of path measures. This limitation becomes critical in the context of Diffusion Bridge Models (DBMs), where the dynamic aspects of the forward and reverse processes play fundamental role. To better motivate the need for our KL-based objective, we next discuss the conceptual differences between KL divergence of path measures and Fisher divergence, illustrating why Fisher-based objectives like those used in SiD (Zhou et al., 2024b) are insufficient in the general setting of bridge matching. Consider two reverse-time diffusions D2 and D2 given by the same start- ing distribution p(IT) and SDEs:
3.5. Mulitistep distillation
We also present multi-step modification of our distilla- tion technique if a one-step generator struggles to distill the models, e.g., in inpainting setups, where the corrupted image IT contains less information. Our multi-step tech- nique is inspired by similar approaches used in diffusion distillation methods (Yin et al., 2024a, DMD) and aims to avoid training/inference distribution mismatch.
Dj dr. v(ax,, l)dt + 9 (1)dvre,,
or p(oT),
We choose N timesteps {0 < t1 < t2 <.. < = =tN==1 T}
D2 : dx, = Mare, t)dt + g² (I)dir,,
IT ~ p(PT)
and add additional time input to our generator Go 2, t). For the conditional Bridge Matching case, we also add conditions on ØT and use Gø(ze, 2, 2T). To perform in- ference, we alternate between getting prediction from the generator To Ga 2, and using posterior sampling ((Ft,, , 1o .i.,,) given by the diffusion bridge. To train the generator in the multi-step regime, we use the same procedure as in one step except that to get input o, for in- termediate times in < tNy we first perform inference of our generator to get Zo and then use bridge q (.))
Let p(æz) and p(are) be the corresponding marginals. Then the KL divergence and Fisher divergence are given by:
(L(Di||D2) Eep(())) 22 (1( -
a(or,)))
(I. )lF(TT))
(11)
-Off((r))=
Drisher PrisherDDi110 (2) (:P2)=Fnn -E 12, logp(x)) - a,, log)())))
6

---

Inverse Bridge Matching Distillation
In SiD (Zhou et al., 2024b), Fisher divergence (Drishe)) is averaged over time and compares only marginal distribu- tions p((x) and p(z+) of path measures. However, path measures with the same marginal distributions might not be equal; thus, in general, minimizing Fisher divergence does not guarantee that D1 x D2 as stochastic processes, For classical diffusion models, the forward drift (x,, 1) is fixed, and reverse drifts are fully determined by score functions:
Algorithm Inverse Bridge Matching Distillation (IBMD)
Input : :eacher network z7 IR x [0, x RD + RD;
Bridge g (To, zT) used for training x=;
Generator network Gg : RD x TR - RD:
Bridge network _6 RP x [0. T] x R + RD:
Input distribution p(ær) accessible by samples;
Weights function X(t) [0, + Rt: Batch size N: Number of student iterations K; Number of bridge iterations L.
F( t) f(a,)) (1) 21 log ((ae); v t) (x), - g² v.. log p(z)).
Output: Learmed generator Go of coupling po (20, 2r) for which Bridge Matching outputs drift e x v". // Conditioning on #T is used only for distillation of Condi- tional Bridge Matching models.
Substituting these into the KI. expression (11) shows that in this specific setting - with a fixed forward SDE - KL di- vergence between path measures becomes equivalent (up to a constant) to the time-averaged Fisher divergence between the marginals. This explains why Fisher-based methods like SiD (Zhou et al., 2024b) may succeed in this context.
for ls to K do for L= 1 to L do
Sample batch or ~ p(FT)
Sample batch of noise 2 ~ NO, 1)
20 + -GG(TT,2) 2)
However, this equivalence breaks down in the case of unconditional bridge matching. Here, the forward drift (( is not fixed and depends on the data coupling p(20, ØT). In turn, the forward drift fo(T2, t) for the gen- erated coupling pa (#o, ZT) also depends on 0. As a result, * t(r,,) fo t). and the reverse drifts cannot be ex- pressed solely in terms of marginal scores. Hence, KL divergence in the case of unconditional bridge matching is not equivalent to Fisher divergence between marginals. This difference is expected since, in the case of an uncondi- tional diffusion bridge, one does not have a fixed forward process, which specifies the "dynamic part" of the measure. This highlights the importance of using KI. divergence be- tween path measures as a high-level objective instead of the previously used Fisher Divergence.
Sample time batch t U/0,T]
Sample batch Se ~ ( (00,
L. +- N Σ [1())| (2e,, (T) - x0113)
Update di by using 00 ac,
Sample batch ør ~ p(xT)
Sample batch of noise
(0,
To + -G((Fr 2)
Sample time batch 1. ~ U10, 7]
Sample batch De ~ ( (20 fT)
- X()) (Zi,, L TT) - =00l 8.,
t,,) - zol1 1|
Update by using à0
5. Experiments
I'SB paper (Liu et al., 2023a), specifically (a) two models for the 4x super-resolution with bicubic and pool kernels, (b) two models for JPEG restoration using quality factor QF= 5 and QF= 10, and (c) a model for senter-inpaintio with a center mask of size 128 x 128 all of which were trained on ImageNet 256 x 256 dataset (Deng et al., 2009).
This section highlights the applicability of our IBMD distil- lation method in both unconditional and conditional settings. To demonstrate this, we conducted experiments utilizing pretrained unconditional models used in 1SBB paper (Liu et al., 2023a). Then we evaluated IBMD in conditional settings using DDBM (Zhou et al., 2024a) setup (85.2). For clarity, we denote our models as IBMD-DDBM and IBMD-I'SB indicating that the teacher model is derived from DDBM or I'SB framework, respectively. We provide all the technical details in Appendix B.
For all the setups dataset, which was used to train the used models. For the evaluation we follow the same protocol used in the 1 paper, i.e. use the full validation subset of ImageNet for super-resolution task and the 10'000 subset of validation for other tasks. We report the same FID (Heusel et al., 2017) and Classifier Accuracy (CA) using pre-trained ResNet50 model metrics used in the IS paper. We present results in Table 1, Table 3. Table 2, Table 4 and Table 6. We provide the uncurated samples for all setups in Appendix C.
the same train part of ImageNet
5.1. Distillation of I2SB (5 setups)
Since known distillation and acceleration techniques are designed for the conditional models, there is no clear base- line for comparison. Thus, this section aims to demonstrate that our distillation technique significantly decreases NFF required to obtain the same quality of generation.
Results. For both super-resolution tasks (see Table 1, Ta- ble 3), our 1-step distilled model outperformed teacher model inference using all 1000 steps used in the training.
Experimental Setup. To test our approach for uncondi- tional models, we consider models trained and published in
7

---

Inverse Bridge Matching Distillation
Table Results the image super-resolution task. Baseline re-
Table 3. Results the image super-resolution task. Baseline re-
sults are taken from SB (Liu et al., 2023a).
sults are taken from I'SB (Liu et al.. 2023a).
4x super-resolution (bicubie)
ImageNet (256 x 256)
4x super-resolution (pool)
ImageNet (256 256)
NFE
FID 4
CA f
NFE
FID +
CA
DDRM (Kawar al., 2022)
20
21.3
63.2
DDRM (Kawar et al., 2022)
20
14.8
64.6
DDNM (Wang al., 2023)
100
13.6
65.5
DDNM (Wang et al., 2023)
100
9.9
67.1
IIGDM (Song et al., 2023)
100
3.6
72.1
IIGDM (Song et al., 2023)
100
3.8
72.3
ADM (Dhariwal & Nichol, 2021)
1000
14.8
66.7
ADM (Dhariwal & Nichol, 2021)
1000
3.1
73.4
CDSB (Shi et al., 2022)
50
13.6
61.0
CDSB (Shi et al., 2022)
50
13.0
61.3
PSB (Liu et al., 2023a)
1000
2.8
70.7
18 (Liu et al., 2023a)
1000
2.7
71.0
IBMD-I (Ours)
2.6
72.3
IBMD-I'SB (Ours)
2.5
72.5
Table Results the image JPEG restoration task with QF=3.
Table 4. Results the image JPEG restoration task with QF=10. Baseline results are taken from IBSB (Liu et al., 2023a).
Baseline results are taken from FSB (Liu et al., 2023a).
JPEG restoration, QF= 5.
ImageNet (256
256)
JPEG restoration, QF= 10.
ImageNet (256
256)
NFE
FID +
CA f
NFE
FID
4
CA 1
DDRM (Kawar et al., 2022)
20
28.2
53.9
DDRM (Kawar et al., 2022)
20
16.7
64.7
TGDM (Song et al., 2023)
100
8.6
64.1
IIGDM (Song et al., 2023)
100
6.0
71.0
Palette (Saharia et al., 2022)
1000
8.3
64.2
Palette (Saharia et al., 2022)
1000
5.4
70.7
CDSB (Shi et al., 2022)
50
38.7
45.7
CDSB (Shi et al., 2022)
50
18.6
60.0
I'SB (Liu et al., 2023a)
1000
4.6
67.9
I'SB (Liu et al., 2023a)
1000
3.6
72.1
JSBB (Liu et al., 2023a)
100
5.4
67.5
ISSB (Liu al., 2023a)
100
4.4
71.6
IBMD-I'SB (Ours)
5.2
66.6
IBMD-I'SB (Ours)
3.7
72.4
Note that model does not use the clean training target data p(ro),, only the corrupted p(zr), hence this improve- ment is not due to additional training using paired data. We hypothesize that it is because the teacher model introduces approximation error during many steps of sampling, which may accumulate. For both JPEG restoration (see Table 2, Table 4), our 1-step distilled generator provides the qual- ity of generation close the teacher model and achieves around 100x time accelcration. For the inpainting problem (see Table 6), we present the results for 1, and 4 steps distilled generator. Our 2 and 4-step generators provide a quality similar to the teacher IBSB model, however, there still some gap for the 1-step model. These models pro- vide around 5x time acceleration. We hypothesize that this setup is harder since it requires to generate the entire center fragment from scratch, while in other tasks, there is already some good approximation given by corrupted images.
et al., 2024). For the image-to-image translation task, dataset (Isola et al., 2017)
utilize the Addges>>anddee with a resolution of 64 x 64 pixels and the DIODE-Outdoor dataset (Vasiljevic et al., 2019) with a resolution 256 256 pixels. For these tasks, we report FID and Inception Scores (IS) (Barratt Sharma, 2018). For the image inpainting task,
the same setup of center-inpainting before.
Results. We utilized the same teacher model checkpoints and as in CDBM. We present the quantitative and qualitative results of IBMD on the amage-to-image translation task in Table 5 and in Figures 12, 10 respectively. The compet- ing methods, DBIM (Zhou et al., 2024a, Section 4. and CDBM (He al., 2024, Section 3.4), cannot use single-step inference due the singularity at the starting point ær-
We trained our IBMD with 1 and 2 NFES on the Edges-Haandaags dataset. We surpass CDBM at 2 NFE, outperform the teacher at 100 NFE, and achieve perfor- mance comparable to the teacher at 50 NFE with 1 NFE, resulting 50x acceleration. For the DIODE-Outdoor setup, we trained IBMD with 1 and 2 NFEs. We surpassed CBD in FID at 2 NFE, achieving results comparable CBT with slight drop in performance and maintaining strong performance at 1 NFE with minor quality reductions.
5.2. Distillation of DDBM setups)
This section addresses two primary objectives: (1) demon- strating the feasibility of conditional model distillation within our framework and (2) comparing with the CDBM (He et al., 2024) leading approach in Conditional Bridge Matching distillation, presented into different models: CBD (consistency distillation) and CBT (consistency training).
For image inpainting, we show in Table quantitative results and in Figure 9 the guantitative results. We train IBMD with NFE in this setup. It outperforms CBD and CBT
Experimental Setup. For evaluation, we use the same setups used in competing methods (He et al., 2024; Zheng
8

---

Inverse Bridge Matching Distillation
Table Results the Image-to-Image Translation Task (Training Sets). Methods grouped by NFE (> 2, with the best metrics
bolded in each group. Baselines results are taken from CDBM.
Edges + Handbags (64 x 64)
DIODE-Outdoor (256 x 256)
NFE
FID |
IS f
FID +
1
DDIB (Su al., 2023)
2 40
186.84
2.04
242.3
4.22
SDEdit (Meng et al., 2022)
40
26.5
3.58
31.14
5.70
Rectified Flow (Liu al., 2022a)
2 40
25.3
2.80
77.18
5.87
SB (Liu et al., 2023a)
2 40
7.43
3.40
9.34
5.77
DBIM (Zheng al., 2024)
50
1.14
3.62
3.20
6.08
DBIM (Zheng et al., 2024)
100
0.89
3.62
2.57
6.06
CBD (He et al., 2024)
1.30
3.62
3,66
6,02
CBT (He et al., 2024)
2
0.80
3.65
2.93
6.06
IBMD-DDBM (Ours)
0.67
3.69
3.12
5.92
Pix2Pix (Isola et al., 2017)
74.8
4.24
82.4
4.22
IBMD-DDBM (Ours)
1.26
3.66
4.07
5.89
Table 6. Results on the Image Inpainting Task. Methods are grouped by NFE (> 4, 4, 1), with best metrics bolded in each group. Baselines results are taken from CDBM.
model and teacher model on these sets in Figures 13 and in Appendix C. We also provide the uncurated samples on the train part in Figures 12 and 10 to compare models' behavior on train and test sets. From these results, that the teacher model exhibits overfitting both setups, c.g., it produces exactly the same images corresponding reference images. In turn, on the test sets, teacher models
ImageNet (256 x 256)
Inpainting, Center (128 x 128)
NFE
FID
CA 1
DDRM (Kawar et al., 2022)
20
24.4
62.1
IIGM (Song et al., 2023)
100
7.3
72.6
DDNM (Wang al., 2023)
100
15.1
55.9
work well for the handbag setups, while on the test set of
Palette (Saharia al., 2022)
1000
6.1
63.0
DIODE images, it exhibits mode collapse and produces gray images. Nevertheless, our distilled model shows exactly
SB (Liu et al., 2023a)
10
5.4
65.97
DBIM (Zheng et al., 2024)
50
3.92
72.4
the same behavior in both sets, i.e., our IBMD approach precisely distills the teacher model expected.
DBIM (Zheng et al., 2024)
100
3.88
72.6
CBD (He et al., 2024)
5.34
69.6
6. Discussion
CBT (He et al., 2024)
4.77
70.3
IBMD-I'SB (Ours)
5.1
70.3
Potential impact. DBMs are used for data-to-data trans- lation in different domains, including images, audio, and biological data. Our distillation technique provides univer- sal and efficient way to address the long inference of DBMs, making them more affordable in practice.
IBMD-DDBM (Ours)
4.03
72.2
CBD (He al., 2024)
5.65
69.6
CBT (He al., 2024)
5.34
69.8
2
IBMD-I-SB (Ours)
5.3
65.7
IBMD-DDBM (Ours)
4.23
72.3
Limitations. Our method alternates between learning an additional bridge model and updating the student, which may be computationally expensive. Moreover, the student optimization requires backpropagation through the teacher, additional bridge, and the generator network, making it 3x time more memory expensive than training the teacher.
IBMD-I'SB (Ours)
6.7
65.0
IBMD-DDEM (Ours)
5.87
70.6
NFE with a significant gap, surpassing both at NFE and maintaining strong performance at NFE while achieving teacher-level results (with 50 NFE) with 12.5x speedup.
Concerns regarding the evaluation protocol used in prior works. For Gddges-Handbel and DIODE-Outdon setups, we follow the evaluation protocol originally introduced in DDBM (Zhou et al., 2024a) and later used in works on acceleration of DDBM (Zheng al., 2024; He et al., 2024). For some reason, this protocol implies evaluation of the train set. Furthermore, test sets of these datasets consist of tiny fraction of images (around several hundred), making the usage of standard metrics like FID challenging due to high statistical bias or variance their estimation. Still, to assess the quality of the distilled model on the test sets, we provide the uncurated samples produced by our distill
A Acknowledgements
The work was supported by the grant for research centers in the field of AI provided by the Ministry of Economic Development of the Russian Federation accordance with the agreement 0000003113225P44000 and the agreement with Skoltech 1339-10--0255-03
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences our work, none which we feel must be specifically highlighted here.

---

Inverse Bridge Matching Distillation
References
Li, B., Xue, K., Liu, B., and Lai, Y.-K. Bbdm: Image-to- image translation with brownian bridge diffusion mod- els. In Proceedings of the IEEE/CVF conference on computer vision and pattern Recognition, pp. 1952-1961, 2023.
Barratt, S. and Sharma, R. A note the inception score. arXiv preprint arXiv: 1a1iv1880111 18018111101011 801.01973, 2018.
De Bortoli, V., Liu, G.-H., Chen, T., Theodorou, E. A., and Nie, W. Augmented bridge matching. arXiv preprint NWV:2311.06997. 2023.
Lipman, Y., Chen, R. T. Q.. Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. URL petttps://ppenee//// net / weerrmmmdddPPMMRDC.TT.
Deng, J., Dong. W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp.
leee, 2009.
Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E. A., Nie, W., and Anandkumar, A. Pesb: Image-to-image schr\," " odinger bridge. The Fortieth International Conference on Machine Learning. 2023a.
Dhariwal, and Nichol, A. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, $4:8780-8799 2021.
Doob, J. L. and Doob, J. Classical potential theory and its probabilistic counterpart, volume 262. Springer, 1984.
Liu, X., Gong, C., et al. Flow straight and fast: Learn- ing to generate and transfer data with rectified flow. The Eleventh International Conference on Learning Representations, 2022a.
Gushchin, N., Selikhanovych, D., Kholkin, S., Burnaev, E.,, and Korotin, A. Adversarial schrödinger bridge match- ing. In The Thirth-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: // //pppnrevvie... net um 1d=1ddd31t -edd--nnniniccn
Liu, X., Wu, L.., Ye, M., and qiang liu. Let build bridges: Understanding and extending diffusion generative models. In NeurIPS 2022 Workshop on Score-Based Methods, 2022b. URL http:: /ppeneeviewww net forum? 1d=0e
He, G., Zheng. K., Chen, J., Bao, F., and Zhu, J. Consistency diffusion bridge models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Liu, X., Gong, C., and qiang liu. Flow straight and fast: Learning to generate and transfer data with rec- tified flow. In The Eleventh International Conference on Learning Representations, 2023b. URL https:
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
/ //oppnreeiee.eeteeeeeeee
Meng, C., He, Y., Song, Y., Song, J., Wu, J., Zhu, J.-Y., and Ermon, S. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2022.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems,
2020.
Pavon, M. and Wakolbinger, A. On free energy, stochas- tic control, and schrödinger processes. In Modeling, Estimation and Control of Systems with Uncertainty; Proceedings of Conference held in Sopron, Hungary, September 1990, pp. 334-348. Springer, 1991.
Huang, Z., Geng, Z., Luo, W., and Qi, G.-j. Flow generator matching. arXiv preprint arXiv:2410.19310 2024.
Isola, P., Zhu, J.,Y., Zhou, T., and Efros, A. A. Image-to- image translation with conditional adversarial networks. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 1125-1134, 2017.
Peluchetti, S.
Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023a.
Kawar, B., Elad, M., Ermon, S., and Song, De- noising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593- 23606, 2022.
Peluchetti, S. Non-denoising forward-time diffusions. arXiv preprint arXiv:2312.14588 2023b.
Kong, Z., Shih, K. J., Nie, W. Vahdat, A., Lee, S.- g., Santos, J. F., Jukic, A., Valle, R., and Catanzaro, B. A2sb: Audio-to-audio schrodinger bridges. arXiv preprint arXiv:2501.11311, 2025.
Saharia, C., Chan, W., Chang, H., Lee, C., Ho, J., Salimans, T., Fleet, D., and Norouzi, M. Palette: smage-to-imane diffusion models. In ACM SIGGRAPH 2022 conference proceedings, pp. 1-10, 2022.
10

---

Inverse Bridge Matching Distillation
Shi, Y., De Bortoli, V., Deligiannidis. G., and Doucet, A. Conditional simulation using diffusion schrödinger bridges. In Uncertainty in Artificial Intelligence, pp. 1802. PMLR, 2022.
Wang, Y., Yoon, S., Jin, P., Tivnan, M., Song, S., Chen, Z., Hu, R., Zhang, L.., Chen, Zz, Wu, D., al. Implicit image-to-image schrödinger bridge for image restora- tion. Zhigiang and Wu, Dufan, Implicit Image-to -aagget---mage Schrödinger Bridge for Image Restoration, 2024.
Shi, Y., Bortoli, V. D., Campbell, A., and Doucet, A. Dif- fusion schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:ttp...mmm https////pppeeeeeeeeeeeeeeeeeeee wwt.:netttoouoeeee id=qy070HsJTS.
Yin, T, Gharbi, M., Park, T., Zhang, R., Shechtman, E., Durand, F., and Freeman, W. T. Improved dis- tribution matching distillation for fast image synthe- sis. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a. URL https: / // //oppnrreieww.. net forum?idee wrrumiddd-ukeeDDee..
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. International conference on
Yin, T., Gharbi, M., Zhang, R., Shechtman, E., Durand, F., Freeman, W. T., and Park, T. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6613-6623, 2024b.
machine learning. pp. 2256-2265 PMLR, 2015.
Somnath, V. R., Pariset, M., Hsich, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. Aligned diffusion schrödinger bridges. In Uncertainty in Artificial Intelligence, pp. 1995-1995. PMLR, 2023.
Yue, Z., Wang, J., and Loy, C. C. Resshift: Efficient diffu- sion model for image super-resolutie by residual shifting. Advances in Neural Information Processing Systems, 36, 2024.
Song. J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference Learning Representations, 2021. URL htttps:///ppereeeeeee net/
Zheng, K., He, G., Chen, J., Bao, F., and Zhu, J. Diffusion bridge implicit models. In The Thirteenth International Conference on Learning Representations, 2024.
Song. J., Vahdat, A., Mardani, M., and Kautz, Pseudoinverss-guided diffusion models for inverse problems. In International Conference on Learning Representations, 2023.
Zhou, L., Lou, A., Khanna, S., and Ermon, S. Denoising diffusion bridge models. In The Twelfth International Conference on Learning Representations, 2023.
Su, X., Song, J., Meng, C., and Ermon, S. Dual dif- fusion implicit bridges for image-to-image translation. In The Eleventh International Conference on Learning Representations, 2023. https: //openrevie.....oe/o/ee n
Zhou, I.., Lou, A., Khanna, S., and Ermon, S. De- noising diffusion bridge models. In The Twelfth International Conference on Learning Representations, URL.nttsss//ppeereeiie...nm 2024a. URI.https: http///pppnweeeeeeeeeeeee www.reewwwwwwweeeeeeeeeeeeeeeeee net//orum? id=FKksTayvGo.
Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation- free schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pp. 1279-1287. PMLR, 2024.
Zhou, M., Zheng, H., Wang. Z., Yin, M., and Huang, H. Score identity distillation: Exponentially fast distilla- tion of pretrained diffusion models for one-step genera- tion. In Forty-first International Conference on Machine Learning, 2024b.
Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9):1 134, 2021.
Vasiljevic, I., Kolkin, N., Zhang, S., Luo, R., Wang, H., Dai, F. Z., Daniele, A. F., Mostajabi, M., Basart, S., Walter, M. R., et al. Diode: A dense indoor and outdoor depth dataset. arXiv preprint 1908. 000663 2019.
Wang. Y., Yu, J., and Zhang, J. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference on Learning Representations, 2023. htLhttps: httlhtts///ppeeeviieee net/ forum? d=mRi eggMtNTT.....

---

Inverse Bridge Matching Distillation
A. Proofs
Since all our theorems, propositions and proof' for the inverse Bridge Matching problems which is formulated for the already trained teacher model using some diffusion bridge, we assume all corresponding assumptions used in Bridge Matching. Extensive overview of them can be found in (Shi al., 2023, Appendix C).
Proof of Proposition 3.1. Since both BM( and M given by reverse-time SDE and the same distribution (r(r)) the KL-divergence expressed in the tractable form using the disintegration and Girsanov theorem (Vargas et al., 2021; Pavon & Wakolbinger, 1991):
(L( = Ez, (g)())))00)))))) (*), +)
(#), ~ pa po IT), t (10,, T),, ~ (( æT).
The expectation is taken over the marginal distribution p(z;) of HI since it is the same for BM((Ta) (Shi al., 2023, Proposition 2). In turn, the drift v(ze, 1) is the drift of Bridge Matching using II i.e. BM([Ta):
V = arg min Eo [llo' [1le t) Vo, log (( [o)
(Io, vT) Po (xo, (To, TT), ~ ([0, T),, xa a(ra (ro, TT).
Combining this, the inverse problem can be expressed in more tractable form:
min B [g² (1) () - o" (.z,, 1)] s.t. arg min E [lle (x+;, 1) Va, log (mo) (|? dl, (Fo, TT) ~ (o IT), ((0, TJ), xt ~ (0
(12)
We add positive valued weighting function X(1) for the constraint:
= arg min Ex,, [A(t)lle' t) - V., log ((r)/20)))
of
since it is the MSE regression and its solution is conditional expectation for any weights given by:
o( t) = Banglan i
Vz, s, log (( [x))].
We can add positive valued weighting function X(I) for the main functional:
E [A(t)]lo(2,, t) 6² (I,, +)13].
since it does not change the optimum value (which is equal to 0) and optimal solution, which is the mixture of bridges with the same drift as the teacher model.
Proof of Theorem 3.2. Consider inverse bridge matching optimization problem:
min Ex,, [X(t)]lu(zz, t) v² (Fe, t)]f]. s.t.
(13)
V = arg min Ex,, [1lo' (ze, t) Vz, log q(ι, ((0) [3]
(zo, (T) ~ (00, ,T), t ~~ U(O, x ~ (<
First, note that since v = arg min,, Ex (1o, - V., log (20) 122] i.c. minimizer of MSE functional it is given by conditional expectation as:
u(z,,1) Exol Va, (o0)(( 4].
(14)
Then note that:
min Eg,, (o - Ve, log g( (xo))
12

---

Inverse Bridge Matching Distillation
En a.o (()))(( Va, 108)(
1000 ((())l 011
2F 2,000 [(1)((((11))) 1), Zz, log (22/(oo)))) + Ben,,, [(t))l [xgg(1)))x))1
Ban (())11((2)0111
2. (2(t) e(ze,t) Boolan Boolet) V.., log loga ((<0)200) ) Ex,,, [Att]| log |00)(11))))
Eg,, ((t)]lo(xa, t)[[
E ((0))l((0,, 1)19] 2 (A()]l((a,,)) 0113] + a [(OIIV,, log [x0)]
-Bag [A(1)]l((22, 0)] + Ber [A(1) log (( [x0)]]]
(15)
Hence, we derive that
Exg [A(O)lu(aa, 0]13] t, [ log (() min En, 1, [1(1)le' V. log (x)
Now we use it to reformulate the initial objective:
E E - (x2,
- -
Ez ot (A()))o((>,,) 0)] 2E,, ((11) (0(2),, 1), 0- (x,,, 1)) Ex [A(O)lo (112 t. a. log (( (x0) - min R. ((()lle((are) - V., log (( (ro))
=E,,, [xC)]]
2E,,, (A(1) (o(ar,, (), 0" (1), + E [A(1)llo (vx,, 011 = E E [X(I)V ny log g [mmo)l3] - 2 at [A(1) (1((2, v² (xx, 0)} Buy. [A())le* (1z,, 1)
min t, (())ll(((x1))) Vz, log (122((1)2)))1
Therefore, we get:
E . [A(t)ll((2,,)) ux,, E Es 11..0.00 1,000 [A(t)] V. Ø, log (( (z)) [zo) 2E [A(t) (v(2r, t), o² (Ize, t)) + Exne,o (z)
min Fr.0 [A(t)]lu'(an) Zz. log ((=;)</))111
To complete the proof, use the relation v(x,,t) Eno Ims. Vn, log (2) (20)) from Equation 14. Integrating these components, arrive at the final result:
E [A(1)]lo( - (xr,
Bury,t,rr [X(1)I xi log ( [x]] 2 [A(t) (Ezulee, [Va, log logg(se (r (xx, Br [)(t)llu*<aa (113 t.oo [A(t) 10/ (xt, 1) Va, log log [10
min F
f, [A(t)]I x, log (ogala/ooo) [1 - 2 t. [A(E) (Vz, log (((.e//rr)))))))) of (xx, t)} + Br
min F..,, t. [a(1) 10 (x Vn, log ( (20)
Eg,, [X(t)llu*(a (xx, - Ve, log (( fo) - min Ex, E. (()))l(((z,, V.. logg(ze/20)113)
Proof of Proposition 3.3. Consider the problem from Proposition 3.2:
min Ex tx,0 [At)]loo (æe, t) Va. log (=<r min Ex, ((t)llos 1) - Va, log (( [x001
0
13

---

Inverse Bridge Matching Distillation
For the priors with the drift f())ar the regression target Vx, (1rg(o)))))) -x- م1 of
Hence the parametriza-
tion u( 1) = (. We use reparameterization of both v" and ui given by:
3t (225
2t a (2,,
ex (x1,t)
of
and get:
min to [2(1)]lo* (222 Vx, log (14(((1
min 10 (2)), Vx, log (3)0((1/)0)011]]
t x(1) at (1,, -xoll²]
min E 1.00 A(1) of
min
(x(())
(((())
min t. ((0)) =ull? min E 1..00
((t)) =0.1 -
where (t) is just another positive weighting function.
Proof of Theorem 3.4. In fully analogical way, as for the unconditional case consider the set of the Inverse Bridge Matching problems indexes by zy:
min [KL( jor Ml
where Mfr, is a result of Bridge Matching conditioned on 24 and Hap is Mixture of Bridges for each 2T constructed using bridge q( #T) and coupling po (#ol2r))),,
By employing the same reasoning in the proof of Proposition 3.1, the inverse problem can be reformulated as follows:
min Ea (0)lo((2,) - o (Ix, (x,, [T)] s.t. 1 arg min E Eg,, f. er [|| (x,, 2,) V., log (22; (x0g)=)))l]]
(x), ~ ( U(O, T)), n. q200
Following the proof of Theorem 3.2, obtain tractable formulation incorporating weighting function:
min 8
Ex, L.
[xt) Vz, log [20))
min E 2o (2),, V log g( mo) [=]
Utilizing the reparameterization under additional conditions (52.3), we obtain the following representations:
0" (x2, t, xr) 24
(1 (,(x,,t,rr) It. ag (x TT)
of
of
Consequently, applying the proof technique from Proposition 3.3, we derive the final expression:
min Ez, ( - =0ll -min Er,, an [x(t)] (2,, 2T) moll3]]
(z0, ~ (ou, xT), t ~~ U((O,T]) xs ~ (00//2000 2T).
14

---

Inverse Bridge Matching Distillation
Task
Dataset
Teacher
NFE
L/K ratio
LR
Grad Updates
Noise
4x super-resolution (bicubic) 4x super-resolution (pool)
ImageNet
I'SB
5:1
5c-5
3000
ImageNet
FSBB
5:1
5e-5
3000
JPEG restoration, QF 5
ImageNet
PBB
5:1
5e-5
2000
JPEG restoration, QF 10
ImageNet
PSB
5:1
5e-5
3000
Center-inpainting (128 x 128)
ImageNet
IBSB
5:1
5c-5
2000
Sketch Image
Edges - Handbags
DDBM
2
5:1
le-5
300
Sketch Image
Edges + Handbags
DDBM
5:1
le-5
14000
Normal to Image
DIODE-Outdond
DDBM
5:1
le-5
500
Normal to Image
DIODE-Outdoor
DDBM
5:1
le-5
3700
Center-inpainting (128 x 128)
ImageNet
DDBM
1:1
3e-6
3000
Table Table entries specify experimental configurations: NFE indicates multi-step training (See. 83.5); L/K represents bridge/student gradient iteration ratios (Alg. $3.4); Grad Updates shows student gradient steps; Noise notes stochastic pipeline incorporation.
B. Experimental details
All hyperparameters are listed Table We used batch size 256 and ema decay 0.99 for setups. For each setup, started the student and bridge networks using checkpoints from the teacher models. In setups where the model adapts noise: (1) We added extra layers for noise inputs (set to zero initially), (2) Noise was concatenated with input data before input it to the network. Datasets, code sources, and licenses are included in Table 8.
Training time. We present the training time of each in Table 9. About 75% of this training time used to get the last 10-20% decrease of FID (e.g., drop from 3.6 2.5 FID in pooling SR setup or from 4.3 to 3.8 FID JPEG with), while training for the first 25% of time already provides good-quality model. On Sketch-to-image and Normal-to-image in multistep regime with 2 NFEs, convergence appears faster than in the corresponding single-step version.
Table 8. The used datasets, codes and their licenses.
Name
URL
Citation
License
Edges-+anndee Handbagg
GitHub Link
(Isola et al., 2017)
BSD
DIODE-Outdoor
Dataset Link
(Vasiljevic al., 2019)
MIT
ImageNet
Website Link
(Deng et al., 2009)
Diffus-Di
GitHub Link
(Dhariwal & Nichol, 2021)
MIT
I'SB
GitHub Link
(Liu et al., 2023a)
CC-BY-NC-SA-00
DDBM
GitHub Link
(Zhou et al., 2023)
DBIM
GitHub Link
(Zheng et al., 2024)
Task
Teacher
Dataset
Approximate time on 8xA100
NFE
4x super-resolution (bicubic)
I2SB
Imagenet
40 hours
4x super-resolution (pool)
I2SB
Imagenet
40 hours
JPEG restoration, QF 5
I2SB
Imagenet
40 hours
JPEG restoration, QF = 10
I2SB
Imagenet
40 hours
1
Center-inpainting (128 128)
I2SB
Imagenet
24 hours
4
Center-inpainting (128x 128)
DDBM
Imagenet
12 hours
4
Sketch Image
DDBM
Edges/Handbags
40 hours
1
Sketch Image
DDBM
Edges/Handbags
hour
2
Normal Image
DDBM
OuoDE-Outdoor
48 hours
1
Normal Image
DDBM
7 hours
2
Table 9, Training times and NFE across different tasks, teachers, and datasets.
B.I. Distillation of SB models.
We extended the ISB repository (see Table 8), integrating our distillation framework. The following sections outline the setups, adapted following the SB.
15

---

Inverse Bridge Matching Distillation
Multi-step implementation In this setup, we use the student model's full inference process during multi-step training (Section 3.5). This means that Zip is generated with inferenced of the model Go through all timesteps (T -N,.... ... 1 0) in the multi-step sequence. The generated zu subsequently utilized in the computation of the bridge Es or student Co objective functions, as formalized in Algorithm 1.
4x super-resolution. Our implementation of the degradation operators aligns with the filters implementation proposed in DDRM (Kawar et al., 2022). Firstly, we synthesize images at 64 x 64 resolution, then upsample them to 256 x 256 to ensure dimensional consistency between clean and degraded inputs. For evaluation, we follow established benchmarks (Saharia et al., 2022; Song et al., 2023) by computing the FID on reconstructions from the full ImageNet validation set, with comparisons drawn against the training set statistics.
JPEG restoration. Our JPEG degradation implementation, employing two distinct quality factors (QF=5, QF=10), follows (Kawar al., 2022). FID is evaluated statistics, following baselines (Saharia et al., 2022; Song et al., 2023).
10, 000-image ImageNet validation subset against the full validation set's
Inpainting. For the image inpainting task on ImageNet at 256 x 256 resolution, we utilize fixed 128 x 128 centrally positioned mask, aligning with the methodologies of DBIM (Zheng et al., 2024) and CDBM (He et al., 2024). During training, the model is trained only on the masked regions, while during generation, the unmasked areas are deterministically retained from the initial corrupted image Zr to preserve structural fidelity of unmasked part of images. We trained the model with 4 NFEs via the multi-step method (Section 3.5) and tested with 1, 2, and 4 NFEs.
B.2. Distillation of DDBM models.
We extended the DDBM repository (Table 8) by integrating our distillation framework. Subsequent sections outline the experimental setups, adapted from the DDBM (Zheng et al., 2024).
Multi-step implementation In this setup, the multi-step training (Section 3.5) adopts the methodology of DMD (Yin et al., 2024a), wherein timestep / is uniformly sampled from the predefined sequence (II...., l.w.. The model Go then generates To by iteratively reversing the process from the terminal timestep iN = I' to the sampled intermediate timestep t. This generated my is subsequently used to compute the bridge network's ' loss L. or the student network's 's loss Lg. as detailed in Algorithm 1.
image-to-image translation task (Isola et al., 2017), with the 64 64 resolution images. Two versions were trained under the multi-step regime (Section 3.5), with 2 and 1 NFEs during training. Both models were evaluated using the same NFE to match training settings.
Edges + Handbags The model was trained utilizing the
DIODE-Outdoor Following prior work (Zhou et al., 2023; Zheng et al., 2024; He et al., 2024), we used the DIODE outdoor dataset, preprocessed via the DBIM repository's script for training/test sets (Table 8). Two versions were trained under the multi-step regime (Section 3.5), with and NFEs during training. Both models were evaluated using the same NFE to match training settings.
a CBDM checkpoint (Zheng et al., 2024). This checkpoint is adjusted by the authors to: (1) condition ær and (2) ImageNet class labels as input to guide the model. Also this is the same checkpoint used in both CDBM (He et al., 2024) and DBIM (Zheng et al., 2024) works.
Inpainting All setups matched those in Section B.1 inpainting, except
C. Additional results
16

---

Inverse Bridge Matching Distillation
IBMD (Ours) NFE=1
Teacher
Input
NFE=1000
Reference
Figure 4. Uncurated samples for IBMD-I'SB I'B distillation of 4x- rx-supperresoltutinn with bicubic kernel on ImageNet 256 x 256 images.
17

---

Inverse Bridge Matching Distillation
IBMD (Ours) NFE=1
Teacher
Input
NFE=1000
Reference
Figure 5, Uncurated samples for IBMD-I'SB I'SB distillation of 4x- ressupoerr with pool kernel on ImageNet 256 x 256 images.
18

---

Inverse Bridge Matching Distillation
IBMD (Ours) NFE=1
Teacher NFE=1000
Input
Reference
A311
2001
010
1013
800
130
UINIO
B3 6
31
B B
B
8
200
40
20
60
Figure 6. Uncurated samples for IBMD-I'SB distillation of Jpeg restoration with QF=5 on ImageNet 256 x 256 images.
19

---

Inverse Bridge Matching Distillation
IBMD (Ours)
Teacher NFE=1000
Input
NFE=1
Reference
460
1000
100
1002
800
130
DIGIO
D
B
8
40
211
40
40
100.
Figure 7. Uncurated samples for IBMD-I'SB distillation of Jpeg restoration with QF=10 on ImageNet 256 x 256 images.
20

---

Inverse Bridge Matching Distillation
IBMD (Ours)
IBMD (Ours)
IBMD (Ours)
Teacher
Input
NFE=1
NFE=2
NFE=4
NFE=1000
Reference
FORNA
1
100
0016
10
50
C
50
U 4
160
U 4 O
150
150
4491
100
4
100
ISC
Figure 8. Uncurated samples for IBMD-12SB distillation trained for inpaiting with NFE- 4 and inferenced with different inference NFE ImageNet 256 256 images.
21

---

Inverse Bridge Matching Distillation
(BMD( (MMs()urs) NFE=1
IBMD Ours)
Teacher
Input
NFE= =
NFE= = 4
MFE=500
Reference
Figure 9. Uncurated samples for IBMD-DDBM distillation trained for inpaiting with NFE- 4 and inferenced with different inference NFE ImageNet 256 256 images.
22

---

Inverse Bridge Matching Distillation
(BMD() NFE=1
IBMD(Ours) NFE=2
Teacher
Input
NFE=500
Reference
Figure 10. Uncurated samples from IBMD-DDEM distillation trained the DIODE-Outdod datasct (256 256) with NFE= and
NFE= 1. inferred using the corresponding NFEs on the training sel,

---

Inverse Bridge Matching Distillation
(BMD(Ours) NFE=1
IBMD()urs) NFE=2
Teacher
Input
NFE=500
Reference
Figure 11. Uncurated samples from IBMD-DDBM distillation trained the DIODE-Outdoor datasct (256 256) with NFE= and
NFE= 1. inferred using the corresponding NFEs on the test set. 24

---

Inverse Bridge Matching Distillation
(BMDs) NFE-1
(BMD(Ours) NFE=2
Teacher
Input
NFE=500
Reference
A
Figure Uncurated samples from IBMD-DDBM distillation trained the Edges Handbags dataset (64 64) with NFE= and
NFE= 1. inferred using the corresponding NFEs on the training sels

---

Inverse Bridge Matching Distillation
(BMDs) NFE=1
(BMD()urs) NFE=2
Teacher
Input
NFE=500
Reference
Figure Uncurated samples from IBMD-DDBM distillation trained the Edges Handbags dataset (64 64) with NFE= 2 and
NFE= 1. inferred using the corresponding NFEs on the test set. 26


Reviews:
[{"review_id": "HzCy21JY2Z", "reviewer": "Reviewer_AN8r", "summary": "This paper proposes a new distillation scheme for diffusion bridge models. The main idea is to parameterize the entire formulation based on stochastic generator G. The student must follow input-output pairs produced by G, which constrains the path of diffusion bridge, which must coincide with the teacher path. Optimizing G, one can find a viable student. To make the formulation tractable, the constrained problem is reformulated into an unconstrained problem, resulting in a bilevel optimization problem that is somewhat GAN-like. The generator G becomes the resulting one-step generator, and the difference between the teacher error and the \"student error\" acts as a discriminator. G can also be designed in multi-step fashion. The entire formulation can be applied to both unconditional bridge matching and conditional bridge matching. Experiments show that the proposed method provides state-of-the-art results with much fewer steps than existing methods.\n\n## update after rebuttal\n\nAll the reviewers have rated the paper positively. The additional results provided in the rebuttal are also convincing. I 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": "Dear Reviewer AN8r, Thank you for your comments. Here are the answers to your questions and comments.\n\n**(1)  In my understanding, the second term with \\phi here acts as an \"expander,\" which means, G can become trivial (or can collapse) without this term ... In other words, the proposed method trains another DBM just to help train $G$. I'd like to ask whether this interpretation is right, and I'd like to see a deeper discussion in the paper regarding the role of $\\phi$.**\n\nYes, this interpretation is correct. To show it more formally, note that the minimal value of the inner problem is the averaged variance of $x_0 \\sim p_{\\theta}(x_0|x_t, x_T)$:\n\n$$\n\\min_{\\phi} \\mathbb{E}_{p\\_{\\theta}(x\\_t, t, x\\_0, \\textcolor{MyRed}{x_T})}  \\big[\\lambda(t) \\| \\widehat{x}_0^{\\phi}(x_t, t, \\textcolor{MyRed}{x_T}) - x_0 \\|^2 \\big]\\Big] = \\mathbb{E}\\_{p\\_{\\theta}(x\\_t, t, x\\_0, \\textcolor{MyRed}{x\\_T})}  \\big[\\lambda(t) \\| \\mathbb{E}\\_{p\\_{\\theta}(x\\_0|x\\_t, x\\_T)}[x_0] - x_0 \\|^2 \\big]\\Big] = \n$$\n$$\n=  \\mathbb{E}\\_{p\\_{\\theta}(t, x\\_t, x\\_T)} \\Big( \\lambda(t) \\Big\\[\\underbrace{\\mathbb{E}\\_{p\\_{\\theta}(x\\_0|x\\_t, x\\_T)} \\big[ \\| \\mathbb{E}\\_{p\\_{\\theta}(x\\_0|x\\_t, x\\_T)}[x_0] - x_0 \\|^2 \\big]\\Big]}\\_{\\text{Variance of } p\\_{\\theta}(x\\_0|x\\_t, x\\_T)}\\Big\\] \\Big).\n$$\n\nFor $t=T$, this is directly the variance of the generator $x_0 \\sim p\\_{\\theta}(x_0|x_T)$. Since we are maximizing this part over $\\theta$, it enforces the generator to produce more diverse outputs and avoid collapsing. \n\nFollowing your advice, we will add this discussion on the interpretation of how the auxiliary model $\\phi$ helps to train the Generator $\\theta$ to the final version of the paper.\n\n**(2) How long does it take for training?**\n\nWe present the training time of each model below:\n| Task                                  | Teacher | Dataset        | Approximate time on 8A100 | NFE |\n|---------------------------------------|---------|----------------|---------------|-----|\n| $4 \\times$ super-resolution (bicubic) | I2SB    | Imagenet       | 40 hours      | 1   |\n| $4 \\times$ super-resolution (pool)    | I2SB    | Imagenet       | 40 hours      | 1   |\n| JPEG restoration, QF $=5$             | I2SB    | Imagenet       | 24 hours      | 1   |\n| JPEG restoration, QF $=10$            | I2SB    | Imagenet       | 40 hours      | 1   |\n| Center-inpainting ($128 × 128$)       | I2SB    | Imagenet       | 24 hours      | 4   |\n| Center-inpainting ($128 × 128$)       | DDBM    | Imagenet       | 12 hours      | 4   |\n| Sketch to Image                       | DDBM    | Edges/Handbags | 40 hours      | 1   |\n| Sketch to Image                       | DDBM    | Edges/Handbags | 1 hour        | 2   |\n| Normal to Image                       | DDBM    | DIODE-Outdoor  | 48 hours      | 1   |\n| Normal to Image                       | DDBM    | DIODE-Outdoor  | 7  hours      | 2   |\n\nAbout 75\\% of this training time is used to get the last 10-20\\% decrease of FID (e.g., drop from 3.6 to 2.5 FID in pooling SR setup or from 4.3 to 3.8 FID in JPEG with $QF=5$), while training for the first 25\\% of time already provides a good-quality model. On Sketch-to-image and Normal-to-image in multistep regime with 2 NFEs, convergence appears faster than in the corresponding single-step version. We will add the approximate time used for training to Table 7 of Appendix B (Table with all hyperparameters). \n\n**Concluding remarks**. \nWe would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns and questions about our work. We are also open to discussing any other questions you may have.", "claims_and_evidence": "The proposed formulation based on stochastic generator G, as well as the bilevel reformulation, is sound and novel. I checked the proof and they are correct in my opinion.", "methods_and_evaluation_criteria": "As mentioned above, the method and the proofs are sound. The method was evaluated on popular benchmark datasets.", "theoretical_claims": "As mentioned above, the proofs ((a) the parametrized matching problem becomes constrained optimization (9), (b) it can be reformulated as an unconstrained bilevel optimization, and (c) it can be reparametrized based on the denoisers/samplers) are sound.", "experimental_designs_or_analyses": "The method was evaluated on popular benchmark datasets, and the experiments show most important metrics (NFE, FID, CA). The proposed method shows state-of-the-art performance with much fewer sampling steps.", "relation_to_broader_scientific_literature": "Diffusion bridge models are an important topic in diffusion models, and they can be used in many data-to-data translation problems. Providing a faster sampling method for diffusion bridge models can benefit many related areas.", "essential_references_not_discussed": "I believe the bibliography is thorough.", "other_strengths_and_weaknesses": "The proposed distillation technique based on the G formulation is sound and novel. This particular formulation allows for the handling of both unconstrained and constrained diffusion bridge models. One downside is that the learning procedure can be quite complicated and heavy, as also pointed out in the Discussion section.", "other_comments_or_suggestions": "Please see the below question.", "questions_for_authors": "The main question I have is about the meaning of the bilevel formulation. In my understanding, the second term with \\phi here acts as an \"expander,\" which means, G can become trivial (or can collapse) without this term. Initially, I thought that \\phi is the student and is the resulting model. However, after reading the whole paper, I realized that G is the final model, and even if \\phi looks like a student, it is actually an auxiliary component. In other words, the proposed method trains another DBM just to help train G. (This makes the whole formulation quite heavy because now we have two additional models G and \\phi. How long does it take for training?)\n\nI'd like to ask whether this interpretation is right, and I'd like to see a deeper discussion in the paper regarding the role of \\phi. Currently, there is not much of a discussion about the meaning of the final formulation.", "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": "SAd6qbbEmU", "reviewer": "Reviewer_bxNd", "summary": "The paper introduces IBMD, an inverse bridge matching distillation method for inverse problems. The key idea is to consider bridge matching distillation as an inverse problem and convert the constrained problem into an unconstrained one using the reparameterization trick. Based on the teacher models DDBM and $I^2SB$, distilled models with IBMD achieve low FID with fewer NFEs in super-resolution, image restoration, image inpainting and image-to-image translation compared to other baselines.\n\n## Update after rebuttal\nI have no concerns regarding the submission, therefore, I will maintain my original rating of 4.", "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": "Dear Reviewer bxNd, thank you for your comments. Here are the answers to your questions and comments. \n\n**(1) It would be better to clarify the reasoning behind the statement in Section 3.2, line 240: 'The key difference in the reformulated problem is that it admits clear gradients of the generator .' For example, explaining that all parts are differentiable would help.**\n\nThank you for this suggestion. In the final version, we will clarify the differentiability of all parts of the final objective. \n\n**(2) One curious point is that multi-step distillation is applied to the distillation of DDBM but not to CBD and CBT. What if single-step distillation is applied to the proposed approach? How would the results change in the metric?.**\n\nIf we correctly understand, you asked about single-step distillation applied to the result of the multi-step distillation. Indeed, we applied multi-step distillation to the original teacher model, e.g., I2SB and DDBM, but not to the already distilled models like CBD or CBT. We did so, since our method is designed for distillation of diffusion-bridge models (like I2SB and DDBM), while CBD and CBT are consistency models obtained from DDBM. Result of multi-step IBMD (ours) distillation is also not a diffusion bridge model. Due to that we do not consider one-step distillation of multi-step distilled models. \n\n**(3) It would be better to use the same notation for 'single-step' and 'multi-step' in line 99.**\n\nThank you for this suggestion. We will change it in the final version.\n\n**(4) How much time does it take to distill each model each model I2SB and DDBM?**\n\nWe present the training time of each model below:\n| Task                                  | Teacher | Dataset        | Approximate time on 8A100 | NFE |\n|---------------------------------------|---------|----------------|---------------|-----|\n| $4 \\times$ super-resolution (bicubic) | I2SB    | Imagenet       | 40 hours      | 1   |\n| $4 \\times$ super-resolution (pool)    | I2SB    | Imagenet       | 40 hours      | 1   |\n| JPEG restoration, QF $=5$             | I2SB    | Imagenet       | 24 hours      | 1   |\n| JPEG restoration, QF $=10$            | I2SB    | Imagenet       | 40 hours      | 1   |\n| Center-inpainting ($128 × 128$)       | I2SB    | Imagenet       | 24 hours      | 4   |\n| Center-inpainting ($128 × 128$)       | DDBM    | Imagenet       | 12 hours      | 4   |\n| Sketch to Image                       | DDBM    | Edges/Handbags | 40 hours      | 1   |\n| Sketch to Image                       | DDBM    | Edges/Handbags | 1 hour        | 2   |\n| Normal to Image                       | DDBM    | DIODE-Outdoor  | 48 hours      | 1   |\n| Normal to Image                       | DDBM    | DIODE-Outdoor  | 7  hours      | 2   |\n\nAbout 75\\% of this training time is used to get the last 10-20\\% decrease of FID (e.g., drop from 3.6 to 2.5 FID in pooling SR setup or from 4.3 to 3.8 FID in JPEG with $QF=5$), while training for the first 25\\% of time already provides a good-quality model. On Sketch-to-image and Normal-to-image in multistep regime with 2 NFEs, convergence appears faster than in the corresponding single-step version. We will add the approximate time used for training to Table 7 of Appendix B (Table with all hyperparameters).\n\n**Concluding remarks**. \nWe would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns and questions about our work. We are also open to discussing any other questions you may have.", "claims_and_evidence": "1. The proposed distillation method reduces NFEs in the inverse problem while preserving the teacher model’s performance. As far as I know, this is the first distillation approach for both unconditional and conditional inverse problems. The equations and derivations are solid.\n\n2. It would be better to clarify the reasoning behind the statement in Section 3.2, line 240: 'The key difference in the reformulated problem is that it admits clear gradients of the generator $G_θ$.' For example, explaining that all parts are differentiable would help.", "methods_and_evaluation_criteria": "1. The authors follow the typical setup for inverse problems (super-resolution, image restoration, translation, and inpainting) and select appropriate teacher models (unconditional for $I^2SB$ and conditional for DDBM) along with suitable baselines.\n\n2. One curious point is that multi-step distillation is applied to the distillation of DDBM but not to CBD and CBT. What if single-step distillation is applied to the proposed approach? How would the results change in the metric?", "theoretical_claims": "I have checked Proposition and Theorems 3.1–3.4, along with their derivations in the Appendix. The proofs seem correct.", "experimental_designs_or_analyses": "The experimental designs are valid.", "relation_to_broader_scientific_literature": "Due to the reduced NFEs achieved by the proposed method, it will be more applicable to real-world inverse problem applications.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strengths**\n\n1. Treating bridge matching distillation as an inverse problem is novel.\n2. The derivation for reparameterization is solid.\n3. It enables one-step inference.\n\n**Weaknesses**\n\n1. The qualitative results for inpainting are not satisfactory.", "other_comments_or_suggestions": "It would be better to use the same notation for 'single-step' and 'multi-step' in line 99.", "questions_for_authors": "How much time does it take to distill each model each model $I^2SB$ and DDBM?", "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": "ps90eVmQcx", "reviewer": "Reviewer_kKRF", "summary": "This work incorporates the technique of score distillation in diffusion models to diffusion bridges for accelerated generation. The empirical performance demonstrates that the proposed approach is superior compared to existing baselines under multiple image-to-image tasks.", "questions": "", "limitations": "", "rating": 3, "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": "Dear Reviewer kKRF, thank you for your comments. \n\n**(1) Relation to SiD [1] and FGM [2].**\nFollowing your suggestion, we will extend the discussion of the related work in the main text. Below, we present the preliminary version of the extension:\n\nUnlike SiD [1] and FGM [2], we focus on diffusion-bridge models used for data-to-data translation and not for generation from noise. Furthermore, our high-level objective is the KL divergence $\\text{KL}(\\text{BM}(\\Pi_{\\theta})||M^*)$ between path measures of teacher model $M^*$ and path measure $\\text{BM}(\\Pi_{\\theta})$ given by the generator $G_{\\theta}$, which differs from Fisher Divergence used in SiD [1]. The motivation of our high-level objective is to restore alignment between data pairs $(x_0, x_T)$. We derive our final tractable objective using different techniques (see Appendix A), i.e., we do not use flow product or score product identities as in [1, 2] but hypothesize that analogous identities can be derived for diffusion bridge models. Our final objective can be rewritten similarly to the final objective used in SiD [1, Eq. 23] ($\\alpha=0.5$) and FGM [2, Eqs. 4.11, 4.12]:\n\n$$\n\\mathbb{E}_{t, p\\_{\\theta}(x\\_0, x\\_t, x\\_T)} \\Big[ \\underbrace{\\|\\widehat{x}\\_0^*(x\\_t, t, x\\_T) - \\widehat{x}\\_0^{\\phi}(x\\_t, t, x\\_T) \\|^2}\\_{\\text{quadratic term}} + 2 \\left< x_0^{*}(x_t, t, x_T) - x_0^{\\phi}(x_t, t, x_T), x_0^{\\phi}(x_t, t, x_T) - x_0 \\right> \\Big]\n$$\n\nHowever, in both SiD [1] (for $\\alpha=1.0, 1.2$ used in experiments) and FGM [2], the authors either omitted the quadratic term or even used a negative coefficient for it in image experiments since it introduced instability. Unlike SiD [1] and FGM [2], we do not omit any parts of the original loss function and use it as the theory provides it. \n\n**Relation between KL divergence of path measures and Fisher Divergence.**\nTo highlight the difference between KL divergence of path measures (which we use) and Fisher Divergence (which is used in SiD [1]), consider two reverse-time diffusions $D_1$ and $D_2$ given by the same starting distribution $p(x_T)$ and SDEs:\n$$\n    D_1: dx\\_t = v(x\\_t, t)dt + g^2(t)d\\bar{w}\\_t, \\quad\n    D_2: dx\\_t = \\widehat{v}(x\\_t, t)dt + g^2(t)d\\bar{w}\\_t.\n$$\nDenote marginal densities as $p(x\\_t)$ for $D\\_1$ and $\\widehat{p}(x\\_t)$ for $D\\_2$, then KL divergence and Fisher Divergences are given as:\n$$\n    \\text{KL}(D\\_1||D\\_2) = \\mathbb{E}\\_{t, p\\_t(x\\_t) } \\Big[\\frac{1}{2g^2(t)}\\|v(x\\_t, t) - \\widehat{v}(x\\_t, t) \\|^2 \\Big] + \\underbrace{\\text{KL}(p(x\\_T)||\\widehat{p}(x\\_T))}\\_{= 0 \\text{ if } p(x\\_T) = \\widehat{p}(x\\_T)}\n$$\n$$\n    \\mathcal{L}\\_{\\text{SiD}}(D\\_1, D\\_2) = \\mathbb{E}\\_{t} \\Big[\\underbrace{\\mathbb{E}\\_{p(x\\_t)}\\|\\nabla\\_{x\\_t}\\log p(x\\_t) -  \\nabla\\_{x\\_t}\\log \\widehat{p}(x\\_t)\\|^2}\\_{\\text{Fisher divergence between $p(x\\_t)$ and $\\widehat{p}(x\\_t)$}}\\Big]\n$$\nNote that in SiD [1], the authors use the time average of Fisher Divergence, which compares only marginal distributions $p(x_t)$ and $\\widehat{p}(x_t)$. However, two path measures with the same marginal distributions might not be equal. As a result, the minimization of Fisher Divergence with the teacher does not guarantee that the learned model will transform data in the same way as the teacher model. Nevertheless, Fisher Divergence allows one to build a Generator $x_0 = G_{\\theta}(z)$ to produce data $p_{\\theta}(x_0) \\approx p_{\\text{data}}(x_0)$, by matching marginals of data and generated samples. In contrast, we use KL-divergence between two path measures, which is zero if and only if two path measures are identical, since we need to get generations aligned to data coupling.\n\n**(2) Acceleration of the unconditional diffusion bridge models.**\n\nTo obtain PF-ODE for the diffusion bridge model, one needs to subtract forward and reverse SDE drifts [4, end of Sec. 4]. For the conditional case, the drift of a forward process $p(x_0|x_T) \\rightarrow \\delta(x_T)$ is known analytically from Doob h-transform. In the unconditional case, the drift of a forward unconditional diffusion $p_0(x_0) \\rightarrow p_T(x_T)$ is unknown. Hence, to restore PF-ODE in the unconditional case, one must learn an additional teacher model for forward-time translation, which is time-consuming. \n\nWe will add that I3SB [3] is used to accelerate the unconditional bridge model. Their sampler coincides with the DBIM sampler but replaces the conditional model with an unconditional one to sample $\\widehat{x}_0$. Their sampler provides good quality for moderate NFE (25+) but performs worse than distillation methods in a single NFE regime, e.g., for JPEG-10, their FID is 17, while ours is 3.8 (we will add it to the text).\n\n**Concluding remarks**. \nWe would be grateful if you could let us know if our explanations have been satisfactory. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.\n\nReferences:\n\n[1,2,3] the same.\n\n[4] Shi Y. et al. Diffusion Schrödinger Bridge Matching.", "claims_and_evidence": "Most of the claim is well supported despite some comments on related work that needs more careful examination.", "methods_and_evaluation_criteria": "The proposed method and evaluation make sense for accelerating the sampling for diffusion bridges.", "theoretical_claims": "The theoretical claims are checked in detail and are correct. However, the equivalent derivation has already been proposed in previous work, and its close connection with the proposed method in this work should at least be discussed (see weaknesses).", "experimental_designs_or_analyses": "The experiments are well-executed.", "relation_to_broader_scientific_literature": "This research is related to the broader literature on solving inverse problems.", "essential_references_not_discussed": "All key papers are cited, but they are not discussed thoroughly (which should be) in my opinion (see weaknesses).", "other_strengths_and_weaknesses": "## Strengths\n* This work demonstrates enough empirical significance in accelerating the sampling of diffusion bridges and designing some key techniques (e.g., noise-conditioned one-step generator) tailored for distilling from diffusion bridges.\n\n## Weaknesses\n* In my opinion, although the authors mentioned some related works on distillation for diffusion models [1, 2], these works deserve a much more detailed discussion given their strong relevance — or even the use of essentially identical distillation loss functions — to the method proposed in this paper. For example, the high-level objective function of this paper is essentially the Fisher divergence used in [1], and the tractable training objective derived here is mathematically equivalent to those in [1, 2] (e.g., the loss function in this paper corresponds to the SiD ($\\alpha=0.5$) loss in [1] and to the combined $\\mathcal{L}_1 + \\mathcal{L}_2$ loss in [2]). The derivation should be essentially the same as what has been done in [2] (same high-level objective and same final loss function). Although the focus shifts from diffusion models to diffusion bridges, and the motivation for adopting a Fisher divergence-like objective is different (which I think is a positive contribution if you can elaborate more on the connection between the KL for path measures and Fisher divergence), I still believe that better contextualization of these related works is necessary and would greatly benefit the community’s understanding of this research direction.\n\n* I feel the claim that \"previously proposed samplers and consistency models can not work with unconditional bridges\" needs more justification than provided and should be re-examined. For example, [3] develops DDIM-like samplers for I2SB, which is an unconditional bridge. For consistency models, it only requires that we can simulate the PF-ODE, which I believe is plausible once we know the drift of the unconditional bridge, given it is a Markovian process.\n\nI am happy to raise my score if these are properly addressed.\n\n[1] Zhou, Mingyuan, et al. \"Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation.\" (ICML 2024).\n\n[2] Huang, Zemin, et al. \"Flow generator matching.\" arXiv:2410.19310 (2024).\n\n[3] Wang, Yuang, et al. \"Implicit Image-to-Image Schrodinger Bridge for Image Restoration.\" arXiv:2403.06069 (2024).", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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: This paper focuses on improving the inference efficiency of diffusion bridge models, enabling single-step generation for both conditional and unconditional tasks. All reviewers recommend acceptance of this work, highlighting the novelty of the proposed method and its evaluation across a wide range of experimental settings. AC recommends acceptance. The authors are encouraged to include the reviewer discussions in the final paper.
Decision: Accept (poster)