Feedback Schrödinger Bridge Matching
Authors: Panagiotis Theodoropoulos
Nikolaos Komianos
Vincent Pacelli
Guan-Horng Liu
Evangelos Theodorou
Venue: ICLR 2025 Oral (2025)
Keywords: Diffusion models
Schrödinger bridge
Distribution matching
Semi-Supervised Learning
Primary area: generative models
TL;DR: We introduce Feedback Schrödinger Bridge Matching, a novel semi-supervised framework that uses pre-aligned pairs to guide the matching of non-coupled samples.

Abstract: Recent advancements in diffusion bridges for distribution transport problems have heavily relied on matching frameworks, yet existing methods often face a trade-off between scalability and access to optimal pairings during training. 
Fully unsupervised methods make minimal assumptions but incur high computational costs, limiting their practicality. On the other hand, imposing full supervision of the matching process with optimal pairings improves scalability, however, it can be infeasible in most applications.
To strike a balance between scalability and minimal supervision, we introduce Feedback Schrödinger Bridge Matching (FSBM), a novel semi-supervised matching framework that incorporates a small portion ($<8$% of the entire dataset) of pre-aligned pairs as state feedback to guide the transport map of non-coupled samples, thereby significantly improving efficiency. This is achieved by formulating a static Entropic Optimal Transport (EOT) problem with an additional term capturing the semi-supervised guidance. The generalized EOT objective is then recast into a dynamic formulation to leverage the scalability of matching frameworks. Extensive experiments demonstrate that FSBM accelerates training and enhances generalization by leveraging coupled pairs' guidance, opening new avenues for training matching frameworks with partially aligned datasets.

Published as a conference paper at ICLR 2025
FEEDBACK SCHRÖDINGE BRIDGE MATCHING
Panagiotis Theodoropoulos! Nikos Komianos', Vincent Pacelli', Guan-Horng Liu Evangelos A. Theodorou 'Georgia Institute of Technology ?FAIR, Meta
(ptheodor3, ano3,, vpacelli3, evangelos whw.oorouuooooooouuoouoooo | hawp:/orrooooaaateeeeee..eee ghl ioliuumeta,.. com
ABSTRACT
Recent advancements in diffusson bridges for distribution transport problems have heavily relied matching frameworks, yet existing methods often face a trade- off between scalability and access to optimal pairings during training. Fully unsupervised methods make minimal assumptions but incur high computational costs, limiting their practicality. On the other hand, imposing full supervision of the matching process with optimal pairings improves scalability, however, it can be infeasible in most applications. To strike balance between scalability and minimal supervision, we introduce Feedback Schrödinger Bridge Matching (FSBM), novel semi-supervised matching framework that incorporates small portion (< 8% of the entire dataset) of pre-aligned pairs as state feedback to guide the transport map of non-coupled samples, thereby significantly improving efficiency. This is achieved by formulating a static Entropic Optimal Transport (EOT) problem with an additional term capturing the semi-supervised guidance. The generalized EOT objective is then recast into a dynamic formulation to leverage the scalability of matching frameworks. Extensive experiments demonstrate that FSBM accelerates training and enhances generalization by leveraging coupled pairs' guidance, opening new avenues for training matching frameworks with partially aligned datasets.
INTRODUCTION
Transporting samples between distributions a ubiquitous problem in machine learning. Given the rise of generative modeling, significant progress has been made using static (Goodfellow et al., 2014), deterministic (Chen et al., 2018; Biloš et al., 2021), or stochastic mappings (Ho et al., 2020; Song al., 2020) that transport samples from noise more complex distribution. A prominent approach lies in diffusion models that simulate a Stochastic Differential Equation (SDE) to diffuse the data to noise, and learn the score function (Hyvärinen & Dayan, 2005; Vincent, 2011) to reverse the process (Anderson, 1982). However, these models present several limitations. For instance, the requirement converge during the noising (forward) process to Gaussian noise suggests that these models must run for a sufficient number of time steps to ensure the final distribution approximates Gaussian noise (Chen et al., 2021). Additionally, this suggests that these models begin their generative (backward) processes without any structural information about the data distribution, which implies randomness being introduced in the couplings that emerge from the diffusion models (Liu et al., 2023). Lastly, there is no guarantee that the optimal path interpolating the boundary distribution minimizes the kinetic energy. (Shi et al., 2023)
In an attempt to overcome these shortcomings, principled approaches that stem from Optimal Trans- port (Villani al., 2009) have emerged. The most prominent alternative has been the Schrödinger Bridge (SB; Schrödinger (1931)), which has been shown to be equivalent to mntropy-regulare Optimal Transport (EOT; Cuturi (2013); Léonard (2013); Pavon et al. (202 I)), and can also be framed Stochastic Optimal Control (SOC) Problem (Chen al., 2016; 2021). In particular, SB gained significant popularity in the realm of generative modeling following advancements proposing a train- ing scheme based on the Iterative Proportional Fitting (IPF), a continuous state space extension of the Sinkhorn algorithm to solve the dynamic SB problem (De Bortoli et al., 2021; Vargas et al., 2021). Notably, SB generalizes standard diffusion models transporting data between arbitrary distributions
1 Equal advising. Work done while Guan Georgia Tech.

---

Published as a conference paper at ICLR 2025
DSBM (k = 0%)
FSBM (Ours) 1% of Aligned Pairs
1 (k = 100%)
No pre-aligned Pairs
Every pair is a priori optimally aligned
k96: Percentage of aligned pairs compared to the entire dataset
Figure I: FSBM connecting existing Bridge Matching frameworks at the extremes, where the dataset is comprised of either fully aligned or fully non-aligned pairs
To, Ti with fully nonlinear stochastic processes, seeking the unique path measure that minimizes the kinetic energy. More recently, building on advancements in Bridge Matching methods (Peluchetti, 2023; Liu et al., 2022c), Shi et al. (2023) introduced Diffusion Schrödinger Bridge Matching (DSBM), framework that solves the Schrödinger Bridge problem while being significantly more efficient. Unlike previous methods, DSBM avoids the need to cache the full trajectories of forward and back- ward SDEs, making it more scalable and mitigating the sime-discreiization and "forgetting" issues encountered in carlier DSB techniques. For training these matching frameworks, we start from two datasets that lack pre-aligned pairs-that is, we sample only from the marginal To. ক an- the framework is trained to determine the optimal coupling in the sense of least-energy transportation. This process represents an unsupervised learning approach, as no prior information about the coupling is provided. However, there are other applications- such as image restoration (Liu et al., 2023) and protein docking (Somnath et al., 2023)-in which we have prior knowledge of the coupling 1o between pairs (æo, 21) To,1 drawn presumably from the SB solutions. Consequently, using an unsupervised upproach for such problems is not idcal, as it discards valuable information about the relationships between the samples of the two boundary distributions. An alternative approach suggests reframing these tasks as inverse problems. In this vein, Liu et al. (2023) proposed recently 12SB, building diffusion bridge matching frameworks between coupled data to perform image restoration.
In practice, in most tasks, the available datasets might possess a few limited coupled pairs, however, the high cost of manually labeling large datasets renders fully supervised approaches infeasible. Unfortunately, existing matching frameworks can not effectively leverage the information in partially pre-aligned datasets. An alternative approach from Optimal Transport (OT) literature involves using guided transportation maps, where only few pairs that belong to a Key-Point (KP) set are annotated. These pairs are utilized to guide the transportation mapping of the unpaired samples (Gu et al., 2023a), substantially reducing the need for extensive human labeling or expert guidance by leveraging a small number of ttrree aligned sample pairs for training (Mustafa & Mantiuk, 2020). More specifically, recent approaches employ "relation-preserving" scheme, which maintains the data's relationship to the given pairs from the KP set (Mémoli, 2011; Sato et al., 2020), or a pairwise distance- preseervigreeeeinggggggg constraint (Gu et al., 2023b). However, the feasibility of adapting semi- semi-supervised guidance in a dynamic Schrödinger Bridge or Bridge Matching setting remains an open question. In this vein, drawing inspiration from the guided OT schemes, advocate semi-supervised guided Schrödinger Bridge Matching framework.
In this work, introduce novel semi-supervised matching algorithm, Feedback Schrödinger Bridge Matching (FSBM), designed to integrate information from partially aligned datasets. Draw- ing inspiration from optimal transport (OT) literature (Gu et al., 2022), our analysis begins from static, semi-supervised OT problem. from which we derive a dynamic objective. Following recent advancements in matching frameworks, we adopt an alternating scheme, where the intermediate path and the coupling are optimized in two separate steps, resulting in a novel matching framework that leverages partially aligned data to guide the transport mapping of non-aligned samples. A key aspect of our approach is that the information from aligned samples is encoded as state feedback within the dynamic objective, effectively steering the transport of non-aligned data. This renders our FSBM a bridge between two extremes: unsupervised matching frameworks, which lack pre-aligned couplings (Shi et al., 2023), and fully supervised frameworks, where data is entirely pre-aligned (Liu et al., 2023; Somnath al., 2023) (see Figure 1). Empirical results show that our algorithm generalizes better, is more robust to perturbations in the initial conditions, and exhibits reduced training time. Our contributions are summarized follows:
2

---

Published as a conference paper at ICLR 2025
We introduce FSBM, the first matching framework leveraging information from partially aligned datasets.
We begin analysis from static semi-supervised OT problem, and derive modified dynamic formulation. Notably, our analysis can be extended to any selection of regularization functions.
We introduce an Entropic Lagrangian extension of the variational gap objective used in BM frameworks to match the parameterized drift given a prescribed probability path
Through extensive experimentation, we verify the remarkable capability of FSBM to gen- eralize under a variety of unseen initial conditions, while simultaneously achieving faster training times. These results were consistent across wide range of matching tasks from low-dimension crowd navigation, to high-dimensional opinion depolarization and image translation.
2 PRELIMINARIES
2.1
SCHRÓDINGER BRIDGES
The Schrödinger Bridge (Schrödinger, 1931) in the path measure sense concerned with finding the optimal measure 1' that minimizes the following optimization problem
min (IP)))
Po mu, P1 π
(1)
where is Markovian reference measure. Hence the solution of the dynamic SB 1 is consid- ered to be the closest path measure to Q. A very convenient, and heavily used property of l its decomposition as p J24x (0, ,d 1 (xn, (1), where x the solution of the static SB mitlzel[(<q.s]) L(T Rt), and Qju.1 1 the Bridge of the reference measure for pinned points (No, X1) (Léonard, 2013). Another formulation of the dynamic SB crucially emerges by applying the Girsanov theorem in Eq. (1)
min ["
/.
Ep, Illle
s.t.
Op: - (upp)) + Apt. and
Po = To,
Pi
Me.P:
Ot
2
Finally, note that the static SB is equivalent to the entropy regularized OT formulation (Pavon et al., 2021; Nutz, 2021).
/
min 31) RddRd
|oo - dIllddT(oo.
cKL(TT @ x
(3)
This regularization term enabled efficient solution through the Sinkhorn algorithm and has presented numerous benefits, such as smoothness, and other statistical properties (Ghosal et al., 2022; Léger, 2021; Peyré et al., 2019).
2.2 SCHRÓDINGER BRIDGE MATCHING
As discussed above, SB Matching algorithms (Shi et al., 2023; Liu al., 2024) have emerged recently, following advancements in matching frameworks (Liiu et al., 2022c; Peluchetti, 2023; Lipman et al., 2022; Liu et al., 2022b), rendering the solution of the SB problem significantly more tractable. Specifically, this family of algorithms solves Eq. (2) by separating the training process into two alternating steps. The first step entails relaxing the boundary distributional constraints Eq. (2) into just two endpoints, by fixing the coupling T0,1, and drawing pairs of samples (Io, 21).
min Mejo.1 J
Ep.o. Ill ||
Opaoo1 S.L. Ot
-V 14. + 00 Aptjo, Xo Zo.
X1 21
This optimization problem returns the optimized intermediate bridges between the drawn pairs. Subsequently, the parameterized drift uf is matched given the prescribed marginal path from the previous step, by minimizing the following variational gap
min / Ep. || u 2
/'
(5)
3

---

Published as a conference paper at ICLR 2025
where of is the drift associated with the fixed intermediate path. Propagation of the SDE according to the matched parameterized drift u., R x [0, 1] -
dX, - u di + odW,, Xu To: X1 1I
induces the parameterized coupling ~b,1 In the next iteration of the BM algorithm, this updated coupling used to generate improved pairs of (1) ~ū, employed to refine the marginal path.
METHODOLOGY
3.1 SEMI-SUPERVISED GUIDANCE
In the context of semi-supervised optimal transport, cer- tain aligned couples of data points from the source and target distributions are employed to guide the match- ing of unpaired samples. These pairs are assumed by construction to be given by the optimal solution of the static SB. We denote with IC = (ar xx] ee]] } Ia the set with the aligned pairs, the Key-Point (KP) set, with N being the number of pairs in the set. Usually, N < M, where M is the total number of samples.
Key
Bource Distribttion
Source Rugion Amaction
: Tanget Region Attraction
TTarget Distributtio
Assumption 3.1. All pairs from the KP set consid- ered to be solutions of the static Schrödinger Bridge.
Figure 2: Source and Target Distributions with their KP set pairs
Note that since the KP pairs are sampled from the solu- tion of the static SB, this ensures their stochastic nature, making them well-suited for guiding the transport map. Our analysis begins with the static entropic OT problem Eq. (3), to which an additional regularization term is included capture the interaction among unpaired samples and aligned samples, thereby guiding the transport map. We consider prreencepresseving scheme. We start by sampling a pair of non-coupled samples (Xo, X1). Let us assume that the 7th sample from the KP set zó K. is closest to Xo in the source distribution. Then, the distance d(Xo, xn) is computed and we proceed to compute the distance d(X1, xf).. Therefore, for fixed cou pair (z), 21) K. define the Guidance function for the non-aligned pair (Xo, X1), and as
G(Xo, X1) (d(Xi,=1) d(Xo, x/))²
(7) creates region of attraction, illustrated in Figure 2. This
Intuitively, this regularization attraction around each KP sample defines clusters that guide the neighboring samples from the source distribution to the target distribution, by preserving the relative distance to the respective KP sample.
Adjusting the cost function of the entropic O.T. formulation in 3, to incorporate guidance from the coupled pairs yields
(11X-- x,112 G(Xo, X1)) (20.. v1) + eK [(TT/m H R1)
min ()
(8)
In this work, we start our analysis from the modified expression of Eq. (3) exploit the flexibility of static OT formulation into adding the regularization term that is responsible for the guidance of the non-paired samples. However, our focal point to derive an equivalent dynamic formulation, resembling Eq. (2), which provides a more thorough description of the transport process. From an algorithmic standpoint, our methodology is based on BM frameworks (see Section 2.2) in order to leverage their great efficiency and scalability in solving distribution matching problems.
3.2 FEEDBACK SCHRÖDINGER BRIDGE MATCHING
We propose Feedback Schrodinger Brige Matching (FBSM) a novel matching framework, which leverages the information of partially, optimally paired datasets to guide the matching of the non- aligned samples. Our analysis starts from recasting Eq. (8) in a dynamic formulation in Sec. 3.2.1. Subsequently, we adopt recent advances in dynamic SB frameworks, that propose a decomposition of the dynamic optimization problem into two components (Liu et al., 2024: Shi et al., 2023). Therefore, separate the training phase into two stages: 1) the optimization of the intermediate path, conditioning on the endpoints (zp, F1), and 2) the optimization of the parameterized drift, which results in refining the coupling.
4

