Categorical Schrödinger Bridge Matching
Authors: Grigoriy Ksenofontov
Alexander Korotin
Venue: ICML 2025 poster (2025)
Keywords: Schrödinger Bridge
Entropic Optimal Transport
Optimal transport
Unpaired Learning
Discrete space
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: This paper extends the Schrödinger Bridge problem to work with discrete time and spaces.

Abstract: The Schrödinger Bridge (SB) is a powerful framework for solving generative modeling tasks such as unpaired domain translation. Most SB-related research focuses on continuous data space $\mathbb{R}^{D}$ and leaves open theoretical and algorithmic questions about applying SB methods to discrete data, e.g, on finite spaces $\mathbb{S}^{D}$. Notable examples of such sets $\mathbb{S}$ are codebooks of vector-quantized (VQ) representations of modern autoencoders, tokens in texts, categories of atoms in molecules, etc. In this paper, we provide a theoretical and algorithmic foundation for solving SB in discrete spaces using the recently introduced Iterative Markovian Fitting (IMF) procedure. Specifically, we theoretically justify the convergence of discrete-time IMF (D-IMF) to SB in discrete spaces. This enables us to develop a practical computational algorithm for SB, which we call Categorical Schrödinger Bridge Matching (CSBM). We show the performance of CSBM via a series of experiments with synthetic data and VQ representations of images. The code of CSBM is available at [this repository](https://github.com/gregkseno/csbm).

Categorical Schrödinger Bridge Matching
Grigoriy Ksenofontov
Alexander Korotin
Abstract
et al., 2024b, D-IMF) has shown promising results in cer- tain unpaired learning problems, enabling faster generation (inference) times than its predecessors.
The Schrödinger Bridge (SB) is a powerful frame- work for solving generative modeling tasks such unpaired domain translation. Most SB-related research focuses on continuous data space IR P and leaves open theoretical and algorithmic questions about applying SB methods to discrete data, e.g. on finite spaces SD Notable examples of such sets & are codebooks of vector-quane qusor--uattized (VQ) representations of modern autoencoders, tokens in texts, categories of atoms in molecules, etc. In this paper, we provide theoretical and algorith- mic foundation for solving SB in discrete spaces using the recently introduced Iterative Markovian Fitting (IMF) procedure. Specifically, we theoreti- cally justify the convergence of discrete-time IMF (D-IMF) SB in discrete spaces. This enables develop practical computational algorithm for SB, which we call Categorical Schrödinger Bridge Matching (CSBM). We show the perfor- mance of CSBM via series of experiments with synthetic data and VQ representations of images. The code of CSBM available at this repository.
Unfortunately, the D-IMF procedure heavily relies on cer- tain thcoretical properties of particular SB setups in con- tinuous spaces. At the same time, a vast amount real- world data is either discrete by nature, such as texts (Austin et al., 2021; Gat et al., 2024), molecular graphs (Vignac et al., 2022; Qin et al., 2024; Luo et al., 2024), sequences (Campbell et al., 2024), etc., or discrete by construction like pector-quantized representations of images and audio (Van Den Oord et al., 2017; Esser et al., 2021). I). These cases highlight a fundamental limitation, as D-IMF is not directly applicable such data. In this work, we address this gap by making the following contributions:
Theory. We provide the theoretical grounds for applying the D-IMF to solve the SB problem in discrete spaces.
Practice. We provide a computational algorithm to im- plement the D-IMF in practice for discrete spaces.
Notations. Consider a state space X and a time set (tn [ where = to t <... < iN < I are N 1 time moments. The space X is referred to the path space and represents all possible trajectories
Introduction
def (Fo, Tin, Itv+1) where = (It,...., RiN) corresponds to the intermediate states. Let P(x be the space of probability distributions over paths. Each & E N+) can be interpreted as a discrete in time --vuuued stochastic process. We use q Tin ((Ni1) denote its density at (rn,
The Schrödinger bridge (Schrödinger, 1931, SB) problem has recently attracted the attention of the machine learning community due to its relevance to modern challenges in gen- crative modeling and unpaired learning. Recently, a variety of methods have been proposed to solve SB continuous spaces; see (Gushchin et al., 2023b) for a recent survey.
E N++ and use q(+l-) to denote its con- ditional distributions, e.g., q(21 (ro) q(Fin 21). Finally, we introduce N+) C N+) as the set of all Markov processes g, i.e., those processes which satisfy the equality q(Zo, Fin Dtx+s) = g( I ((0,, (tt,.)).
One modern approach to solving SB is the Iterative Marko- vian Fitting (IMF) framework (Peluchetti, 2023; Shi al., 2023; Gushchin et al., 2024b). Specifically, within this framework, the discrete-time IMF procedure (Gushchin
2 Background and Related Works
'Skoltech, Moscow, Russia 'MIPT, Moscow, Russia AAIRI, Moscow. Russia. Correspondence Grigoriy Ksenofontov ksenoooooooo @swoottehsoo..........o Alexander Korotin <a.korotin@e skoltech.ru..
In this section, review the formulation and existing ap- proaches to the Schrödinger Bridge (SB) problem, with focus on its generative applications. We begin with the static SB problem (82.1). Next, we highlight the challenges of extending SB methods from continuous to discrete state spaces (52.3). We proceed to the dynamic SB formulation,
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Categorical Schrödinger Bridge Matching
motivating its importance in practice ($2.4). This leads to the Iterative Markovian Fitring (IMF) procedure and its discrete-time variant D-IMF ($2.5). Finally, we summarize the known characterizations of SB (Table 1) and identify the object of our study, namely, establishing theoretical guarantees for the discrete state and time setting (32.6).
In the related literature, this setup is mainly explored in the context of umpaired (unsupervised) domain translation. In this task, the datasets consist of samples from two different data distributions (domains), and the goal is to learn a trans- formation from one domain to the other (Zhu et al., 2017, Figure 2). The problem is inherently ill-posed because, the- oretically, there may be multiple possible transformations. In many applications of unpaired learning, it is crucial to preserve semantic information during the translation, for example, the image content in image-to-image translation. Therefore, SB and EOT are suitable tools for this task as they allow controlling the properties of the learned trans- lation by selecting the reference distribution gee in (1) or the transport cost in (2). Over the last several years, many such SB/EOT methods for unpaired learning have been de- veloped; see (Gushchin et al., 2023b) for a survey.
2.1 The Static Schrödinger Bridge Problem
Consider two distributions po: E P((2) and all distribu- tions q E P(X²) whose marginal distributions are po: Pl, re- spectively. The set of such distributions [1(po. Pi) C P(x") called the set of transport plans. In addition, suppose we given a reference distribution qet E P(
The Static Schrödinger Bridge (SB) problem (Schrödinger, 1931; Léonard, 2013) consists of finding the transport plan q E [ppo.p)) closest to qrt in terms of the Kull- back-Leibler (KL) divergence:
2.3 Discrete and Continuous State Space X in SB
(1) = argmin (I =1)ll ( x1),, (1) (ETT(po.P1)
Most methods (Mokrov et al., 2024; De Bortoli et al., 2021; Vargas al., 2021; Gushchin al., 2023a; 2024b; Korotin et al., 2024; Gushchin al., 2024a; Shi et al., 2023; Liu et al., 2022a; Chen et al., 2022) use neural networks to approximate qº and specifically focus on solving SB in continuous state spaces, e.g., X' = DR. This allows us to apply SB to many unpaired translation problems, e.g., the saage-to-image translation or biological tasks related the analysis and modeling of single-cell data (Pariset et al., 2023; Tong et al., 2024).
With mild assumptions on components of the problem (X' Po: Pi, ) the solution q² to this problem uniquely exists; it is called the static SB.
Notably, the static SB problem is equivalent to another well- celebrated problem- - the Entropic Optimal Transport (Cu- turi, 2013, HOT). Indeed, (1) can be written as
min gEll(po.Pr) En( log (10, (20 x1)
Despite advances in computational SB methods, signifi- cant challenges remain when adapting these generative ap- proaches to discrete state spaces X::
min gEll(po.P1)
(Eg(zo.)1) log gett x1) H())}
min (ET((P.P1))
[ (1)
(2)
1. Their underlying methodological principles are mostly incompatible with discrete spaces X'. For example, (Shi et al., 2023; Gushchin et al., 2023a; Vargas et al., 2021; Liu et al., 2022a) use stochastic differential equations (SDE) which are not straightforward to generalize and use in discrete spaces; (Mokrov et al., 2024) heavily re- lies on MCMC sampling from unnormalized density which is also a separate challenge for large discrete spaces X'; (Gushchin et al., 2024a; Korotin et al., 2024; Gushchin et al., 2024b) theoretically work only for the EOT problem with the quadratic cost on X' R.D etc.
where (q) denotes the entropy of transport plan g( 11)
and o(2o, 21) is a transport cost function.
2.2 Practical Learning Setup of SB
Over the last decade, researchers have approached SB/EOT problems in various studies because of their relevance to real-world tasks (Peyré et al., 2019; Gushchin et al., 2023b). In our paper, we consider the following learning setup, which is usually called the generative setup.
2. Extending any generative modeling techniques to dis- crete data is usually challenge. For example, models such as GANs (Goodfellow et al., 2014) require back- propagation through the generator- - for discrete data is usually done via heuristics related to the Gumbel trick (Jang et al., 2017); flow matching methods (Liu et al., 2022b) can be used for discrete data (Gat et al., 2024) but require numerous methodological changes, etc.
We assume that a learner is given empirical datasets {ag MM C X' and {at] C X, which are i.i.d. samples from unknown data distributions Po and pl. respectively. The goal is to leverage these samples to find solution q A q to the SB problem (2) between the distributions po. P1- The solution should permit the out-of-sample estimation, i.e., for any 20 should be able to generate new 2:!e" (1(x).
At the same time, a significant portion of modern data
2

---

Categorical Schrödinger Bridge Matching
inherently discrete, as discussed in S1. Despite its preva- lence, the Schrödinger Bridge framework for discrete spaces remains underdeveloped, motivating our focus.
distributions are Po and p1- In turn, the solution q' this itself becomes an -'vvalued stochastic process. Note that:
KL(g(ao, Fin, x1)llq" Tin 21)) =
We assume that the state space X' is discrete and rep- resented as X' = =S Here 3 is a finite set, and for convenience, we say that it is the space of cate- gories, e.g., S = f1. 53. One may also consider X' Si x Sp for D categorical sets. This does not make any principal difference, we use S = E to keep the paper's 's exposition simple.
(L I1 +1)1l1"0"" "llg x1)) + By(2o,z1) [g to, x lll ( (zin (20, x1)]] (4)
Since conditional distributions qaain (o, can be chosen independently of g 21), we can consider [ 21) (rmm æ1). It follows that the sec- ond term becomes 0 for every zo. 21. As a result, we see that the joint distribution q ("n, x1) for time I = 0, of the dynamic SB (3) is the solution to the static SB (1) for the reference distribution given by the (xu, z1).
Discrete EOT Methods. We would like to mention, for the sake of completeness, that there is a broad area of research known as discrete EOT, which might appear to be closely related to our work. It includes, e.g., the well- celebrated Sinkhorn algorithm (Cuturi, 2013) and gradient- based methods (Dvurechensky et al., 2018; Dvurechenskii et al., 2018). However, such algorithms are not relevant to our work, as they consider a different setting from the gener- ative one ($2.2) and target different problems. Specifically, discrete EOT assumes that the available data samples are themselves discrete distributions, i.e., Po = M7 Em= 8ay P1 k Σ Gat (the weights may be non-uniform), and the goal is to find a bi-stochastic matrix E MxX K (a.k.a. the discrete BOT plan) which optimally matches the given sam- ples. Since this matrix a discrete object, such methods are called discrete. Works (Hütter & Rigollet, 2021; Pooladian & Niles- 2021; Manole et al., 2024; Deb et al., 2021) aim to advance discrete EOT methods to be used in genera- tive setups by providing out-of-sample estimators. However, they work only for continuous state space X' = RR.. It remains an open question whether discrete solvers can be used for generative scenarios in discrete space X' ="
At this point, reader may naturally wonder: why does one consider the more complicated Dynamic SB, especially considering that it boils down to simpler Static SB?
In short, the dynamic solution allows for leveraging the so- called reciprocal and Markov properties of q" (it is discussed below), which can be effectively utilized in developing com- putational algorithms for SB (Liu et al., 2023; Shi al., 2023; Peluchetti, 2023). In fact, most of the computational methods listed at the beginning of 92.3 operate with the dynamic SB formulation. While some methods (De Bortoli et al., 2021; Gushchin et al., 2024b) consider formulation (3) with discrete time and finite amount N of time moments, (Shi et al., 2023; Chen et al., 2022; Gushchin et al., 2024a) work with continuous time E [0, 1]. Informally, one may identify it with discrete time but N cc. In discussions, we will refer to the continuous time case this way in the rest of the paper to avoid unnecessary objects and notations.
The scope of our paper is exclusively the discrete-time in dynamic SB (N < oo) as it is more transparent and feasible to analyze.
2.4 From Static to Dynamic SB Problems
The static SB problem (1) can be thought of as a problem of finding a stochastic process acting at times Usually, one considers an extension of this problem by incorporating additional time moments (De Bortoli et al., 2021; Gushchin et al., 2024b). Let us introduce N 2 1 intermediate time points 0- = to < t1 <... < IN < NN+1 = 1, extending to these moments. Consequently, q becomes a process over the states at all time steps, i.e., q E P(x Similarly to the static formulation (1), let us be given marginal distributions Po.P1 E P(X) with reference process gre E N22). Then the dynamic Schrödinger Bridge problem is
To conclude this section, we introduce an important defini- tion that is specifically relevant to the dynamic SB.
0, 1.
Reciprocal Processes. A process r E P(x is called a reciprocal process with respect to the reference process net if its conditional distributions given the endpoints zn, 31 match those of the reference process, i.e.:
r(z To, .I1) = q g Zo. 21)
The set of all reciprocal processes for the reference process ge is denoted by Rio( c P(2N+2))
min (Ell((Po.PP)
KL(q(Zo,
Tin, 21) lf (zo, Iin #1)),
(3)
2.5 Iterative Markovian Fitting (IMF) Procedure
where II (po. P1) C 22 is a set of all discrete-time stochastic processes in which initial and terminal marginal
In practice, the most commonly considered case of dynamic SB is when gee E MM c P((NN+2), i.e. qe is a
3

---

Categorical Schrödinger Bridge Matching
Continuous time
Discrete time
(N = 00)
(N < <0)
Theory
Practice (SB algorithm)
Theory
Practice
(SB characterization)
(SB characterization)
(SB algorithm)
Continuous space
DSBM 54
Theorem 3.1 (Gushchin et al., 2024b)
ASBM $3.5 (Gushchin et al., 2024b)
X' R2
Theorem 3.2 (Léonard al., 2014)
(Shi et al., 2023)
Discrete space
DDSBM $3.1
Our work ($3)
X' gD
(Kim et al., 2024)
Table 1. A summary of SB problem setups and existing (D-)IMF-reltte) results. The table lists theoretical statements characterizing the SB solution (as the unique both Markovian and reciprocal process between two given distributions) which allows to apply the (D-)IMP procedure to provably get the SB solution q², see (Shi et al., 2023, Theorem 8). The table also lists related computational algorithms.
Markovian process. In this case, the solution q" to SB is also known to be a Markovian process. This feature motivated the researchers to develop the Iterative Markovian Fitring (IMF) procedure for solving SB based on Markovian and reciprocal projections of stochastic processes.
reciprocal projR and Markovian projM projections:
g2t1 = projR (q²t), 2122 =projm (q 1g24+1).
(6)
Originally, the IMF procedure (Peluchetti, 2023; Shi al., 2023) was considered the continuous time (N o0), but recently, it has been extended to the finite amount of time moments (Gushchin et al., 2024b), i.e., N < co. We recall their definitions of the projections for finite N. In this case, the procedure is called the D-IMF (discrete-time IMF).
Since both the Markovian and reciprocal projections pre- serve marginals po, P1 at times f = 0, 1, respectively, we have that each g' E II (Po. P1). In certain configurations of N, X, gTT IMF provably converges to the dynamic SB x" in KL, i.c., lim KL (4'll9") = O. Specifically, the con- vergence easily follows from the generic proof argument in (Shi et al., 2023, Theorem 8) as soon as it known that q* is the unique process in x and reciprocal. We provide Table 1, summarizing the con- figurations for which this characterization of SB is known. We also list the related practical algorithms.
Reciprocal Projection. Consider process q +2) Then the reciprocal projection proj (4) with respect to the reference process qrdl is process given by:
that is both Markovian
[projι=((]] (Io, Tin; 21) ge (xin to. (2)
Markovian Projection. Consider q N2) Then
Finally, we would like to emphasize that the convergence rate of the (D-)IMF procedure notably depends on the num- ber N of time steps. In fact, for each N it its own separate procedure with different Markovian projection (5), (Gushchin et al., 2024b, Figure 6a).
the Markovian projection (q)) is given by:
[projи(9)] (zg. Finy x1)=
N+1
= ((2))
II
q( -=o--)==
2.6
Object of Study
n=
farward representation
As it is clear from Table 1, for the setup with the discrete space X' = p and finite amount of time moments N < co, there is still no theoretical guarantee that the SB is the unique Markovian and reciprocal process. This leaves a large gap in D-IMF usage in this case, and we close in our paper.
N+1
a(x))
II n=1
(5)
backward sepresentation
At the same time, we note that there is a very recent IMF- based algorithm DDSBM (Kim al., 2024) for the discrete state space X' but continuous time (N oc). However, since working with continuous time is infeasible in practice, the authors discretize the time grid to large finite N. Due to this, the authors apply the D-IMF procedure, although it still lacks any theoretical ground in this case. In contrast, our work shows that theoretically even N is enough.
The reciprocal projection obviously preserves the joint distri- bution ((o,, 41) of a process at time moments t = 0, 1. The Markovian projection, in general, alters 4 .1) but pre- serves the joint distributions (g(at., -t,-1 )) 1 at neigh- boring time moments and the marginal distributions q())....
The D-IMF procedure is initialized with any process S E Pt). Then the procedure alternates between

---

Categorical Schrödinger Bridge Matching
3 Categorical Schrödinger Bridge Matching
(DTMC). A DTMC is defined by N + transition matrices On of size |x| 12), where [Qn] Ft. represents the probability of transitioning from state 2t,-1 to state "
We start by establishing the convergence of the D-IMF framework (NV oc) to the SB under general Markov reference process ($3.1) with the proof's in Appendix B. Then we provide a practical optimization procedure and implementation details of the proposed method ($3.2).
g( Le. ,) = 1QnIe,,-.
Thus, in theory, one can model any such DTMC q explicitly, However, in practice, the size |X] may be large. In particular, we consider the case X' = SD, where S is a categorical space leading to exponential amount S of elements in X.
3.1 Theoretical Foundation
The result of (Gushchin et al., 2024b, Theorem 3.6) char- acterizes the SB solution in X' - _ and N < 3o as the unique Markovian and Reciprocal process which allows the usage of D-IMF procedure. However, that proof assumes a specific reference process qet = q" induced by the Wiener process W (EOT with the quadratic cost) and thus cannot handle general Markov qee or discrete X.
This raises two natural questions: (a) how to choose refer- ence process gr and work with it? and (b) how to parame- terize and update the process g during D-IMF steps? Both these questions will be answered in the following generic discussion about the parameterization and implementation of reciprocal and Markovian projections.
3.2.1 Implementing the Reciprocal Projection. The reciprocal projection is rather straightforward if we can draw samples from our current process ((Fo, 21) and the reference bridge ( (z0, 21). Indeed, sampling 11 (prji)rrpj--0000 (q) is just merging these two.
Below we provide our main theoretical result for the discrete space X' and general Markov reference process get which characterizes SB and immediately allows the usage of D- IMF (N < oo) procedure to get it.'
Theorem 3.1 (Characterization of the solution for the dynamic SB problem on a discrete space X with a Markovian reference g'cf). Let X be a finite discrete space and let po. P1 E (X)) be distributions with full support. Let qet E (M be reference Markev process with full support on N22 If E P((NN+2) satisfies the following conditions:
3.2.2 Choosing Reference Process. As is clear from the paragraph above, it is reasonable to consider reference processes get E M(x for which sampling from their bridge (z |zo, z1) is easy. We give two popular ex- amples of qed which appear in related work (Austin et al., 2021) that lead to practically meaningful cost c for EOT (2). For both examples, we start with dimension = 1.
Case 1 (Uniform Reference q In this case, we as- sume that the set of categories S is unordered, e.g., atom types, text tokens, latent variables, etc. Define a process where the state stays in the current category Zt,-, with high probability, while the remaining probability is distributed uniformly among all other categories. This process qmmil is called uniform and has transitions matrices Qn:
1. q" (xro) pu(aa) and q" = (1(1)). i.e., (xO,II) is transport plan from II(po, PI);
2. q" E EM)NN+2) and q² E (2 (2),, iee. 9"'
is both the reciprocal and Markov,
then o' is the unique solution of the dynamic SB (3).
Our theorem immediately yields the following corollary.
ar if It., = 11,-11
[Qn)a Fin
Corollary 3.2 (Convergence of D-IMF on discrete spaces).
if st, # Ttn-17
S-1
The sequence {q' ) produced by the D-IMF procedure on a discrete space X and for Markov reference process from the theorem above converges to q² in KL:
where a E [0, 1] is the stochasticity parameter that controls the probability of transitioning to a different category.
lim KL (g (dll = O.
Case 2 (Gaussian Reference g If we know that the categories are ordered, specifically, S = (1, 2,....,.) S), and two neighboring categories are assumed to be related, the transitions may be chosen to reflect this. Consider the Gaus- siae-like reference process geauss with [Qn] 112
3.2 Practical Implementation
In this subsection, discuss our computational algorithm to implement D-IMF and get the SB problem solution q*
4/s -nn-1 )2 (cra)d
exp
Since we consider a finite amount N of time steps, the processes g E +2) are discrete-time Markov chains
> exp
462
I1, # Ita-it
(8)
1 Σx., f0en
[2 P Rto
Ft, Ita-it
In fact, our proof argument can be applied to any i.e. not only discrete, thus, the ASBM algorithm (Gushchin et al., 2024b) for continuous Rd" can he applied for general Markov ddd
where a > 0 is an analog of the variance parameter, and
A S 1 is a maximum distance between categories.
5

---

Categorical Schrüdinger Bridge Matching
Dimension D 1. The construction of quaí (or y gen- eralizes to higher D by combining several such independent processes (one per dimension). The bridges ntd (xin ( 21) can be easily derived analytically and sampled thanks to the Markov property and the Bayes' formula.
Algorithm 1 Categorical SB matching (CSBM)
Input: number of intermediate time steps N;
number of outer iterations L E N;
initial coupling q ( (1);; reference process gtat.
Output: forward model qe (1t,, [Tt,- backward model In (Ft,.- (rte.)
For more details on the construction and selection of refer- ence processes qtaf, please refer to Appendix D.I.
for to L do
3.2.3 Parameterization of the Learnable Process. There are |S*] = 5 possible states x (z ,,'....,""" x²)) in the space, where S is the number of categories for each variable. Consequently, each transition matrix Qn is of size SD x SS i.e., it grows exponentially in dimension D. Due to this, explicit modeling of the transition matrices of the process that we learn is computationally infeasible. We follow the standard practice in discrete generative models (Hoogeboom et al., 2021; Austin et al., 2021; Gat et al., 2024; Campbell et al., 2024) and model the transition prob- ability via combining two popular techniques: posterior sampling and factorization over the dimensions. Firstly, we parameterize the transitions (I (t...) ) as follows:
Forward step (repeat until convergence):
Sample n ~ U N + 1];
Sample x1) ~ pr (21) I
N+1
On (F (=c.):
Sample -t,,-1
(3 100221)
Train qe by minimizing Lg (21);
Backward step (repeat until convergence):
Sample n ~ U N + 1];
Sample (Io. 21) Po (x())) I 4o (2t, lee,
N+1
Sample It,, ~ (1x,, [ 11);
Train qn, by minimizing Ln (22);
end for
N
(0( [It_..) (311 [ (It, (2t_,, =.)]
L(m) del
Egtuo.er)
Σ
E Eer mm
where go (3i [ir,, .) is learnable distribution. This param- eterization assumes that sampling of zt,, given Ttn-1 can be done by first sampling some "endpoint" 21 ~ To (71 (Z).))).)) and then sampling from the bridge (ste,, [at,, 1 T.). Sec- ond, the parameterization for T (r)...) is factorized:
KL (1 lor, i (1)(,, Tor, 1))
- -Baaleey moom)) [log m( [
(10)
D (o [Tt, 1) R II d=1
among the Markov processes m E M() Fur- thermore, this objective is also equivalent to optimizing NN+1 1 Be(s, 1) KL. (g(xs, [et,.-)]mm(e (xt,,)))
To
(M) pe., J)-
In this case, for each Tt we just need to predict a row- stochastic D S matrix of probabilities To( ). See Appendix A for a discussion of the limitations of this ap- proach. Following the common practices, we employ neu- ral network SD + D x S which outputs row-stochastic matrix for each input 2 -1.-1. Typically, predicting endpoints at each time step n - would require N + 1 distinct mod- els for each To (51 ..). Instead, we use single neural network with an additional input indicating the timestep.
Note that the key distinction from standard losses in dif- fusion models, such as (Austin et al., 2021, Equation 1), lies in the sampling of rin-1. Instead of drawing from the noising process of 1 it is sampled from the ref- erence bridge distribution nrd (I [10, As a result, with the proposed parametrization and Markovian projec- tion representation, we can effectively apply the learning methodology from D3PM (Austin et al., 2021). The explicit loss formulation is provided in Appendix D.2.
3.2.4 Implementing the Markovian Projection. The Markovian projection is little bit more complex than the reciprocal one and requires learning a process. From $2.5, the goal of the projection is to find a Markov process whose transition probabilities match those of the given reciprocal process q. Fortunately, we show that this can be achieved by minimizing an objective that closely resembles the optimiza- tion of the variational bound used in diffusion models (Ho et al., 2020; Austin et al., 2021; Hoogeboom et al., 2021). Proposition 3.3. Let q E ( N+2) be a given recipro- cal process. Then, the Markovian projection proju((g) E M +2)) can be obtained by minimizing:
3.2.5 Practical Implementation of the D-IMF Proce- dure. With the reciprocal and Markovian projections fully established, we now proceed to the implementation of the D-IMF procedure. This method is conventionally applied in a bidirectional manner (Shi al., 2023; Gushchin et al., 2024b), incorporating both forward and backward represen- tations (5). This is because training in unidirectional man- ner has been shown to introduce an error in IMF (De Bortoli et al., 2024, Appendix I). Therefore, we follow a bidirec- tional approach. which naturally leads to the Categorical Schrödinger Bridge Matching (CSBM) Algorithm 1.
6

---

Categorical Schrödinger Bridge Matching
4 Experimental Illustrations
KL vs D-IMF iteration (S = 50. 10)
40:
0 0.65
We evaluate our CSBM algorithm across several setups. First, we analyze the convergence of D-IMF on discrete data (54.1). Then, we demonstrate how CSBM performs with different reference processes in 2D experiments ($4.2). Next, we test CSBM's ability translate images using the colored MNIST dataset ($4.3), varying the number of steps N. We then present an experiment on the CelebA dataset (64.4), showcasing CSBM's performance in a latent space. Finally, we explore the text domain by solving sentiment transfer on the Amazon Reviews dataset (Appendix C.4). Experimental details are provided in Appendix D.3 and ad- ditional immages in Appendix D.4.
In-
oenito 0.1
ttredtadd 1ed0
10
D-IMF iterationn U)
(a) Dependence the stochastisity parameter
KL vs D-IMF iteration (S = 50, 0.05)
40
W 120
10
iteretions (1)
4.1 Convergence of D-IMF on Discrete Spaces
(b) Dependence on the number time steps N with geass
KLldl xcilee ((ouxs))
D-IMF iteration 50, 0.01)
In this section, we derive analytical expressions for D-IMF and compare its convergence on discrete data under sev-
in- 1n
-
eral setups. As noted in 32.5, the Markovian projection preserves the one-step transition probabilities of the given process 2 Thus, our task reduces to replicating:
10-
10
21+2 ( ( For each D-IMF iteration, these transition matrices can be extracted from the joint distribution:
Vn E [1].
D-IMF Iterations (1
(c) Dependence on the number of time steps N with q"tir,
Figure I. Dependence of convergence of D-IMF procedure on dis- crete data under different N, a and grd
2 (It,, Pta-1)==
Σ
21+ (2tt(x0, 21)
q (xr,, (ro, If (Ite, 1 mi)
the Sinkhorn algorithm (Cuturi, 2013). As cost matrix, we use the negative logarithm of cumulative transition matrix I Qn. The resulting convergence curves, shown in Figure I, indicate notably fast convergence of KL (g'lla").
where ne (xt, (o. 21) and grae( (I (1t,_1 could be de- rived using Markov property and Bayes' formula.
Given (x Lt, )), we obtain the desired transi- tion distribution q2' (x,,, (=t,,)) = !Q I 1 St.,
4.2 Illustrative 2D Experiments
by normalizing the joint distribution over the marginal (2 which is computed by summing over all st, E ESD in g2l+1 (Jt,,, Ft., stt,-1). )). We then get the conditional distribution (2 (20) by multiplying the transition ma- trices Q21+2, i.e., 22 (20)
In this experiment, we take the initial distribution Po as 2D Gaussian and the target distribution P1 as a Swiss Roll. Both are discretized into 5 = 50 categories, resulting in a 2-dimensional categorical space with | = S = 50 x 50 points. Compared to the previous experiment, this setup in- volves working with N matrices of size 500 x 2 500, mak- ing it significantly more demanding computational task. Therefore, from now on, we solve the SB problem using our proposed Algorithm 1. The goal of this experiment is to ex- amine the impact of the reference processes ganss and q'iif, Thus, we train CSBM with N = 10 intermediate steps with different a and qe For sauss, we test a E (0.02, 0.05). In the case of q qiii we use a E {0.01, 0.005].
N+1 Q
Finally, we reweight this conditional distribution with Po(au) to obtain new coupling ( 21) po(zo) N+1 Q21+2] of the next iteration.
20.01
All of these equations tractable and can be efficiently computed for small values of S and D. Therefore, in our experiment, we solve the SB problem with S = 50 and
D D=1 1 between the following marginals:
Figure 2 demonstrates that increasing the parameter in- creases the number of jumps. In the case of g the jumps mostly happen only to neighboring categories (Fig- ures 2c and 2d). In the case of qin the jumps happen to all categories (Figures 2e and 2f). This is aligned with the construction of the reference processes.
Po(za) - 5 (a)) =
ES
To assess convergence as in Corollary 3.2, we also required to have the ground-truti bridge d², which we compute via
7

---

Categorical Schrödinger Bridge Matching
8
3
2
3
2
2
2
2
2
2 2
3
2
2
2
0
3
2
2
2
2
2
2
2
(a) a~
(b)
N
(c)
N
(a) x (b)
8
2
2
2
2
2
2
2
2
3
}
2
2
2
2
2
2
2 2
3
2
2
2
2
2
2
2
3
2
2
2
2
2
2
2
2
HEte and
(d)
(e)
N
10
(l) N
25
Tragecior staet
Trajectry Mart Tiuedtory (inadi
Pajectory
3
2
2
2
2
2
2
2
2
(c) Low stochasticity, grues
(d) High stochasticity, g a: = 0.05.
3
2
2
2
2
2
2
2
2
= 0.02.
3
2
2
2
2
2
2
2 2
2 2
3
2
2
2
2
2 2
(g) ao-po
(h) 50
(1 100
Figure 3. Results of colored digits unpaired translation "3" - "2" learned by our CSBM algorithm with reference process guurs and varying number of time moments N.
Risd Traject y RON Thaddrre (1ll))
Trajensory rtain Ripoctar and Hntaed
profictions
(e) Low stochasticity, uit 0.005.
(I) High stochasticity, g 0.01.
Specifically, each pixel is represented using three 8-bit chan- nels (RGB), i.c., S 256, and the data space is of size 256 where D = 32 x 32 x 3. The goal of this experiment is to evaluate the capability of CSBM to perform unpaired translation with different numbers of intermediate steps N. Since each color channel values have an inherent order, we utilize the Gaussian reference process geauss with a 0.01.
Figure 2. SB between 2D Gaussian and Wwiss-Roll distributions learned by our CSBM algorithm different reference processes niill and ooue with varying parameters
Remark. Beyond the theoretical objectives established match the distributions using alternative loss functions, such as MSE, or through adver- sarial methods, as in ASBM (Gushchin et al., 2024b). For completeness, we conducted additional experiments using the MSE loss and observed results comparable to those obtained with KL. Details on the experimental setup and loss generalization are provided Appendix C.I.
in Proposition 3.3,
The results in Figure 3 suggest that even with low N 2, the generated outputs maintain decent visual quality and preserve the color. However, some pixelation appears in the samples, which is likely due to the factorization of the learned process (recall $3.2.3). The effect declines slightly as N increases, reflecting a trade-off between model sim- plicity and the ability to capture inter-feature dependencies. Moreover,
be observed that similarity reduces propor- tionally to N. We hypothesize that this issue is related to underfitting, since all models were trained with the same number of gradient updates. Presumably, larger N requires proportionally more updates to adequately train all transition probabilities (9). Additionally, we experiment with qaair with details provided in Appendix C.2.
4.3 Unpaired Translation on Colored MNIST
Here, we work with the MNIST dataset with randomly col- ored digits. Inspired by (Gushchin et al., 2024b, Appendix C.3), we consider an unpaired translation problem between classes "2" and "3" of digits. In our case, we work in the discrete space of images,
continuous space.
8

---

Categorical Schrödinger Bridge Matching
Low stochasticity
(a)
(b) CSBM (ours)
(c) ASBM (Gushchin et al., 2024b)
(d) DSBM (Shi et al., 2023)
High stochasticity
(e) ~
(I) CSBM (ours)
(g) ASBM (Gushchin et al., 2024b)
(h) DSBM (Shi et al., 2023)
Figure 4. Comparison of male + female translation on the CelebA 128 x 128 dataset using CSBM (ours), ASBM, and DSBM. ASBM and DSBM operate in continuous pixel space, whereas CSBM operates in discrete latent space of VQ-GAN (Esser al., 2021). The low-stochasticity setting for CSBM corresponds to a = 0.005, while the high-stochasticity setting corresponds to a = 0.01 of the reference process g The images for ASBM and DSBM taken from (Gushchin et al., 2024b).
4.4 Unpaired Translation of CelebA Faces
identical across all images for CSBM, which not the case for all other methods, especially in high stochasticity setups.
Here, we present an unpaired image-to-image translation experiment on the CelebA dataset using vector quantiza- tion. Specifically, we focus on translating images from the male to the female domain. We train VQ-GAN autoen- coder (Esser et al., 2021) to represent 128 x 128 images as D = 256 features with S = 1024 categories (a.k.a. the codebook). This formulation reduces complexity, as the data to modeled has dimensionality of SD 1024256 Indeed, this smaller than the raw colored MNIST image space (54.3) and considerably smaller than the raw pixel space of CelebA. As there is no clear relation between the elements of the codebook, we use uniform reference grt
Table 2. Metries comparison CSBM (ours), (Gushchin et al., 2024b, ASBM), and (Shi et al., 2023, DSBM) for unpaired male
female translation on CelebA 128 x 128 dataset.
Low stochusticity
High stochasticity
Mettic
CSBM
ASBM
DSBM
CSBM
ASBM
DSBM
a
0.005
==
6=
a 0.01
==
FID
10.60
16.86
24.06
14.68
17.44
92.15
CMMD (I)
0.165
0,216
0.365
0,212
0.231
1.140
LPIPS (4)
0.175
0.242
0.246
0.170
0.294
0.386
The standard FID (Heusel et al., 2017), CMMD (Jaya- sumana al., 2024), and LPIPS (Zhang et al., 2018) metrics comparison in Table 2 quantitatively demonstrates that our approach achieves better results than the other methods on the test set. Still, it is important to note that our experiments are conducted with N = 100 in D-IMF, which is higher than the N = 3 used in continuous-space D-IMF in ASBM, i.e., the trade-off between the number of time steps N and the generation quality should be taken into account.
We test E {0.005, 0.01) and N = 100.
For completeness, we compare our CSBM method with ASBM (Gushchin et al., 2024b) and DSBM (Shi et al., 2023), which operate in the continuous pixel space. For the rationale behind not training them in the latent space, see Appendix C.3. We take their results from (Gushchin al., 2024b, 54.2). Qualitatively, we achieve comparable visual
results (Figure 4). Notably, the background remains nearly

---

Categorical Schrödinger Bridge Matching
Acknowledgements
Deb, N., Ghosal, P., and Sen, B. Rates of estimation of opti- mal transport maps using plug-in estimators via barycen- tric projections. Advances in Neural Information Process- ing Systems, 34:29736 29753,, 2021.
The work was supported by the grant for research centers in the field of AI provided by the Ministry of Economic Development of the Russian Federation in accordance with the agreement 0000000C31392554F0000 and the agreement with Skoltech 10139-10-2025-03..
Dvurechenski,, P., Dvinskikh, D., Gasnikov, A., Uribe, C., and Nedich, A. Decentralize and randomize: Faster algo- rithm for Wasserstein barycenters. In Advances in Neural Information Processing Systems, pp. 10760-10770, 2018.
Impact Statement
Dvurechensky, P., Gasnikov, A., and Kroshnin, A. Com- putational optimal transport: Complexity by accelerated gradient descent is better than by Sinkhorn's 's algorithm. In International conference on machine learning, pp. 1367- 1376. PMLR, 2018.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
References
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981-17991, 2021.
Gat, I., Remez, T., Shaul, N., Kreuk, E., Chen, R. T., Syn- naeve, G., Adi, Y., and Lipman, Y. Discrete flow match- ing. In The Thirty-eighth Annuul Conference on Neural Information Processing Systems, 2024.
Banerjee, A., Merugu, S., Dhillon, I. S., and Ghosh, J. Clus- tering with Bregman divergences. Journal of machine learning research, ((0ct):1755-1744 2005.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley. D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances neural information processing systems, pp. 2672-2688 2014.
Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, T. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 5453-5512. PMLR, 21- 27 Jul 2024. URL https: http://oooceeee/ge mlr. press /v.ess/v2355 / campbe 124a. 124ma html.
Gu, S.. Chen, D.. Bao, J., Wen, F., Zhang, B., Chen, D., Yuan, L., and Guo, B. Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10696-10706, 2022.
Gushchin, N., Kolesov, A., Korotin, A., Vetrov, D., and Bur- naev, E. Entropic neural optimal transport via diffusion processes. In Advances Neural Information Processing Systems, 2023a.
Chen, T., Liu, G.-H., and Theodorou, Likelihood train- ing of Schrödinger bridge using forward-backwalt SDEs
theory. In International Conference Learning Repre- sentations, 2022.
Gushchin, N., Kolesov, A., Mokrov, P., Karpikova, P., Spiri- donov, A., Burnaev, E., and Korotin, A. Building the bridge of Schrödinger: A continuous entropic optimal transport benchmark. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023b.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. Diffu- sion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:116655-1777 17709, 2021.
Gushchin, N., Kholkin, S., Burnaev, E., and Korotin, A. Light and optimal Schrödinger bridge matching. In Forty-first International Conference Machine Learn- ing, 2024a.
De Bortoli, V., Korshunova, I., Mnih, A., and Doucet, A. Schrödinger bridge flow for unpaired data transla- tion. In The eiihty- Annual Conference on Neural Information Processing Systems, 2024. URL https: //openreview wwt./ee//eeeee wwwp//er///eeeeeeeee ftt.formmmmd1121111cc
Gushchin, N., Selikhanovych, D., Kholkin, S., Burnaev, E., and Korotin, A. Adversarial Schrödinger bridge match- ing. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL https: / oppnreviivveeo ew, net/ forum? d=1.3knn gic..
10

---

Categorical Schrödinger Bridge Matching
He, J., Wang. X., Neubig, G., and Berg-Kirkpatrick, T. A probabilistic formulation of unsupervised text style trans- fer. In International Conference on Learning Representa- tions, 2020.
Léonard, C., Relly, S., and Zambrini, J.-C. Reciprocal pro- cesses. meassure-hheoretiaal point of view. Probability Surveys, 11:237-269, 2014.
Li, J., Jia, R.,, He, H.,, and Liang, P. Delete, retrieve, generate: simple approach to sentiment and style transfer. In Pro- ceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pp. 1865-1874, 2018.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. GANs trained by a two time-scale update rule converge local Nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems, 33:6840 6851, 2020.
Liu, A., Broadrick, O., Niepert, M., and Broeck, G. V. Dis- crete copula diffusion. arXiv preprint urXiv:2410.01949 2024.
Hoogeboom, E., Nielsen, D., Jaini, P., Forré, P., and Welling, M. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34:12454-12465, 2021.
Liu, G.-H., Chen, T., So, O., and Theodorou, E. Deep generalized Schrödinger bridge. Advances in Neural Information Processing Systems, 35:9374-93880 2022a.
Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E., Nie, W., and Anandkumar, A. 1 sb: Image-to-image Schrödinger bridge. In International Conference on Ma- chine Learning pp.
Hütter, J.-C. and Rigollet, P. Minimax estimation of smooth optimal transport maps. 2021.
PMLR, 2023.
Jang, E., Gu, S., and Poole, B. Categorical reparameteriza- tion with softmax. In International Conference on Learning Reprexentations, 2017.
Liu, X., Gong, C., et al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Repre- sentations, 2022b.
Jayasumana, S., Ramalingam, S., Veit, A., Glasner, D., Chakrabarti, A., and Kumar, S. Rethinking FID: To- wards a better evaluation metric for image generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 9307-9315, 2024.
Luo, F., Li, P., Yang, P., Zhou, J., Tan, Y., Chang, B., Sui, Z.,, and Sun, X. Towards fine-grained text sentiment transfer. Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2013- 33 2022, 2019.
Kholkin, S., Ksenofontov, G., Li, D., Kornilov, N., Gushchin, N., Burnaev, E., and Korotin, A. Diffusion & adversarial Schrödinger bridges via iterative proportional Markovian fitting. arXiv preprint arXiv:2410.02611, 2024.
Luo, X., Wang, Z., Lv, J., Wang, L., Wang, Y., and Ma, Y. CrystalFlow: flow-based generative model for crys- talline materials. arXiv preprint arXiv:2412. /193,, 2024.
Kim, J. H., Kim, S., Moon, S., Kim, H., Woo, J., and Kim, W. Y. Discrete diffusion Schrödinger bridge matching for graph transformation. arXiv preprint arXiv:2410.01500. 2024.
Manole, T., Balakrishnan, S., Niles-Weed, J., and Wasser- man, L. Plugin estimation of smooth optimal transport maps. The Annals of Statistics, 52(3):966-995 2024.
Korotin, A., Gushchin, N., and Burnaev, E. Light Schrödinger bridge. In The Twelfth International Confer- ence on Learning Representations, 2024.
Mokrov, P., Korotin, A., Kolesov, A., Gushchin, N., and Bur- naev, E. Energy-guided entropic neural optimal transport. In The Twelfth International Conference Learning Representations, 2024. URL
Kudo, and Richardson, SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, 2018.
net
Mukherjee, S., Kasner, Z., and Dušek, O. Balancing the style-content trade-off sentiment transfer using polarity-awand denoising. In International Conference on Texz, Speech, and Dialogue, pp. 172 186. Springer, 2022.
Léonard, C. A survey of the Schrödinger problem and some of its connections with optimal transport. arXiv preprint arXiv: 300 15, 2013.
Ni, J., Li, J., and McAuley, J. Justifying recommendations using distantly-labeles reviews and fine-grained aspects.

---

Categorical Schrödinger Bridge Matching
In Proceedings of the 2019 conference on empirical meth- ods in natural language processing and the 9th interna- tional joint conference on natural language processing (EMNLP-IJCNL)), pp. 188-197, 197,19972999 2019.
Salimans, T., Goodfellow; Zaremba, W. Cheung, V., Rad-
ford, A., and Chen, X. Improved techniques for training GANs. In Advances in neural information processing systems, Pp. 2234 2242, 2016.
Papineni, K., Roukos, S., Ward, T., and Zhu, W.-. BLEU: method for automatic evaluation of machine transla- tion. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pp. 311-318, 2002.
Schrödinger, E. Über die Umkehrung der Naturgesetze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u. Company, 1931.
Shen, T., Lei, T., Barzilay, R., and Jaakkola, Style transfer from non-parallel text by cross-alignment. Advances in neural information processing systems, 30, 2017.
Pariset, M., Hsich, Y.-P., Bunne, C., Krause, A., and De Bor- toli, V. Unbalanced diffusion Schrödinger bridge. arXiv preprint XxXv:2306.00099, 2023.
Shi, Y., Bortoli, V. D., Campbell, A., and Doucet, A. Dif- fusion Schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023, URL https: oppnrevieneeeeeeeeeeeeeeeeeeeeee net/ forum? id=qyC70HsJIz
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF interna- tional conference computer vision, pp. 4195-4205, 2023.
Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L.,, Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation tfeeee Schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pp. 1279-1287. PMLR, 2024.
Peluchetti, S. Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
Peyré, G., Cuturi, M., et al. Computational optimal transport. Foundations and Trends@ in Machine Learning. 11(5-6); 2019.
Van Den Oord. A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
Pooladian, A. and Niles-Weed, J. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109. 12004, 2021.
Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
Prabhumoye, S., Tsvetkov, Y., Salakhutdinov, R., and Black, A. W. Style transfer through back-translation. In Pro- ceedings of the 56th Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers),
Vignac, C., Krawezuk, I., Siraudin, A., Wang. B., Cevher, V., and Frossard, P. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209. 14734, 2022.
pp. 866-876, 2018.
Qin, Y., Madeira, M., Thanou, D., and Frossard, P. De- FoG: Discrete flow matching for graph generation. arXiv preprint ar K41:22410002633 2024.
Wang, K., Hua, H., and Wan, X. Controllable unsupervised text attribute transfer via editing entangled latent repre- sentation. Advances in Neural Information Processing Systems, 32, 2019.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Xu, M., Geffner, T., Kreis, K., Nie, W., Xu, Y., Leskovec, J., Ermon, S., and Vahdat, A. Energy-based diffusion language models for text generation. arXiv preprint arXiv:2410.21357, 2024.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffuuson models. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Zhang, R., Isola, P., Efros, A. A.. Shechtman, E., and Wang. O. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted intervention-MICCA. 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part 18, pp. 231. Springer, 2015.
Zhu, J.-Y., Park, T., Isola, P., and Efros, A. A. Unpaired image-to-image translation using cycle-consistent adver- sarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232, 2017.
12

---

Categorical Schrödinger Bridge Matching
A Limitations
One limitation of the proposed algorithm stems from the factorization of the transitional probabilities (see 93.2.3). This simplification comes at the cost of losing some information, as dependencies between features the same step are not explicitly accounted for. However, it should be taken into account that this limitation is inherent the most modern flow-based (Campbell et al., 2024; Gat et al., 2024) and diffusion- based (Hoogeboom al., 2021; Austin et al., 2021) methods for discrete data. Recent approaches aim to address this issue by modeling the transition joint distribution using copulas (Liu et al., 2024) or energy functions (Xu et al., 2024). Additionally, in the Colored MNIST experiments ($4.3) the reference process varies slightly with N, due to implementation specifics. The MSE between cumulative transition matrices bounded by 10-4, confirming that the induced discrepancies are statistically insignificant. Thus, these differences are negligible and do not impact the experiment's goals or broader implications.
B Proofs
Proof of Theorem 3.1. As stated in the theorem, we consider a process g( Tin, 21) (po. p1) with N 2 intermediate time steps that is both Markov and reciprocal and a reference Markov process aed E M(X We focus on the joint distribution of the boundary elements Tig 31, and a selected intermediate state It,, where n E [1, NJ. This distribution, p(zo, It,, T1). can be expressed in two equivalent ways using the Markov or the reciprocal properties:
g( a5 11)44 ( (i,,, I vi) = g( Ft,, 21) = p(1oo)((ar (xn.)))))
by reciprocal property
by Markov pruperty
Rearranging this equation and applying the logarithm thus we get:
log (x21 [x)) = log ( (20) + log q( (xx,) - log (2e, x1).
Note that all the probability terms are strictly positive by the theorem's assumption. The knowledge that the last term log (2,, (zu 21) is Markov leads following equation:
q (2o) ( (oo)) (201( (ou, x1)
log (( (x)) log (mo) + log g(<)) (zt,.) )- - log
= log ( (x)) log of gfo (2,, (20) - log of (x0)) + log ((21 (Tt,)) - log a (21 (111 + log nd (2o.
fo(
# (1))
Thus, we get:
f (xo, def log (x) - log (2o, x1) = fo( + fi (T),
(11)
Notably, f(ro 1 can be represented as sum of two single-variable functions, (o( and g1 (z1). This could be observed by setting = xt in (11), where æt X is some fixed point in the state space. Indeed, we have:
f(zo, x1) f(xo, x1) = Falvorn,.) + fi (ft,,, - falrorar..) fi fi(2e, fi (122, (Itt,, x!).
f(zu, Likewise, with fixed 2t,,, the Fixing 21 = xt makes f depend only on Fu. so, we define 90(2v) def difference f f 21) ( f(r,, x1) depends only on 21. Thus, we set (x1) def f 21) -f---f(aap xf). Finally, we obtain:
log q( (2o) go (zo) + (1())) + log g n00 (z).
Exponentiating both sides and multiplying by p(xo), we derive:
q v1) eo(vo) (xr (ro) (()))
v(me)
((m))
According to (Léonard, 2013, Theorem 2.8), this formulation describes the optimal transport plan 9" for the Static Schrödinger Bridge problem between Po and P1- Alternatively, this can be derived as in (Gushchin et al., 2024b). Given that the assumption of the theorem ensures q( 21) = (2 (xo, x1), it follows that g Dmg 21) is a dynamic Schrödinger Bridge qº (xo, Pin, 21).
13

---

Categorical Schrödinger Bridge Matching
Proof of Proposition 3.3. Thanks (Gushchin et al., 2024b, Proposition 3.5), it is known that
(projm(a) Vin x1) argmin
KL (g(200 Zin x #1)lm Fin 21)) MEM(XN+3)
(12)
where q (2) is a reciprocal process. Thus, we can decompose this KL divergence as follows:
g( Win, x1) m Tins F1)
KL Fin, Fin 21)) mi) log
po (2 (2))g ( ( v2)
E (g( log
(13)
(m(z))((z (x1 I m( (ta--1)
Here, the denominator can represented this way because m a Markov process, while the numerator expressed using the reciprocal property of q. Next, we separate the corresponding colored terms, leading to:
(13) = - E 21) [log (211[(()) + (1) log I qeel (I (2 +
II m( (
KL (po(zz))lm((z) + E ((m [zog((r][00]] (14)
Rewriting the product inside the logarithm (violet term) sum of KL divergences, we obtain the following equation:
N Eg(s Lt
(14) = Li +
KL (1t (rt, 11 2. (mm)m,, Ir,, .)) + Lo + Ci.
(15)
1==1
We observe that, by construction, the Markov process preserves the terminal distribution when represented in forward manner (5), i... m(xo) po(z)). Consequently, Lo can be omitted since KL O, which completes the proof:
(15) L1 +
B , (I(a"((I (x1,, F1 (m)))lm(m((,,. (zm_.)) + C1.
(16)
1=1
Additionally, because the Markovian projection (5) leaves the neighbouring time joint distribution q( vt..) unchanged, can train with the alternative objective:
KL (g(2), Zin D1 (l Zin. 21)) Bx(2o:=mma1) log g(zo; Tins x1) m(Io, Tin, x1)
N+1
E q ( KL (g(2e, (e ( + KL (po(2))]lm(zo)) (17)
KL
Lg
Similarly, we discard Lu, leaving us with an objective that minimizes the divergence between one-step transition probabilities of the given process q and the desired Markov process m.
Additional Experiments
C.1 Alternative Losses
Proposition 3.3 shows that two equivalent KL-based training objectives yield the same optimal solution. This naturally suggests a generalization broader class of divergences D.
14

---

Categorical Schrödinger Bridge Matching
The Original Objective. First, let consider the original objective function given in (16). To ensure that substituting an alternative divergence does not alter its minima, the replacement must be equivalent in this context. Specifically, the T1 term can be reformulated as the KL divergence between Kronecker delta distribution and the transition distribution of m, i.e.;
L1 E :1 [log m(zi( E 21) Bo,, (E1) log Se, (T1)
m (are))
E ay ) KL (8xr (31 (( = Byy mr (0:))
KL (q(21
Consequently, the L1 term can be moved under the sum of the violet term, leading to:
N+1 Σ
(16)
B B KL (are, let, 2n (area (et,, .) C1.
By restricting the choice of divergences to the Bregman family, we ensure that the minimum is attained at the same value, namely, Ey(es (=)..1) [g lire, 21: g(>r,,, (rr, ,) (Banerjee al., 2005). Thus, any Bregman divergence can be used the objective. an example, we consider the MSE loss as an alternative the KL divergence:
N+1 E B ,) KL (g ( (T (ztn-)))
argmin MEM((N+22
n=1
N+1
argmin mEM(X++2)
Σ
Bate 1) (ort (0t - m( [
(18)
Applying this loss parametrization from 33.2.3 and repeating the derivation leads to the following objectives:
N+1
IMSE (0)
E By
(2o.a) (ae,, (er, 11) - Bioo (xr, l0,, 1 Tilll]
(19)
n=l
N+1
IMSE (1)
E m1) But (w ) ( 1.t. (pen.a0 vo) - (20l / (x Ieen mo)l
(20)
n=1
for forward and backward parametrization, respectively. To test the MSE loss, repeat the 2D domain translation experiment between the Gaussian and Swiss-Roll distributions. It could be observed that the generated samples and trajectories with the MSE loss in Figure 5 appear visually similar to those obtained using the KL loss shown in Figure 2.
Treestary teckkh rewcrs and
praiciions
(a) ~ Po. Pr.
(b) Low stochasticity, puuss with 0.02.
(c) High stochasticity, ains with = 0.05.
(d) Low stochasticity. unit with 0.005.
(e) High stochasticity. umi with 0.01.
Figure 5. SB between 2D Gaussian and Swiss-Roll distributions learned by our CSBM algorithm with MSE loss in Equations (19) and (20) different reference processes q and y gass with varying parameters a.
The Alternative Objective. Analogous reasoning extends to the alternative objective in (17). Although the conditional distribution q (or, ,) is generally unavailable in closed form, it can be sampled. This property suggests employing an adversarial training strategy, following the approach (Gushchin et al., 2024b).
15

---

Categorical Schrödinger Bridge Matching
C.2 Unpaired Translation Colored MNIST with q quii
We perform an additional Colored-MNIS experiment using uniform reference process qnnif. Here we set N 25 and test (0.01, 0.05). Mini-batch OT is not applied at the D-IMF iteration. The samples in Figure demonstrate the failure to match the digit colors, showing that a uniform transition matrix is not suitable for this domain.
8
2
2
2
2
2
221
2
2
2
3
3
3
3
3
3
3
3
3
2
2
2
2
2
2
2
2
2
3
3
3
3
3
3
3
3
3
2
2
2
2
2
2
2
2
2
3
3
3
3
3
3
3
3
3
2
2
2
2
2
2
2
2
2
3
3
3
3
3
3
3
3
(a) ~ Po-
(b) a = 0.01.
(c) a 0.05.
(d) ~ Pl.
(e) a = 0.01.
(f) a 0.05.
Figure 6. Results of colored digits unpaired translation learned by CSBM algorithm with reference process q and varying stochasticity parameter a.
C.3
Continuous Methods in Latent Space
For completeness. we also trained DSBM the latent space. For a fair comparison, we train DSBM on the same latent space used for CSBM, following the approach in (Rombach et al., 2022, Appendix G). Concretely, because the decoder expects discrete tokens, pipeline proceeds as follows: (1) map the images to their continuous latent representations, (2) apply DSBM in this continuous space, (3) vector-quantize the resulting latents, and (4) pass the quantized tokens through the decoder. Unfortunately, the results are not satisfactory, as the model tended to collapse to the identity mapping with e = 1 and = 10 (see Figure 7). Due to these limitations, we do not proceed with training ASBM and choose not to compare both methods with CSBM in such settings. One may ask why CSBM performs better in this setting. We hypothesize that this due to the choice of the reference process unif" which is better suited to the VQ-GAN latent space.
Po-
(b) 1.
(c) 10.
(d) Pl-
c=)==1 1.
(l) 6 = 10.
Figure 7. Results of training DSBM (Shi et al., 2023) on VQ-GAN lantent space of CelebA. The VQ-GAN model is the same as in the main experiments (54.4).
16

---

Categorical Schrödinger Bridge Matching
Table Metries comparison of CSBM (ours), CAAE (Shen al., 2017), Rel.&Ret. (Li et al., 2018), Seq2SentiSeq (Luo al., 2019), BST (Prabhumoye et al., 2018), FGIM (Wang et al., 2019), PST (He et al., 2020) and SCT, (Mukherjee et al., 2022) for unpaired negative + positive style transfer on the Amazon Reviews dataset. Bold denotes the best value, and underline the second best. Metries of baseline methods are taken from (Mukherjee et al., 2022) and marked with superscript f.
CSBM
CSBM
Metric
CAAE
Del.&Ret.
Seq2SentiSeq'
BST
FGIM
PST
SCT,
a
0.005
a
0.01
Accuracy (t)
79.3
76.5
88.6
69.9
92.4
93.5
79.3
91.5
82.0
NLL (1)
5.4
5.4
74.0
85.1
42.0
61.0
116.8
65.9
79.6
BLEUU (t)
72.5
74.8
3.2
14.7
0.0
0.9
10.6
9.5
13.7
Table 4. Style transfers of CSBM (ours), Ret. (Li ct al., 2018), BST (Prabhumoye et al., 2018), FGIM (Wang et al., 2019), PST (He al., 2020) and SCT, (Mukherjee et al., 2022) on Amazon Reviews dataset. Samples of baseline methods are taken from (Mukherjee et al.. 2022) and marked with a superseript f.
negative + positive
positive + negative
Source
movie was waste of money: this movie totally sucks.
my daughter loves them
CSBM
movie was great value for the money: this movie totally wass.
my daughter hatex them of
0.005
CSBM
movie was great value for the money: : this movie totally superb.
my daughter hates them of
0.01
Ret.
our favorite thing was movie story: : the dream class roll
my doughter said was still not acknowledged.
BST
stan is always a great place to get the food.
do i't be going here.
FGIM
movie is a delicious atmosphere of: this movie totally sucks movie ! this theater was a great place, we movie totally amazing. -
i should not send dress after me more than she would said not
PST+
yup daughter has left ourselves.
SCT,1
movie: great deal of money: this movie is absolutely perfect.
my daughter hates it my daughter.
Source
nothing truly interesting happens in this book.
best fit for my baby: this product is wonderful !!
CSBM
everything truly interesting happens in this book.
not fit for my baby: : this product is junk !!
0.00G
CSBM
everything truly interesting happens in this book
not fit for my baby : this product is had !!
0.01
Ret. t
nothing truly interesting happens in this book.
my mom was annoyed with my health service notice.
BST
very good for the best.
bad customer service say the food. and n't.
FGIM
nothing truly interesting happens in this book make it casual and spot
do not buy my phone: this bod crop was worst than uninspired. .
PSTt
haho truly interesting happens in this book.
SCT,
this book is truly a really great book.
not good for my baby: : this product is great! ! !!!!!!!
C.4 Unpaired Text Style Transfer of Amazon Reviews
This section examines the text domain, focusing style transfer in the Amazon Reviews corpus (Ni et al., 2019). The task is to convert reviews with negative sentiment into ones with positive sentiment and vice versa. We adopt the filtered, pre-processed split of (Mukherjee et al., 2022). Reviews are tokenized with a unigram SentencePiece model (Kudo & Richardson, 2018) that has a vocabulary size set to S = 8 192. Each review is then padded or truncated to fixed length of D = 100. We evaluate the uniform reference process gre for a E (0.005, 0.013. The reported scores are averaged over both transfer directions negative <+ positive and compared with baselines, using the metrics from (Mukherjee et al., 2022).
To mirror the image-domain protocol, we select analogous text metrics. Target aligmment is measured with the Hugging Face pipeline's s default sentiment classifier, complemented by the negative logelikelihood (NLL) under GPT-2 Large (Radford et al., 2019). Similarity between the transferred text and its source is measured with BLEU (Papineni et al., 2002). Quantitative metrics appear in Table 3, while representative samples are shown in Table 4.
CSBM excels at content preservation, achieving the highest BLEU score and the lowest NLL, indicating fluent, meaning- faithful rewrites. Its sentiment- aatimment-transfer accuracy is lower than half of the methods, yet manual inspection of the samples in Table 4 suggests that most generations convey the correct polarity.
17

---

Categorical Schrödinger Bridge Matching
D Practical Details
D.1 Construction and Selection of Reference Processes 9te
Construction of yTd, The article touches only briefty on how the reference processes yet are built. the current scheme, qtd assembled by chaining intermediate transition probabilities qrd (2,, ( Consequently, the full end-to-ent transition (23 (ro) varies with the choice of a, the transition matrix Qn. and the discretizations level N, rather than remaining fixed across settings. Due to this, for example, the increasing number of steps N forces us to choose smaller a. If a remains too large, the overall transition probability q (x, (zo) converges to the stationary distribution, making every start state equally likely to reach every end state. A uniform distribution is not inherently wrong, but it defeats our aim, as we want a to control the overall stochasticity in the process. Thus, building non-uniform, non- non-Gaussian qee is considerably more challenging, prompting us to explore new construction strategies the future.
Selection of at Across many experiments, observed pattern for choosing Overall, the general idea follows the same intuition choosing in continuous SB methods (Shi et al., 2023; Gushchin et al., 2024b). Specifically, lower values of lead to less stochasticity in the trajectories, resulting in higher similarity to the input data but lower-quality approximation of the target distribution. At very low values, the model may collapse due to insufficient stochasticity. Conversely, higher values a introduce more variability, reducing similarity to the initial data. Beyond certain point, excessively large values a make the model difficult to train, leading to a drop both quality and similarity. Unfortunately, the effective range of these behaviors highly dependent on the dataset and the chosen reference process. Nonetheless, we provide reasonable baseline values from which can begin and adjust needed.
D.2 Loss Function of CSBM
In this section, we outline the optimization procedure for the parameterization (9), obtained by substituting m = qo into (10). Following (Austin al., 2021), we parameterize the model predict the terminal point 2 or zo for the forward or backward reparameterization. respectively, and adopt hybrid loss that sums the base loss with the loss Lsimple scaled by weighting factor A. The resulting training objective is therefore given by:
N
L(0) = -E (a)...n)))
E (ns, (m
Liiok
KL ( (ig 110 (tt,,
- - log qo (2i (x
E mm [log To (wn (F(
(21)
Since the backward decomposition of m also holds for Proposition 10, we can apply similar parametrization. In this case, we use neural network with parameters to predict zo:
N+1 Σ
((n)) Eeg(e
(xn (m0.01) n=2
Liaegte
KL (It, 11se,
BB) )l "1a ,1,,,, Foo)) Alog Tr (T (Tol
Epri(z (m0 [log T, (2ol )]
(22)
For further details on the training process, we refer the reader to (Austin et al., 2021).
D.3 Training Aspects
For the implementation of the training logic, we use the official D3PM repository (Austin et al., 2021) as a reference:
https://gitte cttp..co//ggg/// http:://hhhhooooogererrrrrrreree
18

---

Categorical Schrödinger Bridge Matching
Table Hyperparameters for experiments. Ir denotes the learning rate, and represents millions. Params indicate the number of model parameters, where for the CelebA dataset, the first value corresponds to the model and the second to the VQ-GAN.
Initial
D-IMF
D-IMF=1
D-IMF
Batch size
Experiment
Lr
Params
coupling
outer iterations
grad updates
grad updates
2D
Ind
10
400
40000
10
512
0,0004
46588
2. 4,
Colored MNIST
MB
3
000
40
10, 25,
128
0.0002
34m
50, 100
93m
CelebA
Ind
4 800 000
40 000
100
32
0.0004
+
70m
Amazon Reviews
Ind
5 800 000
40 000
100
32
0.0004
100m
Shared Aspects. For all experiments, we use the AdamW optimizer with fixed betas of 0.95 and 0.99. Additionally, apply Exponential Moving Average (EMA) smoothing stabilize training and enhance final model performance. The EMA decay rate is consistently tuned across all experiments and set to 0.999, except for the Colored MNIST experiment, where it is set to 0.9999. For all experiments, we set the weighting factor of .ssimple to 0.001.
For the 2D and colored MNIST experiment, we follow the preprocessing approach from (Austin et al., 2021), where the logits of qo ( ..) are modeled directly as the output of neural network.
Notably, various previous works have introduced different initial couplings ( #1), such the standard indepen- dent coupling (o(2)) (x1) (Shi et al., 2023; Gushchin et al., 2024b), couplings derived from a reference process, e.g., Po po(zo)g (eog (x) (mr) (Shi et al., 2023), and mini-batch OT couplings referred as MB, i.e., discrete Optimal Transport solved on mini-batch samples (Tong et al., 2024). For a more comprehensive overview of coupling strategies, see (Kholkin et al., 2024). In this work, we focus exclusively on the independent and mini-batch initial coupling.
Experiment-speeific Aspects. For the 2D experiment (84.2), we use a simple MLP model with hidden layers of size [128, 128, 128] and ReLU activations. To condition on time, we use simple lookup table, i.e., an embedding layer of size 2.
For the colored MNIST experiment (84.3), follow (Austin al., 2021) and use an architecture based PixeICNN++ backbone (Salimans et al., 2016), utilizing a U-Net (Ronneberger et al., 2015) with ResNet-like structure. The model operates at four feature map resolutions, with two convolutional residual blocks per resolution level and a channel multiplier of (1, 2, 2, 2). At the 16 x 16 resolution level, a self-attention block is incorporated between the convolutional blocks. For time encoding, we apply Transformer sinusoidal position embeddings to each residual block. We train the model on a training subset of size 60 000 and generate images from the hold-out set.
For the CelebA experiment (54.4), we employ VQ-Diffusion (Gu al., 2022), which consists of two models: VQ-GAN (Esser et al., 2021) and transformer- based diffusion model. The VQ-GAN component trained using the official repository:
https: / //tt/ttitttg thub. com Comp/i / taw.iamttransfooomoee-t transformers..
We slightly modify the experimental setup of unconditional generation for CelebA-HQ from (Esser et al., 2021) by reducing the number of resolution levels to three, with scaling factors of (1, 2, 4). This adjustment accounts for our use of CelebA 128 x 128 resolution, compared to 256 x 256 in Ce. The discrete diffusion model is adopted from:
https: /tte::///ittttbt httpggitthhu.coooo
ffusion
Our diffusion model consists of multiple transformer blocks, each incorporating full attention and a feed-forwardd network (FFN). We follow the small model configuration from (Gu et al., 2022), which consists of 18 transformer blocks with an increased channel size of 256. The FFN is implemented using two convolutional layers with a kemnl size of 3, and the channel expansion rate is set to 2. Additionally, we inject time step information through the Adal.N operator.
We train the model on 162 770 pre-quantized images of celebrities. For evaluation, we compute FID and CMMD using 11 816 hold-out images to ensure consistency with the evaluation protocol from (Gushchin et al., 2024b). Likewise, the images presented in the main text of the paper are generated using this hold-out set.
19

---

Categorical Schrödinger Bridge Matching
For the Amazon experiment (Appendix C.4), we train unigram SentencePiece tokenizer (Kudo & Richardson, 2018) that includes explicit start-of-sentence (
) and padding (<pad>) tokens, following the procedure of (Austin et al., 2021). The backbone is the DiT model (Peebles & Xie, 2023), with the implementation available at:
We employ the "small" variant with 12 transformer blocks, each with hidden size of 768 and 12 attention heads. Every block contains multi-head self-attention, rotary positional embeddings, and an MLP with a dropout rate of 0.1. Noise-level information is injected via a 28-dimensiioal AdaLN modulation vector. The model is trained on 104 000 pre-tokenized reviews and evaluated on 000 reviews from the held-out test set. The rest hyperparameters are presented Table 5.
Computational Time. Training the 2D experiment requires several hours single A100 GPU. The colored MNIST experiment takes approximately days train using A100 GPUs. The most computationally demanding task, the CelebA and Amazon Reviews experiments, requires around five days of training on four A100 GPUs.
D.4 Additional Images
2
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
(a) # ~ P1-
(b) 2.
(c) = 4.
(d) N = 10.
2
3
3
3
3
3
33
3
3
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
3 3 3
3
3
3
3
3
(e) Pl-
0 N = 25.
(g) N = 50.
(h) N 100.
Figure 8. Results of colored digits unpaired translation "2" + "3" learned by our CSBM algorithm with reference process qauss and varying number of time moments N.
20

---

Categorical Schrödinger Bridge Matching
Low stochasticity
(a) æ ~ Pt-
(b) CSBM (ours)
(c) ASBM (Gushchin al., 2024b)
(d) DSBM (Shi et al., 2023)
High stochasticity
x ~ Pt
(f) CSBM (ours)
(g) ASBM (Gushchin al., 2024b)
(h) DSBM (Shi et al., 2023)
Figure Comparison of female male translation the CelebA 128 128 dataset using CSBM (ours), ASBM, and DSBM. The low-stochasticity setting for CSBM corresponds to a 0.005, while the high-stochasticity setting corresponds to 0.01. The stochasticity parameters for ASBM and DSBM taken from (Gushchin al., 2024b).
21

---

Categorical Schrödinger Bridge Matching
D
N
00
UU
U0
UO
UO
nn
llts
in
Figure male - female translation trajectories on the CelebA 128 dataset using CSBM 0.01. Each column corresponds to time moments 0, 10, 25, 50, 75, 90, and 101.
22

---

Categorical Schrödinger Bridge Matching
D
U
W
BR
UO
00
U
nn
mn
n
n
n
n
Figure male + female translation trajectories the CelebA 128 128 datasct using CSBM with corresponds to time moments 0, 10, 25, 50, 75, 90, and 101.
0.005. Each column
23

---

Categorical Schrödinger Bridge Matching
Figure 12. femfemale male translation trajectories on the CelebA 128 dataset using CSBM 0.01. Each column corresponds to time moments 0, 10, 25, 50, 75, 90, and 101.
24

---

Categorical Schrödinger Bridge Matching
Figure 13. female male translation trajectories the CelebA 128 128 datasct using CSBM with 0.005. Each column corresponds to time moments 0, 10, 25, 50, 75, 90, and 101.
25


Reviews:
[{"review_id": "FJm17HbJsP", "reviewer": "Reviewer_McJt", "summary": "The paper does what the title says: it establishes the basic framework for the version of Schrodinger bridge diffusion models, for the case of discrete (categorical) spaces. This means that it has a theoretical result describing why and how an iterated projection method for finite-steps markov processes can be made to converge to the optimum bridge process, and then it describes how to implement this iteration in practice using neural networks, and shows a few examples to illustrate the phenomena.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer McJt, thank you for pointing out the unclear parts you encountered. We will revise the text where you have highlighted, as much as possible. \n\n**[R.1] Line 110 column 2: \"additional properties\" sounds too vague, maybe state some examples**\n\nIn line 110, by \"additional properties\", we refer to the reciprocal and Markovian properties, which allow us to use diffusion models (specifically, bridge matching models) as the backbone of our domain translation framework. Without these properties, we would be limited to training one-step models such as GANs, which have since been largely replaced by modern diffusion-based approaches.\n\n**[R.2] Line 235 column 2: $D=1$ seems a weird way to put it, and it confused me because I had forgotten what $D$ was.. maybe just say that you work with $\\mathbb{S}$ and it'll be clear that $D=1$**\n\nIn line 235, the purpose of $D = 1$ is to simplify the transition matrices defined in Equations (7, 8). Since we are using factorization, we consider feature-wise rather than data-point-wise transition matrices. So, to maintain compactness of the text, we omit introducing $Q_n$ for arbitrary $D$, as it is not used in our approach, again, due to the factorization.\n\n**[R.3] Line289 column 1: when you introduce this factorization, maybe briefly write about the limitation that this imposes (I know that this has been written in the \"limitations\" part, but I think it's worth putting it here too)**\n\nWe agree that it is indeed important to clarify the factorization in line 289. We will include a reference to the limitations section in order not to distract the reader from the flow of the main text.\n\n**[R.4] Line 290-299 column 1 : \"Since, in fact, we need N+1 neural nets to do the prediction of endpoints at each time step, we simply use a single neural network with an extra input n\" this sentence is unclear to me...**\n\nIn lines 290–299, we intended to explain that we use a single neural network for all time steps rather than separate networks for each step, which is a common practice. To simulate the stochastic process, one would typically require $N+1$ distinct functions $q(x_1 | x_{t_{n-1}})$ for each $n$. However, training $N+1$ neural networks is computationally expensive. Instead, we use a single neural network for all transition steps (i.e., the transition function $q(x_1 | x_{t_{n-1}})$), with additional time conditioning $q_\\theta(x_1 | x_{t_{n-1}}, t_{n-1})$ for all $n \\in [1, N+1]$. If you are also interested in the sampling procedure using the trained model, please refer to our response to the reviewer [atNH](https://openreview.net/forum?id=RBly0nOr2h&noteId=0v01XUlBGp) [Q. 1].\n\n**[R.5] Line 354-355 column 2: \"discrete space of images but not continuous\" this looks weird.. yeah it's discrete and not continuous, what did you want to say with the \"but\" part? Please correct/erase whatever is needed to make this right.**\n\nIn lines 354–355, we intended to highlight that image spaces are typically treated as continuous rather than discrete, i.e., the space of image pixels is commonly represented as the interval $[0, 1]^D$ rather than the discrete set {$0, 1, \\dots, 255^D$}. \n\n**[Q. 1] I'd like to know more precisely some examples of what kind of dependencies will be lost in the dimensional factorization, and also if the authors have any ideas on how to \"do better than\" the factorization $S^D \\to D\\times S$ metioned in the un-numbered formula following (9). I mean, how would you increase a bit the complexity of the models to try to make it \"forget less\" about the dependencies between dimensions? In the case of a discrete space this is easier than in general, so it's worth giving some pointers.**\n\nRegarding factorization, to the best of our knowledge, the work [1] is the only existing work addressing this issue. Authors introduce an additional generative model that models the copula of the factorized distributions, which is, informally, the part of the joint distribution $q_{\\theta}(x_1 | x_{t_{n-1}}) = q_{\\theta}(x^1_1, ..., x^D_1 | x_{t_{n-1}})$ that is lost due to the factorization.\n\nHowever, practically, it seems that the issue of factorizayion can also be mitigated just by increasing the number of steps, as demonstrated, for example, in our experiments with C-MNIST. The more steps that are taken during sampling, the more expressive the resulting composition of distributions becomes. As a result, repeatedly incorporating the full previous state leads to more correlated features. Even though factorization has implications for our practical implementation, it is important to recall, as stated in the article, that this issue is inherent to all discrete diffusion models.\n\n[1] Liu, Anji, et al. \"Discrete Copula Diffusion.\" arXiv preprint arXiv:2410.01949 (2024).", "claims_and_evidence": "Yes I think that the claims are all carefully argued and convincing.", "methods_and_evaluation_criteria": "I think that the examples are toy model examples, even the VQ-VAE one which is the hardest is a toy model. I don't see a big issue with that, given that this is really the first formulation of SB's for categorical data, toy model experiments are to be awaited. Nevertheless, I may be biased towards theory, a referee interested in practical applications may require stronger evidence that the model helps to advance applications. (The fact that this model would be better than any competitor, is not part of the experimental validation.)", "theoretical_claims": "Yes, I checked all the proofs. Furthermore the main results are in line (similar hypotheses and theses) with well established versions as indicated from Table 1. This means that there is little surprise that the results are true, and on the positive side, it gives good reason to trust that the results are valid, even for people who would not have read the proofs.", "experimental_designs_or_analyses": "I did not check in detail the experiment implementation (i.e. I did not run the codes from the supplementary material), but I believe that the experiment outcomes are realistic and that the declared setup is sound.", "relation_to_broader_scientific_literature": "As said before, the SB framework was previously restricted to continuous spaces, and no version for discrete spaces was available. Thus this paper fills a gap. \n\nThe gap was not hard to fill: it was sufficient to build similar versions of proofs as in cited reference [Guschin et al 2024b], but for the discrete space case, and no surprises appeared anywhere. Some simplifying tricks like the one for passing from S^D to S\\times D for practical purposes, were inherited from previous literature.  \n\nEven if the goal of the paper does not face strong new difficulties, it is important that the gap in the literature has been filled. Also, the new theoretical result (grey box in the paper) is elegant and simple to state so it's worth publishing.", "essential_references_not_discussed": "I don't know of any.", "other_strengths_and_weaknesses": "A strength of the paper is its simplicity. Some referees may shun that, saying that the work is somehow minor since it is not technically involved, but I disagree.", "other_comments_or_suggestions": "Here are a few places where some rewriting may help:\n\nLine 110 column 2: \"additional properties\" sounds too vague, maybe state some examples\n\nLine 235 column 2: \"D=1\" seems a weird way to put it, and it confused me because I had forgotten what D was.. maybe just say that you work with \\mathbb S and it'll be clear that D=1\n\nLine289 column 1: when you introduce this factorization, maybe briefly write about the limitation that this imposes (I know that this has been written in the \"limitations\" part, but I think it's worth putting it here too)\n\nLine 290-299 column 1 : \"Since, in fact, we need N+1 neural nets to do the prediction of endpoints at each time step, we simply use a single neural network with an extra input n\" this sentence is unclear to me. I don't follow the implication beyond the \"since [...]\", and I don't see how \"n\" is an input of a neural network. I feel that this sentence tries to abbreviate stuff too much and it became unintelligible. Can you expand and state this clearly please?\n\nLine 354-355 column 2: \"discrete space of images but not continuous\" this looks weird.. yeah it's discrete and not continuous, what did you want to say with the \"but\" part? Please correct/erase whatever is needed to make this right.", "questions_for_authors": "I'd like to know more precisely some examples of what kind of dependencies will be lost in the dimensional factorization, and also if the authors have any ideas on how to \"do better than\" the factorization S^D \\to D\\times S metioned in the un-numbered formula following (9). I mean, how would you increase a bit the complexity of the models to try to make it \"forget less\" about the dependencies between dimensions? In the case of a discrete space this is easier than in general, so it's worth giving some pointers.", "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": "20eaTHc0ic", "reviewer": "Reviewer_MBG2", "summary": "This paper introduces Categorical Schrödinger Bridge Matching, an approach that extends the Schrödinger Bridge (SB) framework to discrete spaces. While SB has gained traction in generative modeling and domain translation, most prior work has been confined to continuous spaces. The paper addresses this gap by developing a theoretical foundation and a computational algorithm tailored for discrete data.\n\n## update after rebuttal\n\nI have decided to maintain my score, as two main concerns remain:\n\n1. It remains unclear how the discrete-time formulation offers meaningful insights to the field. While the authors argue that prior work relying on continuous-time models is theoretically unsound, I find that the improvements in mathematical rigor on such a fine-grained detail do not constitute a substantial breakthrough from a research perspective.  In particular, I do not perceive a theoretical barrier to translating the results from continuous to discrete time.\n\n2. The similarity to DDSBM is also a concern, as the proposed framework appears to have significant overlap and does not clearly offer novel contributions beyond existing approaches.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer MBG2, thank you for raising important questions regarding our paper. \n\n**[W. 1] Questionable Novelty of the Theoretical Contribution**\n\nWe respectfully disagree. First, as highlighted in Table 1, the continuous-time frameworks DDSBM [1] and DSBM [2] rely heavily on the theoretical foundation established in [3], which does not extend to discrete time. Therefore, our work should not be viewed as a trivial extension of these frameworks but rather as a distinct theoretical foundation that shares certain similarities. Furthermore, the core theoretical contribution of our paper goes beyond simply generalizing the D-IMF procedure from [4] to discrete data: the discrete case emerges as a consequence of a broader generalization of the D-IMF procedure to **arbitrary reference processes**. As noted in the article (see footnote 1, page 5), our generalization enables ASBM [4] to operate with any Markov process $q^{ref}$.\n\n**[W. 2] Moreover, there is a deeper concern regarding the algorithm itself..., one must question: Is CSBM simply DDSBM with a different motivation?... For example, DDSBM has been proposed specifically for graph-structured discrete data, yet it is not included in the evaluation... In what ways does the derived algorithm differ from DDSBM?**\n\nWe agree that adding practical distinctions will strengthen the comparison. Thus, we will include them in the revision. \n\nFirst, let us clarify how DDSBM derives its loss. In theory, it matches the generator matrices $A$ (see [1, Equations (5, 6)]). However, in practice, authors of [1] discretize time, which leads to minimization of $D_{KL}(q^{ref}(x_{t_n}|x_{t_{n-1}}, x_1)|| q_{\\theta}(x_{t_n}|x_{t_{n-1}}))$ (see the derivation in [1, Appendix E.1]). Thus, as it could be mentioned, it is indeed the same loss as presented in our article. However, since we have theoretically alternative objective $D_{KL}(q^{ref}(x_{t_n}|x_{t_{n-1}})||q_{\\theta}(x_{t_n}|x_{t_{n-1}}))$ (see Prop. 3.3), we can match distributions directly by employing **various loss functions** such as MSE or even adversarial training as in ASBM [4]. We conducted extra experiments using only an MSE loss and observed comparable to KL (**LINK**: see [figures](https://anonymous.4open.science/r/images-64B3) with prefix 'toy'). Thus, the similarity in practical implementation reflects our design choice to use this particular parametrization. Alternative approaches could be easily used. This also explains why we did not include a comparison with DDSBM.\n\n**[W. 3] A missing baseline discrete OT with [5]**\n\nIf we correctly understand, you propose first using classical discrete OT methods to get aligned data followed by the application of [5]. However, there is one crucial problem with this setup. The SB problem that lies behind [5] is still considered to be continuous, i.e., continuous Brownian motion is used as a reference process. This breaks the assumptions on discrete data. We will cite [5] in the final revision and discuss it, but we think that a comparison with it is impossible due to the reasons mentioned above.\n\n**[W. 4] Additionally, while the paper compares CSBM to ASBM and DSBM, these are methods designed for continuous spaces, meaning they are not necessarily the most appropriate baselines...**\n\nTo our knowledge, there are no other discrete domain translation models for unpaired data. Thus, the only methods we compare are continuous-space methods ASBM and DSBM. \n\n**[W. 5] ...all experiments focus on image-based tasks, which, despite using vector quantization, are inherently less discrete than the originally mentioned data types...**\n\nWe apologize, but the meaning of \"less discrete\" is unclear to us. Still, for further clarification regarding the choice of experiments (especially texts), please refer to our response to the reviewer [atNH](https://openreview.net/forum?id=RBly0nOr2h&noteId=0v01XUlBGp) [W. 1].\n\n**[Q. 1] What values does the discrete-time perspective bring to the SB framework for categorical data?**\n\nThe main theoretical advantage is the ability to consider CSBM with $N = 1$, which is guaranteed to converge to the SB. In contrast, continuous-time setups typically require to assume $N=\\infty$ to achieve convergence. From a practical side, our framework also enables the flexible selection of loss functions for matching the transition distributions, which we mentioned in answers to your previous questions.\n\n**Concluding remarks.** We hope that, with the above clarifications, you will kindly reevaluate our work and find it deserving of a higher rating.\n\n[1] Kim, Jun Hyeong, et al. \"Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation.\"\n\n[2] Shi, Yuyang, et al. \"Diffusion Schrödinger bridge matching.\"\n\n[3] Léonard, Christian. \"A survey of the Schrödinger problem and some of its connections with optimal transport.\"\n\n[4] Gushchin, Nikita, et al. \"Adversarial Schrödinger Bridge Matching.\"\n\n[5] Somnath, Vignesh Ram, et al. \"Aligned diffusion schrödinger bridges.\"", "claims_and_evidence": "**Limited Empirical Support for Discrete Claims:** \n\nWhile the paper is motivated by discrete data applications (e.g., text, molecular graphs), all experiments focus on image-based tasks, which, despite using vector quantization, are inherently less discrete than the originally mentioned data types. A stronger demonstration on truly discrete datasets (e.g., molecular graphs, categorical sequences, or text-based data) would significantly bolster the claims.", "methods_and_evaluation_criteria": "A central issue with this paper is the lack of comprehensive comparison to existing methods. While the authors position CSBM as an advancement for discrete-state Schrödinger Bridge problems, several relevant baselines are missing from the experiments.\n\nFor example, DDSBM (Discrete Diffusion Schrödinger Bridge Matching) has been proposed specifically for graph-structured discrete data, yet it is not included in the evaluation. Given that DDSBM also addresses discrete Schrödinger Bridge problems, a direct comparison would be necessary to assess whether CSBM provides meaningful improvements or is simply an alternative formulation. The omission of such a baseline weakens the empirical claims of the paper. In particular, this paper would greatly benefit by replicating the experimental setup in the DDSBM paper and demonstrating improvement.\n\nAdditionally, while the paper compares CSBM to ASBM and DSBM, these are methods designed for continuous spaces, meaning they are not necessarily the most appropriate baselines for a method explicitly aimed at discrete settings. A fairer assessment would include methods developed for discrete generative modeling, such as categorical diffusion models.", "theoretical_claims": "**Questionable Novelty of the Theoretical Contribution:**\n\nThe discrete-time theory presented in this paper appears to be a rather trivial extension of the existing continuous-time Schrödinger Bridge theory. The results largely follow from prior work and do not introduce fundamentally new insights beyond what has already been established in the continuous-time setting. While the authors claim to provide a theoretical foundation for discrete-time setting, it is unclear whether this contribution is substantially novel or merely a straightforward adaptation of known results.\n\nMoreover, there is a deeper concern regarding the algorithm itself. The submission presents CSBM as a new computational approach, but it is unclear whether this is substantively different from DDSBM. DDSBM is explicitly motivated by continuous-time considerations, yet it is naturally discretized for implementation. Since the proposed CSBM method operates in discrete time, one must question: Is CSBM simply DDSBM with a different motivation? The paper does not make a clear distinction between the two, and without this clarification, the novelty of the algorithm is questionable.\n\nIf the authors aim to claim CSBM as a distinct method, they must explicitly differentiate it from DDSBM and explain what aspects of their approach do not follow directly from the standard discretization of the continuous-time Schrödinger Bridge framework.", "experimental_designs_or_analyses": "The experimental design does not seem to substantiate the advantages of the proposed algorithm; see my comments in **Methods And Evaluation Criteria**.", "relation_to_broader_scientific_literature": "The problem studied in this paper is **important and valuable**, as extending **Schrödinger Bridge methods to discrete spaces** is highly relevant for applications like **molecular generation, text modeling, and vector-quantized representations**.", "essential_references_not_discussed": "**Missing Baseline: Discrete Static OT Solver:**  \n\nA missing baseline in this paper is a **discrete static optimal transport (OT) solver** built on top of **Reference [1]**. Since the proposed approach is fundamentally an extension of Schrödinger Bridge methods to discrete spaces, it is essential to test whether a **simpler discrete OT solver**—one leveraging existing **static OT formulations**—could serve as an effective alternative within the **bridge matching framework** along side ASBM and DSBM. \n\n\n\n[1] Somnath et al., Aligned Diffusion Schrödinger Bridges, 2023.", "other_strengths_and_weaknesses": "The paper is well-written and effectively conveys its ideas.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- In what ways does the derived algorithm differ from DDSBM?\n\n- What values does the discrete-time perspective bring to the SB framework for categorical data?", "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": "0v01XUlBGp", "reviewer": "Reviewer_atNH", "summary": "The paper addresses the Schrodinger Bridge (SB) problem for discrete spaces (categorical data). It proposes CSBM: a method that extends IMF (actually D-IMF) to discrete categorical spaces, proving theoretical convergence, propoing a concrete implementation and showing experimental evidence with two practical reference processes.  Evaluations on synthetic datasets, Colored MNIST, and CelebA demonstrate competitive or superior performance in unpaired image-to-image translation tasks compared to baseline methods (ASBM and DSBM).\n\n## update after rebuttal\nthe authors addressed my concerns and i maintain my recommendation", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer atNH thank you for your questions and commentaries.\n\n**[R. 1] Hoogeboom et al., and Gat et al, are mentioned but could use further discussion and perhaps comparison**\n\nRegarding the references [1, 2] you mentioned, we believe they are not suitable for comparison in our setting. For [1], the practical differences from D3PM [3] (which is our backbone method) are minimal, as [3] extends the discrete diffusion framework introduced in [1]. While the second method [2] could, in principle, be used as a backbone for our method, it is well known that lower values of $\\epsilon$ (or $\\alpha$ in our work) make it significantly harder for the model to approximate the target distribution (see [4], Figure 7). Current discussion will be added to the revised version. \n\n[1] Hoogeboom, Emiel, et al. \"Argmax flows and multinomial diffusion: Learning categorical distributions.\" Advances in neural information processing systems 34 (2021): 12454-12465.\n\n[2] Gat, Itai, et al. \"Discrete flow matching.\" Advances in Neural Information Processing Systems 37 (2024): 133345-133385.\n\n[3] Austin, Jacob, et al. \"Structured denoising diffusion models in discrete state-spaces.\" Advances in neural information processing systems 34 (2021): 17981-17993.\n\n[4] Shi, Yuyang, et al. \"Diffusion Schrödinger bridge matching.\" Advances in Neural Information Processing Systems 36 (2023): 62183-62223.\n\n**[W. 1] Could benefit from additional experiments on other categorical datasets (e.g., discrete text tokens or molecules) and perhaps more comparison with baselines.**\n\nWe believe the current set of experiments is sufficient to support our claims. In particular, the practical implementation of the VQ-GAN setup does not differ significantly from potential setups with text data. The model and reference process will remain the same. While graphs are indeed more challenging due to their structural complexity, this domain has already been addressed by DDSBM. As explained in our response to reviewer [MBG2](https://openreview.net/forum?id=RBly0nOr2h&noteId=20eaTHc0ic) [W. 2], we do not include a direct comparison with DDSBM and instead focus on exploring other experimental settings.\n\n**[Q. 1] Clarification: In Alg 1, when sampling $(x_0,x_1)$ using $q_\\eta$ or $q_\\theta(x_1|x_0)$, do the authors mean $x_1$ is simulated from $x_0$ with $N$ steps? or estimated directly from the $x_1$ timestep? if it is the latter -- could you provide more details on how this is done in practice?**\n\nYour first suggestion is correct. To generate samples, we follow a standard diffusion sampling procedure. After the $l$-th step of D-IMF, we obtain $q^l_\\theta(x_1 | x_{t_{n-1}})$. Moving to $(l+1)$-th D-IMF iteration, we first apply the trained model to generate $x_1$ taking as an input a dataset point $x_0$. We then perform posterior sampling using $q^{\\text{ref}}(x_{t_n} | x_{t_{n-1}}, x_1)$ to obtain $x_{t_n}$. This procedure is repeated for $N+1$ steps to obtain samples from coupling $q^l_\\theta(x_0, x_1)$. For the backward parametrization, we use the same scheme.\n\n**[Q. 2] Can the authors discuss ways to reduce the information loss arising from the factorization of the conditional distributions?**\n\nPlease refer to our response to reviewer [McJt](https://openreview.net/forum?id=RBly0nOr2h&noteId=FJm17HbJsP) [Q. 1].\n\n**[R. 2] However, the chosen parameters for stochasticity level (alpha) appear somewhat ad-hoc. further explanation of the choice of particular values would be helpful**\n\nThe pattern of selecting $\\alpha$ follows the same intuition as choosing $\\epsilon$ in continuous SB methods. Specifically, lower values of $\\alpha$ lead to less stochasticity in the trajectories, resulting in higher similarity to the input data but a lower-quality approximation of the target distribution. At very low values, the model may collapse due to insufficient stochasticity. Conversely, higher values of $\\alpha$ introduce more variability, improving the quality of the approximation but reducing similarity to the initial data. Beyond a certain point, excessively large $\\alpha$ values make the model difficult to train, leading to a drop in both quality and consistency. Unfortunately, the effective range of these behaviors is highly dependent on the dataset and the chosen reference process. Nonetheless, we provide reasonable baseline values from which one can begin and adjust as needed.", "claims_and_evidence": "(1) Theoretical Claim: The uniqueness and convergence of the discrete-time IMF procedure for categorical spaces. \nEvidence: A formal theorem is provided, clearly proving the convergence under stated conditions. The proof appears rigorous to the best of my judgement\n(2) practical algorithm: The proposed CSBM is claimed to be effective in practice for categorical SB problems.\nEvidence:supported by experimental results on several datasets: gaussian to swiss roll, colored MNIST and VQ CelebA, demonstrating visually good translations and quantitative improvements  compared to baseline methods.", "methods_and_evaluation_criteria": "The methods and evaluation criteria (FID, CMMD) make sense given the problem of discrete unpaired translation. I liked the use of VQ representations of the celeb-A images as it aligns  with common practices in generative modeling. The select reference processes (Uniform and Gaussian-like) are well suited for common scenarios (unordered vs ordered categories).", "theoretical_claims": "Theorem 3.1's proof appears correct and rigorously executed", "experimental_designs_or_analyses": "The experimental designs appear solid, with clear evaluation metrics (FID, CMMD) and visual evidence to assess qualitative performance. However, the chosen parameters for stochasticity level (alpha)  appear somewhat ad-hoc. further explanation of the choice of particular values would be helpful", "relation_to_broader_scientific_literature": "The authors do a really good job positioning their work within existing literature on SB, optimal transport, and diffusion models. They  distinguish their contributions from related methods -- continuous-space IMF (e.g., Shi et al., 2023; Gushchin et al., 2024) and the part clarifying difference from discrete optimal transport methods (Sinkhorn algorithm, gradient-based approaches) was also a good addition (even though i felt the distinction was clear, it was nice to read and to be stated explicitly).", "essential_references_not_discussed": "Hoogeboom et al., and Gat et al,  are mentioned but could use further discussion and perhaps comparison", "other_strengths_and_weaknesses": "Strengths:\n\n+ Clear motivation and justification for extending IMF to discrete spaces.\n+ Thorough theoretical support and convincing experiments showing effectiveness across tasks and settings\n+ Very clear presentation -- the paper is fun to read\n\nWeaknesses:\n- Could benefit from additional experiments on other categorical datasets (e.g., discrete text tokens or molecules) and perhaps more comparison with baselines.", "other_comments_or_suggestions": "none", "questions_for_authors": "* Clarification: In Alg 1, when sampling (x0,x1) using q_eta or q_theta (x1|x0), do the authors mean x1 is simulated from x0 with N steps? or estimated directly from the x1 timestep? if it is the latter -- could you provide more details on how this is done in practice?\n\n* Can the authors discuss ways to reduce the information loss arising from the factorization of the conditional distributions?", "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": "t6PlUy5yDv", "reviewer": "Reviewer_1AGf", "summary": "The paper proposes an algorithm based on Iterative Markovian Fitting (IMF) for solving Schrödinger Bridge (SB) in discrete (categorical) space. The contribution of the paper therefore lies in the extension of SB, originally constructed in continuous state spaces, and its data-driven learning-based algorithm to discrete setup. Experiments are conducted on 2D synthetic dataset and latent-space images.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer 1AGf thank you for your questions and commentaries. \n\n**[Q. 1] Can the author clarify the difference between proposed method to DDSBM, which does present theoretical results for continuous-time IMF in discrete spaces? I'm fairly familiar with DDSBM and, since both are based on IMF, the methods seem to collapse in practice when time is discretized.**\n\nThe answer to this question can be found in our response to reviewer [MBG2](https://openreview.net/forum?id=RBly0nOr2h&noteId=20eaTHc0ic) [W. 1] and [W. 2].\n\n**[W. 1] I'm not convinced by Table 2. The author should report ASBM and DSBM in a similar GAN-based continuous latent spaces for a fair comparison.**\n\nRegarding the setup on the CelebA dataset, we agree with your concerns. We did attempt to train DSBM in the latent space. For a fair comparison, we ran DSBM on the same latent space used for CSBM, following the approach in [1, Appendix G]. However, the results were not satisfactory, as the model tended to collapse to the identity mapping with $\\epsilon = 1$ and $\\epsilon =  3$ (**LINK:** see [figures](https://anonymous.4open.science/r/images-64B3/) with prefix 'latent'). Due to these limitations, we did not proceed with training ASBM and chose not to compare both methods with CSBM in such settings. One may ask why CSBM performs better in this setting. We hypothesize that this is due to the choice of the reference process, with $q^{\\text{unif}}$ being more suitable for the latent space of VQ-GAN.\n\n[1] Rombach, Robin, et al. \"High-resolution image synthesis with latent diffusion models.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n**[W. 2] It seems faulty to claim the dimensionality of CelebA Faces to be $1024^{256}$. Given the factorized parametrization, the dimension that the proposed method handles should be $1024*256$.**\n\nThis is not a mistake but rather a point that could be clarified more explicitly. When we refer to $1024^{256}$, we are indicating the complexity of the data, not the complexity of the model parametrization. The quantity that you mentioned, $1024 \\times 256$, corresponds instead to the complexity of the generated samples under $q_\\theta(x_1 | x_{t_{n-1}})$.\n\nIt is also worth noting that as the number of sampling steps increases, the complexity of the resulting composition of distributions also grows. As mentioned in our response to reviewer [McJt](https://openreview.net/forum?id=RBly0nOr2h&noteId=FJm17HbJsP) [Q. 1], this increase in sampling steps can lead to higher-quality samples and help mitigate issues related to factorization.\n\n**[R. 1] Fig 3's caption should clarify that images are generated in VQ-GAN latent spaces. I think it's boarderline misleading to not mention specifically in the caption that it's a GAN-based latent-space image experiments. GAN's latent spaces are not only of lower dimension but also much structural.**\n\nWe will specify this aspect of training in the caption, as you have suggested.\n\n**Concluding remarks.** We hope that, with the above clarifications, you will kindly reevaluate our work and find it deserving of a higher rating.", "claims_and_evidence": "- Can the author clarify the difference between proposed method to DDSBM, which does present theoretical results for continuous-time IMF in discrete spaces? I'm fairly familiar with DDSBM and, since both are based on IMF, the methods seem to collapse in practice when time is discretized.", "methods_and_evaluation_criteria": "Y", "theoretical_claims": "Y", "experimental_designs_or_analyses": "- I'm not convinced by Table 2. The author should report ASBM and DSBM in a similar GAN-based continuous latent spaces for a fair comparison.\n\n- It seems faulty to claim the dimensionality of CelebA Faces to be 1024^256. Given the factorized parametrization, the dimension that the proposed method handles should be 1024*256.\n\n- Fig 3's caption should clarify that images are generated in VQ-GAN latent spaces. I think it's boarderline misleading to not mention specifically in the caption that it's a GAN-based latent-space image experiments. GAN's latent spaces are not only of lower dimension but also much structural.", "relation_to_broader_scientific_literature": "Discrete space is an important extension of data-driven methods. Thm 3.1 could potentially handle other data types/spaces, which may also be of independent interests.", "essential_references_not_discussed": "N", "other_strengths_and_weaknesses": "N", "other_comments_or_suggestions": "N", "questions_for_authors": "My main questions and concerns, as listed above, are the similarity to DDSBM (ICLR'25) and experiment setups. The only comparison to prior SB works is Table 2 which is not conducted fairly.", "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": "iMEjT1Djt0", "reviewer": "Reviewer_hwSE", "summary": "The authors:\n- provide a proof for the convergence of discrete-time IMF in discrete-state spaces.\n- develop an algorithm called \"Categorical SBM\" that approximates a solution to the SB problem for discrete-state spaces.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and positive evaluation. We will correct the typos you mentioned. If you're interested in the topic, we recommend checking the other reviews and these works on the Schrödinger Bridge Problem [1, 2, 3].\n\n[1] Kim, Jun Hyeong, et al. \"Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation.\" arXiv preprint arXiv:2410.01500 (2024).\n\n[2] Shi, Yuyang, et al. \"Diffusion Schrödinger bridge matching.\" Advances in Neural Information Processing Systems 36 (2023): 62183-62223.\n\n[3] Gushchin, Nikita, et al. \"Adversarial Schrödinger Bridge Matching.\" The Thirty-eighth Annual Conference on Neural Information Processing Systems.", "claims_and_evidence": "I'm not well equipped to answer this question.", "methods_and_evaluation_criteria": "The experiments make sense to me (inter-domain translation of images), and the method performs on par with what was compared.\n\nI don't know the literature enough to fully appreciate these results.", "theoretical_claims": "I checked the proofs to the best of my ability, and they look correct.\n\nHowever, I can't provide any merit on how relevant it is to the current understanding of the topic.", "experimental_designs_or_analyses": "These experiments are standard.", "relation_to_broader_scientific_literature": "I don't know.", "essential_references_not_discussed": "I don't know.", "other_strengths_and_weaknesses": "The paper is very well written.", "other_comments_or_suggestions": "There are many typos regarding word order (see lines 112 and 116) as well as the use of the word \"the\".\n\nThey don't affect the understanding.", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper extends the Schrödinger Bridge (SB) framework to discrete spaces by developing a theoretically grounded and practically implementable algorithm based on Iterative Markovian Fitting (IMF). The proposed method is validated through rigorous theoretical analysis and empirical studies on synthetic data and quantized image representations.
While some reviewers expressed concerns regarding novelty and similarity to recent works, such as DDSBM, the authors provided clear justifications, including theoretical distinctions and flexible algorithmic design. In contrast, several reviewers found the contribution meaningful, particularly for filling an important gap in the literature and providing a solid foundation for discrete SB. The experiments, while limited to image data, are carefully conducted and demonstrate competitive performance.
Given that some negative comments remain, for example, especially on empirical scope and relation to prior work, the paper makes a significant contribution to the field by formalizing the discrete Schrödinger Bridge problem and providing a viable algorithmic solution. Therefore, I recommend acceptance.
Decision: Accept (poster)