---

Published as a conference paper at ICLR 2025
3.2.1 DYNAMIC OBJECTIVE
The theorem below provides general framework recast the entropy regularized semi-supervised O.T. problem in Eq. (8) in a dynamic setting, regardless of the selection for the guidance function G. Theorem 3.2. Assume Xg, X1 E Rd with Xg To. and X, 2i and consider stochastic random variable Xx connecting Xo and X1, whose law is the continuous marginal probability path Pt. joining To, and my Additionally consider the KP pairs {a% a1) n=1 with their fixed interpolating paths. Starting from the static semi- sene-super guided entropic OT problem in Eq. (8), we derive the following relaxed dynamic formulation.
min {
o² (2 (fluell2 + Bpois [lu] VG..ol + AGr,o) pidrdt,
)
Pe Rd ((.1]
(9)
s.t. Öp.
Ōt = = (uspp) + 22 Api, and po = To, =
Notice that our dynamic formulation is a modified and generalized version of Eq. (2) capture the guidance from the KP samples. For degenerate guidance term (G = 0), one readily obtains the standard entropic analogue of the fluid dynamic formulation (Benamou & Brenier, 2000; Gentil et al., 2017) in Eq. (2) which also corresponds to the objective of SB objective (De Bortoli et al., 2021; Shi et al., 2023).
3.2.2 INTERMEDIATE PATH OPTIMIZATION
As discussed above, the initial step of our methodology entails the optimization of the inter- mediate path. Let the marginal be expressed as a mixture of conditional probability paths Pt = To More explicitly, we sample non-coupled pairs (xo, œ1) and fix the coupling at the boundaries No.1- This results in relaxing the distributional constraints in Eq. (9) to just two endpoints. The optimization of the intermediate path yields the optimal marginal conditioned on the two endpoints 1 of the uncoupled samples, which crucially preserves their relative distance to the KPs throughout the entire trajectory.
Proposition 3.3. Let the continuous marginal path pt satisfy the following decomposition pt 1 To, dx For optimizing with respect to the intermediate path, the coupling no,1 is frozen which results in Eq. (9) being recasted
min /
Gllu 0.1 +l VG(X AG(X,, dzd,,
1
P1jo.1 Rax (.4x(u,) 0,1]
(10)
Öpro.1 - -V..(pp)) 1 20² Aprj..1 Xo= =
s.l.
X1 = x
Ot
Importantly, the conditioning on the unpaired sample zi acts as an anchor. Our guidance function G groups the unpaired data into clusters around the KPs in the source distribution, and penalize the deviation of the unpaired samples from the trajectory that their assigned keypoint sample follows, preserving structural and distance information in the dynamic transport map. This forces unpaired samples to preserve their relative distance from the closest KP sample in the source distribution at each time-step.
Remark 3.4. Notice that in the extreme case in which the entire dataset is comprised of only pre- aligned couples, G would be trivial and Eq. (10) would retrieve the diffusion bridges between optimal pairs, similarly to SB (Liu et al., 2023). Conversely, the other extreme case, where the algorithm retrieves DSBM.
KP set is the empty
Feedback Interestingly, in stochastic optimal control literature (TTheodorou et al., 2010)) it has been reported that the inner product of the drift with the state cost acts state feedback. Hence, the term VG 1 1 is interpreted as state-feedback that steers the non-coupled samples, by projecting the conditioned drift on the gradient of Gyjo Figure 3 illustrates the gradient field of the guidance function, and how it pushes the optimal trajectory towards the path of the aligned data. Note that the trajectories of the KP samples along with their coupling assumed to be given and fixed for each KP pair. Furthermore, the absolute value implies that the angle projection always acute, which greatly improves stability during training.
5

---

Published as a conference paper at ICLR 2025
0.33
0.677
T
Vector Held of Gradienn Optimal Trajectory
Keypoint Sample Irajectory
Figure 3: Gradient Field of VG in the S-tunnel
Intemediate Path Recall that we have relaxed the distributional constraints in Eq. (9) to just two endpoints. We solve Eq. (10) for each pair (xo, æ1) ~ Pb,1 and later marginalize to construct pt- The conditioned probability path is approximated in simulation-free manner Gaussian path Prjo.1 R (1 o2 Ia). where Io = [tt, To. 21) is the interpolant function between the pinned endpoints Tq. 21, and is the standard deviation with op σi This enables to compute a closed-form expression for the conditional drift of the uncoupled pairs (Särkkä & Solin, 2019; Albergo et al., 2023).
0g -0 (
0,00
22
)
1+)
(11)
20
where v by definition of o vv/t(1 - t). More details are left for Appendix B.1.
3.2.3 COUPLING OPTIMIZATION
Subsequently, we will match the parameterized drift uf, given the prescribed path P:jo,1 of the uncoupled samples from the previous step, which will further improve the intermediate path in the next iteration. The more general form of the Lagrangian in Eq. (10) prompts to express the variational gap through the Bregman divergence, following recent advancements in bridging optimality gaps for general Lagrangian costs Neklyudov et al. (2023). However, before deriving the expression of the variational gap. we need the convex conjugate of the Lagrangian in Eq, (10).
Proposition 3.5. The convex conjugate of the Lagrangian in Eq. (10) is defined as the Hamiltonian H (Z+, ar, f) = u (uajo,1, at) - L(x,, ue)o,1,t). t). The optimization with respect to the drift yields 4rjo1 de g(a (Go) VG o. where VG.jo G(X,, and g(-) R TR. is given by
TvGnall if |a Va]VGull S [|V
=
sgn(a? VG..)
else
Finally, we obtain the Hamiltonian associated with the Lagrangian in Eq. (10)
1 H(s,, tt, 221aa. VaT[VGr)) VGrool2= AG(X,,
o²
(12)
At this point, introduce our Entropic Lagrangian Bridge Matching objective through the Bregman divergence, using the Hamiltonian in Eq. (12).
Proposition 3.6. Consider the optimized conditional drift f the optimization in Eq. (10). The parameterized drift uf is matched to u 0,1 given the prescribed marginal path by minimizing the variational gap expressed through the Bregman divergence: Dr, H L(Xe, ue, + H (X,, ax, (a,, uz)
min I'
E.po. E., o. |1 uijo, - g(a g(aTVGro)) VGrjol
(13)
From Proposition 3.5, the parameterized drift can be expressed as uf af g(a. VGrjo- Therefore, the matching objective of Eq. (13) can be rewritten as the objective Eq. (5)
min 1'
Bp. Blllou -u2l1 -ufl1aa
(14)

---

Published as a conference paper at ICLR 2025
Algorithm 1 Foodback Schrödinger Bridge Matching (FSBM)
1: Input: Boundary distributions no. and and the pairs from the key-point (KP) set {#] 24] [ ] and their fixed trajectories
2: Initialize ug. and To, as independent coupling xo & m-
3: repeat
4: Sample 2o, 21, Xt, forward propagating using u/. on time steps < t <... tx <1
5: Find closest aligned sample un to ro, and calculate dl(0o, au)
6: Using the same 't data point from the KP set, calculate d(Xt, 24p)
7:
G(Xt; zo) from Eq. (7)
8:
Fix the zo, 41, and obtain Po using Eq. (10)
9:
Calculate 1 from Eq. (11)
10:
Update ur from Eq. (14)
11:
until converges
3.3 TRAINING SCHEME
summary of our training procedure is presented in Alg. 1. Our algorithm finds a solution (p², už) to Eq. (9) by iteratively optimizing Eq. (10), and Eq. (14). First, we draw pairs (r0,, ). and fix the parameterized coupling ~ Subsequently, to compute the guidance function, let the 7th data point from the KPs set z E IC be closest to zo in the source distribution, the distance d(æo. 25) between the two samples is evaluated. Next, for every sampled time-step 4x, we compute the distance d(Xta, X between Nta and Xi_ and evaluate Gi zo). Then, optimizing Eq. (10), between the endpoints (xo, 21) from the first step, yields the optimal conditional path P/ whose corresponding drift un 0.1 obtained from Eq. (11). This drift ensures that the uncoupled data follow the trajectory of the KPs. Lastly, we fix the marginal path and optimize the coupling, by matching the parameterized drift uf to the optimized drift of the uncoupled samples "/j0, using Eq. (14). The optimized parameterized drift ug induces an improved coupling ~θ,1 by propagating the dynamics through the SDE in Eq. (6).
4 EXPERIMENTS
We demonstrate the efficacy of our FSBM in a variety of distribution matching tasks, such as Crowd Navigation, Opinion Depolarization, and Unpaired Image Translation, compared against other state- of-the-art distribution matching frameworks such as GSBM (Liu et al., 2024), DSBM (Shi et al., 2023), and Light and Optimal Schrodinger Bridge Matching (LOOSBM; Gushchin et (2024)).
4.1 CROWD NAVIGATION
We first showcase the efficacy of our FSBM in solving complex crowd navigation tasks and its superior capability to generalize under variety of initial conditions, listed as: i) Vanilla: initial distribution is the same as the one the model was trained on, ii) Perturbed Mean: shift the mean of Vanilla, iii) Perturbed STD: shift the standard deviation (STD) of Vanilla, iv) Uniform Distribution: the initial distribution is defined as a uniform distribution. We compare our FSBM with the GSBM framework (Liu al., 2024). The aligned pairs used to guide the transport map were less than 4% of the uncoupled pairs for both tasks. More details are left for Appendix C.I.
Table 1 shows that FSBM generated more accurately the targeted distribution in both tasks. Addition- ally. Figure 4 depicts our FSBM consistently producing better trajectories, successfully avoiding the obstacles under all initial conditions, and generating final distributions pf. that more closely match the target distribution T1, also shown by the smaller values of the Wasserstein distance W2 m)). In contrast, under GSBM, several particles were observed to diverge completely in the S-tunnel task under shifted initial conditions, causing the increase of the W2 distance. Additionally, notice that the final distributions generated by GSBM (see navy blue particles in the upper row Figure 4), even for the samples that did not diverge, were substantially more inaccurate when compared to FSBM. Lastly, Fig. illustrates that our FSBM required significantly less wall-clock time for training. This is attributed to the faster convergence of FSBM with regards to the number of epochs, but also due to GSBM requiring the computation the entropy and congestion cost, which according to our experiments has been shown increase training duration significantly.
7

---

Published as a conference paper at ICLR 2025
GSBM Perturbed Mean
GSBM Uniform Dist.
GSBM Perturbed Mean GSBM Perturbed STD
FSBM Perturbed Mean
FSBM Uniform Dist.
FSBM Perturbed Mean FSBM Perturbed STD
Figure 4: Crowd Navigation comparison between GSBM and FSBM the S-tunnel, and V-neck under perturbed mean and uniform initial distribution, and in the V-neck under perturbed mean and perturbed STD. The color of particles means: i) Yellow: initial conditions, ii) Green and Cyan: intermediate trajectory, iii) Red: target distribution, iv) Navy Blue: generated distribution
Table 1: W2 distance between generated and ground truth distributions for GSBM and FSBM in (Lower is better)
Parcentage of training zime of FSBM with respec: to GSBM
S-tunnel
V-neck
GSBM
FSBM
GSBM
FSBM
Vanilla
0.02
0.02
0.01
0.01
Perturbed
Mean
8.87e+5
0.67
0.68
0.34
CSBM
PSBM
Vinwkk
Perturbed
2.43e+10
0.34
3.03
0.11
Figure 5: Training time percentage comparison between GSBM and FSBM in Crowd Navigation tasks
STD
Uniform
6.65e+9
0.49
0.10
0.10
Distribution
4.2 OPINION DYNAMICS
Subsequently, we consider high-dimensional opinion depolarization, where each particle possesses a high-dimensional opinion X; E 100 that evolves through interaction with the other agents under polarizing dynamics which tend to segregate them into two groups of diametrically opposing opinions (see Figure 6: first row, second column). More information about the polarization base drift left in Appendix C.2. To mitigate the segregation, we apply our FSBM method employing the KPs samples guide the rest of the population to more uniform opinion, closer to the target distribution (see Figure 6: second row third column). We compare against DeepGSB and GSBM (Liu et al., 2022a; 2024), which employ state cost to mitigate polarizing dynamics. Lastly, it worth noting that the size of the KP set utilized for the opinion depolarization
approximately 2%, of the entire population.
The results shown demonstrate that the utilization certain opinion leaders, that guide the overall opinion of the population to depolarization renders our FSBM highly efficient model to mitigate opinion segregation. The metrics used to evaluate the proximity of the generated and target distribu- tions (c.g. pf. respectively) Table 2 and Figure indicates that our FSBM retrieves a solution that is slightly closer to the targeted distribution compared to GSBM, albeit needing half of the training time.
the (2f. Ti) distance and the KL divergence. Observation of
8

---

Published as a conference paper at ICLR 2025
Initial
Polarization Dynamics
Table 2: Comparison of W2 TT) distance and KL divergence between target and generated distribution in High Dimensional (X R R:10)) Opinion Dynamies among Deep GSB, GSBM and FSBM (ours) (Lower is better in all metries)
Condition
Targeted it
PCA vs.
g
Time Optimizer (hours:min)
W2 Distance
KI. divergence
Deep GSB
GSBM
FSBM
PCA vs.
DeepGSB
4:37
58.16
0.096
GSBM
5:29
43.41
0.034
DS
FSBM
2:43
42.91
0.034
Figure 6: Opinion depolarization from our FSBM and GSBM, and DeepGSB
4.3 IMAGE TRANSLATION
Finally, present the results of our FSBM in two types of image translations: i) Gender translation, ii) Age translation. We follow the setup of (Korotin al., 2023) with the pre-trained ALAA autoencoder FFHQ dataset (Karras et al., 2019) to perform the translation in the latent space of dimensions 512 1, enabling more efficient training and sampling. Notably, the number of aligned images for the age translation. Our method is compared with other state-of-the-eee aht matching frameworks, such DSBM (Shi et al., 2023) and Light and Optimal SBM (LOSBM)(Gushchin al., 2024). Additional details are left for Appendix C.3.
(Pidhorskyi et al., 2020) the 1024 x x024
4% of the total dataset for the gender translation and 8%
Table 3: FID values in 4 translation tasks for DSBM, LOSBM, and our FSBM (Lower is better)
Input
FSBM
LOSBM
DSBM
W > M
Man Woman
Woman Man
Old to Young
Young Old
Optimizer
Y
DSBM
15.90
17.02
16.83
17.76
LOSBM
18.89
17.24
17.89
19.32
Figure 7: Translation comparison be- tween our FSBM, LOSBM and DSBM
FSBM
14.78
16.12
15.52
17.14
Table shows that FSBM consistently achieves lower FID value, than the other methods. Additionally, it is demonstrated that our FSBM recovers better couplings through the qualitative comparison in Figure and quantitatively through the LPIPS values in Table 5. We can see that DSBM and FSBM return very close LPIPS values, along with very similar images. which suggests that the two methods converge to close solutions. Conversely, the translations of L.SSBM in the validation set were losing some of the features of the input image. Figure 8 presents additional examples of the couplings obtained from our FSBM. Our FSBM demonstrates highly accurate translations preserving key features identical, such smile, accessories, angle of the face, and background color. Notably, Table 4 demonstrates that FSBM required only an additional 180 MB of VRAM compared to DSBM, manage the aligned latent vectors, which is tractable for most modern GPUS, while training DSBM and FSBM for the same number of epochs requires virtually identical training duration.
4.4 LIMITIONS AND DISCUSSIONS
Our FSBM has demonstrated remarkable results in leveraging partially aligned datasets for distribution matching. However, framworks relies heavily the availability and the quality of the available KP pairs. The manner in which we sample from solution of the static SB does not ensure that all modes equally and satisfactorily represented. This implies that in cases where minimal number

---

Published as a conference paper at ICLR 2025
Woman to Man
Man to Woman
Old to Young
Young to Old
Figure 8: Image Translations using our FSBM
Table 4: Epoch duration and memory usage (GB) of different stages in Alg. I, measured the age translation task
Table 5: LPIPS values translation tasks for DSBM, LOSBM, and our FSBM (Lower is better)
Epoch Time
Memory
Man to Woman
Woman to Man
Old to Young
Young to Old
Optimizer
(hours:min) (GB)
Guidance (lines 5-8)
DSBM
0.47
0.48
0.53
0.49
00:01
0.18
LOSBM
0.70
0.69
0.71
0.70
Matching (lines 9-11) 02:56
FSBM
0.44
0.43
0.48
0.47
1.51
of these pairs available, the regions of attraction formed around these may not sufficiently cover the diversity the source distribution. This could lead to wrongful guidance, suboptimal matching and poor generalizabilty.
In future work, we have identified several directions which could improve our work even further. For instance, we aim to devise a more intricate mechanism to sample from the solution the static SB, while guarantecing that all modes in highly diverse datasets are covered and represented through the sampled KPs. Additionally, we aim further explore the capabilities of the guidance function and devise more task-specifee function which furhter improve perfomance generalizability by optimizing the information offered by the KP samples. Additionally, we intend to improve the way we impose guidance, by allowing the aligned pairs to preserve the relative distance to k-nearest neighboring KPs, instead of simply the nearest one.
CONCLUSION
In this work, we developed FSBM, novel matching algorithm capable of using the information from pre-aligned data pairs, guide the transport mapping of samples from the source to the target distribution. We demonstrated our algorithm's enhanced generalizability and reduced training over prior methods through extensive experimentation. In future work, wish delve deeper into the effect of the KP samples on the matching procedure, along with using more intricate and task-specific guidance functions for further improvements. In summary, our work paves new ways for training diffusion models with partially aligned data.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The authors would like to thank Tianrong Chen and Kevin Rojas for the useful and fruitful discussions.
REFERENCES
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions, 2023.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Marin Biloš, Johanna Sommer, Syama Sundar Rangapuram, Tim Januschowski, and Stephan Gün- nemann. Neural flows: Efficient alternative to neural odes. Advances in neural information processing systems, 34:21325-213377 2021.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianrong Chen, Guan-Horng Liu, and Evangelos A Theodorou. Likelihood training of schr\." odinger bridge using forward backward sdes theory. arXiv preprint arXiv:2110.11291, 2021.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 1-911, 2016.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamoniy. Joint distribution optimal transportation for domain adaptation. Advances in neural information processing systems, 30, 2017.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Valentin De Bortoli, James Thortoo,, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to based generative modeling. Advances in Neural Information Processing Systems, 14:176911,7799, 2021.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. URL http: jml org/ http:/aaperrsseeeeeee /22/ 20122/20-455 4500 html.
Jason Gaitonde, Jon Kleinberg, and Éva Tardos. Polarization in geometric opinion dynamics. In Proceedings of the 22nd ACM Conference on Economics and Computation, pp. 499-519, 2021.
Ivan Gentil, Christian Léonard, and Luigia Ripani. About the analogy between optimal transport and minimal entropy. In Annales de Faculté des sciences de Toulouse: Muthématiques, volume 26, pp. 569-600, 2017.
Promit Ghosal, Marcel Nutz, and Espen Bernton. Stability of entropic optimal transport and schrödinger bridges. Journal of Functional Analysis, 289(9) 2022.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farles Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural information processing systems, 27, 2014.
11

---

Published as a conference paper at ICLR 2025
Xiang Gu, Yucheng Yang, Wei Zeng, Jian Sun, and Zongben Xu. Keypoint-guided optimal transport with applications in heterogeneous domain adaptation. Advances in Neural Information Processing Systems, 191144492-989555 2022.
Xiang Gu, Liwei Yang. Jian Sun, and Zongben Xu. Optimal transport-guide guided conditional score-based diffusion model. Advances in Neural Information Processing Systems, (6:365500-65555 2023a.
Xiang Gu, Yucheng Yang, Wei Zeng, Jian Sun, and Zongben Xu. Keypoint-guided optimal transport, 2023b.
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 13:6840-6854. 2020.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4). 2005.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schr\" odinger bridge. arXiv preprint arXiv:2310.01174, 2023.
Flavien Léger. A gradient descent perspective on sinkhorn. Applied Mathematics & Optimization, 84 (2):1843-1855, 2021.
Christian Léonard. A survey of the schr\ odinger problem and some its connections with optimal transport. arXiv preprint 100.:0330.00115 2013.
Chi-Heng Lin, Mehdi Azabou, and Eva I. Dyer. Making transport more robust and interpretable by moving data through small number of anchor points. Proceedings of machine learning research, 139:6631, 2021.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.027477 2022.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schrödinger bridge. Advances in Neural Information Processing Systems, 35:9374-9388, 2022a.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. RSB: Image-to-lmage Schrödinger bridge. 2023.
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A Theodorou, and Ricky TQ Chen. Generalized Schrödinger bridge matching. In International Conference on Learning Representationss 2024.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 030003, 2022b.
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models, 2022c.
Loshchilov. Decoupled weight decay regularization, arXiv preprint arXiv: 2017.
Facundo Mémoli. Gromov distances and the metric approach to object matching. Foundations of computational mathematics. 11:417-487, 2011.
Aamir Mustafa and Rafat K Mantiuk. Transformation consistency regularization semi-supervised
paradigm for smage-to-mmage translation. In Computer Vision-ECCT 2020: 16th European Conference, Glasgow, UK, August 28-28, 2020, Proceedings, Part XVIII 16, pp. 999-615. Springer, 2020.
12

---

Published as a conference paper at ICLR 2025
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learning stochastic dynamics from samples. In International conference on machine learning. pp. 25858- 25889. PMLR, 2023.
Marcel Nutz. Introduction entropic optimal transport. Lecture notes, Columbia University, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, pp. 8026-8037, 2019.
Michele Pavon, Giulio Trigila, and Esteban G Tabak. The data-driven schrödinger bridge. Communi- cations on Pure and Applied Mathematics, 74(7):1545-1573 2021.
Stefano Peluchetti. Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51. 2023.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications to data science. Foundations and Trends@ in Muchine Learning, 1(5-6)):355-607, 2019.
Stanislav Pidhorskyi, Donald A Adjeroh, and Gianfranco Doretto. Adversarial latent autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 14104-14113, 2020.
Simo Särkkä and Amn Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Ryoma Sato, Marco Cuturi, Makoto Yamada, and Hisashi Kashima. Fast and robust comparison of probability measures heterogencous spaces. arXiv preprint arXiv:2002.016/5, 2020.
Erwin Schrödinger. Über die unkehrung der naturgesetze. Verlag der Akademic der Wissenschaften in Kommission bei Walter De Gruyter u..., 1931.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching, 2023.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne, Aligned diffusion schrödinger bridges. Uncertainty Artificial Intelli- gence, pp. 1985-1995. PMLR, 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. 2020.
Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. A generalized path integral control approach to reinforcement learning. The Journal of Machine Learning Research, 11:3137-3181, 2010.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrüdinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compula- tion, 1367111 2011.
Yuguang Yan, Wen Li, Hanrui Wu, Huaqing Min, Mingkui Tan, and Qingyao Wu. Semi-supervised optimal transport for heterogeneous domain adaptation. In I/CAI, volume 7, pp. 2018.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.I PROOF OF THEOREM 3.2
Assume Xo. X1 Rd, with Xo ~ no, and X Ti, and consider a stochastic random variable X connecting Xo and X1, whose law is the continuous marginal probability path pt. joining no, and n Starting from the static remi-supervisee appropriately and obtain the following relaxed dynamic formulation.
guided entropic OT problem in Eq. (8), we modify
min /
(Glluell + Erore [lu AGool)
Pe. 22,1
(15)
Pp. Ot
o² Apt. and No. T1
s.t.
(uep)) 2
Before the proof, we provide sketch of the proof overviewing the key steps towards the derivation of our dynamic formulation.
Sketch of Proof. Let no: 7I be the two given boundary laws, recall the static formulation: min, f ||x X1 + X1) + ( TT @ π1). The first step of our analysis re- arrangement of the terms of the static semi-superivsed EOT problem in Eq. (8). This rearrangement allows us to rewrite the modified EOT as regularized static Schrodinger Bridge problem.
min KL + G(Xo: 1doddr
Starting from the EOT problem, from 40 and IN 1, we can consider any intermediate point and find the optimal paths pi. p5 between the endpoints in the two newly formed subproblems (i.e.. from to to f1. and t to tx). The union of the optimal subpaths forms the optimal policy of the initial problem, and the converse is also true Villani al. (2009). This can be extended for any number N of intermediate points. As N -> 0o. this leads us to to consider 'local' versions of optimal transport problems between infinitesimally close distributions. This implies that the local path Pt.,t++: converges to the continuous marginal Pt, for arbitrarily large number of steps in [0, 1] Villani et al. (2009). Therefore, of local EOT problems.
next step, we express the guidance constraint EOT problem as a sequence
N min min
N K L / li-i + Bpg E G(X.. G(4r)
2
Subsequently, we rewrite the G Rd x Rd + R function between Xt, and Xt... and as a sum of telescopic series with constant reference point Xo from the source distribution. This induces the difference G 2(X,,, 1 Xp) G,,,,, Xo). Although, this term is positive at the optimal path, in order to ensure positivity of this term throughout the optimization process, upper bound it with its absolute value.
N-1
Σ Ep Ep G(Xe,, Xo) G(Xiun, Xo) + Σ Xr
min
Finally, we apply the Girsanov theorem in the KL divergence and Ito's Lemma in the difference G(X Xx) G(X,,,X, Xo) introduced in step 3, where the stochastic terms dropped assuming zero mean Wiener noise.
Assuming now that the number of intermediate steps N co, this yields our continuous dynamic formulation
min Pejo.
1 Ep. [
M 11+ VGyol+ -²
AGHoo] s.t.
Opeo.1
V..440.1P0001) Apro.1
öt
14

---

Published as a conference paper at ICLR 2025
Proof. First, let c(Xo, X1) be the cost from the static optimal transport expression, and let no. R1 be the two given boundary laws. Recall the static formulation: min, |x - X1ll + G( X1) + eK L (o & T.). Our objective is to derive a dynamic formulation for the static EOT problem with the cost function X1)
||Xo X1 11 + (1). We rewrite the divergence ( Too & TT), and the quadratic cost | x1, as ( ( assuming gaussian prior measure Q (Nutz, 2021). More specifically, it holds that log T x Xo -xlll// /ell hence n o gllXo-Xill//e /e, We rewrite the objective in Eq. (16), as [ + G( x mo drgdæ1. Therefore, the static objective can be rewritten
min KL (L + XI)oo, dxodæı
(16)
In this work, our focal point is the derivation of dynamic expression for the semi-supervised regularized cost function in 8. Through the prism of entropic interpolation Gentil et al. (2017), the solution of dynamical OT formulations correspond to solving 'local' versions of optimal transport problems between infinitesimally close distributions Villani et al. (2009). Let stochastic random variable path (Xt)oste joining Xo ~ To to X1 ~ T1, namely X₂ is an interpolation from Xo to X1. This random variable is assumed to be the solution of the following SDE
dX. u.di + adW,
(17)
Equivalent we can consider a marginal density induced by 17 (pr)ogesi which joins To, and π Villani et al. (2009) admits the Fokker Plank Equation (FPE):
1 Öp: - (u:pi 2 0 Apt Ot
(18)
Now, application of [Villani et al. (2009), Theorem 7.21] suggests that if we consider Pt be the solution to the infimum above, namely a minimizing curve. The endpoints of the random curve X whose law admits an optimal transference plan, are optimal coupling (Xt,) Then from the definition of the optimal cost and the minimizing property of the X₁
N
N 1
min min
KL (41) + Ep.
Σ X(xx) 1=0
(19)
for with Po To. and pN N. It generally holds Ep, c(r;, xi+1) 1) 2 Ep, c( xi41) ) Ep, c(zo. xi). with the equality being true for the minimizing path. Therefore, we can say that
min min Ep, o(,,, X.-1) Pe
min min Ep, (Xo,,X++1) Ep, c(Xo,
(20)
Pt X,
X,
This yields the modified version of the dynamic analogue
min min
(L(Po...110..41) K KL (40,))
(21)
E E [G(X Xo)l Ep,, [G(XX,,,
However, note that K LL 190 K KL 1qu,)) S KL 19)-144)) Subsequently, note for any coupling between Pt,, and 1t, 11 it holds that
Ep,,y [c(X Xo)l - Eep.,.oce
Xo)l < E [o(X,, Xo) - ((X,....
Finally, in order to ensure positivity of the difference G(X,,, Xo) - G(X Xo). we upper bound with its absolute value, and finally obtain the relaxed dynamic objective, which admits the following formulation
N
min X
E E
G(Xi, Xo) G 1(X
Σ
K
(22)
1=0
15

---

Published as a conference paper at ICLR 2025
For the first term, assuming the Markovian property of the base measure, apply the Girsanov theorem in the infinitesimally small time-step i - +
K (p)
dt |uc,112Ax
(23)
For the second term, have
E 1. (G Xo) G(X,,,
For discretized random variable X obeying the following discretized SDE: X4411 Xt, u1, Ali Tt, AW,, application of the discrete version of the Ito's Lemma yields
aG,, d1
Epg
VGG, + 2 oo + (or, VG,))WWJ)l
(24)
8G 0, since does not explicitly depend on t. and also assume that the Wiener holds that St, process increment has zero expectation, hence the equation above is rewritten as
Vu, 2
Tr(VaG.))AA.)
(25)
Now, make of the fact that [[f(x)]] S [ to take the expectation out of the absolute value.
d²
o
N-1
(u VG o² of TV(V
E u
V
2
(26)
Note that G is convex, i.e., e' ,'TH(V?G) 2 0, thus we conclude that the expression in Eq. (33) is upper bound by
[[u[VG] + d² THV
E
Σ 1=
G(X,, Xten)
E Ep
(27)
Finally, let us substitute back in the expression above the function G, and consider the continuous analogue of Eq. (23) and 27
( [1u Epr.o (Iu]
d² de
min min where Gi, G(X,, Xo). Finally, conclude the relaxed dynamic formulation is given by
AG.,))
(28)
2
min [' E luell Epo [lu[V o² AG.oo
min min
(29)
subjected to the FPE: and p1 ni.
Sp: a1
= (upp)) + 00 222A Ap,, and under the boundary constraints: Po = -To
A.2 PROOF OF PROPOSITION 3.3
Let the continuous marginal path satisfy the following decomposition p. = Pejo, To, dxodxx For optimization with respect to the intermediate path, we freeze the joint distribution at the boundaries Tu,1, which results in Eq. (9) being recasted as
min ['
a² 2
P.jo.1 J 10 VG(X,, + AG(X,, xo) deot,
(30)
Optj.1
s.t.
(PHo, ..)++ + 0² Aprio,
Øt
16

---

Published as a conference paper at ICLR 2025
Proof.
Initially, let us deal with the dynamic objective in Eq. (29)
o²
min min Pc
Ep, 122 Ill 12 Epoj- [luf [ vG + 2
AGroo
(31)
2he
The marginal is constructed as a mixture of conditional probability paths = J Peo. To, To. dxodx)- Now, sample pairs of so, 21, and fix the couplings. This is equivalent to fixing the vector, which induces the couplings. Consequently, fixing the coupling results in recasting Eq. (29) as
+ VGuol + AGro]
min Psj0.1
Epp
(32)
where 110,11 is used emphasize the drift that corresponds to the conditional probability path, Gejo = (o). Additionally, notice that by fixing the coupling, the expectation Epol. is dropped.
N E X 2=
min min Xx
(33)
q0.1
Now, we want to prove that the conditional path preserves the marginal. Our process to show that is isolate from every in the Fokker Plank the joint distribution 50, and preserve the Fokker Plank formulation after freezing the coupling.
dt f 7o, =
Oppc..
at no, 1xxodx Eo
Opna..1 86
(P)) V (uejo,1 f P.jo. 1x J /V (1 E E V (u a.1Pja,1))
Ap: = V Vp: = V v V dxxodk = -V ddxoox = E 1
where in every term we used Fubini's Theorem, to change the order between the integral from the definition of the marginal and the differentations from the FPE. Therefore, conclude that by freezing the coupling the Eq. (29) is recasted as
|1 11 [1u VG + a² AGrol] 2
min 0.1
Ep,
(34)
0p -V (u + 0² s.t.
2 Aprj,1
01.
A.3 PROOF OF PROPOSITION 3.5
The conver conjugate of the Lagrangian in Eq. (10) is defined as the Hamiltonian H(x,, ar,t) sup a ("ijo,1 - L(xe, tt). The optimization with respect to the drift 4too1 yields uejo.1 =
a, g(a (oo) where VG G(Xe, zu),, and
aTVG.p TvG if Va. < |
=
sgn(a, V else
Finally, we obtain the Hamiltonian associated with the Lagrangian in Eq. (10)
o² H( ar, t) 111 g(a VGoo) VG
1
AG(X,, 2
(35)
Proof. We begin with the definition of the Hamiltonian, dependent the gradient field de-
H H ac, t) sup (u (.jjo,:aa) ae) L(xx, t))
(36)
#4j0.1
17

---

Published as a conference paper at ICLR 2025
A more detailed analysis the gradient field is given in Neklyudov al. (2023). Given the convexity of Eq. (10), we compute the Hamiltonian through the subgradient of the absolute value ([u VG|). Optimizing with respect to the conditioned drift 4 yields the following
at - (we + (u, (V. = 2/j0,1 - ar - V (V..)
(37)
where a(.) denotes the subgradient. Note that
in VG, > 0: u = a. (VG,) = (at - (VG.) 2 20 2 /VG,1
" VG, < 0: = =a + (VG,) = (at + (VG,) 5 2 o.VG - <
u, VG, 0: a(. = 1-. from which we infer a. VG. a VG, , for [-1, 1].
This essentially means that the inner products u. VG, and a VG, share the same sign, namely (gn(u[VG,) = sgn(a, VG,,, when u. VG, # O, otherwise a. VG, al|V,,11
(Gt) VG:) = sa VG.,, O(I[VGr = of VG.
for Va > |
V.G. for |a S VG,|
We define function g R + R. be equal to:
(gg for Va 2 |VG,
saa VG.)
VG, G.ll
for [o s ||V
and substitute it in Eq. (36), which yields expression for u, function of the gradient field
1rj0,1 a, g(a]VG,) VG,
(38)
Now, we will provide expression for the Hamiltonian dependent only on the gradient field and the feedback term (X,,
= (a= V., - - VG. (a VG, VOi) - 0² AG. (39)
Therefore, we conclude that the Hamiltonian can be expressed as:
o² AG 1 = - vG,
H(I,, t)
(40)
2
A.4 PROOF OF PROPOSITION 3.6
Given the optimal conditional drift 1., match the parameterized drift uf, by minimizing the optimality gap given by our Entropic Lagrangian Bridge Matching objective as
min (' Ep. aa u/ - g((a]GG) VG
(41)
From Proposition 3.5, find that the parameterized drift can be expressed as up = ag VG.) G Therefore, the matching objective of Eq. (13) can be rewritten as follows
min #
Ep. 1/7 adl
(42)
Proof. Given the prescribed path, we seek to learn a parameterized drift field that matches the optimal conditioned paths, acquired from the previous step, which minimized the Lagrangian in Eq. (30). Following advancements in matching frameworks, we extend the notion of bridging the variational gap for entropy regularized optimal transport problem with general langriangian cost. We will first prove the following proposition which shows that we can express the variational gap that we need to bridge is expressed through the Bregman divergence.
18

---

Published as a conference paper at ICLR 2025
Proposition A.1. For Lagrangian function L(X,, ue t). the convex conjugate is given by the Hamiltonian: H (zt. ae t) = sup(a,, ut) L( d,,t), dependent on the gradient field or. The variational gap between the optimal drift and the parameterized gradient field is expressed as the following Bregman divergence.
S(¹ Di, (hooalProo.dd dddd
L - fr -
(43)
Proof. Starting from a general minimization problem for Lagrangian strictly convex with ua.
min ( / L(2,, 1 t)p dzdtt subj.: Ppro. Majo,a
(P 0, 02 Apro. o and Po To. P1
Ōt
21
(44)
and expressing it Lagrange optimization problem yields, with 8 is the Lagrange multiplier enforcing the Fokker Plank equation constraint
( J Laxe, 1 i
Pprj.1
(())=
min "ejo,1
subj.:
-
(Pujo,12eoo,1) + 00 Ap.o.
Ot
('' /op. [x u.oo1 p +
S
Pp.j..1
(Po10.1 00²² 12p
sup min Ue 0,5
81
Ot
Pp drdl
sup min
Lax, " () +
0,3
0 Jad
Rd
OL
9t V.V. (Pejo, 1 d
drdt [/' [
St
0² Apoo, dadt
2
Rd
la
(45)
We separate each one of the integrals and perfrom integration by parts
--=oooo See 81 @pga. drdt [sl So J' P ot
ans drdt
13 = Se V.V a²
-So Sas (Vor u)t,
T So f se (VP dn)dt
$ Sc dn)dt So Sgg siApoo, dtdr So (Voc. Vp.o,)d
-S' S&c As Pajo dx + o² & sp
Substituting them back to Eq. (45), we obtain
S' / (((x1,)e0111))
Ost Ot
L. = sup min
(Vst, u) Ass) Pao. - Epo [zo] + Ep, [x]]
81 1110.1
o²
f 8t (pejo.1,
2
('
Os:
σ² / o st se(e)o,,,
sup
(-H(Ir, Vs,,t)
Ast) Peoo Eyy [wo] Ep. [x1]]
Rd
at
min 0 Bd
/
(H Vs,, t) +
3s:
+ Ast) Prjo, dxdt + Epp [20] Ep, [x1]]
2 /
{ St (e dn)
at
(46)
where, the second equation stems from defining the Hamiltonian of the Lagrangian L(x, u, t) with respect the Lagrange multiplier 8, H (ze, Ss. 1) sup(u,, Vs.) T( 44 f). Now, consider
19

---

Published as a conference paper at ICLR 2025
the optimality gap between parameterized gradient field, and optimal gradient field st. which corresponds to a loss function L* We compute the optimality gap as
_ - -- !, (H Vst, +
a²
/ f 81 sp
Ast) 1zjot + Epo [xo] - Ep, [221]
0t
2
J'¹ jad
(H( 1) + dsj
a² / o (Pto,,1 dn)
+ As;) Pejo, + R [vo] Epp [mr]
jt.
1/' J (H( Vs
(0,, Vs, () Poo + f Jee
dst
Ot Prio, drdt + Epo (oo) Epi [21]]
So Sas 0p.jo.1 sydade
isi δt
Pejo, 1 dzdt + Ep [Fo] Ep, [21] S"' / ne
o² 2
As
st se(pHo,, dn)
Rd
So Sad Opejo.1 dfddt
So Sas 1
As Ps dxddt o² 2
So Sea 1s
(H Vs (3 dddtt + S J/e (s, st)( ppri,11 + Aprio, drdt
1-' [
J&c (H(De, Vs.,t) (o,, Vs, P) dzdt + S. [ (st s (Proo 17 10
1 drdt
I, [ (H Va.,)) (ma, Vst. ()) (V(s: +1), 1io..ippo.ddddd
(47)
At this point, we define the gradient field function az (X.)) Rd IR,, which is equal Vs, = de which finally yields
L C- [(¹ [ H(I,, a),t) L(xe, u/o 1. t)) (ar, PP drdt L'
(48)
= -f0 [ Di,uraa : u/o, l...
where the first equality stems from the fact that at the optimal point (aj, uf), holds that H(s,, at, = (at, u;) L(xt, ut,t) f) (Neklyudov et al., 2023; Boyd & Vandenberghe, 2004). This allows us to substitute the expression in the last line of Eq. (47), (at, (at,))) - H(,,, at, f) with L(x,, ut f), and express the variational gap as the Bregman Divergence, which emerges by the definition in the last equality finishing the proof of the proposition.
Therefore, to minimize the optimality gap, we sample sufficiently many uncoupled pairs (ro, 2n and parameterize the gradient field af. and try to match the parameterized gradient field on the optimized drift u,. We express our matching objective an expectation over the sampled pairs of optimality gap derived from the previous proposition
min Ero. [L - L*] = min BrO.. [./0 [ H(ae, ai + L(ae, u/o,1: 6) - (af, ut uu) dzds
(49)
Substituting for the expressions we have for the functions L, and H, we obtain
-- = . / (Elld g: vG,l1
o²
EnO,, [C L-
AAG
2
(50)
1
x²
+ 2 |u/o,111 +luijo1 TV
+
AG; (af. u.o..) P.jo. no, dedt ioud
2
20

---

Published as a conference paper at ICLR 2025
At this point, we take cases for the different values of gr. First: sgn(a, sgn(u, VV,) > and g(a, VG2) = 1:
/
1
Bmo.1 [ - -6- J Rd
Glla?
vG,12 + 11 112 + u/o TGG - (a 1) Prdddt
! JRd
+ = 2 1 (1u aE Prdxrtt
2 ala
f R2
1
(
VG. -Wjo. pp)dd
(51)
Second: gn((aVGi) sgn(ug VG,) < 0, and
-1:
B [ - C*] / / Gllae vG.11
1
1 112 47j0,1 TG. (al, uio, (pddrrt
2
/-¹ gd
( vG,1? + l1
1
(uijo. af VGi)
/. /
1 [ flo + VG. - -ujo.1 11 pp)dd
(52)
Third: VG.) 0,, and Va. A with A E [-1,1]:
/ 100 -V +
1
E [L" L*]
ア |ui
(aa u 1))
R
/ ggd
1
lof 2V
+ |ui
+ tu TVG, (al u/o, didddt
2
VG,
+ 2 11 (u at AVGi))
/ gd 2 la AVG. -ujj. Pidddd
(53)
where the second equality comes from the fact that we added the quantity Au VG, 0. Finally, we can succinctly express the optimality gap combining Eq. (51), and 52:
min Ex., [C [-] = min / /
RR= Gll g(a,, VG,) u/jo, ppdddt
(54)
But notice that by the definition of uf in Eq. (38), we have uf = of g(a. VG VG.,,. hence the minimization problem of the variational gap in Eq. (54) is expressed
min 1' f |1 af pdddtt
(55)
B DISCUSSIONS
B. GAUSSIAN PATH APPROXIMATION
Finally, following recent advancements matching frameworks (Liu et al., 2024; Albergo al., 2023), the conditioned probability path is approximated in simulation-free manner, as a Gaussian path
Pri011 R (NT
(56)
where the process is given from
X,
Tooo1 + 0,7
21

---

Published as a conference paper at ICLR 2025
Algorithm 2 Spline Optimization
1: Input: 21, (Xt,) where 0<4, <(x <1, and the pairs from the key-point (KP) set {u0 =i Ja=I.
and obtain their trajectories.
2: Initialize J, Fz
3: repeat
4: Sample X, ~ o
Evaluate the objective GiX,,, 2o) using KP pairs
Estimate the conditional drift (e(X, (X, zo, Eq. (11)
7: Take gradient step w.r.t control pts. (Xt,: Tt
8:
until converges
9:
Return Prjo,1 parametrized by optimized 4., d2.
where Tojo1 T(t, To, vi) is the interpolant function between the pinned endpoints, and Z ~ (O, Ia) is the Brownian motion. evaluate 4, and or through the following equations for the vector field and the score function (Särkkä & Solin, 2019)
X,X, = Ör, + 0,01 (X,- - Tt). Vlogpt(X.) = - 1,)
Lastly, using these expressions, obtain the following expression for the conditional drift
u (X,,((o, m1) = a. 0.0r (X, It)
20 (X: Tt)
at
where v given by the definition of o, = vv We can see that substitution of 4. := (1 - t) + 1æ and or == / (iv/((1t)) t), to the conditional drift expression indeed yields the desired drift of the Brownian bridge 21=X1
Spline optimization. Efficient optimization is achieved by parameterizing 4, o respectively as d- and 1-D splines using several knots (i.e., control points) sampled sparsely and uniformly along the time steps 0 < 4 <... < lx < 1: X1, E Rd and ot, E R. Among those knots, we fit the dimendependent interpolant function 4, and the standard deviation oz
I Spline(t; zo, (Xx, 21)
0e :=== oo=0, )1=0)
(57)
Notice that the parameterization in Eq. (57) satisfies the boundary in Eq. (56), hence remains as feasible solution Eq. (10). The number of control points K is much smaller than the discretization steps <30 for all experiments). This significantly reduces the memory complexity, compared to prior methods caching entire discretized SDEs (e.g. (De Bortoli al., 2021: Chen et al., 2021)).
We follow the algorithm proposed in (Liu et al., 2024), shown in Alg. 2, which, crucially, involves no simulation of an SDE. This is attributed the utilization of independent samples from Prjo. which are known in closed form. Finally, recall that since we solve Eq. (10) for each pair (Io, x1) ~ Pb.1 and later marginalize to construct pt-
C ADDITIONAL DETAILS ON THE EXPERIMENTS
In this section, provide further information about the experiments to validate the efficacy of our FSBM.
Baselines We compare the efficacy of our FSBM in the crowd navigation tasks against the state- of-the-art matching framework GSBM (Liu et al., 2024), with the default hyperparamters. For the opinion depolarization task, we compare against the GSBM, and the DeepGiB (Liu al., 2022a) also run with their default hyperparameters. For the DeepGSB, we adopt the "actor-critic" parameterization which yields better performance. Finally, the image translation experiments, we compare against with DSBM (Shi et al., 2023), and Light and Optimal SBM (LOSBM) (Korotin et al., 2023). In our experiments, DSBM was implemented by considering trivial guidance returning the analytic solution of non- guided Brownian bridges. All other frameworks were using their official implementation, and default hyperparameters. All methods, including our FSBM, are implemented in Torch (Paszke al., 2019).
non-guided
22

---

Published as a conference paper at ICLR 2025
t= 0.25T
= 0.5T
t= 0.75T
T=T
t-0
0-2T57
t- 0.5T
t- 0.75T
bridiin Rudd depport sanple Vipecarr
Figure 9: Comparison of gradient fields ensuing after training with UP: distance- preeervvir G,
DOWN: cellation-prseervie G
Guidance Function Recall that analysis begins with the static entropic OT problem in Eq. (3), to which an additional regularization term is included to capture the interaction among unpaired samples and aligned samples, thereby guiding the transport map. In practice, the selected Guidance Function was a slightly modified version of Eq. (7)
(58) where a E IR is the guidance regularization hyperparameter responsible to amplify G, and bypass the minimization of the kinetic energy of the non-paired samples enough to satisfy the feasibility of the solution. where for our experiments, the selected distance between the non-coupled samples and the fixed KP sample was the L1 norm x((,,, ||X; æill L1. The th aligned data point was the closest keypoint sample to non-aligned particle X, at time O. To compute the guidance function, we need the trajectories of the KP samples. These are assumed be given only through their samples and fixed for each KP pair. In practice in the next sections, we discuss how we acquired those trajectories for each of our experiments.
G(Xo, X.) = a (d(X,, d( X(.o. xf))²
Cecaatio-Prreeevvng vs Distance Preserving Although the relation- pettrion-preervinng scheme was indeed the more effective choice in the static setting in [1]. this was not the case in our applications. This could be attributed to the fact that the softmax averaging relation preserving scheme works more effectively in static frameworks. However, in dynamic formulation, the softmax averaging in relation glatton-preserving schemes (similar to [1]) failed to generate a gradient field capable of providing effective guidance to the unpaired samples. In contrast, the distance-preservingg scheme offered significantly better guidance. Figure presents a comparison of the ensuing gradient fields after training with the pratrion-rreerving guidance function similar (Gu al., 2022), and our distance- distance-preserving guidance function.
observed empirically that
Related Works Super-sue OT) Recently, the field of Optimal Transport (OT) has seen advancements in developing partially supervised frameworks aimed reducing the reliance on source-target aligned image pairs for training, thereby alleviating the high cost of labeling supervised datasets (Mustafa Mantiuk, 2020). In particular, the semi-supervised domain adaptation framework (Sato et al., 2020) demonstrates how limited labeled target domain data can act as keypoints, enabling the adaptation of unlabeled data through iterative optimization. These approaches generally share common theme: the way partial supervision is integrated into the OT formulation. For many applications, annotated keypoint pairs are pre-determined, and guidance is incorporated into the OT problem regularization terms that constraint the transport map (Gu et al., 2023a). For example, (Courty et al., 2017) explored regularizing the OT cost using functions that encode class label information. Their regularization term used to preserve the data structure while encouraging the alignment of labeled data with shared class labels across source and target domains for domain adaptation. Similarly, (Yan et al., 2018) proposed method that uses the Gromov- Wassersttinn (GW)
23

---

Published as a conference paper at ICLR 2025
model to transport source samples to the target domain. Their approach induced constraints in the transport process by minimizing the distance between the centers of transported source samples and labeled target samples with the same class labels. Furthermore. (Lin al., 2021) regularized the OT cost to capture the group structure between source and target distributions and their representative anchors. In this framework, the anchors promote clustering, similar our approach, and are used to enhance robustness to outliers by imposing rank constraints on the transport plan. However, unlike the concept of keypoints in our approach, the anchors do not represent pairs of annotated points explicitly used to guide the transport map. More closely aligned with our approach, (Gu et al., 2023b) proposed a relation- preseerving keypoint- guiddd model that steers OT matching by preserving both the correspondence between keypoint pairs and the relationships of each data point to the keypoints. In their framework, this guidance introduced as a regularization term the OT formulation. While the utilization of keypoitns by Gu et al. 2023 is similar to ours, their choice of the guidance function differs. Specifically, they modeled the distance as a softmax- softmax-avcraaed relationship between unlabeled samples and keypoints, with the guidance expressed as the -aasnn-Shannon divergence between the softmax averages of the source and target distributions. Lastly, in (Gu et al., 2023a) it was shown that this partially supervised OT framework could be applied to effectively guide score based diffusion model.
Network Architectures For all experiments, we adopted the same architectures from GSBM, which consists of 4 to 5 residual blocks with sinusoidal time embedding. All networks are trained from scratch, without utilizing any pretrained checkpoint, and optimized with AdamW (Loshchilov, 2017)
Forward and Backward Scheme Training of our FSBM entails a originally proposed in DSBM (Shi et al., 2023). This necessitates the parameterization of two drifts, one for the forward SDE and another for the backward. During odd epochs, we simulate the coupling from the forward drift, solve the corresponding conditional lagrangian problem in Eq. (10), and then match the resulting Pt with the backward drift. On the other hand, during even epochs, we perform the reverse procedure, matching the forward drift with the marginal pt derived from the backward drift. As discussed in Liu et al. (2024), this alternating scheme improves the performance of the matching algorithm, averting the matching framework from bias accumulation, as the forward drift always matches the terminal distribution Ti and the backward drift is ensured to match the source distribution To-
scheme
C.1 CROWD NAVIGATION
We revisit the efficacy of our FSBM in solving complex crowd navigation tasks and its superior capability to generalize under a variety of initial conditions. More explicitly, the two compared frameworks were trained on transporting the samples Xo To to X1 TT In the evaluation phase, the task is always to transport the samples to the same target distribution as in the training phase, however, the initial conditions change throughout the different tasks. For the first task (Vanilla), we draw new samples from the same distribution Ty = -N(uo on), for the second task we shift the mean of the distribution, for the third task increase the standard deviation, and in the fourth task sample from a uniform distribution. The initial conditions for each task are the following:
Vanilla: AV ((1 [1,11,,11,,00
Perturbed Mean: /V [1-11,
/ -111,-4].0.33 -4].
Perturbed ATD:W//
STTD:N((-11,, (111,-11),3
Uniform Distribution with geometric mean [-11,-1]
where ((Jui,Ha) represents a Gaussian with mean the point [2], and standard deviation equal to a. Note that for the S-tunnel crowd navigation, we used 100 KPs, which is approximately 4%, of the entire population, and in the V-neck task only 20 KP samples were employed, which is approximately 1% of the population. In experiments, for the trajectories of the aligned data, utilized the POT python package (Flamary et al., 2021), to acquire optimal pairings, and GSBM for one epoch to obtain the trajectories between the endpoints. Figures 10, and 11 illustrate the
24

---

Published as a conference paper at ICLR 2025
trajectories of the coupled samples, that guide the matching of the unpaired particles. Finally, Figure 13 presents additional plots of FSBM in the crowd navigation tasks examined under the variety of initial conditions mentioned above.
S-tunnel Keypoints
V-neck Keypoints
Keypoint Trejectories
Keypoint Trajectories
Source Keypoints
Source Keypoints
Target Keypoints
Target Keypaints
Figure 10: Interpolating path of the KP samples for the S-tunnel
Figure 11: Interpolating path of the KP samples for the V-neck
FSBM Vanilla
FSBM Perturbed Mean FSBM Uniform Dist.
FSBM Perturbed STD
Figure 12: Additional Figures on S-tunnel using our FSBM. The color of particles means: i) Yellow: initial conditions, ii) Green and Cyan: intermediate trajectory, iii) Red: target distribution, iv) Navy Blue: generated distribution
FSBM Vanilla
FSBM Perturbed Mean
FSBM Uniform Dist.
FSBM Perturbed STD
Figure 13: Additional Figures V-neck using FSBM. The color of particles means: i) Yellows initial conditions, ii) Green and Cyan: intermediate trajectory, iii) Red: target distribution, iv) Navy Blue: generated distribution
25

---

Published as a conference paper at ICLR 2025
Epochs
Table 6: The W Distance under Vanilla and Per- turbed Mean in the S-tunnel for varying the number of KP samples
25
Percentage KPs
W Vanilla
W
w Distance
1.
Pert. Mean
0,034
0.77
Perturbed Meen
2.5
0.027
0.55
Vanilia
66
5
0.018
0.79
53
Percentage of Aligned Data
23
10
20
8
0.015
0.82
10
0,013
0.81
Figure 14: UP: Training Epochs, DOWN: The 8 Distance under Vanilla and Per- turbed Mean in the S-tunnel for varying the number of KP samples
12
0.010
1.01
15
0.008
1.23
20
0.008
1.48
C.I.1 IMPACT OF KP PAIRS THE S-TUNNEL
In our methodology, becomes clear that the number of KP pairs used to guide the rest of the samples affects the effectiveness of our algorithm. In Fig. 14, and Table ??, we see that the vanilla performance of our framework continues increase the number of aligned pairs increases. As depicted in Figure 14a, although an increase of the aligned pairs results in reduced training time. Additionally, in Fig. 14, we see that the vanilla performance our framework continues to increase as the number of aligned pairs increases. However, it is shown that the empirical generalization to unseen initial conditions (e.g. when the mean of the initial distribution perturbed) does not necessarily improve by adding more keypoints, suggesting overfitting like phenomenon. There appears to a "sweet spot", which suggests that using too many KP pairs hinders the ability of our algorithm generalize effectively akin to overfitting- note that the hyperparameters of the framework were tuned to maximize the performance while using few keypoints. As result, increasing the number of aligned pairs may require re-tuning the hyperparameters to maintain optimal performance. Lastly, it should be mentioned that even this decrease in generalizability understanding the effect of the aligned data points in the efficacy, and generalizability of the model is a topic for future work.
effect. Lastly, it is important
still very small compared the benchmark methods. Deeper
C.2 OPINION DYNAMICS
Subsequently, we revisit the high-dimensional opinion depolarization of population. Each particle possesses a siigh-imensional opinion X, 100 that evolves under polarizing dynamics
dX, = fpolariee (X,, t)dt + odW,
(59)
which tend to segregate them into two groups of diametrically opposing opinions.
Polarization drift We the polarization drift from DeepGSB and GSBM Liu al. (2022a; 2024), based on the party model (Gaitonde et al., 2021). At each time step t, all agents receive the same random information & Rdd sampled independently of pt, then react this information according to
1 if (g)(((,,1)) sign((v. (60)
((): = Ryy- (a(x, ys E,) a(s,y. v.
otherwise
where y are opinions and ū = 1 and a(x, {,) is the agreement function indicating whether the two opinions agree the information & This suggests that the agents tend reject opinions. with which they disagree, while they accept opinions closer to theirs. As the dynamics of this model evolve
26

---

Published as a conference paper at ICLR 2025
Woman to Man
Man to Woman
Old to Young
Young to Old
Figure 15: Additional Translation Examples
Recall that the size of the key-point set utilized for the opinion depolarization was approximately 2%, of the entire population. In our opinion depolarization experiments, the optimal pairings were obtained through the POT python package, however, in contrast to the crowd navigation tasks, in this case, the trajectories of the opinion guide particles were obtained by linear interpolating between initial and target positions.
C.3 IMAGE TRANSLATIONS
Finally, we review the experiments in the two types of image translations: Gender translation, ii) Age translation. We follow the setup of (Gushchin et al. (2024)) with the pre-trained ALAE autoencoder (Pidhorskyi al. (2020)) on 1024 x1024 FFHQ dataset (Karras et al. (2019)) perform the translation in the latent space of dimension 512 x enable more efficient training and sampling.
Recall that the number of aligned images was 4% of the total number of images for the gender translation and 8% for the age translation. In our translation experiments, the aligned pairings, used guide the rest of the matching algorithm, were obtained by employing lightweight SB algorithm using Gaussian mixture parameterization (Korotin et al., 2023). Theoretically. this algorithm provides guarantees to solve the SB, and empirically the quality the images generated in the training set of this algorithm was deemed satisfactory to constitute our KP set. Tables 7 and 11 present additional results establishing the superiority of the coupling obtained by our FSBM. More specifically, Table presents the L2-norm of the difference between the pixel values of the input images and the generated images for each the translation tasks. Note that the values have also been divided by the total number of pixels (1024 1024) derive mean 7.2-norm difference. Furthermore, Table 11 presents
27

---

Published as a conference paper at ICLR 2025
the Structural Similarity Index Measure (SSIM) between the input and the generated images in the translation tasks. Finally, Figure 15 presents additional translation instances using our FSBM.
Table 7: L2 - norm/ (1024 x 1024) values between input and generated images in 4 translation tasks for DSBM, LOSBM, and our FSBM (Lower better). Note the values are divided by 220, which is the total number of pixels to get a mean L2-norm over all the pixels
Man to Woman
Woman to Man
Old Young
Young to Old
Optimizer
DSBM
0.87
0.84
0.88
0.87
LOSBM
1.42
1.47
1.41
1.35
FSBM
0.79
0.81
0.83
0.88
Table 8: SSIM values between input and generated images in 4 translation tasks for DSBM, LOSBM, and our FSBM (Higher is better)
Man to Woman
Woman to Man
Old to Young
Young to Old
Optimizer
DSBM
0.46
0.51
0.50
0.43
LOSBM
0.34
0.35
0.35
0.36
FSBM
0.49
0.47
0.46
0.45
Finally, we also used ChatGPT to assess the image couplings between 100 pairs of input and generated images each translation task. Given that the coupling from LOSBM was considerably worse than the coupling generated by the other two methods, we only performed comparison between DSBM [1] and our FSBM. ChatGPT was instructed give score on the coupling quality on a scale 0-5, with 5 being the best. The coupling quality was assessed using the following criteria: i) Background Consistency: consistency terms of color, lighting, background objects, ii) Identity Preservation: maintain the identity of the original face in terms of facial features and overall resemblance iii) Structural Integrity: preserve structural integrity of the input image terms of accessories worn, pose, orientation. The results in 4 translation tasks are shown in the Tables below.
Table 9: Background Consistency score between input and generated images in 4 translation tasks for DSBM, and our FSBM (Higher is better), assigned by ChatGPT
Man to Woman
Woman to Man
Old Young
Young to Old
Optimizer
DSBM
4.05
4.04
3.99
4.08
FSBM
4.58
4.21
4.05
4.10
Table 10: Identity Preservation score between input and generated images in 4 translation tasks for DSBM, and our FSBM (Higher is better), assigned by ChatGPT
Man to Woman
Woman to Man
Old to Young
Young to Old
Optimizer
DSBM
3.89
3.09
3.15
3.49
FSBM
4.02
3.48
3.72
3.60
28

---

Published as a conference paper at ICLR 2025
Table 11: Structural Integrity score between input and generated images in 4 translation tasks for DSBM, and our FSBM (Higher is better), assigned by ChatGPT
Man to Woman
Woman to Man
Old to Young
Young to Old
Optimizer
DSBM
3.38
3.27
3.11
3.30
FSBM
3.69
3.61
3,55
3.33
29


Reviews:
[{"review_id": "u9RGigrv3M", "reviewer": "Reviewer_T3hw", "summary": "Drawing inspiration from guided optimal transport schemes, this paper introduces Feedback Schrödinger Bridge Matching, a semi-supervised matching framework that uses a small set of aligned pairs to guide the transport map of non-coupled samples.\nThe proposed approach is evaluated on various distribution matching tasks, such as crowd navigation, opinion depolarization, and unpaired image translation, and compared against other state-of-the-art matching frameworks.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed semi-supervised approach uses only a small percentage of paired data and achieves faster training time.\n\nExperimental validation across a diverse set of tasks demonstrates the effectiveness of the proposed approach.", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "6aZ2vQTBcK", "reviewer": "Reviewer_GeTG", "summary": "This paper tackles a semi-paired matching for generative models, based on Schrödinger bridge. In the semi-paired matching problems, along with non-coupled data, a small portion of pre-aligned data pairs are provided. In methodology, the paper is inspired by (Gu et al., 2023a) which investigates the guided optimal transport using a few paired data in static OT setting. The authors utilize a \"distance-preserving\" scheme to develop the dynamic OT model with entropy regularization, i.e, the Schrödinger bridge. Building on this, the matching algorithm is developed. Experiments on three settings are conducted to evaluate the proposed methods.", "questions": "- How about including the metrics, e.g., LPIPS distance, perceptual similarity loss, SSIM, $L_2$ distance, etc. between the input and output images for the image-to-image translation? \n- For the user study, could you provide the generated images by different methods to distinct persons, and ask which one is better for the desired translation (e.g., better preserving the identity, background, stricture, etc.)? Meanwhile, could you ask the large language models (LLMs), e.g., GPT, to answer the question, as LLMs show a strong ability to understand text and images?\n- For the other experiments, could you train a classifier on the target domain data and apply it to the generated data, to see if the accuracy could be improved with the keypoints as guidance?\n- Could you try the relation-preserving scheme? At least, discuss the challenges faced. Meanwhile, could you provide more analysis on the effect of the number of keypoints in-depth, to figure out why a larger number of paired data do not bring performance gain?\n- How about selecting $i$ such that $x_1^i$ is closest to $X_1$ or a weighted combination of selecting based on $X_0$ and $X_1$. Which strategy is optimal or most reasonable?\n\nI am happy to increase the score if the questions are well addressed.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is interesting to explore the guidance of a small fraction of paired data to data translation in flow/bridge matching.\n- It is reasonable to deduce the dynamic version of keypoint-guided optimal transport with a few paired data as guidance for developing the corresponding matching models (I did not check the proof and mathematical deductions).\n- The paper writing is clear with clear organization and presentation.", "weaknesses": "- The main weakness to me is the unclear performance gain by introducing the keypoints, from the experiments. The authors reported the W-distance or FID between generated data and real data. However, to my understanding, the tackled problem in this paper belongs to the controllable generation, in which the transported data should preserve certain structures or information of the original source data. For example, in the image-to-image translation (Woman->Man) as in the experiments of this paper, the transported images should preserve as much information as possible of the input image as long as the transported image is a \"man\". The authors show some transported images, but I can not find the corresponding metric. Maybe a user study is helpful for this issue, e.g., ask the GPT to evaluate the generated images. Without such a metric, the importance of the keypoints can not be illustrated (the same issue exists in the other two experiments).\n- It is quite confusing that adding in more keypoints worsens the performance. Is this because the distance-preserving is not effective for such cases? The authors mentioned that Gu et al., 2022 presented the relation-preserving scheme, which seems to be more effective than the distance-preserving scheme as shown in (Gu et al., 2022). Why not consider the relation-preserving scheme instead of the distance-preserving scheme? Unfortunately, both the above issues are left as future work.\n- In Eq.(7), the $i$ is selected such that $x_0^i$ is closest to $X_0$. What if selecting $i$ such that $x_1^i$ is closest to $X_1$. The authors should prove that the designed strategy is optimal or justified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "6xuaQrg52I", "reviewer": "Reviewer_Pfgj", "summary": "Summary of the paper \n\nThe paper proposed a novel semi-supervised matching framework, termed as  Feedback Schrödinger Bridge Matching (FSBM), which was designed to improve the efficiency of distribution transport problems by incorporating a small percentage of pre-aligned pairs as state feedback. This framework was developed in order to minimalize  supervision on basis of scalability. It also addressed the limitations of existing methods, which often face trade-offs between computational efficiency and optimal pairing access during training. The authors claim that FSBM enhances generalization and accelerates training by reformulating the static Entropic Optimal Transport (EOT) problem into a dynamic one, thereby facilitating the matching of non-coupled samples.\n\nContributions \n+ Analysis begins with a static semi-supervised OT problems and derivation of a modified dynamic formulation out of it.\n+ Introduction of FSBM, the first matching framework leveraging information from partially aligned datasets and Entropic Lagrangian extension of the variational gap objective used in BM frameworks to match the parameterized drift given a prescribed probability path.\n+ Verification of remarkable capability of FSBM through extensive experiments to generalize under a variety of unseen initial conditions, while simultaneously achieving faster training times. These results were consistent across a wide range of matching tasks, from low-dimensional crowd navigation to high-dimensional opinion depolarization and image translation.\n\nMerits\n+  Proposal of a semi-supervised framework. This framework leverages minimal aligned pairs, which is a significant contribution to the field, because it effectively addresses the practical challenges posed by fully unsupervised and fully supervised methods. This method is novel and efficient.\n+ The paper offers a solid theoretical basis mathematically  by deriving a generalized EOT objective and demonstrating its advantages in terms of efficiency and generalization capability.\n+  The extensive experiments conducted across various tasks, including crowd navigation, opinion dynamics, and image translation, provide compelling evidence of the effectiveness of FSBM compared to state-of-the-art methods. Additionally  comparative experiments and their results were also shown to compare results with methods proposed in this paper and baseline methods.\n+ The inclusion of detailed performance metrics and visual comparisons strengthens the claims made, illustrating the superior performance of FSBM in practical scenarios. The results were shown vividly.\n\nWeaknesses:\n+ While the theoretical aspects are robust, the manuscript could benefit from clearer explanations and simplifications, particularly in sections discussing the derivation of dynamic formulations and coupling optimization. A few formula faults are discovered in the paper.\n+ The manuscript lacks a thorough discussion regarding the limitations of FSBM, particularly in relation to scenarios with highly variable datasets or when the fraction of aligned pairs is insufficient for effective guidance.\n\nSuggestions for Improvement:\n+ Simplify the presentation of theoretical concepts to enhance readability and comprehension for a broader audience.\n+ Expand the comparative analysis to include more specific scenarios where existing methods excel, providing a more nuanced view of FSBM’s performance.\n+ Discuss the limitations of FSBM more explicitly, particularly regarding the trade-offs associated with the number of aligned pairs used.\n+ Include suggestions for future research that could build on the findings and address any identified limitations.", "questions": "See Summary", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "See Summary", "weaknesses": "See Summary", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "n7tdeZm5u6", "reviewer": "Reviewer_LqYG", "summary": "This paper proposed a novel approach, which introduces Feedback Schrodinger Bridge Matching (FSBM) as a semi-supervised matching framework for addressing distribution transfer issues. The authors initiate their study with the static semi-supervised Optimal Transport (OT) problem. They judiciously incorporate an additional regularization term to capture the interaction between uncoupled and aligned samples, thereby enhancing the transfer mapping process.\n\nThe paper's strength lies in its innovative use of a state feedback term to encode information from aligned samples, effectively guiding uncoupled samples' transfer. Additionally, the authors optimize the drift field using divergence, supported by a variational objective.\n\nStarting from the Static Euclidean OT problem, the proposed term $G$ is seamlessly integrated into the cost function $c$, and feasible optimization algorithms are derived using Lagrangian and Hamiltonian formulations. Overall, the paper's motivation, methodology, and experimental results are partially well-presented and contribute to the existing literature.", "questions": "Minor issues:\n1. The resolution of figures should be improved.\n2. Some typos. \n3. What are the advantages of the current FSBM framework's dynamic target design through state feedback? Is there any theoretical evidence that makes it more transparent?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "+ Introducing Feedback Schrödinger Bridge Matching (FSBM): FSBM is the first matching framework capable of leveraging information on partially aligned data pairs, significantly improving efficiency.\n+ Semi-supervised guidance: Accelerate training and enhance generalization capabilities by using a small number of pre-aligned pairs as state feedback to guide the transfer mapping of uncoupled samples.\n+ Dynamic Target Rewriting: Reformulate the static entropic optimal transport (EOT) problem into a dynamic form to exploit the scalability of the matching framework.\n+ Experimental validation: Extensive experiments show that FSBM has better generalization ability under various unseen initial conditions and is faster in training time.\n+ Multiple task applications: FSBM performs well in various matching tasks, such as low-dimensional crowd navigation, high-dimensional opinion depolarization, and image translation.\n\nIn particular, I think this paper's theoretical derivation and experiments are good, and utilizing partially paired samples is also novel and exciting.", "weaknesses": "1. The key contribution is the proposal of the distance function $G$. However, the motivation of $G$ is not clearly illustrated. Can authors further illustrate why they give such a constraint to paired samples? What is the meaning of $G$? What is the meaning of $d(X_1,x_1^i)$ and $d(X_0,x_0^i)$? Is it relative to clustering? Authors can illustrate the behaviors of paired samples and unpaired samples more specifically. Is the behavior of FSBM related to the \"clustering\" of trajectories? More discussions, examples, and theoretical derivations to illuminate the motivation of the term $G$ are expected to be shown in the paper.\n\n2. Paired samples and stochastic properties of SDE. As we know, randomness is an important property of SDE/Diffusion/Bridges. From $x_0$ to $x_1$, if $x_0$ is sampled from $\\pi_0$, $x_1$ can be stocastic. However, paired samples seem to be deterministic. What is the relation between paired samples and the stochastic property of SDE/Diffusion/Bridges? Please address this concern explicitly in the methodology section.\n\n3. Figure 3 directly shows the effect of $G$. However, Figure 3 is qualitative and not rigorous. The term $u^{\\top}\\nabla{G}$ forces the trajectory of unpaired samples to the trajectory of paired samples. What is the reason? Please provide a more rigorous analysis of how G affects the trajectories, perhaps through a theoretical derivation or a quantitative experiment that demonstrates the effect.\n\n4. The proof Theorem 3.2. In Eq. 20, authors discretize $G(X_0,X_1)$ into a sum sequence $G(X_{t_i},X_{t_{i+1}})$. The authors seem to assume $\\delta t \\rightarrow 0$ and transform the sequence of $G$ into $G(X_0, X_1)$. Is the application rigorous? Authors and other reviewers can confirm that. Please discuss under which conditions the transformation (i.e., from the sum sequence to the time interval $[0,1]$) is valid. I suggest that authors provide a proof sketch to make the paper more readable.\n\n5. The contradiction of formulations of Hamiltonian function. As for the formulation of Hamiltonian function, in line 860, authors use $\\sup_{u} \\langle u,\\nabla{a} \\rangle-\\mathcal{L}(x,u,t)$. However, in Eq. 37, authors $\\sup_{u} \\langle u,a\\rangle-\\mathcal{L}(x,u,t)$, which is distinct from line 860. In the following derivations, authors use the formulation in Eq. 37, where the Lagrangian multiplier $a$ is not with the gradient operator. However, in Eq. 47 of Proposition A.1, from line 998 to line 1000, authors eliminate $\\langle  \\nabla{a^*, u^*}\\rangle$ using the product term in Hamiltonian function $H^*(x,a^*,t)$. Then, the authors use the Hamiltonian formulation with $\\nabla{a^*}$. Thus, in the proof of Proposition 3.5 and the proof of Proposition A.1, the formulations of Hamiltonian functions are contradictory. I am not familiar with Hamiltonian function, but the contradiction in the proofs should be claimed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Oral)