Tight Bounds for Schrodinger Potential Estimation in Unpaired Data Translation
Authors: Nikita Puchkin
Denis Suchkov
Alexey Naumov
Denis Belomestny
Venue: ICLR 2026 Poster (2026)
Keywords: Learning theory
stochastic optimal control
Schrodinger potential
non-asymptotic bounds
Primary area: learning theory
TL;DR: Sharper bounds on the accuracy of Schrodinger potential estimation with applications to generative modelling and unpaired data translation.

Abstract: Modern methods of generative modelling and unpaired data translation based on Schrodinger bridges and stochastic optimal control theory aim to transform an initial density to a target one in an optimal way. In the present paper, we assume that we only have access to i.i.d. samples from the initial and final distributions. This makes our setup suitable for both generative modelling and unpaired data translation. Relying on the stochastic optimal control approach, we choose an Ornstein-Uhlenbeck process as the reference one and estimate the corresponding Schrodinger potential. Introducing a risk function as the Kullback-Leibler divergence between couplings, we derive tight bounds on the generalization ability of an empirical risk minimizer over a class of Schrodinger potentials, including Gaussian mixtures. Thanks to the mixing properties of the Ornstein-Uhlenbeck process, we almost achieve fast rates of convergence, up to some logarithmic factors, in favourable scenarios. We also illustrate the performance of the suggested approach with numerical experiments.

Published as a conference paper at ICLR 2026
TIGHT BOUNDS FOR SCHRÖDINGER POTENTIAL ESTIMATION IN UNPAIRED DATA TRANSLATION
Nikita Puchkin' & Denis Suchkov*
HSE University, Faculty of Computer Science (npuchkin, d. such.ovv
Alexey Naumov
HSE University and Steklov Mathematical Institute of Russian Academy of Sciences anaumovehse.
Denis Belomestny
Duisburg-Esse Essen University and HSE University denis be 1 Toooomtttttt mymeestn@unn ny@uni ceiidduuedee
ABSTRACT
Modern methods of generative modelling and unpaired data translation based on Schrödinger bridges and stochastic optimal control theory aim transform an initial density to target one in an optimal way. In the present paper, we assume that we only have access i.i.d. samples from the initial and final distributions. This makes our setup suitable for both generative modelling and unpaired data translation. Relying on stochastic optimal control approach, we choose an Saastiin-hllebecck process the reference and estimate the corresponding Schrödinger potential. Introducing a risk function as the Kullback-Leillrr diver- gence between couplings, derive tight bounds on the generalization ability of an empirical risk minimizer over a class of Schrödinger potentials, including Gaussian mixtures. Thanks the mixing properties of the Ornstein-Uhlenbern process, we almost achieve fast rates of convergence, up to some logarithmic fac- tors, in favourable scenarios. We also illustrate the performance of the suggested approach with numerical experiments.
INTRODUCTION
The Schrödinger bridge problem has emerged as a powerful framework in modern generative mod- elling and unpaired data translation (in particular, image-to-ine to-ii translation). It formulates the task of transforming an initial distribution into a target one via a probabilistic coupling that is as close as possible prescribed reference dynamics in the sense of relative entropy. In the present paper, we assume that both the initial and target distributions are absolutely continuous and denote their densities with respect to the Lebesgue measure by Po and pr. respectively.
Let qe(y) w), process X (typically diffusion), and let T² be the joint distribution of xX) Rd Rd given by the initial distribution Po and the forward dynamics of the reference process. Then, in the Schrödinger bridge problem one looks for coupling (I. y) between Pn and pr that minimizes the Kullback-Lein Leibler (KI.) divergence
[0, T), denote the transition density of time-homogeneous reference Markov
inf xell(po.Pr)
KL(T, x²),
where [p Pr) denotes the set of all couplings with the prescribed marginals. This marginal formulation of the Schrödinger problem avoids dealing with full path measures and is particularly amenable to sample-based estimation.
"qqual contribution

---

Published as a conference paper at ICLR 2026
Under mild regularity conditions, the minimizer T" of the variational problem 1 admits a specific form (see Leonard 22044. Theorem 2.4)):
x <² N (I) dr ()) (1))
(2)
where vi and v are positive functions referred to as Schrödinger potentials. These scaling func- tions modify the reference transition structure so that the resulting coupling satisfies the marginal constraints. From this coupling, one can construct the corresponding Schrödinger bridge process X", which time-inhomogeneous Markov process with dynamics governed by drift-modified version of the reference process.
Usually, researchers consider reference process of the form
dxp = a dW,, O
(3)
where W. is a standard d-dimensional Wiener process and 0. The motivation for such choice is that the problem with the reference process (3) is equivalent to the entropic optimal transport problem
/
[ 2 2
inf TE .T) Rd Rd
ull (2 y) dz dy + PLT,L(TT Po @ PT),
which is studied quite well. However, there are two drawbacks. First, the entropic optimal transport formulation deals with the scalar parameter a and does take into account possible anisotropy of the data. Second, if the reference dynamics obeys (3). then the correlation between X and Xp decays at slow polynomial rate. In the context of that initial images have an overabundant and possibly negative influence on the final ones. As consequence, learner has take larger values of a or T to mitigate this effect. In this work, we take the Ceustein-Uhlenber (OU) process as the reference dynamics. To be more specific, we will consider the base process X solving the SDE
translation, this could mean
dxP = b (m- x!) di + aW/ 0 EETT,
(4)
where b > 0 controls the drift rate, m E Fd represents the mean-reversion level, and Σ E d positive definite symmetric matrix. Analytic tractability and exponential mixing properties of the OU dynamics make it especially suitable for both theoretical analysis and practical computation. The corresponding Schrödinger bridge process X* evolves under the modified drift
ax; = 3d (t, X2) + I1/2WW,
Xi Po-
(5)
where
B (1.z) = 1m x) EV log uf (x), and 1/ (I) denotes the time-evolved Schrödinger potential given by
(6)
v (x) = 1 (x) qr- x)) dy.
This yields controlled diffusion that transitions from Po to pr in way that remains close to the reference process. We also would like to note that the controlled diffusion (5) solves the stochastic optimal control problem (see Dai Pra (191))..
The expression for the optimal coupling n" suggests us look for the solution of the problem (1) of the form
(I, = v08 (z) (8l x)) (()) Th E I[po. Pr), where & belongs to a closed parametric class &. The best log-potential op minimizes the Kullback- Leibler (KL) divergence between x and rP:
/ Rd Rd
x (I. v) IP (z. v)
6eth E argmin KL(T"" ;),
where KL(T", (L(T)TT)=-
log
" y) dx dy.
PEΦ
For any function g: Rd IR and any t> >>0 0, let
[z)) /
g(v)) qr(v) x))
Rd
2

---

Published as a conference paper at ICLR 2026
Table 1: Frequently used notations
Notation
Eq.
Meaning
Source (initial) distribution
Pr
Target (terminal) distribution
9r x))
Transition density of the OU process, density of N (mme (a)) Ex)
my (x)
(1 e )m + e-
Dt
11-6-26t)>) )z/
vū, Pt 60t
Left and right Schrödinger potentials
Target Schrödinger log-potential
7
Target coupling, solution of the variational problem
TP
Coupling corresponding to the log-potential candidate +p
Log-potential estimate, empirical risk minimizer
Coupling estimate, corresponding to Ø
stand for the Ornstein-Uhlenheed operator applied to Let us note that the condition
Po(z) x (xx y) dy
yields that vz) (z) For this reason, we can rewrite If in the following form:
x& (z, v) Po(x) qr (y [x) et())
(8)
Tr ( (x)
Similarly, the optimal coupling satisfies the equality
(x)) qr (y) x)) (x) (v)
xt" (1r, v)
80 log(u)))
Tr Trler (xx))
Hence, we can represent the KL-divergence between x" and I' as the difference
KL(T" , = [(p)) C(l"),
(10)
where
((p) Ez log Tr (Z) Ey Ey~pr
(11)
Our focus lies on statistical aspects of Schrödinger potential estimation. We assume that Po and pr are accessed through i.i.d. samples Zi..... NN and Vi...., Yn, respectively. In what follows, we assume that N = n for simplicity. The objective L(p) has a remarkable property that it includes only marginal distributions Po and pr, rather than the coupling n². Hence, we do not have to make any presumptions about joint distribution of Zj 's and Y,'s
A natural strategy to estimate '²" and 4p is replace the integrals in I1]) by empirical means and consider the estimates
c = and
T E argmin ((4)) 4E#
(12)
where
Z(p) 1 n
Σ p(Y)).
1 n
log T (7,))
(13)
In the present paper, we are interested in generalization ability of the empirical risk minimizer
Contribution.
The main contribution of this paper is a non-asymptotic high-probability generalization eтror bound for the empirical risk minimizer (12) (Theorem Its proof relies on concen- tration arguments and ergodic properties of the OU process.
When the approximation emror is small, we establish nearly fast convergence rates (up to logarithmic factors) for the population KL risk. These results highlight the statistical effi- ciency of the Schrödinger bridge approach in the finite-sample regime.
We conduct numerical experiments on synthetic and real data illustrating that a proper choice of the reference process can improve data generation quality.

---

Published as a conference paper at ICLR 2026
Notations. Throughout the paper, qr(y] x) denotes the transition density of the OU process which is Gaussian with mean m2)r) = m + e e= m) and variance Σg = (1 - e 2bt | (26).
Based on this, we define the Ornstein Caieein-Uhhebecck operator
Tel](()) ((v))((y)l Ra
The notation f S 9 g 2 f means that f Ogg). Besides, we often replace max(a, b) and min(a,b) b) by shorter expressions a V b and a A b, respectively. For any 8 3 1, the Orlicz ths- norm of random variable ç ~ p is defined as |kEll,. = inf {u [u>0: 0: Eelil"/v" /ue <2). Sometimes, we use the notation (l) to emphasize that & is drawn according to p. Finally, given p > 1 and a probability density p. the weighted Ly-norm of function f defined as IIfIle,(t)) Given two probability densities & q, the Kullhack-Leebble divergence between them is defined as KL(p, q) = Ee-p log (p())/q(E)) Some other frequently used notations are collected in Table
2 RELATED WORK
The empirical success of denoising diffusion models (Ho ct al. 2020 Song al. 2021) and flow matching (Lipman et al,, 2023) made researchers study more general frameworks for mapping an initial distribution to target one, such stochastic interpolants (Albergo & Vanden-Eijnden,
2023 Albergo et al. (025), optimal transport (Tong et al. 2020: 2024b[a) and Schrödinger bridges (De Bortoli al. 2021: [Shi al., 20233] Liu al. 2023). In the context of the Schrödinger bridge problem, a lot of attention was paid to development of solvers based on iterative proportional fit-
ting (De Bortoli al. 2021) and iterative Markovian fitting (Shi al.
2023 Liu al.) 2023
Gushchin et al.) 2024b). Some recent papers (Rapakoulias et al. 2025
Domingo-Enrin et al.
2025) suggest using stochastic optimal control theory (see Dai Pra
(1911] for background) con- struct Schrödinger bridges. Besides, a lot of efforts were made to construct efficient algorithms for solving the Schrödinger bridge problem and accelerate statistical inference (see, for instance, (Taan al. (2024); Gushhchi et al. (2024a b)). Despite numerous methodological advances in application of Schrödinger bridge framework to generative modelling, its theoretical justification often quite challenging and largely missing. For instance, only recently Conforti al. (2024) (see also Eckstein (2025) extended results on geometric convergence of IPF from compactly supported pu and pr to strongly log-concave marginals with unbounded supports. We would like to emphasize that Con- forti et al. (2024) consider an idealized setting ignoring sampling error. A comprehensive analysis of the statistical guarantees for the Schrödinger potential estimation is still to be done. We provide a detailed discussion for few contributions in this direction below.
In Rigollet & Stromme 2025), the authors considered the following empirical dual objective:
1
12
((P,, = n1
log vo (2)) + logut (Y)
mo(2)) qr q 2,) vT (Y)) + 2
72
nT
;=1
This objective arises variational representation of the entropic optimal transport problem. where the regularization term penalizes deviations from the empirical product measure po @ 8on with regularization parameter o² > 0. Note since Po is fixed in our setting, the empirical KL. objective
13 reduces to
L(vo, vT)
Σ log vo (Z,)
Σ
log Pr (Yi).
i=1
which corresponds to the log-likelihood of the coupling n (x; v) vo(x:) qr(y ()) vr(y) evaluated on the data, up to an additive constant.
The dual objective Š thus coincides with 7 up to a constant when the marginal penalty term vanishes, ie., when x² O. More generally, we may write:
s = 2. marginal penalty constant.
4

---

Published as a conference paper at ICLR 2026
This formulation highlights the interpretation of æ trade-off parameter controlling how strictly the marginal constraints are enforced: in the limit E - 0. the dual objective recovers the pure like- lihood form. It was shown in (Rigollet & Stromme, 2025) that the empirical objective S(vo, Pr) converges its population counterpart at the rate O(I/n) in expectation. However, tant to note that no theoretical guarantees were provided for the convergence of the coupling-level Bellback-Leibler divergence KL(T*, R). where the estimated coupling is given by ((,, v) Do (2) cr(y (x) (v))
impor-
and π* denotes the true optimal coupling. As a result, while the dual potentials converge in the objective sense, this does not immediately imply convergence in distribution or divergence between the associated couplings.
An empirical objective similar to formulation 13) was investigated in (Korotin et al. 2024).
The authors studied the practical performance of the objective E and also established theoretical guarantees on its statistical behavior. In particular, they proved that 2 converges to its population counterpart L in expectation at the rate ((1/vn) when the right Schrödinger potential is a Gaus- sian mixture. This result implies the convergence of the induced coupling (1,, y) expected KI. divergence with the same rate. However, simple numerical experiment indicates that learner can hope for much more optimistic rates in the case of Gaussian marginals Po and pr (see Figure| []. In contrast, our work provides significantly stronger guarantees: we prove high-probability upper bounds on the KL divergence KL(TT", #), where π² denotes the optimal solution and the coupling स is defined in [22. Our result holds for general class &b of log- potentials, including the class of logarithms of Gaussian mixtures considered in (Korotin et al., (022)). Under appropriate assump- tions on the function class 4, we obtain a fast convergence rate of order (loo"(n)/mn This sharp control of the statistical error reflects the actual dependence of EKL(m", #) the sample size in the case of Gaussian Pa and pr much better (see Figure [T and makes our approach particularly attractive for generative modelling and style transfer tasks with limited data.
Figure 1: Dependence of EKL(x*, т) the sample size in the case of two- dimensional Gaussian marginals Po and pr. The expected KL-divergence was es- timated from 10 runs. The purple line was fitted with linear regression (I²² = 0.94).
Let us also mention the work (Pooladian & Niles- Weed 2025). where the authors show that the Schrödinger potentials obtained from solving the static entropic optimal transport problem between source and target samples can be modified to yield natural plug-in estimator of the time-dependent drift defining the Schrödinger bridge between the two distributions. The authors show that, if both Pu and pr are supported on compact & k-dimensional submanifolds in Rd then it holds that
1 Vn (T- O<T<T.
CTV2(P% Plo.r)) S
(14)
Here Po.-)) and P(o.r) are the path measures corresponding to the optimally controlled process (5) and the process with estimated drift. Unfortunately, the right-hand side of (14) blows up as ap- proaches -| This highlights a fundamental challenge in learning Schrödinger bridges near the boundary of the time interval.
Let us stress that Korotin et al.) (2024), Pooladian & Niles Weed 2025). and Rigollet & Stromme 2025 restrict their attention to the where the reference process is Brownian motion. While this
Probably, this is due to singularity of the target distribution with respect the Lehesgue measure in Rdd
5

---

Published as a conference paper at ICLR 2026
choice offers analytical convenience and aligns with the classical formulation of the Schrödinger problem, it limits modelling flexibility, especially in contexts where prior structure or time- inhomogeneity is important. In contrast, framework accommodates more general reference processes, such as the Omstein-Uhlenbeec process, enabling better control of prior dynamics and improved statistical regularity through mixing properties.
MAIN RESULT
In this section, present high-probabillty upper bound on the KL-divergence between the optimal coupling x² and the empirical risk minimizer ń defined in (12). It holds under mild assumptions listed below. First, mentioned in the introduction, assume that the reference process X obeys the Ornstinn-Ulleneeck dynamics.
Assumption 1. The base process
X solves the following SDE
dxp = b (m- X?) di + 12 aw,, 0 OELET.
where b m E R Σ a positive definite symmetric matrix of size d x d, and W is a d- dimensional Brownian motion.
Second, we assume that the initial distribution has bounded support. Assumption 2. The initial density po has bounded support and |E-1/2(x m)ll < R for some R > O and all E supp(po)-
We do not require Po to be bounded away from zero infinity supp(po). For this reason, Assumption 2 is satisfied for a large class of densities. It plays crucial role the proof of Lemma C.I which helps us verify Bernstein-type condition (29). Extension of this result the of unbounded supp(po) will require significant efforts. A requirement on the target distribution is even milder, we just assume that it sub-Gaussian.
The target distribution time T is absolutely continuous with the density pr and sub-Gaussian with variance proxy v², that is,
Assumption 3.
E Ex Y & /2 /2 for any E Rdd
(15)
We would like to note that (15) yields that Ey = This is just a normalization condition, which does not diminish generality of the setup we consider. In Korotin et al. (2024), the authors assumed that both Po and pr have compact supports. contrast, we allow supp(er) to be un- bounded. We proceed with two assumptions properties of the log-potential yo associated with the optimal coupling ाt (see (+) and the reference class &.
Assumption 4. There exist positive constants and M such that every yp E & [U(() satisfies the inequalities
-11E-1/2(x 11)12 M < p(x) & M for all E Rd. Moreover, for any p 1 u { it holds that Txip = 0.
for all ip + {1-} because of the fact that Schrödinger log-potentials are defined up to an additive constant. Thus, the condition Toop = helps avoid ambiguity. Assumption 4 together with Assumptions and 3 ensures that the random variables (Y,), < i & n. and log Tz [e ( (Z,). < < n, are sub-exponential. This allows us quantify deviations of the empirical risk [3 from its expectation using concentration inequalities. In PPuhhiin et al. (2025), the authors provided several examples of reference classes satisfying Assumption 4 including classes ReL.U neural networks and classes of logarithms of Gaussian mixtures. We will use the last ones in Section 4 with numerical experiments. According to Korotin et al. (2024) Theorem 3.4), Gaussian mixtures possess a uniform approximation property. For this reason, proper choice of Gaussian mixtures class may lead to convincing practical results. Further discussion of Assumption
We have to introduce the requirement Too
deferred to Appendix A
Finally, assume that & is parametrized by diiteedimensional parameter 0
R +
{spo : f E O), where θ is a subset of a D-dimensional cube [-1, 1] and each function Po maps Rdd onto RR. We suppose that the parametrization is sufficiently smooth in the following sense.

---

Published as a conference paper at ICLR 2026
Assumption 5. The class &b is parametric: + {4p8 0 b) where θ S [-1,1] Moreover, there exists A 3 0 such that
[po(2)) - yor (x)( < A (1+ ||=]|²) 110 0' d'll for all of e e and all x RR!
A similar condition appeared in Puchkin et al. (2025), where the authors argued that Assumption 5 holds for the aforementioned classes of neural networks with ReLU activations and classes of logarithms of Gaussian mixtures. We refer reader to Puchkin et al. (2025) Section 4) for the details.
We are ready to formulate the main result of this section.
Theorem 1. Grant Assumptions i 4 and For any n E N and ò E (0, 1), introduce
T(n 6) = D(M ( ) (M + log ((L A)nd) + log(1/6)) log
Then there exists
Ld +
+ M log log
To - log log T(n, 1) V log d (d+LR*
LR
b
such that for any E (0, 1) and any T E To, with probability at least (1 - 5), the coupling गे defined in (12) satisfies the inequality
(bV L)(M d)
KL(x", - T(n,8)
Vog
bA
(bV L) + d)
+ T(n
1 V
(16)
bA
where A = inf KL(T*" :
Remark 2. A reader can note that the quantity ((n, 8) introduced in Theorem
Sccales as
T(n, 8) 2 (log + log d log(1/5)) Dd logn
n
If the class rich enough in sense that A x ((n, 8), then the upper bound T6 simplifies to
KL(T", (t) S T(n, 8) logn & (logn - log d log(1/8)) Dd
To our knowledge. Theorem provides the first non-asymptotic highh-rrobbbiity upper bound on the KL-divergence between the optimal coupling n² and the empirical risk minimizer म. The previous result of Korotin et al. (2024) (see Theorem A. their paper) concerned only the expected excess risk EKL(TT, क) A. We also would like to note that our rate of convergence sharper than the OLI/Vá) bound of Korotin et al. (2024) when the approximation error A is small. This is the case when there exists & approximating jo" with a good accuracy with respect to L1 (pr)- and L, (N norms.. Indeed, according to Lemma E.2 and the representation (10), we have
KL(T*, 11p 40 (( + (Txlpp 1)) == (T)
Quantitative bounds on 1 ----==11 -0,l, (o(rr) and Taclyp | require the log pog-potential yp to belong to class of smooth functions, like Hölder or Sobolev classes. Unfortunately, we are not aware of results of this kind. For this reason, present paper. However, we can observe that in a favourable scenario, when S T(m, 8), the rate of convergence Theorem I becomes (log (n)/n), which is much better than ((1/ (n). Besides, Theorem 1 exhibits nearly linear dependence O(d d) on the dimension d, which makes the result suitable for high- nigh-dimensiona settings.
leave the study of the approximation error out of the scope of the
Theorem requires bT to satisfy the bound BT > log log It will become clear from Section B.I that it plays an essential role in the proof of Theorem [ and is necessary for verification a Bernstein-type condition. The assumption 1T 2 log log n seems to be necessary for the fast rates,

---

Published as a conference paper at ICLR 2026
because of the loss function (see (3) In the limiting case T 0 the loss becomes linear in ip and has no curvature. Nevertheless, we would like to emphasize the following. First, the inequality BT 2 log log n should be considered as condition on 6T, rather than on the time horizon T alone. This means that we can take 7 = (as is usually done in Schrodinger bridge problems) and b Z log log n. Second, a violation of the condition iT log log does mean that the ERM o becomes inconsistent. Using the standard approach based on Rademacher complexities, one can show that the KL-divergence between π and will still tend zero but slower rate O vn). Finally, the OU process is exponentially ergodic while the condition on bT exhibits the doubly logarithmic dependence. Hence, there is still a significant correlation between the initial and final points of the SDE-governed stochastic process. In the context of image-to-ing tonge-to translation, this means that the style of initial images will be successfully transferred to the final ones.
In generative diffusion models one is often interested in accuracy of the optimal drift estimation. According to Dai Pra 1991. Theorem 3.2), the optimally controlled process X². < t & T, obeys the SDE
X Po.
dx; = - m)dt + u (X,, t)dt / 0<t<T,
(17)
where u² (x = VV log -/ (xx). With the estimate @ at hand (see 12) for the definition), we can consider the corresponding process
No Po:
dd. -b( m)dt + ū ū(,, t)ds dw, <<i<T,
(18)
with ū(x, t) EV log TT- (e
(xr). This brings to natural question whether ū(r, t) is close
u (x.,, t). According to Girsanov's 's theorem (see, for example, Domingo-Enrich et al. 2025
Theo-
rem 2 and eq. (138)) it holds that
E 12-1/2 ( u (X; |) dt (L R).
Ex-
The expectation in the left-hand side is taken with respect to the trajectory of X 0 < <<< T, con- ditioned on data Yi,..., Yn upper bound on the control estimation.
and Zz...., Hence, Theorem also establishes a high-probability
NUMERICAL EXPERIMENTS
This section is devoted numerical experiments illustrating the advantage of the Ornstein- Uhlenbeck dynamics (4) as the reference process. For this purpose. we adopt ideas behind the light Schrödinger bridge (LightSB) algorithm from Korotin et al. (2024), where the authors suggested a simple yet quite efficient approach for generative modelling. We do not pursue the goal of devel- oping statc-of-the-ae like to illustrate how proper choice of the base process improves data generation. Throughout this bt x (1 -
method, it goes beyond the scope of the present paper. However, we would
section, we use the notations of (1 - 2bt /2bb and mx)x) In what follows, we consider the process (4) with T and Σ eId. where
bt )m, 0 < < T.
0. The light Schrödinger bridge considers a reference class of log-potentials = {saa : E θ}. such that the functions ve (v)) = (*o(w) I1all ((2of), referred to adjusted potentials, are nothing but Gaussian mixtures:
K
(v( Σ ar p(y: So,
(19)
Here, for any h E (1....,) : cofS)) is the density of the Gaussian distribution with mean Tk and covariance sx. In this case, the parameter = {(a,) Tk. Sk) : < k & K) is just the set of triplets. The main advantage of such parametrization is that the Ornstein-UnI -uaieene- transform (z))))(<z) admits a closed-form expression. This significantly saves computational resources. At the same time, according to Korotin al. (2024 Theorem 3.4), choosing proper class of Gaussian mixtures, one can approximate quite complex potentials. With the described (pe, the corresponding coupling Se and the risk L(yp)) the (see and (T),, respectively) reduce to
90 (x. v) po())) e² TT mi(())((eo)) (o)) (v)/co(v))
(20)
8

---

Published as a conference paper at ICLR 2026
and [(po) Bz~pp log cZ(Z) By mpr log (Y(Y), where
K
co ca(x) = 2d (mmm())(soo vo()
Σ (x)
(21)
Ra
km 1
and
rhe (x) m1 (x))
||5k
mn
xx(2) = exp
+
(22)
200
The log- oog-potential estimate 0 constructed through minimization of the empirical risk
N
E(po) 1 N
Σ j=1
log co (o
1 log 4o (Y;)
Σ 1=1
over all admissible values of We call the described algorithm LightSB-O OU and provide its pseu- docode in Appendix F2 The standard LightSB algorithm (Korotin al. 2024) performs the same
1ly steps but replaces qı (y) æ)) with the Gaussian kernel (2TE) d/2e - VE.
corresponding the
transition density of the scaled Wiener process (3) with
Below we thoroughly evaluate our solver on setups with both synthetic (Section and real data sets (Sections 4.2 and (.3). The method demonstrates stable increase in quality compared to the original LightSB and shows itself well in comparison with other Schrödinger bridge solvers. The source code is available at G
4.1 EVALUATION ON THE GAUSSIAN MIXTURE
We start with evaluation of the modified LightSB algorithm a Gaussian mixture problem. For this experiment, we used mixtures of 25 Gaussians under three configurations: uniform grid with standard covariance matrices. random grid locations with standard covariance matrices, uniform grid with anisotropic random covariance matrices, with the sliced approximation of the Wasserstein distance W, as the metrie, see Appendix F2 for definition of metrics.
Table 2 summarizes the performance for K 30 potential approximation components. With sufficient number of potential approximation components, our approach demonstrates consistent improvements in both W1, MMD and Mode Coverage (see Appendix F.3 for detailed plots and Appendix for for hyperparameters values used).
Table 3: The quality of inter- mediate distribution restora- tion for the single-cell data
Table Comparison of the LightSB and LightSB-OU performance for K 30 across experiments with Gaussian mix- tures with 25 components
Solver
W;
OT-CFM Tong al. 2024a
0.790 0.068
SF M-Exact (Tng al. 2024b
0.793 0.066
LightSB-OU (ours)
0.812 0 0.020
LighSB (Korotin et al. 2024
0.823 0.017
Experiment
LighiSh
OUU
Reg. CNF (Funlay et
0.825 N/.
Sliced W
0.260 1 0.016
0.156 0.018
T.. Net (Tong al. 2020)
0.848 +: NJA
Standard
MMD
0.0001 0.0001
0.0003
0.0001
DSB De Bottoli ct all 2021)
0.862 =: 0.023
Mode Coverage
24.2 0.4
25.0 1
Slized W9
0.525 1 0.024
0214 0.028
I-CFM [1og Long cl 20245
0.872 0.087
Irregular
MMD
0.0024 + 0.0003
0.0004 0.0001
[SF (long et al. 2024b
0.879 0.148
Mode Coverage
21.8 0.4
25.0 1 0.0
NI.S (Koshizuka Salo 2024
0.970 N/A
Shoeed
(1225 0.017
[SE]-M-SinS (MMM-Smk(oose (Tong all 202AI
1.198
0.342
Anisotropic
MMD
0.0007 0.0001
0.0006 = 0.0002
SB-CFM ( (Tong et al. 2024g
1.221 0.380
Mode Coverage
2.5
DSBM (Shi ct al. 2023
1.775 0.429
"The authors not report the standard devi-
ation.
Translation

---

Published as a conference paper at ICLR 2026
4.2 EVALUATION ON THE SINGLE CELL DATA
To comprchensively evaluate LightSB-OU, we conducted experiments biological data. Following Tong et al. 2020), we considered the problem of transporting a cell distribution from time 4 to time 44+1 for i (1.2.3). The results of Tong al. (2024b) allow us compare the performance of LightSB-OU not only with LightSB but also with many other methods (see Table 3}. We predicted the cell distribution the intermediate time & and computed the Wasserstein distance W, between the predicted distribution and the ground truth. The results were averaged across three setups (i = repeated the experiment five times. Using the default
2,2, 3). To ensure statistical robustness, LightSB parameters from the original paper (Korotin et al. (2024), the baseline achieved W 0.823 + 0.017 and was outperformed by our OU-modified version (W1 0.810 + 0.020). The hyperparameter values of LightSB-Or are reported in Appendix F.6
4.3 UNPAIRED IMAGE-TO- IMAGE TRANSLATION
(a) LightSB: FID 24.1
(b) LightSB-OU: FID 24.0
Figure 2: Translation results from the latent space of ALAE using LightSB (left) and
Top: Adult to Child - skin color and face shape preserved. Middle: Male to Female glasses and face shape are preserved. Bottom: Male to Female skin color and face shape preserved.
To further illustrate the importance of the reference process for data generation, we consider task of unpaired image-to-image translation. Following the experimental setup of Korotin et al. (2024), we employ the ALAE autoencoder (Pidhorskyi et al. (020) and focus on the Adult-to-Chill translation task using the full 1024 x 1024 FFHQ dataset (Karras et al. 2019).
Below, we present comparative translation results for selected samples in Figure While the base- line LightSB successfully performs the translation, sometimes fails to preserve key attributes the source distribution such as skin tone, eye color, and facial structure leading noticeable deviations. In contrast, LightSB-OU achieves more accurate translations while better retaining the distinctive features of the original distribution. We emphasize that these examples universal superiority but rather to showcase specific strengths of our method in certain cases (the hyperparameter values are specified in Appendix F.6).
intended suggest
CONCLUSION
In the present paper, we established that the empirical risk minimizer defined in 122 satisfies a Bernstein-type bound (see Theorem []. This is a consequence of exponential ergodicity of the OU process and of some properties of the empirical risk (13) exhibiting a kind curvature when 5T Z log log The importance of the reference process was also illustrated with numerical experiments on artificial and world data. Further research may include an extension of Theorem | to the case when Pe has unbounded support and to heavy-tailed distributions pr. It is also of particular interest whether the result of Theorem Sicooremhola holds one considers a larger class of exponentially ergodic reference processes.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
discuss some insights behind the proof of Theorem I and then perform rigorous derivations in Appendix B.2 The proofs of all the auxiliary statements are collected in Appendix.
The provide the proof of our main result, Theorem [ Appendix Appendix B.I
we provide additional details for the numerical experiments described in Section In Appendix FF)
including the pseudocode of the LightSB-OU algorithm (see Algorithm IT. The code for the numerical experiments is available in the supplementary material and at G
We have not used large language models (LLMs) for retrieval, discovery or research ideation. All the mathematical derivations and numerical experiments were performed solely by the authors. We used LLMs just to polish writing.
ACKNOWLEDGEMENTS
The work supported by the grant for research centers in the field of Al provided by the Ministry of Economic Development of the Russian Federation in accordance with the agreement 100000C31199255400002 and the agreement with HSE University X1139-15-2025-009
REFERENCES
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. In The 25th ACM SIGKDD Interna- fional Conference on Knowledge Discovery & Data Mining, pp. 2623 2019.
Michael Albergo, Nicholas M. Boffi, and Eric Vanden-Eijndens Stochastic interpolants: A unifying framework for flows and diffusions. Journal of Machine Learning Research, 26(209):1-80, 2025.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
Giovanni Conforti, Alain Durmus, and Giacomo Greco. Quantitative contraction rates for Sinkhorn algorithm: beyond bounded costs and compact marginals. Preprint. ArXiv:2304.04451. 2024.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied Mathmmatccs and Optimization, 23(1):313-323 1991.
Valentin De Bortoli, James Thortoo,, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances Neural Information Processing Systems, 34:17695-177099 2021.
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen. Adjoint matching: Fine-tuning flow and diffuusin generative models with memoryless stochastic optimal control. In The Thirteenth International Conference Learning Representations, 2025.
Stephan Eckstein. Hilbert's projective metric for functions of bounded growth and exponential convergence of Sinkhorn's algorithm. Probability Theory and Related Fields, pp. 1-37, 2025.
Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. In International conference machine learning, pp. 3154-3164. PMLR, 2020.
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal Schrödinger bridge matching. In Proceedings of the 41st International Conference Machine Learning, volume 235 Proceedings of Machine Learning Research, pp. 17100-17122. PMLR, 2024a.
Nikita Gushchin, Daniil Selikhanovych, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Adversarial Schrödinger bridge matching. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Translation
11

---

Published as a conference paper at ICLR 2026
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 1840-6851. Curran Associates, Inc., 2020.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light Schrödinger bridge. The Twelfth International Conference on Learning Representations, 2024.
Takeshi Koshizuka and Issci Sato. Neural Lagrangian Schrödinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2023.
Guillaume Lecué and Charles Mitchell. Oracle inequalities for cross-validation type procedures. Electronic Journal of Statistics, 6:1803-1837, 2012
Christian Leonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems, 34(4):153311574 2014.
Yaron Lipman, Ricky Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Guan-Homg Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. SB amage-to-imagg Schrödinger bridge. In Fortieth International Conference Machine Learning, 2023.
Stanislav Pidhorskyi, Donald A. Adjeroh, and Gianfranco Doretto. Adversarial latent autoencoders. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 14104 14114, 2020.
Aram-Alexandre Pooladian and Jonathan Niles-Weed Plug-in estimation of Schrödinger bridges. SIAM Journal on Mathematics of Data Science, 7(3):1115-136,, 2025.
Nikita Puchkin, lurii Pustovalov, Yuri Sapronov, Denis Suchkov, Alexey Naumov, and Denis Be- lomestny. Sample complexity of Schrödinger potential estimation. Preprint. AtXiv:2506.0304.
2025.
George Rapakoulias, Ali Reza Pedram, Fengjiao Liu, Lingjiong Zhu, and Panagiotis Tsiotras. Go with the flow: Fast diffusion for gaussian mixture models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Philippe Rigollet and Jan-Christian Hütter. ArX1v:2310.19244, 2023.
High-dimensional
statistics.
Preprint.
Philippe Rigollet and Austin J. Stromme. On the sample complexity of entropic optimal transport. The Annals of Staristics, 53(1):61-99) 2025.
Yuyang Shi, Valentin De Bortoli, Andrew Camphell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems, volume 36, pp. 62183- 62223. Curran Associates, Inc., 2023.
Yang Song, Jascha Sohl-Dickstein Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021.
Zhicong Tang, Tiankai Hang, Shuyang Gu, Dong Chen, and Baining Guo. Simplified diffusion Schrödinger bridge. Preprint. AXXi:24433114233 2024.
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In Proceedings of the 37th International Conference Machine Learning, 2020.
12

---

Published as a conference paper at ICLR 2026
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing with minibatch optimal transport. Transactions on Machine Learning Research, 2024a.
generative models
Alexander Y. Tong. Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free Schrödinger bridges via score and flow matching. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 1279-1287. PMLR, 2024b.
Konstantin Yakovlev and Nikita Puchkin. Generalization error bound for denoising score matching under relaxed manifold assumption. In Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pp. 5824-5891, PMLR, 2025,
A ON ASSUMPTION 4
In Section 3 we suppose that the pogg-ootenial 6 associated with the optimal coupling ' meets As- sumption While the requirement 4p (xx)) M is mild, the lower bound (>(xx) > -1//--1/2(
m)l M needs a discussion. We are going to provide a sufficient condition for Assumption 4 to
hold. Let define Piax sup (x) +oo. view of (9). we have
aERd
pr(y) = f
qr x)) e
x (I. dr =
po(x) dr.
Tr Te ((x)
Rd
Rd
This yields that
(r x)) Tr lev0 ((x) Po da.
66 (1) = log (T()) - log
Rd
Using Lemma B.3 from Puchkin al., 2025) and taking into account yp (x) < Pmaxx we obtain that
- 1/21/2(x - m)ll?
+ (Pinax V vIvva eb
log
S
Tr Trler" [(x)
vd
for all x E supp(pu). Since, log (3 x)) S d for all y, Rdd, holds that
R
+ (Piax V viva e bT
(v) $ - (ogpr(y) + d
Vd
In the last line, we used Assumption 2 about the support of pa- Hence, we observe that, then .0t satisfies the inequality
-log $ ||5-1/2(y (1)]) + d,
(23)
R Va
+ (Finax v vIlva -bT or
4" (v) S 1/2 +d+ d+
and, as consequence, fulfils Assumption nonparametric Gaussian mixture model
The inequality
(3) is satisfied, for instance, for
/ - (ll)) 1/(2)) du
pr (T = (V2mro)-
[0.1]P
recently considered by Yakovlev & Puchkin 2025) in the context of generative diffusion models. Here 9: [0, 1 + d is smooth map defined on low-dimensional latent space. In (Yakovlev & Puchkin 2025), the authors discussed that such a model is quite general and suitable for description of high-dimensional data occupying vicinity of smooth low-dimensional manifold.
13

---

Published as a conference paper at ICLR 2026
B PROOF OF THEOREM 1
B. PROOF INSIGHTS
Finally, would like discuss some intuition behind the proof of Theorem 1 Rigorous derivations are deferred to Appendix B As we discussed in Section 3 Assumptions 3 and 4 us to use concentration inequalities for sub-exponential random variables and quantity the deviation of the empirical risks [((p). ip E & from its expectations. However, to get rates convergence possibly sharper than ((1/vn) we must verify Bernstein-type condition. That is, we have to show that both
Var(z* - -p(i)) and Var log Tr(e0((A)
TT TT(e [(2)
be controlled by subquadratic function of KL(#" TP). This is the most challenging part of the proof. We would like to note that in Puchkin et allll (2025). the authors used properties of log-densities with subquadratic growth for similar purpose (see their Lemmata C.2 and C.3). Unfortunately, we cannot apply them in our setup, because the joint distribution of Y;'s and Zj 's may not follow x². In other words, we have access samples from the marginals of πT but from x² itself. This makes us to use another strategy. Let us observe that for any 9 E +
KL(T" x)) B ((p(((()) (x(+)) + E log E [ex(x)) Xo]
where X is the endpoint the Schrödinger bridge process defined in (5). Due the tower rule, KL(T", x) can be considered as expectation of the conditional cumulant generating function
((A, Elog E (((())) (x(+)) Xo AE [p(XX)
(X+) Xo!
at 1, that is, KL(T"" xT) = Ex Go). In Appendix B.3, we prove the following general result relating variance and the cumulant generating function of sub-expee cobbeexpoential random variable.
Lemma B.1. Assume that a random variable & has finite 1b norm. Then, for any satisfying the inequality
52115 Egl
4c/A| 5 + log
IlE - -EEl < 1,
(24)
Var(E)
it holds that Var(g) & 4A -2 log EeA(E-RE)
Using this lemma, we show that (see Lemma
Expxpo Var((((77) 4 (x(7)) Xo) S x (K" TT) (Ivlog
KL (TT", TT) (T*,TT)
where x S. (LVL + d). However, this bound not enough, because Var(p" (Yi) (Yi)) is always not smaller than the conditional variance. Fortunately, we can use properties of the Ornstein- Uhlenbeck process to derive the following result.
Lemma B.2. Under the assumptions of Theorem
7 for any w > 0 there exists
log log log d d Ldf + (MM
To S
such that for any T > To and any E & satisfying the inequality Var Var(p(X()) - (0(()))) > it holds that
Var (E (x(xt) - x0 (X+) Xol) < Var (((()) - (X())
< (Varr(p(X)) -p" Xo).
The proof of Lemma B.2 is postponed to Appendix B.4 It immediately yields that if 7 satisfies the conditions of Theorem then either the variance sufficiently small, so that Var (p((X)) sp (X+)) & T(n,6), or
Var(p(X))) (*(x+) & xXL (T", (7) 1 (1vlog log KL (TT xT)
14

---

Published as a conference paper at ICLR 2026
desired. On the other hand, according Lemma C.4 it holds that
(21)
X Var((((X)) 40 (X?))) + x KL (T" (,",))))) (IV Volo log
Var log (e(e [(1]
KL ( TT)
This reveals a dichotomy. Under the assumptions of Theorem if for some E Φ the variance Var((>((X) (Y{)) is of order (n, then the corresponding difference [(p) [(p), is
of order ( (n, 6)) as well. Otherwise, we have Bernstein-type condition allowing us to derive faster rates of convergence using standard techniques from learning theory.
B.2 RIGOROUS PROOF
The proof of Theorem 1 consists of three major steps. On the first one, establish uniform centration bounds for the sample means
Tr fef (Zz) Tr @ ((2,)
(p((.) - (V)())) and
log
n
72
around their expectations. Then we show that the loss functions of interest satisfy a Bernstein-type condition. Finally, we convert the results of the first two steps into the final large deviation bound.
Step 1: uniform concentration. imply that (see Lemma
To start with, we would like to note that Assumptions and
1| (X+) ((x)))l (1V
L b
(M + d)-
On the other hand, according to Lemma D.2 under Assumptions and have
Tr (X log Tr [e(ee"] (Xx))
S M
for all y E Φ. This means that we can exploit properties of sub-exponential random variables (including concentration bounds) during analysis of the empirical risk minimizer C. In particular, Lemma E.I states that, under the assumptions of Theorem [ there exists
To S log d LR + Ld
+ M
such that for any E (0, 1) and any T > To we have
(o(Y) - -x((
x(x)))
Var (Yi) - (((i)) (D log(And) + log (1/))
1 L (M + ()))( (Dglo(()))))) + log(1/5)) log
n
and
n
Tr 0 (Z,) E log (Z1)) TrleP ((2;)
log
Tr(es [(21)
1=1
Tr (Zi)
(Dlog(And) + log(1 1/5)
Var
log
Tr (e (((1)
(M + LD (And) + log(1/0)) logn
15

---

Published as a conference paper at ICLR 2026
with probability at least (1 8) simultaneously for all (p Φ. We provide the proof of Lemma E.1 in Appendix E.1 relies on Bernstein's ' inequality for random variables with finite 4/1 -/r-norms (sec, for example, (Lecué & Mitchell| 2012, Proposition 5.2)) and the E-net argument. This approach is quite standard. similar strategy was used by |Puchkin et al. 2025).
(M (M+ + 1 (Iv ())
log log ((LV A)nd) log(1/6) log n
((n, = = DMM d)
for brevity, we can rewrite the uniform concentration bounds in the following form: for any å (0, 1), with probability at least (1 6) it holds that
(p(()) ((x)) - E(p((i)) - ,(Xi))
Var (x())) (M ((,
(25)
and
Tr Tr(e((Zz)) E Eog
Tr 1 (Z1)
log
Tr Trlee-" ((4,)
T e ((21)
Var log (21)
T(n,
+ (n,
(26)
Tr |r( (r(r))((Z)) (/b) + +d)
simultaneously for all y $
Step 2: variance bounds. show that both
The most challenging and technically involved part of the proof is to
and Var log (r(e [e
Tr Tr( [(21)
can be controlled by a sub- uubqquarratic function of KL(T", TT). This is the key ingredient derivation of rates of convergence possibly faster than (nn-1/2). We would like to note that Puchkin et al. 2025), the authors used properties of log-densities with subquadratic growth for similar purpose (see their Lemmata C.2 and C.3). Unfortunately. we cannot apply them in our setup, because we do not have access to samples from x², we are given only samples drawn from the marginals of n² This makes us to use another strategy. Let us observe that for any
KL(T" xx)== = =B(P ((())) + log [ex(x;)
Xo
(27)
Due to the tower rule, KL( TT) can considered as expectation of the conditional cumulant generating function
G(A, Xo) = Elog [(((())) Xo
AE [(( - 40 (X+)) Xo]
1, that is, KL(T", TTF Xx G(I, In Appendix B.3, we prove the following general result relating variance and the cumulant generating function of a -ubeexxonettaal random variable.
Lemma B.3 (restatement of Lemma B.1. Then, for any satisfying the inequality
Assume that random variable & has a finite 1 norm.
521 E
4c/A|
5 log
- -EEll, <
Var(E)
holds that
Var(E) < 12
log A(E-EE)
16

---

Published as a conference paper at ICLR 2026
Based on this result, in Appendix C.I prove Lemma C.1 claiming that
E X (ar((((7)) (x()) Xo)
iv + ( (M + d) KL (a TT) 1 V log (1 LMb)(M + d)
KL (T)
However, this bound is not enough, because Var( - (x()) is always not smaller than the conditional variance in general. Fortunately, cess (in particular, Lemma B.3 from Puchkin et al. (2025)) to derive the following result.
use properties of the Omstein-Uhlenbeel pro-
Lemma B.4 (restatement of Lemma there exists
B.2). Under the assumptions of Theorem 7 for any w >
1 v log L.R - log log
Ld
b + M) ))
To
such that for any T > To and any yp o satisfying the inequality Var(p)(X))) (xx))) > holds that
Var (E [x((7) (X++ Xo])
(x+)) & (x+))
Var
EVar(p((X)) 40 (X7) Xo)
2
The proof of Lemma B.2 is postponed to Appendix B.4 Taking w T(n, 1), we immediately obtain that there is
log d LRR + Lef b + M) M)
To S
log log w
such that for any T To and any vp E satisfying the inequality Var it holds that
> T(n 1)
Var (E [x(( - 4 (x+)) Xo)) <
- (xxxx)))>(((x)))))) (X++) & EVar
- x (X+) Xo).
In other words, if T satisfies the conditions of Theorem | then either the variance of yp ((x7)
xp* (Xt) is sufficiently small, so that
-40 < T(n, 1) & T(n, 8),
or it holds that
- -=-0==(()) < 2RVar((((77) xx(7)) 1Xo)
4Iv (M + d) KL (a xT)) LV
(1 V L/b)(M +
KL ;T")
Summing up these two bounds, we obtain that
+ (M + +d)
Var(y - --====())))) & (m, 8) + ( 5
(1 v LM +d)
KL (T",, ())),)))))))))
1 V
(28)
KL (TT,
On the other hand, according to Lemma
C.4 it holds that
Var log Tr [r(e)
Trles
I Var (((()))))))))))))))))) - 40" (X;))
(IV b
(M + d) KL (TT ,P) (1Vlog (1V L/b)(M +
(29)
KL (LL" (T,)
iv I. (M + KL (T"",, TP) (Iviog
(1 (L/D)(M +
S T(n, +
KL (TT", xP)
This reveals dichotomy. If is sufficiently large, then either the hight-hand- sides of 25 and 26 are of order T(n, 6) or we have that the loss functions of interest meet Bernstein-tyee conditions 28 and (29) allowing to derive faster rates of convergence.
17

---

Published as a conference paper at ICLR 2026
Step 3: generalization error bound. Bringing together (25)
(26
28
and 29
obtain that,
for any 6 E (0, 1), with probability at least (1 - 5) it holds that
(p((Y)) p(Xi)) E(((())) - (((ii)
(Iv log (1 V L/ + d)
KL (T* )
T(n, 6L,, KL KOAKKLL(T") (TT,
+ T(n.
and
Tr Tr( (Z,)
Tr (21))
log
-Elog log
Ti 1or (Zzz)
Tr (ev0 [(211
(1 V L/b)(M + +d)
T(n, 8) KL (T", TP)
1 log
T(n,
KL (TT TT)
simultaneously for all & E Φ. With these inequalities at hand, a generalization error bound is almost straightforward. For any tp Φ, let
KL(TT क = (+()) - ((V)) + 1 log
(Zl)) Tr Tr(ee"((z))
72 j=1
i=1
Then, due to
10 and the triangle inequality, it holds that
[(, 6) KL xT)) 1 V log (1 (1V L/b)(M + d)
KL (('"""))))))))))))) KL( x
(n,
KL (TT (T9)
with probability at least (1 5) simultaneously for all yp E $. Let us define
8p argmin KL(x", m²) and A = inf KL( n pe#
Note that, due the definition of the empirical risk minimizer, we have KL(m= क) KL TT" の < KL (m',*²) TP) for all 1 Φb
Then, with probability least (1 5), it holds that
KL (T",, m000) KL(T", R) A < KL( (L",, KL (T(=T" (m'; + KL T=",, xt²²)
( (1 L/W/ d)
$ IT(n, 6) + T(n, (T*,
KL (T*
+ T(m, V (1V + +d)
A
(1 L/b)(M + d)
< T(), + T(n, (T*,) 1V
A
V T(m, 8), (I (Vvlog (VVL/))(M + d)
A
In view of the inequality va+bs Va+ Vb. which holds for all non-negative and b, obtain that
log (1V L/b) d)
KL(T", .) - T(n, + y I(m, (KL (m+, A)
1 VVlog
A
+ M(n 0) 1 log (I L/b)(M d)
A
'If the minimum not attained, should consider minimizing sequence.
18

---

Published as a conference paper at ICLR 2026
on the same event. Solving the quadratic inequality with respect to KL(TT", T)) claim of the theorem: it holds that
A, deduce the
(1V L/b)(M (1VL + d)
KL(T" A
V T(mm0]
1 V
log
+ T(n,8)) VVllo (1V + +d)
A
on an event of probability at least (1 - 6).
B.3 PROOF OF LEMMA B.
Let note that it is enough check the theorem statement of 0. In the case X 0 one should introduce 5 = -f and consider its cumulant generating function. Let g(A) stand for the cumulant generating function of g, that is,
(A)) log Ee (E--EE -E)
Fix an arbitrary À 0 satisfying 24
For Borel function f R - I. let define
(C)( E (f({)e>(5-R))
IE =5)
We would like to note that Pf)( can be considered as an expectation with respect to a probability measure Px. Let
E (EE)) = inf { 0: Pxele E-BE/
BE/t
< 2}
stand for the corresponding Orlicz norm. Then, according to Lemma cumulant generating function g(A) can be expressed as follows:
E.5
the derivatives of the
(A())= = (E - EE), g"(A) = PA (5 PxE)² and gdd (A) = Px PxE)
straightforward to observe that (0) 0 and that g"(0) Var(E) < 4 Ef due to Lemma E.8 Let us elaborate on (A). For this purpose, let us introduce
(30)
52115 - gD(
log (3)
+ log
Applying Hölder's inequality, obtain that (A)) < Px - Px < (Px (E PPA))) 1-8
= ( (Px (PAIS P
(31)
Consider the second factor in the right-hand side. Due to the triangle inequality, it holds that
[ E - Pxc s/1+2ec)
PA E5 2 e/(1+2c)
+ PA E Eg|.
In view of Lemma E.8 and Lemma E.6 we have
2115 - IBE||
Px E - Eg| < 2|15 BEller(Ps) < 1 A/L RE
and
Px 15 1E/2+166 e/(1+22)
21/(1426)
3
||g (P
|| - EEll+
21(+))
3 ē
All |
19

---

Published as a conference paper at ICLR 2026
This yields that
PAx - Px5/-11/)) e/(1+26)
1
115 - Bgllo'
2/15 -EEl
< 2
3 - - - BEll
1 - 1-1115 EEE v/1
11E 15E11
2=/(1+2))))
20+t)/((+22) +3 +
2/15 1
< 2
5+
1
115 - B d/
1 - AL - 11
Since, according to the condition of the lemma, satisfies the inequality
1 || - -EEll A &
52lls - E5I
E
5 log
4e
g0(0)
4e'
we have
e/(1+z))
< 25/(1+2))
(5+ 1 (1
8
11E - EEE ab
4e
Combining this bound with (3.. obtain that
(5 5
1
1+2c
-1-2e
[X)l < 2
I|E - B 10
1+22 (o"(A)))
6
4e
= -(1-//(0e)) =6())) )
50
1 +
1 22
(1-
E -1
5
;)
1 E
||E EEC 11 (a"(
5e
4e
The inequality 8 < (5 + log
implies that
50
(1 [1E/(4))) < 52.
Moreover, taking into account the inequality
1 u E See u R, we obtain that
(1 (1-e/(4)))² )
50
1 5e
2x
(1 4e
< 52 2/5 (1-
1 4e
1
(5 + log 13)
< 22
and then
(X) E. 2. 52 (5 (5+2 ||5 TRE (o"(
1
This inequality yields that d da = (((( 2- + 5E)llE - 11
=
and, therefore, for any À satisfying 24 we have
(A)) > (d"(
2| 52 + 5c) - EEl1
5211 - BEl1
= ())
1 2/ + 5:) - -BE
(0)
52 115 IBS 123 ))
(d)
1 2/AllE EEllV
g²(
Taking into account
30 obtain that
5 - (0)
("(A) > (0) (1 2e|le - EEIl., (a
for all À fulfilling 244)
Hence, for any À such that
52115 E (0)
RElE
4e/2/ 5 log
- - < 1,
20

---

Published as a conference paper at ICLR 2026
it holds that
g'()) > g' (i
5+ BEl /s"(0)
1
g"
2 log(52
In the last line used the fact that (1 1xx/2) > 2- for all (0, 1/2) and that
52|lc B
5 + log
> 5 log 13 > 5.
g²(0)
Then, using Taylor's expansion with the Lagrange remainder term, we deduce that
(AA g'(0)1 for any X such that 4e/A|
4 log 52/lc Egll g²(0)
||E - REll at $ 1.
Hence, we proved that
4g(A)
4
Var(E) (0) &
log Ee AEE-6-28
X²
for any À such that
4e/A/ 5 log 52/15 - Eg
11E -BElla & 1.
Var(E)
B.4 PROOF OF LEMMA B.2
Let us fix an arbitrary
0 and consider a log-potential yp E & satisfying the inequality
Var (x(()) 4 (X7)) > W. Also, let T z log d)/b b be such that K(T) < /2, where the
function IC(t) defined in (39). The proof of Lemma B.2 relies on Lemma C.2 which provides
an upper bound on the variance of E [yp( (p(()) (X;, Xo]- Due to the mixing of the diffusion
process X", it not surprising that Var (E [x(())) - -p" (Xt.) Xo]) tends to zero as grows. To
be more precise, Lemma C.2 claims that
Var (IE [x(()) - x0 (x7) Xoo)
de--2bT (IM2 v (p()))((()))) (XG()) 1/12(2)) +
(/1()11))))))
.exp (oovv [ (VVev- log M LR (Ld (Ld /b)
(M2vv)dee 20T (w- (W-1+1/X2(T) wf 242 Var( - (X7))
-epp vop(oVa (VQ log (M LR (Ld /b)
where M is given by
32
Let us note that the definition of M implies that
M LT- L' (xrXxx))(())) L 1/2 -m)
m)]" +4 S Ld LR? + M.
The last inequality yields that Var (E [x(+)) - 40 (X+) |x,])
I (M 1)de- 20T (L
1-1/K(²TT)
-20
Var()((;7- (x++)
w2
exp fo (VQ(MeIVe WT log LR (Ld (LdP16))))
Since, due to (39)
log d
log (C(T) = 2e² V arcsin(e-")) 552 552 Vdlool (1 6 c-2T)) S Vdc-bt for all >
b
21

---

Published as a conference paper at ICLR 2026
obtain that
Var (E [
- (X+) Xo])
x (M2 V 1) de Var((s( 40" (Xi)) exp {0 ( (Vde (T log (1v )}
exp explo (o(VV(MVV))
M LRZ V (LdP/b))
Let To > 0 be such that
(M² v 1) de exp for Vide-ir log (Iv +w
exp fo (vā (MVIVe- log M LR v (Ld /6) (Ldd/)))))} < 1/2
Note that one can take To satisfying the inequality
[og log v log d LR? + Ld² M )
To S b
b
Then, for any T > To. it holds that
Var (E [x((7) - (x(())) Xol) < Var(p((7)) - 10' (X7)).
Since
Var -xxx))(-((xx))) EVar (V()))(()7)))))))))))))))))))))) 80 (x7) Xo) + Var (E (x(7)) 40' xx77) Xoo)
this yields that
Var (E (x -p (X+) Xxo) } Var(+(X7)) (x7)) < EVar(<(XF) (x77)) Xx).
C LOG-POTENTIAL VARIANCE AND CONDITIONAL VARIANCE BOUNDS
In this section, we elaborate on variance bounds. The presented results is key argument in deriva- tion of the rates of convergence, which possibly faster than O(n 1/2). We use Lemma C.2 in the proof of Lemma B.2 and then combine the results of Lemma B.2 and Lemma C.1 to derive an important intermediate result in the proof of Theorem
[ the inequality 28
Lemma C.I. Grant Assumptions
2 and
4
Then it holds that
E Ex Var()((7)) 4" (x()) Xo)
S KL (TT", x))
1
(d+ M R2c-20T KL (L(xx,, 77)
(LV + M R8e- 25T)
log
1+
KL (T))
The hidden constants behind S are absolute.
Lemma C.2. Grant Assumptions 1 and
4 Let T be large enough in a sense that K² & 2
with K(t) defined in (39). Then it holds that
Var (T. [<(XG) 40" (X+) Xol)
I de -2TT (IM2 v 1)Var (p(x)))())))-))))))))))))()))) (X.))) 112010) + Var(p(X)) (x))) (/1
exp Vā (MVVVVVVo M (M
MVVLRR LR (Ldd/bb)) e }
where K(tt) is defined in (39..
M = LTN + L //// ( - mil' + +M,
(32)
and the hidden constants behind S and O()) are absolute.
22

---

Published as a conference paper at ICLR 2026
The proofs of Lemma C.1 and Lemma C.2 are moved to Appendices C.1 and C.2 respectively.
While Lemma C.I is an almost immediate consequence of Lemma B. and the representation 27
the proof of Lemma C.2 relies on the following property of the Ornstein- -hitein-Uhheebcck operator.
Lemma C.3. Let s¹ Rd - I be such that vp(x)
< M for all x E Idd and T
2 1. Let
f: Rd + R be an arbitrary function such that
fx) < A A1|=-1/20
1/2
m)]l +B for all x R
for some non-negative constants A and B. Then, under Assumption 2 for any T > 0, it holds that
(Tr(fe*](Xo))
Tax [S= 2/10(1)
Var
F2 max
FmaxxTx Tooles")
2M 1) + (A((T) (C(T) + x(T)
exp
11) (((T) K(T) )²
2
M
+ +2(A(T))² ? (X(T) + K(TT)
F2 +
1
)²
2
Tool [ool Fmax Too (eo(e)))
2/1(1)-22(11)
max
K(T) K(TT)
6²
1 } 2M (K(T) - 1) + 2A(T) (K(T) + I((T)
exp
where
2
F Be
R² 2Ac M
(400d d²),
b
the function K((t) is defined in 39 for any t > 0
bee R22
A(1)=
+ 46 Vd Vd
arcsin(e (t) 10e² vee2Voloo log (1 11-e" 20d).
(33)
Vd
We provide the proof of Lemma in Appendix following counterpart of Lemma C.I for log (Tr(e?) (Xo)/
Finally, in Appendix C.4 we prove the
/TT (T(e (el [(Xo)).
Lemma C.4. Let us fix an arbitrary up E & and denote
[(Xo) - E [x((})) - (x+)) Xo] [e00 [(Xo)
<= log Tr(
Then, for any T > O, holds that
(Xo))
Var
log Tr [r(e ((Xo) (XXo)
< 2ar (RE [x((7)) - x0 (xx)) Xx])
+ 3/2 KL(T"" x +
(1 (1+5 log
2e 2e
KL( TTP
Similarly to Lemmata C.I and C.2
we apply Lemma C.4 on the second step of the proof of Theorem
I to derive the inequality (29
C.I PROOF OF LEMMA C.I
Due to the definition of the KL-divergence, we have
" (20, xT) "" (wr, IT) drader log
KL (m²,T²)-
T (0(00 vT)
Rd
/ (xT) - (((()(())))))))))))))) ((((T)))(((TT)) dxy +
Tr (r(e:((20)
log T [60 ((ro) po (200 dao
Rd
Ra
Expopr (x+))) :((+)) + Exo log E [xx(x;)- xo]
23

---

Published as a conference paper at ICLR 2026
Introducing the conditional cumulant generating function
Go), log E exp -
-p (Xx))) Xo zо
AE [xx(X)) (x+) Xo mo]
observe that
KL (x", ,T)) = Exo~po Go).
For any zo, let us denote
(zx) - B|x(X7) (x7) Xg =o po ((
where 7² (=) is the conditional distribution of X3 given Xo In. and
1
52 (I.)) 5 log Var(+((;)
A(zo)=
(20)
40 (x+) No xo
Then, according to Lemma B.1
4G
Var(p(Xi) - (X(:) Xo x20) <
for any zo and any |\| & A(zo).
X
If (zro) & 1, then
Var(p((F7) (x7())) No = mo) < 4G
Otherwise, we have
Var(p((7)) - 10" (Xt.) Xo zo) 4G(A(z)). Fo)
4G zo)
A2
A(zo)
Here used the fact that, for any zo. the map À + G (x,, zo)/A / is non-decreasing on R.. Thus, we obtain that
1)
Var
-x0 (XF Xo = xO) < IV A(zo)
((1,
or, equivalently,
13e 132 (200
(1v 1
G(1,
Var( - 4r (X7) Xo = x0)
A(za)
Var(+p(XT) - x0" (Xr) No = x0)
( (e log
526 (20)
Var
- 40 (X7:) Xo = =xo)
Then Lemma E.9 implies that
(T20)(((0))
(e log 13c5))2((0))
>
Var(s((;;) 40" (X+) No x00)
G(1,20)
G(1,
Hence, we have
136 (x0)
-(((((())-'"<pp)))) = xo) < 20)
(VVe log
G(1,
4e+) log (1+ 13c V 2 (xo)
< 4G(1,
G(I,
Applying Lemma D.3 we obtain that
(xx) $ 1 1 ) (d+ M + 26T
|| //-1/2(00 - mnl²²) (mo=-
for any Do E supp(po).
24

---

Published as a conference paper at ICLR 2026
In view of Assumption we have
(1ro) (Iv!) dd M -
sup z Euppp((p)
This yields that
Exo=po Var(p (x(x)) Xxo)
S E Ex G(1, Xo) + (Iv
(d+ M + R2e -2T)
(1V L/b)(d - M + (PPe-81))
((),
1+
Go),
Let us note that the map
L/b + M 18 - 26T
u+ u log 1
u
is concave R.. Then, due to Jensen's 's inequality, obtain that
Po Var (ar(p(()))
(x7))) (Xo)
LV L 5.
X Ex Xo),
(d (d+ RR---226 Exorpo Xo)
1 (1 L/b)(dd M (20T)
.log
Ex Go)
KK ++) + (Iv
+) (d M R2e -25T) KL (TT,, x))
-log 1+ (1V L/b)(d + M 1-2--221
KL "
C.2 PROOF OF LEMMA
Let us introduce
((x) p(x) Exp())) and (x) = (x) Epp (X7).
Both (xx) and nr (z) exhibit sub-quadratic growth. Indeed, due to Assumption 4 holds that
-M < - Ex(x)) LE 1 m)| + +M
L + L M 1/2 (EX; - -m)] + M
and, similarly,
-M < (Xpp((()))) < LE LE||--11/ 1-1 - -m)l + .M
-T1 (ar((--Va(XXx))) + +1|12 11/2-11/(EXX; (EX; m)ll M.
Recalling that (see 32
M = LTT( Var
(TT- (x+)) + (,))))+(+++))>))))) 1/2 - -m)l' 4M
and taking into account Assumption 4 we obtain that
[m(z) x (x)) < < x- m) + +M.
25

---

Published as a conference paper at ICLR 2026
At the same time, it is straightforward to observe that
Var (B (x(x) (x Xol) Var (B -
Xol)
and
Var(p(X;)) 4- (X7) Var (ar()(())) - -(2
For this reason, is enough to check that
Var (E [Ø((7) - (X+) Xo])
I d2- (IM2 v (p(((((()))))))))))))))))))))))) - 10" (X7))) 1/1=((T)
(xx)) (1))))))(000))
exp {o (oaa (va (M (MVVV
M V LR (Ld2 /b1)) DT
Let us note that the conditional density of X given Xo In is equal to "t ><(rr| (20)) = ex (or) x0 (IT Fo)
Tr le (300)
This yields that
E
( -x-7) Xo]
1
(0(0T) - -xT() (xr)) ( (o) der
Tr Trree" (300)
12
Tr (+5 "*))+"" (Xo) (Xo)
Tr [er"" (Xp)
Let us introduce
(6 - Φ = V (G -6- )) and (--))=-v((--
Then we can represent (ō "0*) as difference of two non neg-tieet functions:
/ ~" (3 )*)+ (5 )")
As a consequence, we have
Tr (8 - 3-) Le0" Tr EeP") (Xo)
(Xo)
Tr [ --)__e"] (Xo)
E [+(;;) G (x+) X1]
Tr Te("" (Xo)
Applying Young's inequality, we obtain that
Tr Tr --=+ [Xoo)
Var E [x(7]) - - (x+) [xo)) < Var
Tr Tr(EP" (Xoo
Trr (6p ) (xxo)
+ 2ar
Tr [ew0" (Xo)
An upper bound on the terms in the right-hand side easily follows from Lemma Indeed, intro- ducing
M
Pmax Me 2Le M n2
(400d + df)
b
and applying Lemma C.3 we obtain that
Tool(p + 2/13(T)
Tri(ā- - (*) ((Xo))
Var
< Puass
Tr Tr(er"](X [Xx)
Pmax Too Too(et0")
exp (M(T)) 1) + (A(T) (KIT) + x
M2 (X(T)) K(T)) )² 2
+2(AA())² (KIT+ 1 2
x(M) -07)+((7)
}
1 )
Tooll - (0) + Fmax Toc (EF
2Kx()))
K(TT) K((T)
(M((((()) -) 2A(T) (K(T)) <((T) }
}
exp
6²
26

---

Published as a conference paper at ICLR 2026
where A(t) and (t) are defined in (33) and 39). respectively. Similarly, we have
TrilE )-)) 1(Xo)
Var
Tr (xXo)
T )
2/X((T)
2M - +2((T) (K(T) + x(()) +X(T)) }
< Pmax
Pmax T0 (ESP))
exp
1 2
M
(((T) K(TT)
+ ((A(T)))
(K(T)+ + K(TT)
P
(Tel1(S _0-)
4/1((T)-2X(7)
+
K(T) - K((T)
e2
Xmax Too [EFF)
(K(TT) 1) + 2A(T) (X(IT) + K((T)
exp
This yields that
Var (E [x(}7) - 55 (X7) Xxo)
Taol - 10*1
2/K((T
(K(T) x(T)
2pmax
exp (M(((()) 11 + 2A(T)
Pmax Tac [etp"]
2
M
(((T)- K((T)
1
+ +2(A()))) (KIT)- K(TT)
2(X(T))-X((T)
2ppmmxx (K(T)- x(T))
1
+
ProxxToc Pmax Too [eF
e2
exp (M(()( -) + 2A(T)
1 (K(TT) + K((T)
1
Here we used the fact that both Tox 1( - -=*)+ and Tax [(F - _=""= do not exceed T 1| 6 "e*"] and that 4/K((T) 2K((T) > 0 due to the conditions of the lemma. On the other hand, it holds that
Var (x)()))))))))))))))))))))))))( 40 (X?)) Var(p(X))- -((A)) (X-))²
Due to the tower rule, we have
EE -xx))) X.] =E Tri - 3 (((o) Tx Tree [(Xo)
Applying Lemma E.10)ana and using the bound
(7(z) (²(()) (1.11=-1/218x3 (100> -m)/ m)]l
) + .M
< 2T. Σ (Ex; -mil' +227
we obtain that
Var(p((77)) - 40 (X7))
Tri 3- ) 2 ((Xo)
Tr Fr 1(xo)
Tac 10 - (0 K(T)
Tooleit" [ep0
E (H(X)))² 1-X(TT)
exp -M (1- K(T))
A(Xo,, T) (K(TT) K(TT)
}
where we introduced
32I22
H(x) = - + M
1-1/2(
m)
(2560000d + +dd).
+ b²
27

---

Published as a conference paper at ICLR 2026
Let
32L2
H maxx EM + +16L2eMR" R*
(2560000d + d+t)
6²
and note that
H(z) Hmax and A(r,t) & A(t)
for all x E supp(po).
This yields that
X(T) (T<lll0- (ol1(-00)PPEt P K(T)
- Var -xx (X77)
Too [et')
exp { -M (M-
(1-x(7)
1
A(T) (C(T) + 1
x((T)+<())
Applying the Cauchy-Schwee Schharr-- inequality, we obtain that
T ("
2K(T))
Too Too 2+ A(T)
Tao [et0)]
Too [et0]
((((T)-1 Var Var(p(())) - -x7())
max
M (1 x(T)) + + (K(T)+ + x(T)) }
exp
This allows us to conclude that
Var
E [x() - ~ (X+)
2-2/((() maxx
(22/K())21/X2(1) (x)())(X))))))))))))))))))))))) "0" (X;:)) 1/12(0)
exp M 2+ K3(TT)
(1) + A(T) (2 x>>)))
-1)
(K(T) + k(T)) +
1
)
1 +2((((T)) (K(T)- + K(T)
1
2
K(T)-
K(TT)
2xpimxx
24/2(1)--/K(11)
H 11/11)-1
6²
(KITT)
1
Var(o(XF) (x(x))) 2/12
K(T)
exp { (M 2 (T)
2
x(T)) (KC(T) 1) + A(T)
1
[+
1 (+x2() (K(ET) + x(T)
It remains to simplify the expression in the right-hand side relying on the definitions of Pmax, Hmaxs A(I), and KC(I). First, note that
max (log Pmax, log Hmax} S M + log (M V LR v (Ld2 (LdP/b))
Second, the relations
R Vdl
0 < K(TT) 1 S Vde- -T and 0 A(T)= S
-VT
(id
yield that
Var (E ((}
(X+) Xol) S
Oivde "T) "T) Holv DT
exp (0 ( (vd(M 1)e 1T)
Fma
(d(M/ V 1)e 20T) Var(p(X7)) - (x7))) 1/K2(T)
+ OlvTe- -r Holvde
exp {0 (Vd(M (()
max
de -20T Var(p(X))) (x+)) (T--
28

---

Published as a conference paper at ICLR 2026
Applying the inequality
& Olivee b)) HOlvie exp (V((
max
< exp (o (VV[Mvv (M LRL (Ldd((b))) ( o-
finally obtain that
Var E - - Xol)
S de de- (IM2 V VI)Var(p(XX)) (x((() (X+)) + Var - (x(+) (x+3))2000))))
exp fo (Vd (V(VV(MVIVV) (MVIVIOs (M LR (Ld (LdP/bi)) e )
The proof is finished.
C.3 PROOF OF LEMMA C.3
Let us define
F(x) = Be M + 2A ||=-//(00 - m)]² + b
2A
(400d + (())
Then, due to Lemma E.10 for any Rdd, it holds that
Fr[se("]xx) Tz(er"](2)
< (F(x)) (F(x)))-1/<<7) ( (T=l5ee7)) 1/10(1)
Too [EFt]
exp { ((C(T) - + A(xx T) (C(T) + x(T)
and
Trifer" Tr|ep" (xx)
(F(x)) (F(x))1-x()) 1-K(T)
Toolfe K(TT)
Too (evp
exp { -M (1 (1- (-xxm)
1
A(I,T) (K(T) xim) )
where the functions A(x, 1) and K.(I) are defined in tion 2 for any x E supp(po), it holds that
38 and
39
respectively. In view of Assump-
A(z, 8) < A(t) and F(x) F Faax.
This yields that
Trifet (Xxp) & (Fmax)) Tr (eve ((x))
(Tol 2-1/X(1)))))))) Too (es.)
1/K(T)
(M(TTT - 1) + A(T)
(X(IT= + 1
+x(T)
exp
and
Tr(fe*") (x(o) Tr(e"" ((Xo)
Tool K(T)
> (F maxx
1-X(()
Too [eS-]]
-M (1
(1-x(T))
A(T) (K(T) )
exp
K((T)
29

---

Published as a conference paper at ICLR 2026
almost surely. Then the variance of Tr Iff ((Xo))/T ((Xo) does not exceed
Tr(fe ((Xo) Trlee" ((Xo)
Too I 1/K((T
1-1/K(()) (Fmax)1-1/(())
Var
(Fmax)
Too [e8-]
exp
(M((2) -1) + A(T)
K(T)) x(())
(Fmaxx)
1-x(T)
(To
k(())
Too [eF]
M M1-1-11 (<(m)) (1-(((()) A(T) (K(T)+ + ((t)t(t)) )²)
exp {M
1
Here we used the fact that if a random variable & takes values in aa. b] almost surely, then its variance not greater than (b - aa2 2//4 Simplifying the expression in the right-hand side, we obtain that
Trifet (XXo) Trlep") ((Xo)
Var
2/K((T)
F2 Faxx
exp (K(T) 1) + 2A(T) (KCTT) + )
4
Fiaxx Toe [etF*]
(1 (1-exp exp {
(C(T) - 111
Fmax T Toolet0")
-A((T) (K(T) + x(T)) )}) 2
M + log
Toa Tealser"")
K(T))
Taking into account that 0 & 1 - - & v for all v > and the inequalities
Talfee & Fmax: Txle 3 1,
and A(T) > 0,
we deduce that
Tr( (T(()e(((Xo))
Var
Tr(e""
F2 F
[Set
111 1 ) M + log Fmax FmaxTTolef Taxle91) + +2A(T) (KIT) + x(T))
2/K(TT)
exp 2M(((()) 1) + +2A(T) (KIT) + IC(T)
4
Frmax Too (e100)
(KIT) - K(TT)
Txclfer
The Cauchy-Schweer inequality ba+ 2)2& & 402 +41 + 2c implies that
Tr(ffe" ((Xo)
Var
TT (XXo)
Too Toet)
2/K(TT)
2M 2M (K((T) - -1) + (K(T) }
F2 max
F Fmaxx Too[es")
exp
1 Finax ( 1) +2( (C(T)+ K(T)) "
(((T)) K(TT)
2 + log"
Toolfer']
Finally, taking into account that
FunxTToo Toolet")
>1
Too [fet"
and using the inequality
log u
1
<
for all u >> 1,
u
e
we obtain that
o (K(T)
log
F Fiax Txo
<
1
Fiax Too Telee") K(T)-1/1((T)
K(T)
Too fep"
é
Txclfer'
30

---

Published as a conference paper at ICLR 2026
Hence, it holds that
Var [TrSSeP") ((Xo) Trler (Xx)
Tox [[
2/K(TT)
F max
F Fmux Too (es0*)
2M(((()) -1) + 2A((T) (<(T) + 1 k(n))
exp
2 (K(T)- -K(T)
1 )² +2(A(T))) (K(T))+ )(m)
M
F2
x(T)) ) 2
(
Tx [
4/((7)-2((2))
(K(T) -
+ e²
Fiax Too (e0(e
(M(()( - 1) + 2A(T) (K(T) + I((T) }
+x(7))
exp
C.4
PROOF OF LEMMA C.4
Due to Young's s inequality, have
Tr(e (Xoo)
Var log
Var (5+ E [x( 40 (XF) xxo)
Ty [e0" 1(Xo)
2Var(() 2Va (IE [x((7) x0" =x((F)) xx])
Thus, it remains to prove that
1 2 log 2e|1511vr
Var(<) < 2/ KL(T",L""' "orr",PP)le0 115111 1+
KL TT)
For this purpose, let recall that
(xr(
x (xT x0)
π (xT [x0)
Tr jes" [(xo]
is the density of X2 given Xo To. This allows to rewrite log (Xo))=((X))TTrlee" ((xo) in the following form:
Tr (Xo) log (x(x)-00 Xo
log
Tr [e (Xo))
This yields that ( > 0 almost surely. Indeed, according to Jensen's inequality, it holds that
= log (x( (X;) X E (x(x+) - =0" (x;) Xo] log e F((((;)) (x+) Xo) E [x((X) 10" (x:) Xo| 0 almost surely.
Furthermore, is it straightforward to check that
EÇ = Elog Tr (Xx)) B((((7) - (x7)) = KL(T"K T.). Ty ee ((x)
Let us take
=1(108
[og
2e 2e KL(T** (TT"",, [TT)
E E=
We would like to note that, according to Lemma E.8 KL(TT, TT) = ECC < 21/s This implies that e (0, 1). Then, applying Hölder's inequality, we obtain that
Var(ç) < E = -E(C =, (1+) < (E<)²- (e)/+)) 8 = KL(T", (L(T",TT)1-- (BC())))))*
31

---

Published as a conference paper at ICLR 2026
Let us elaborate Eç(1+e)/5, Using Lemma E.8 again, we observe that the expectation of ((++))/o does not exceed
(1+0)/=
1+1/e
B [((1-
< (2
2
||s || (2te)/e = 22+1/6
(1+ 2
1-1/
E
This yields that
(Be(a+ove)
< 22+e
1
2e
1 1-8
Ik
1
2E )
1 1)² . (1 (1
2E
1
(41kslle)) )
Since + a E e" for all R. we obtain that
(Bc(Lto)lE)
Ve
1+
1
(411s1le)) 1+5
2
2s
Hence, it holds that
Var(S) < 2
Ve -9)1-- ,L(TT,=7))--
KL(T",
(BC
=
1+
(4115110)) 1
2e
=2VeKL(T" x 1llllaa (1
1
4/15|1e,
25
KL( TT))
Since
1
2E|1S11os > log 411511
log KL(TT" ,
8
[_TTTT,7)))
we deduce that
4/|511 KL(T", TTF
Se.
This yields that
Var(S) & 2/ KL(x", ( llslll, (1 1 log KL TT)
2el1S1ler
and we finally obtain that
log (Xoo) [r(et") [elee"]][[" (Xxn)
Var
< 2Va (E [x((}]) - 30 (x+) xx])
+ 3/2 KL(T" x 111111 (1 (1+ 2 log 2e|1511
KL( " TT)
D ORLICZ NORM BOUNDS
In this section, we present upper bounds on the 1 -oom of the log- log-potential difference (Xx)) sp" (Xi.)) well as of log (Xo)) [T""""""") (Tr"e""(((Xx)) [(Xo)). This is starting point in our analysis of the empirical risk minimizer & defined in (12). These results play a crucial role in the proof of Lemma
about uniform concentration of the empirical means
Tr Tr(e:p10 (Z,) log TT @ [(2g)
(Y.) x and
j=1
around their expectation as they allow to use Bernstein's incquality for sub-exponential random variables (Lecué & Mitchell 2012, Proposition 5.2).
Lemma D.1. Under Assumptions
2 and 4 for any T Z 1/b it holds that
(x+)( - ((*7)), $ (1v ( (M+d+ d de ub
R²
+ Vd ++V v(++br))
Vd
32

---

Published as a conference paper at ICLR 2026
Lemma D.2. Assume that log-potentials yp + and y" satisfy Assumption 4 Pu. and Po fulfils Assumption Then, for any T 0, it holds that
Suppose that Xo
(Xo))
(be" It arcsin(e bT)
2 log llggK(1) K(T)
+ K(T)
MK(())
log TT Tr| ((xo)
+ log 2
Vdlog2
log2
KC(T)
where K(T)= is defined in (39
The proofs of Lemmata D.I and D.2 are deferred to Appendices
D.1 and D.2 respectively. We
also provide an upper bound on the Orlicz norm log [ (X Zo x (X7 mo)) with respect to the conditional density (- (o)). Unlike Lemmata D.1 and D.2 we use Lemma D.3 in ver-
ification of Bernstein-type condition for the empirical KL-loss. To be more precise, we com- bine this result with Lemma B.I to obtain an upper bound on the expected conditional variance
EVar (Ø² (X+) ((X)) Xo) (see Lemma C.I for the details).
Lemma D.3. Grant Assumption and suppose that the log-potential s² corresponding to the opti- mal coupling satisfies
Tooip" and ("(x)) & M for all Rd.
Let us take arbitrary y satisfying Assumption Then, for any zu supp(pu). it holds that
log x (X² xo)
()) (a+ ee 20T Σ E-1/2(x0 m) m)]l")
1V (Vv))
T 20)
(- (-1
where S stands for inequality up to an absolute multiplicative constant.
The proof Lemma D.3 is postponed to Appendix D33
D.1 PROOF OF LEMMA
+(x+) 11 S (Iv!) (M (M+d+e-t ( Vd +aa+va+bR? bR
[lx)((;)
Let us note that Assumption Hensures that
(p(})) (*+))) < Σ (X m) 1| M.
For this reason, the Orlicz norm of ((x))) 40 (X7.) does exceed
1112-112 mi]²
(1x(X;) - 40" (xx))l
M
n/
L 1 - m2) mm)l/|
M
+
log 2
In the last inequality, we used the fact that ||e||s < |el/ /og 2 for any constant R. The rest of the proof is devoted to the study the Orlicz norm ||-1/2 - m) According to Lemma E.3 for any Zo supp(pu), it holds that
E bI (2(77-m3))
Xo 20] < 24/2 exp {MK(TT + T)K((T)
Σ E-1/2 m) 11]]
2bc 267
exp
In view of Assumption
and the definition of A(x, t) (see (88).. have
be R²
A(so, T) &
4e2 Vd
aresin(e (T) 10e² Villog (1 c- -26T)
Vd
and
3/ 1/2 1/2 - m)) < RR2
33

---

Published as a conference paper at ICLR 2026
for any zo supp(po). Thus, due to the tower rule, it holds that
Ee (21
EE - -1 -m)l
Xo
2d/2
(1 267) 10c2vá
( R²
MK((T) + K(T)
+ +e22 va)
arcsin(e-bT) + 26e-267
26T R2
exp
vi
Let us denote the expression in the right-hand side by and take
1 logg B
v v.++bbe) br²) ( (va +va
b b M + +d++ vr
Then, applying the Lyapunov inequality, we obtain that
Ee BellE (Xp-=1))////// Hence, the 2 of /120 m)l does not exceed u, and we conclude that
< EE (161)-11(((--) 1226 (xx=-m)])))
1/(be)
<
B/L(vol B) < 2.
[|p(X)) ('(x))l1 L ||| 1/-
m) m)|l|
M log 2
41
M
< Lo + log 2
V 1)
M d e
(R² R2
Vd vVd+ bEF v(++bRe)
(Vd
D.2 PROOF OF LEMMA D.2
According to Lemma E.10 for any x E Rdd, holds that
Toc [e
[x) eM
A(x,T)
Too l
-Axx,T)K(( + KC() log
< log
log e.
M
K(TT)
K((T)
and, similarly,
-Axx, ()k(T)))())))) KK(T) log Too [eo M
Tx(e
A(x, T) K(T)
1 K((T)
Too Tooles") eM
<
log
+
log
eM
where the functions A(x, 1) and K((1) are defined in 38 and the inequalities
39 respectively. This, together with
M eo > Tax and M Tooe
> TTup"
=1,
implies that
(Xx)) Tr [r(e"" (XXo)
< A((o. T) K(T) + +X(T) MK(T).
log
Thus, the Orlicz norm of Tr(e*)(Xo)
(Xo)/TT(e* 1(Xo) satisfies the inequality
log (Xoo)
< <|(XX ( (xITT + +x(27) + MMKCE
Trlee" [(Xo)
(Kim) + +X(T)) |IA(Xo, T) + K(T))|MM)
1
<
K(T)
34

---

Published as a conference paper at ICLR 2026
Let us note that any constant R has finite 4h1 norm cc/logg2 This helps to deduce that
(Xo) Tr|e [Xx))
(x(2) k(T) IA((
MMC(T)
log
log 2
Let us elaborate the first term in the right-hand side. Taking the definitions of A(x, and K((t)
into account (see (38) and (9)). one can observe that
be
A(Xo, T)
M (Xo m)) arcsin(e (T) + log K(TT).
In view of Assumption 2 the first term in the right-hand side does not exceed
bea
Σ E-1/2 - m)]l arcsin(e-
be arcsin(e bT) vd
almost surely.
Vd
This yields that ||A(XX, 1)
be? R² arcsin(e "T)
log X(T)
be? R² arcsin(e bT)
2 log K(T) log 2
+
Vd
valog
and, finally,
log [x))) (elee=((Xo) 11
be R2 arcsin(e bT)
2 log llgg((T)) log 2
( (((T) + -(())
MK(T)
vdlog?
log 2
D.3 PROOF OF LEMMA D.3
Let us take À = min(1/2, 6/(2))) and consider the conditional exponential moment
ा' (X To) x (X2 xo))
E exp
log
Xo xo
According the Cauchy-Schward inequality, it holds that
x
mn) T
x xX77 x0)
exp
log
Xo xo
< B
exp
21
log
Xo 30
Xi オタ
TP (Xū xo
Let us note that the expression the right-han hand side does not exceed
π (X;
Eu)) (X7
+ E exp 2X TP (Xp 1o
E exp
2X
log
Xo zo
Xo =
xo)
(X² xo)
Using Hölder's 's inequality and taking into account that A < 1/2, we obtain that the latter term satisfies the bound
(X7 2o) n (X7. xo) }
T (Xi mn)
22
exp
2A
Xo = 20
< E
Xo = =20
n (X² xo)
2A
mp (xT xx)) x rxrr vb)
n (ET xO) dor
21
T (xT 2o) det
=1.
Thus, we proved that
n" (XG xo)
E
exp
>
log
Xo Do
x (x+ xn
n TX²r x0)
E
exp
2A log
Xo .0
Tテ TT (X xn)
35

---

Published as a conference paper at ICLR 2026
In the rest of the proof, we going to show that
' (Xi xo)
{ Xo 200
E exp
2A log
T (Xq. xo)
d d'log2
+ (K((T) +3) + A(2o, (2KK(T) +1)
< exp
2
exp
2be
2be
-26T
26T
2-1/2 (10mmmm) -m)l
For this purpose, let note that
T' (X7 xa)
- x((;)) -xxx)) + log (x))
log
x (Xi zo)
=
Tr e |r|e** (
Since, due to the conditions of the lemma, (x) < M and pp(x) > - all Rd, holds that
m)]1² M for
4 ('( x( & L 1/2 m)]l + 2M.
(34)
Moreover, according to Lemma E.10 we have
(zo)) log
A(Io. T)
1
+ K(TT)
log Toole*] + A(zo - KC(T) Too [e
Tr| Tr(er" ((x0)
K(T)
M
< T)
(K(T)+ + K(T)) KC(T)
(35)
In the last line, we used the fact that [oo & eM and
Toa [e > exp [Too =
Summing up
34 and
35
we obtain that
"" (Xi x0) log
(11 M + A(zo. T)
1 (K(T) +
L L 1/2 (XT m)ll= 2 x(T)
x(T))
T& (X 2o)
1/12-1/2 -X m)ll + 3M + A(To, 2) (IC(2) +
This immediately implies that
π (Xa. xо)
E exp
2x log
Xo x0
TF (X xo))
GAM+2AA(E0 (X(T)++ E.
2AL (x;;
Xo = m0
Applying Lemma that
E.3
and taking into account that \ = min min (1/2, < b/(2).. we conclude
x" (XG zo0
E exp
2A
x (X+ zn) }
Xoo xo
dlog2 2 + M (K(TT) + 62) + A(ro, (K((T) + 2AK(T) + +21) }
< exp
4e 25T
AL 1 AL
exp
dlog2 2 + M (K(T) + 3) + A(zo. T) (2K(()) +1)
<
exp
x-1/2(50 m) 2 |?}
2bc 200
eexp
36

---

Published as a conference paper at ICLR 2026
Hence, it holds that
x² (X; 200 TP (XG xn) |
E
exp
x
log
Xo 20
(X² ro)
1+ EE
exp
2x log
Xo = 20
TT (x xo)
+ M (K((T) +3) A(so, ()( }
2 2/1+1)/2 exp
exp { be 26T 11111 m)
This
inequality yields the desired upper bound on the Orlicz log (Xx
norm of
/ (X2 wo)). Indeed, let v be the maximum of log(2)/X and
log? 2
d +1
M
(K(T) + 3) + A(ro. (I((T) + 1/2) + 26T Σ 1/2 1/2 m) m))l")
+
A
2
2
Then it holds that
1
I (Xt. zo) } Xo = _o]
n (Xi [no)
exp
log
v
X log
T (X; xx0) } Xo = zo)
x² xX² [00
1/(2n)
E
exp
(K((T) +3) +A(1oo ()((T) ((((T))+1/2)) 1/21 }) 1/(20)
1(141)/2 exp
M
2
be -26T
|>- 1/21 -m]l²)
exp
Au
2.
Thus, we obtain that
x" *X+ xz) log
lv+)
+e 20T |--1/2
-m))²)
mi)l²).
TP x ((((())
where S stands for inequality up to an absolute multiplicative constant. consantt.
E AUXILIARY RESULTS
This section collects auxiliary results. We provide their proofs (except for Lemma E. in Appen- dices E.T
E.9
For the proof of Lemma 10, reader is referred to (Puchkin al. 2025
Lemma
B.3).
Lemma E.1. Grant Assumptions
1 2
and Then there exists
Ld
To 5 b
log d LR
+ M
such that for any 8 (0, 1) and any T 3 To we have
(v(Y)) - (X:))
n
Var - ((()) (D (Dlog(And) 1og(1/))
n
1v (v/1 (1) (M + d) loogggg()dd))) + log(1/5)) log n
L
37

---

Published as a conference paper at ICLR 2026
and
Tí (2,)
E Boo (Z1))
log
Tr ((2g)
Tr [es ((21)
$V Var (log Ts (Z)
(D logd) (1/6)
2
(T(eP" ((z)
Tr
(M + () logllod(nn + log(1/0)) logn
with probability at least (1 5) simultaneously for all E Φ. Lemma E.2. Grant Assumption - Let fu Rd + and fi R + R be arbitrary functions satisfying Assumption 4 Then, for any supp(po), holds that
log Tr (eft (x) log Ty ( (x)
< /
LR 2
L(50d + d2/8) 1-1/K(T)
M
(Txlfi - fol)
Lemma E.3. Grant Assumption
Then, for any À & 6/(2)), it holds that
E 2A
Xo xo] < 2 exp (MK(T) + A(ro, T)K((T}
exp f 4e 4c-20T 2bT AL 1/- (20 -m) m)]
Lemma E.4. Under Assumption 1 it holds that
log (et xx- m)) Xo x7 < MMC() + A(ro, T)K(T) + e 5T u
- 1/2 1/2 m) + (1 e -262 )|(c|
4b
where A(x, t) and K(t), are defined in 38 and (39, respectively.
Lemma E.5. Let E be a random variable with a finite Orlicz norm [Ell < +00. For any such that 0 & X & E BElls, let us define
g(A) = log (E-EE).
Then for all (0,1/ll Egllo)) holds that = - IEE), (A(A) P -EE, and (A) = PA(E PAE)3
where, for any Borel function f IR R.
Pxf(E) = E (EE
BeA(E-E5 35)
Lemma E.6. Let E be random variable with finite Orlicz norm llEllv, < +00. For any Borel function f + I. and any X E (0, 1/1E BEll.,). let us define
E (((
Pxf(E)
BeA(s-E9 -20)
and let
I (Ps))..PP) = inf {+> 0: Px RE < <2)
be the weighted Orlicz norm corresponding to the measure Px. Then it holds that
|1g REl. < ||5 REllo1
for all 0 & À <
A|lg BElly,
[[-
38

---

Published as a conference paper at ICLR 2026
Lemma E.7. For any p it holds that I'p + 1) & (p+ 1)2, where
to0
T(p + 1) /
u'e u'e"" du
stands for the gamma function.
Lemma E.8. Let be arbitrary random variable with a finite Orlicz norm llEll +00. Then, for any p > O, it holds that
E & (p + 1)I & 2(pp + 1)
Lemma E.9. Let positive numbers u, and x be such that
u11 V (xx >
Then it holds that
a
u >
1 (x log
Lemma E.10 (Puchkin al. (025)). Let f RRd TR and g Rd - R. Assume that there exists
M E R and some non-negative constants A, B, and such that
f(x) & M, o g(x) < 1- m)
+ B for all x Rd.
(36)
Let us fix arbitrary x E Rdd and t > 0 and introduce
G(x) = Be + M Σ 1-1/2(x m) || + Ac (2b)-a/2 -a/2 (100va)a++ (2))_a((((1ee +d))
(37)
1be2
1/1 m)l + 462 +eva arcsin(e - 10c² 10e22 (1 e
A(x,t) = Vd
(38)
and
K(t)= = e1 - Vd
exp { 26 Se2vdl arcsin(e }
(39)
Then (a) & G(z)
for all x E RRd for any E [0, too), and, moreover, it holds that
(Tolged)) G(x)
K(I)
[x)) ((A2))//1())))
Too isee [gd)
1/K(()
(((,,)))(()
G(x) (G(z)
G(x)
E. PROOF OF LEMMA
The proof relies concentration inequalities for sub-exponential random variables, smooth parametrization of Pe guaranteed by Assumption and the standard 6-net argument. Let E (0, 1) be a parameter be specified a bit later. Let θ2 stand for the minimal e-net of O with respect the loo -orm and introduce
Φ = : ++.).
Since C [-1, PP it is known that
(2)
( D
c < 10 &
The rest of the proof proceeds three steps. First, we ensure that for any ip E Φ the random variables
Tr (Z1) log (2n)) Tr fe ([z)
(11 (KK)) (Y.)) (VY)) and log
Ty lev- (Z,,
sub-exponential. Second, apply (Lecué & Mitchell 2012 Proposition 5.2) quantify large
deviation bounds simultaneously for all yp E 4.. Finally, we extend these bounds to concentration inequalities holding uniformly for all yo E $.
39

---

Published as a conference paper at ICLR 2026
Step 1: Orlicz norm bounds.
The goal of this step is to establish upper bounds on 2/1 -ooms of
Tr 1 6 (Yi - (Y,) (YY))) and log "T(ee' (Z)) ................................ log
Tr (Zn)
TT (T(,)
They follow from Assumption 2 3 and Indeed, according to Lemma D.I under Assumptions and 4 it holds that
Re ( Vd
11 (Vi) (VV))l
x
1V b
M M+d+e
vā +a++v++bR))
L
(M + d)
for all <iin
(40)
On the other hand, Lemma D.2 implies that Tr(e Tr Tr 1(2g)
bee R2 arcsin(e
2 log K(T)
((()) + K(T))
MK(()
log
Vdlog2
log
log2
j & n,, where KC(T) defined
(39..
Let us note that there exists To S dVog R)/b
such that for all T > To
R
d
br $ and log S vde br 21.
Vd
Hence, if T is large enough, then
log Tr (r(p)((Z)) Tr [(2
R2 Va v
bT +M 2 M M+1 for any 1 < (41) vdde
Step concentration of subexxponential random variables. The 인o -norm bounds obtained on the previous step allow us to use Bernstein's 's inequality for unbounded random variables. Let us fix an arbitrary i E &. and apply (Lecué & Mitchell, 2012, Proposition 5.2)). Then, for any ô (0, 1), with probability at least 6/(4/021 (), it holds that
(o(()) - +())) Elx ((1))
Var(p"((2) - (o(()) (o 1215))
(42)
log ( (/4..118)
max (1²((Y)) (x,))) Isign
n
Similarly, with probability at least 1 - 8/(4/(b2)) we have
Tr (Z(,) Tx (er ((2))
Tx Tr(e+1 (21)
log
F. log
Tr [e(e0 (Zz)
Tr (e
log( (4 (10)
Var
log
(43)
Tx let ((Zz)
n
Tr (r(,) Tr - [(2z)
log 4 (142//5)
max log
16
Due to the the union bound, there cxists an event E, of probability at least 1 - 8/2 such that the inequalities (42) and (43) hold simultancously for all y E 4. on this event. Moreover, using Pisier's inequality (see, for example, (Lecué & Mitchell 2012 1827)), we conclude that
Var (1- (Yi) - (x())) log( (4/22
(5' - -(x)) EX)'(Y)) v(ri) $
log( log(4[4be/40 log
+ (Yi) ((i))l
40

---

Published as a conference paper at ICLR 2026
and
Ty Tr (Z,)
Tz Te (Zz)
log
E log
Ti lee (((,)
Tt [ee (Zz)
Tr Te( (c))(()))
log( 4/////
Var log
ler ((z))
n
Tr [r( (2,)
log(4 Φ (21/6) log
+ log
Tr Tr(eP" (T(1)
on the event & simultancously for all yp Dg. In view of 40 and
41
we have
(V) - - -
(()) +((1)
Var(o* - ( log( 4
(44)
n
1 ( (M d) log 2 166) log
+
n
and
Tr [e (Z,)
Tr ( Tr (es Tr(e ((Z))
log
log
Tr ((2;)
j=1
Tr [T{e)((Z)) log (/1)
Var log
(45)
Ta [e [(Z1)
(M 1) log de 1/6) log
n
on the same event simultaneously for all o E 4
Step 3: from e-nets to uniform bounds. The goal of this step is to transform this upper bound to one holding uniformly for all E Φ According to (Rigollet & Hütter 2023 Theorem 1.19), the norm of a uub-Gaussian random vector satisfies the inequality
{
u²
P ( u) & 6d exp
for all u > 0.
8v2
Then, due to the union bound there is an event & of probability at least (1 - 6/2) such that
max
|
6 + 8v² log(2n/6) on E
(46)
We are going show that the desired uniform bounds hold on an event Ey U E₁ of probability at least (1 5).
Let us fix an arbitrary 0 E O and let 0g be an element of O such that
0 0211,0 <s.
The existence of such 0g follows from the definition of the e-net. Let denote the corresponding to and 0g functions by &p E and ya E 4 respectively. Due to Assumption it holds that
(x()) (Y:)) < As (1 [|Y]]??)
almost surely for all 1 < <is
This implies that
[B - (v())) < (Vi) - (P)()) + (
(V)) (x())) + AA
(1+
(47)
(X)) (P())) +As (1+v'd)
41

---

Published as a conference paper at ICLR 2026
and
& 2Var (x
(K)) + 2 (e(()) -
< 2Var(p((K)) - +'(M)) + 2A (1 +||1111))
(48)
x Var( -00 + A + v
Moreover, in view of (46). we have
1 a
(p(()) - (x)())) &
(Pc(Y) (x(())
(p(()) (x.)
72
As Σ (1 [1
(Y) -((()))
(49)
2n Ae 1 log + 8v log
(vo(YY)) - (Vi))
8
Combining the bounds (77. (48). and
with the inequality (44). we obtain that
(v(Y)) - (V.)) (Vi) - (P(i)
n
Varlx> =V()) log 4/00///)
S Aed +
(50)
(M + d) log(4 क5 1 logn
+ i
n
on the event & simultaneously for all E tb
The analysis of
1 log (Z))) Elog (z(1) Tr(er" ((Z,) )
Tr [r ((2,)
n
carried out in a similar way. As before, let fix an arbitrary θ let 0g be the closest element of O2 and denote the log-potentials corresponding to 0 and 0g by and pss respectively. According to Lemma E.2 holds that
log Tr [e+] (z) log Tr [x))
LR L(50d + dP/8) 1-1/K(T)
AAATTIIEKTTTH//XCTTH+ (M+ +
(Txclp - sp 1/K(T)
2
for all i E supp(po), where the functions A(x, 1) and KC(t) are defined in 38 and 39 respectively.
Assumption ensures that
Tolp - sp5 (V)N((,)) (p()) (x)(
< Ey (Ey~N(m,))
Aee + [||||3)]
As(1+ ||]]] + Tr(E)) S Ae.
Moreover, due to 38
33 and Assumption 2 we have
bee
A(x,T) A(T)
4ea Vd arcsin(e-tT) bT) 10c Vdlog (1 - e -20T)
Vd
for all x E supp(po). Note that the expression in the right-hand side is of order
o ((
R2
+
vd
e -br 6T
whenever
87 2 log R V log d.
Vd
42

---

Published as a conference paper at ICLR 2026
Introducing
LR
L(50d ///8) 1-1/K(()
H(T) = A(T)
M [M
2 b
for brevity, we obtain that
log Tr [efe - log Tr ( [ 2 H(T)
(51)
almost surely for all j E ,l,...,n] This yields that
n2..
E log Tr
e (Z1) log TT [ett [et] (Z1)
E log Tr [e*] (21) - log Tr (T(e (e ][ + Elog Tr (Zi) - log Ti [e( [T(e)"((i))
(52)
E log Tr [e (2z) log Tr e [Z1) H(T) (Ae)
and
Var [lg Tr [e
[[]e):]](1) (Z1) - log Tr [ep {(z))
Var (log Tr [e+] (Z1) log Tr [e-"] (21)
2ar (log Tr (e (Z1) - log Tr [et] (Z1))
(53)
Var (log Tr [ [ (Zz) log Tr [e (21) (T) (AE)
The inequalities 51
(2) and (53) combined with (45) yield that
1
Tr Tr(e>>(Z)) Tr [00' (Zz) log
Tr (
log
Tr lev (Zz)
$ V Var log Tr|e (Z1)
(=1/5) log( l
(54)
Tr|es ((1)
n + H(T) 1/K(T) + (M + 1) log /// 00) log
simultaneously for all o on the event E.
Step 4: choice of and final bounds. It remains to choose a proper e > to finish the proof. The inequalities (50) (54). and |$s] < (2/s) imply that
Σ 1=1
(vr( = (Yt) E(s*(Y)) -((()))))
(D log(1/8))
S Aed +
(55)
I.)
(M + (Dlogg(()))))) + log(1/5)
log n
1 [v()
n
and
1 n
Tr [e (Z,) Trle: ((2,)
EElog 1 (Zz)
log
Tr [ef [(21)
11
1=1
Tr Trr (2z) Trlep" ((21)
(Dlog(1/e) + log(1/5))
<
Var
log
(56)
n
(M + ()loog())))) + log(( logn
+ H(T) 1/K(T)
43

---

Published as a conference paper at ICLR 2026
simultaneously for all to on the event & of probability least 1 5. Let take
K(TT)
1 A
nin
(TH(T)
nd
Such choice ensures that
max Aed, H(T) (Ac) (
1
n
and that
log < log A + (log(nd) V (C(T) log log(n?((T)))
1
< log A + log(nd) + A(T)
(((T) + K(TT)
+ MK(T) + (1- K(TT)
log(M LR² Ld
Then, due (33 and (39. there exists
To S log d LR + Ld
b + M
such that for any T > To have log S lnd(And) +
Ri2
+ vd
e bT vde- bT (M + log( M + LR² + IBP) $ log(And).
Vd
Substituting this bound into
55) and (56). obtain that
(v(()) 40 (Y)) -
n
Var(pp- (Vi) ((r)) (D log(And) log (1/6))
n
(Iv4) 4
(M + (D + log(1/8)) log n
and
1
Tx (Z,)
Tr 1
log
E Elog
TT (op' ((Zg)
Ty [ev" ((Z))
Tr (eve)((Z)) ( log(dAd)) 1og(1/5)
$ Var
log
Tr [epp ((Z)) (M + (og(DoogAAn)) + log(1/0)) log
n
simultaneously for all & on E. The proof is finished.
E.2 PROOF LEMMA
Let us fix arbitrary E supp(po). For any E [0, 1]. we introduce fa(y) and
F(s) log Tr (cf (z))
By the mean value theorem, we have
dF((s
log (x) log Tx (2)](()) = [F F(0) < sup
sE(0,1]
ds
44

---

Published as a conference paper at ICLR 2026
For this reason, the rest of the proof is devoted to the study of (FF() (d)/ds. Let fix arbitrary s E [0, 1] and note that the absolute value of the derivative does not exceed
dF(s) ds
TT [(f1 fo) (x) Trlee.]
Tr [1 fo e (v)
Tr (x)
An upper bound on [dF(ss/ds) will follow from Lemma E.10 Indeed, due to the conditions of the lemma, for any y E Rd. it holds that
f.(v) = sv) + (1--))oll)) & sM + (1 - 8)M M
and
[s(0i) < (w) + +(1 +))s(()) L1|E-1/((y - m)]1 + M.
Applying Lemma E.10
we obtain that
Tr)
(Toesh X(T))
2 A(x,T))((( e.M
(57)
cM
and
Too [li Soleds] 1/K(()
Tr [11 fo Solefa] (x)
(58)
((z)
g(z)
where the functions A(I,t) t) and KC(t) are defined in 38 and
39
respectively, and
((x) = Me Le
E-1/2(x (x mm)
2 +
Le
(400d + d²).
2
8b
The inequalities 57 and 58 yield that
(s) ds
((x))
Too [Ifi - Solesa]
1/K(T)
X(T))
< (LTE)
e.
G(x)
Tx
< (lcx)))00((())11)))
((x)
e Too lo1 - -fol
1/K(T)
eM
K(T))
eM
G(I))
Too
((x) 1-1/K(()
e
K(T)
(Txc ful) fu)) 1/K(T)
eM
Too
The expression in the right-hand side can be simplified, if one takes into account the normalization conditions To To = O. According to Jensen's 's inequality, we have
e 2 Tacl -1 =
Moreover, in view of Assumption
it holds that
((x) cM = M
L 2
M
1/-1/2(x= m) |
L
(400d+ < M+
LR
L(50d + d2/8)
8b
2
b
Hence, we finally obtain that
dF((s)
< exp A(x,T) [K(T)+ + K(T))
1
+ +MK(T)
ds
LR L(50d + d (8)
1-1/K(()
(Toc /K(T)
M +
b
for any s E [0, 1]. This yields the desired bound.
45

---

Published as a conference paper at ICLR 2026
E.3 PROOF OF LEMMA E3
Let (M),.... mTa" ~ (N, be Gaussian random vector, which is independent of Xi and Xo- Then it holds that
2ALIE
m))
Xo = 20 B,j [2avinn' 1/2(x7-m) 7). Xo xb
E
Applying Lemma
we obtain that
R R
e 2VALn x- (x7-(X
17, Xo = x0
exp (MK(T) + A(ro, T)K({}}
E, exp 20 oT VAL E E-1/2 m) +
AL(1 c 26T |lml
c-267
b
exp (M)((T) + A(To, T)K(()}
Bug exp 2e-br -I VALng (5-1/2(50 (10- AL(1 - 26T )n²
II 1=1
b
The expectations with respect to 's can be computed explicitly. Indeed, it is straightforward to check that
(E-1/2(20 -0mmm)))))) m))) AL - e--66-26 267 262)) )m V
Bn exp
2e 2e-bl bT VAL
AL(1 - e-2NT)
-1/2
2e 2e-26T AL 101112((0) m)) ?
26T
exp
1- -L(1 2-267)/0
This yields that
E
e 2XL112 1/3 2X21I1-1/91x3-
m)
0o = x0
= EE,E. [22VLn 1/²
17, Xo = Do
AL(1 - e 2
-d/2
exp (MICTT +
b
20T
AL Σ 1/2 (D: - m)
exp
1- NL11 E -6T
6/2
AL.
cxp (MK(() + A(zo. T)C(T)}
2e 26T T 12-1/2(00 -m)l
exp
1 1-A e -2TT 366
Since & min (1/, < b/(2)). we conclude that
E m))1? Xo xoo mo] < 2/2 exp (MK(T) + A(zo:T)A((T)
exp { 4 -20T 26T XL 1-12 (m)
E.4 PROOF LEMMA E.4
Let us recall that the conditional density of X given X, = 2o is equal to
(xr())) 7 wo)
(ar)
T 7<(T xо) = T Te (x0)
46

---

Published as a conference paper at ICLR 2026
Applying the change of measure theorem, we observe that
1 u x
Tr(e="](xo)) (x3 Xo 20
E e" " 1x2 (x9-1/(23--))=0000 Xo zо
e.
B ee e (x Xo - 20
Tr Tr-eet"(000 (300)
According to Lemma E.10
we have
M
M Too
K(T)
A(zo,
Trlee"
(1o)
[esp*]
The expression in the right-hand side can be simplified, if one takes into account that
T > exp (Tocp"}=1 =1
Then it holds that
MM
Tr fe (m0)
Furthermore, since {Xp 0 < & < T}, the Ornstein- Uhlenbeck process, the conditional distribu-
tion of XQ. given Xo Fig is Gaussian N (mr( (zp), Σy). This yields that
E e" T 1/3(x--m) Xo øo exp
u Σ 1/2 (mo m) + u -//,
Σ- 1/2 1/2222
{
c-bo bb
ttuuuT " T (z21 m) + (1 e-2bT |1/111
=exp
46
Hence, we obtain that
log E e xx (x- m)
Xo = 2o < (KK(T) + A(5o, T)K(()
6T 1 T Σ 1/2 mm) + (1- - e -20T |
4b
E.5
PROOF OF LEMMA
E.5
First, it straightforward check that
E [[E EGE BE)
g' (A)
Px(E - EE),
BEA 3X)
so let us focus on the second and the third derivatives of (A). Let us note that, for any Borel function f: IR + such that Px)(E) and P (E)( are well-defined, it holds that
aP dPx E. [E(E)(E - RRE)
[[(S))N(S-E3]] (E)(E-EE) yp PA(S(E)(E - -EE] - -PAS(E)PA(E Rg).
TE [(E (())e((S-BE)
Bc -RE)
EEMeS-RA RE))
(59)
Applying this formula to f(E) E Eg,, we obtain that
g"()) PA - Eg)² - -PPx( -EE))²
(60)
On the other hand, it holds that
P - 2xc)² = PA - Eg PA(E - EEE)2
-PX(E Eg)² 2Px [(E EEE E{)] +PXXE -EE))²
= PAXE EŞ)² - (PA(E - B
Hence, we showed that
(A()) PA(E -
47

---

Published as a conference paper at ICLR 2026
It remains to consider (A) to finish the proof. Taking into account (60 and applying (59)
obtain that
(A)
yd (Px(E
d
BE)² BE)2 - (Px(E - BE))²)
(PA(E E Px(5 IEE) PX(E EE))
2Px(E B8) (Px(E - EE)² (Px(E - -BS))²
(61)
Px(E Eg)² 3Px EEE) - (EE) + 2(Px EE))*
Let us note that
PA - PAE) = PA(5 - Eg - PA(E - RE)) 3
PA(E - EE) 3Px - PA - EE)
+P - RBE) (Px -RE))² - (PA(E - IRE))*
= PA(E - -E 3P EC) - EEE) + 2(( -EE))
Thus, the right-hand side of (61 simplifies to
g²" = PX -PX)
as required. The proof is finished.
E.6 PROOF LEMMA E.6
Due to Jensen's inequality, for any À IR holds that
Re X6--5) exp (XTEE RB)) = 1.
This yields that
PAelS-BE1/C
Eell Bals-Es//+AA(SBBB
5 5 Egl + A(5 - EEE)
E exp
EEMEE-B BE)
+ (
TEE }
E exp
Taking
1|g EEEL v
t 1 A/EE - 10E
we obtain that
{(
1 A/|E REll
Eexp [ 15 E[
PAcle-BE Eexp
+ 2) [S-E REI|
< 2.
115 - BE
- 1
The last inequality implies that
11E - EElly,
1
E (Ell
for all 0 S
X/E5 Bgll
|| |
E.7 PROOF LEMMA E.7
Let us fix an arbitrary p > 0 and introduce
h(u) p log u.
48

---

Published as a conference paper at ICLR 2026
Since the function h(u) is convex on (0, +oo), for any 0 we have
h(w)) > h(p + 1) + h' 1)( p- 1)
4-p
= + 1 -
p+1
= p - p-plog(p + +)
p+1
This yields that
p p p
p+1 p
T(p 1)
e (--h(u) du &
(0(((p))))))))))) du (p +1)
0
Applying the inequality E ep, which holds for all p R. obtain that
p p
< (p+1)) 1)P.
E.8
PROOF OF LEMMA
E.8
Since [C] is a non-negative random variable, we can compute its expectation according to the for- mula
too
+oo
E|E|
/
P ((E)P tt /
P ( 811)) dt.
0
Making a substitution t//r u. E (0, +oo), we obtain that
ELEP
uP-iP (15) Te u) du.
Due to the Markov inequality and the definition of 4 -ormm, the probability of the event {(E m} does not exceed
P (IE) > u) w/ilelle Eelll// < 2e w/llell,,
This implies that
+00
EK
2p
up-le //llllls du 2p 11
v 2pI(p)lElll,, 2r 1)
Applying Lemma
E.7
we conclude that
EKE < 2(p 1(p++))
E.9
PROOF OF LEMMA E.9
Let us show that the inequality < a/ (zlog a)) yields u(1 V (xlog u)) < a. Indeed,
u < 1 ax)
then holds that
ax loga a(1 V (xloga)) a.
m11 v (xlog
<
a
V lo
IV (xloga)
1 V (x log a)
Hence, the bound u(1 (zelog u)) > a implies that
u >
V x log a)
49

---

Published as a conference paper at ICLR 2026
F FURTHER DETAILS OF NUMERICAL EXPERIMENTS
In this section, we elaborate on details of numerical experiments presented in Section The section organized as follows. In Appendix F.I we discuss general implementation details. Appendix presents information about the metrics used and describes the algorithm performance, Appendix F5 provides additional details about the Gaussian mixture experiment. Appendix [F.5 provides addi- tional information about the 121 translation experiment in the ALAE latent space. Finally, Appendix F.6 presents all final hyperparameter values used for the experiments.
F.1 GENERAL IMPLEMENTATION DETAILS
In our experiments, we followed the official LightSB implementation from the repository of the authors. For comparison, the author's parameters preserved where possible.
Our modification provides additional parameters of the -hiiebbe- process and m. For the experiments, these parameters were iterated over using grid search and Bayesian optimization in the form C' (1, 1,.., 1), where is some constant. Such optimization is not perfectly comprehensive and does not cover all possible types of parameters. However, is much simpler and cheaper computationally and allows for statistically significant improvement in the performance of the basic LightSB.
Due to the strong restriction of the space in which the parameters are iterated to [-1, 1] for both 6 and m, it was possible to optimize them using Optuna (Akiba et al. 2019) and achieve good convergence of optimization. We used 50 trial optimization for optuna over the sliced Wi metric in the 25 Gaussian and I2I translation experiments and 100 trial optimization for optuna over the sliced W, metric for Single Cell data.
The calculations were performed on the personal computer (AMD Ryzen 5 1600 Six-Core 3.20 GHz CPU, 24 GB RAM) and took reasonable amount of time.
F.2 METRICS USED IN THE EXPERIMENTS
Metrics below are industry standards because they:
(a) provide quantitative, theoretically-gronnde measures of distribution similarity;
(b) are applicable to high-dimensional data;
(c) can detect different failure modes of generative models (mode collapse, memorization, etc.):
(d) have been widely adopted research papers and practical applications.
Metrics used in the empirical analysis of LightSB-OU:
Sliced Wasserstein Distance:
Sliced (X,,
Σ Wi (1 (,) (1v kn))
where (0,, : 1 n N} are i.i.d. random vectors drawn from the uniform distributed on the unit sphere Sd-1
2. Maximum Mean Discrepancy:
MMD Y) E[k( X)) (Y, 2E(((X
with Gaussian kernel k(x,)) exp( exp(->|l= y/
3. Energy Distance:
Energy Distance (x, Y) X Y E| X'I EEJY Y'Il, where X" and Y are independent copies of X and Y, respectively.
50

---

Published as a conference paper at ICLR 2026
Algorithm LightSB-OU Training
1: Input: the number of mixture components K, the number of iterations Nsteps, learning rate ") batch size time horizon 7' and parameters b > 0, m E RR and > 0 of the reference Ornstein-Ullenheck process
x!) dt deWW, OEIET. dxp b(m X
2: Initialize parameters: 0 {(TA: Sk, log 0k)) metric positive definite matrices S SK Rddd and with the positive weights a1.... ak such that 01 + .... ax 1.
K} with vectors : JK R sym-
3: for from 1 to Steep do
4: Receive batches of fresh i.id. observations
~ Po and yi..... Vig ~ pr-
5: Define the adjusted Schrödinger potential estimate
vo(y) Σ ak so VER,
where, for any E {I,.... KJ. rk, S07 the density of the Gaussian distribution N(rk,eof.S Sx). 6: Define the normalizing constant
Fk mT((2)
|| mT
(=)=))²
co(2)
exp
2007
k=1
where
1 -
-26T
26T
(= bT + (1
e- -7)m and 07
6T
2b
7: Update the parameters via gradient descent:
1 BB
1 nB
0 0
Σ
No (Y
Σ
log ca (Zj)
nB
nB
8: end for
9: Return the adjusted potential (v(v)=
ai SO,
4. Number of Covered Modes (for Gaussian Mixture Experiment):
Covered Modes
Covered
where
Covered =T min(x yEX
(4))
" (x: 1xx) X2.0
For the Sliced Wasserstein Distance, 100 random projections were used, and for the number of covered modes, a confidence interval of 90% was used.
We also provide a description of the training algorithm for LightSB-OU (Algorithm D. The pseu- docode provided allows a reader to visually evaluate how our approach differs from regular LightSB.
F.3 DETAILS EVALUATION ON THE GAUSSIAN MIXTURE
To evaluate the algorithm's 's performance Gaussian mixtures, the parameters for the basic LightSB and our modification were chosen based on reasonable considerations for such a problem and the parameters for the 2D Swiss Roll given in the original paper.
51

---

Published as a conference paper at ICLR 2026
To find the metrics sliced W1. MMD, covered modes and energy distance, 5 iterations used, including sampling and translation of 10000 random points from the standard normal distribution by trained models. They were used to find the mean values and standard deviations in the table and on the graphs. During the inference, we have used the fact that the conditional density X1 given X has form (see (20)
x (0 y)
x v (u|z)=
ex mu(a)/(cof) 00 (v))/c(())))
Po(x)
where (y) and (x) given by (99) and 21 respectively. When (v(y) is Gaussian mixture, is straightforward to check that the conditional distribution is a Gaussian mixture as well. More specifically, if vy is defined by (I9). then
K
π T" (v|
Σ
ās (()))(
where (xx(x) is given by (22). p(u: TA (x), sofSk) stands for the density of N ((re(E)) and (x) = + (x) for all [I....,,} K}.
The Figure 3 shows the metrics for LightSB-OU for three experimental setups with mixture of 25 Gaussians for different with the interval of one std shown in ctmi-tranpprrent color. The Figure
4 shows joint graphs for LightSB and LightSB-OU, which allow
evaluate the comparative behavior of the algorithms for different setups and different K, and finally the Figure 5 shows examples of the results of both algorithms for different setups with K 25.
Dhark SESIAL
TM STROA
Figure 3: Translation results from standard normal distribution using LightSB. Top: uniform grid and standard covariance matrices. Middle: random location on the grid and standard covariance matrices. Bottom: uniform grid and anisotropic random co- variance matrix.
F.4 DETAILS OF EVALUATION ON THE SINGLE CELL DATA
We analyzed cell differentiation across five stages from day 0 to 27 (i.e. to day 0 to t : day to 9, la day 12 to 15, fa day 18 to 21, f4 day 24 to 27), processing the scRNA-seq data through quality filters and PCA to create distinct feature vectors. In this experiment the goal was to predict the cell distribution at intermediate time point (t,) by "transporting" data from the surrounding intervals (t:-1 and 1.+1), where E [1, 2, 3). To evaluate the model's accuracy, we calculated the Wasserstein-1 (W1 distance how closely our predictions matched the actual observed data.
52

---

Published as a conference paper at ICLR 2026
Pur SERAA
Figure Translation results from standard normal distribution using LightSB and our modified approach. Top: uniform grid and standard covariance matrices. Middle: random location on the grid and standard covariance matrices. Bottom: uniform grid and anisotropic random covariance matrix.
F.5 DETAILS OF EVALUATION ON THE UNPAIRED IMAGE-TO-IMAGE TRANSLATION
www.coo///ee.eeeeeeo and and neural network extracted attributes for the FFHQ dataset from https:: http:///tttu..c//cccco/ wttp:///////oo-eeueeeeeeeeeeeeee -dattts...............
We used the official code and ALAE model from
To evaluate the algorithm's performance in the 121 translation task, the parameters given for the LightSB model in the official implementation were used. From the trained model, 3 face translation results were sampled 100 times, and examples illustrating the best coverage of the latent space by our modification were found using manual selection.
It should be noted that the parameters for LightSB task, while we performed the comparison on both it and the Male to Female translation task. Due the proximity of these tasks and the moderately good behavior of both models, was decided to keep the parameters from the Adult to Child task.
given only for the Adult to Child translation
F.6 FINAL HYPERPARAMETER VALUES
Below are the final hyperparameter values for all experiments. should be noted that in the ex- periments that follow the original LightSB paper, we used the author's parameters, leaving only the Ornstein-Uhlentee Uhlenbeck process parameters for tuning.
Empirical KL Divergence Experiment:
Parameters for experiment: batch size 128, 9.0, Ir 0.002, K 1, diagonal false, b b= 2.0, m 0.0.
ALAE I2I Experiment:
Parameters for experiment: batch size 128, 6 0.1, Ir 0.001, =10, diagonal true, b = 0.02, m = 0.0.
Single Cell experiment:
Parameters for = 1 batch size = 128, E = 0.1, = 0.01, K 100, diagonal
true, b b= -0.2, 4.0,
53

---

Published as a conference paper at ICLR 2026
Figure 5: The translation result from standard normal distribution using LightSB (left) and our modified approach (right). Top: uniform grid and standard covariance matrices. Middle: random location the grid and standard covariance matrices. Bot- tom: uniform grid and anisotropic random covariance matrix. Green; samples from the target Gaussian mixture; orange: LightSB samples; blue: LightSB-OU samples.
Parameters for batch size 128, 0.1, Ir 0.01, K 100, diagonal true, b 0.2, m 4.0,
Parameters for = : batch size 128, E 0.1, 0.01, K 100, diagonal
true, b 0.2, -1.0.
54

---

Published as a conference paper at ICLR 2026
In the experiment with mixture of 25 Gaussians, on the contrary, pre-selected good parameters for LightSB, and then tuned the parameters for the Ornstein- Ornstein-Uhlenbeed process.
It is worth noting that the process parameters b, remained quite stable in all three experiments with a mixture of Gaussians with respect to changing the number of components K for approx- imating the potential, which is in line with expectations. The process should facilitate transport independently of subsequent control, and therefore is likely to be unique and stable with respect to the hyperparameters well-posed problem.
Standard Gaussian mixture with 25 components:
Parameters for
batch size 128, = 0.1, 0.002, K 15, diagonal 0.901,
true, b = -0.125,
Parameters for K = 20 batch size = 128, 0.1, Ir = 0.002, K 20, diagonal = true, b -0.125, m 0.901,
Parameters for K 25 batch size 128, 0.1, Ir 0.002, K 25, diagonal
true, b = 0.232, m = 0.197,
Parameters for = 30 batch size 128, E = 0.1, Ir = 0.002, K = 30, diagonal true, b 0.101, 0.416.
Irregular Gaussian mixture with 25 components:
Parameters for = 15 batch size = 128, 6 = 0.1, Ir = 0.002, K = 15, diagonal = 0.901,
true, b = -0.125,
Parameters for K 20 batch size = 128, 0.1, Ir = 0.002, K 20, diagonal true, b = 0.332, m -0.575,
Parameters for K 25 batch size 128, 0.1, 0.002, K 25, diagonal true, b = 0.332, m -0.575,
Parameters for = 30 batch size = 128, 6 = 0.1, Ir 0.002, K = 30, diagonal = true, b 0.332, m 00.575.
Anisotropic Gaussian mixture with 25 components:
Parameters for K = 15 batch size = 128, = 0.1, Ir = 0.002, K = 15, diagonal = true, b -0.125, m 0.901,
Parameters for K 20 batch size 128, 0.1, Ir 0.002, K 20, diagonal true, b -0.125, m = 0.901,
Parameters for 25 batch size = 128, 5 = 0.1, = 0.002, 25, diagonal true, b 0.232,
0.197,
Parameters for K = 30 batch size = 128, 0.1, Ir -- 0.002, K 30, diagonal = true, b 0.101, m 0.416.
55

---

Published as a conference paper at ICLR 2026
CONTENTS
Introduction
2
Related work
3 Main Result
Numerical Experiments
Evaluation on the Gaussian Mixture
4,2
Evaluation on the Single Cell Data
10
Unpaired Image-to-lmage Translation
10
Conclusion
10
On Assumption
13
Proof of Theorem
14
B. Proof Insights
14
B.2 Rigorous proof
15
B3 Proof of Lemma B.1
19
B.4 Proof of Lemma B.2
21
I
Log-potential variance and conditional variance bounds
22
Proof of Lemmal
23
Proof of Lemma
25
Proof of Lemma
29
4 Proof of Lemma
31
D Orlicz norm bounds
32
D.I Proof of Lemma
33
D.2 Proof of Lemma D.2
34
Proof of Lemma D.3
35
Auxiliary results
37
Proof of Lemma
39
Proof of Lemmal
Proof of Lemma
46
Proof of Lemma
46
Proof of Lemma E.5
47
E,6
Proof of Lemma E.6
48
E.7 7
Proof of Lemma F.7
48
56

---

Published as a conference paper at ICLR 2026
E.8
Proof of Lemma E.8
49
Proof of Lemma
49
Further Details of Numerical Experiments
50
General Implementation Details
50
Metries Used in the Experiments
50
Details of Evaluation on the Gaussian Mixture
F.4
Details of Evaluation the Single Cell Data
52
H5
Details of Evaluation on the Unpaired Image-to toe-tttoommmmm to-lmageeeeam----e-- Translation
53
F.6
Final Hyperparameter Values
53
57


Reviews:
[{"review_id": "grvWh6YNjI", "reviewer": "Reviewer_BNq5", "summary": "This paper provides learning theory style high probability bounds for estimating Schrodinger bridge potentials when the reference dynamics is OU process. This is a theoretical paper where the main result is a nonasymptotic, high probability bound improving prior work.", "questions": "- notation is dense. A short “notation table” early in the paper would help.\n- Can Theorem 1 extend to sub-Gaussian cases? Please discuss Assumption 2 and its applicability in real-world settings.\n- The authors claim the dependence of the bound to dimension is 'nearly linear'. The bound itself, however, is not clearly written. It would be great if the authors provide a corollary, perhaps cleaning up some unnecessary quantities and display the dimension dependence clearly.\n- Please provide a discussion on when $\\Delta$ can be made small (for the favourable convergence rate), i.e., for what classes of initial/target distributions, what families would make it small?\n- The theoretical analysis is elegant, but it’s not clear what it tells practitioners: (i) Does the theorem suggest a way to choose the OU drift parameter b? (ii) Does it explain why LightSB-OU performs better numerically?\n\nIt is this last point I am more keen on seeing explained -- I appreciate the paper's main contribution is theoretical but it would be still good to understand how theoretical bound would give insights for practical improvements.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is mathematically strong, well structured, well motivated on its own domain.\n- There are mathematical novelties, for example, the high-probability bound is novel compared to prior work.\n- The derived convergence rate is fast.", "weaknesses": "- Some assumptions of the paper are probably unrealistic for real-world scenarios\n- The bound, while a nonasymptotic, high-probability bound, is not very clean.\n- Practical implications of the result is not clear, there's a gap between theory and why how it explains the empirical improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LP9xmbQTaw", "reviewer": "Reviewer_ibUA", "summary": "**Disclaimer**\n\n*I do not have specific expertise in quantitative generalization bound estimation. Therefore, while I provide an assessment of the clarity, motivation, and overall contribution of the paper, I will defer to other reviewers with stronger expertise in the theoretical aspects when evaluating the technical soundness of the quantitative rate derivation.*\n\n**Summary**\n\nThis paper studies the finite-sample generalization behavior of Schrödinger bridge (SB) estimation for unpaired data translation. While most prior works adopt Brownian motion as the reference process, this work introduces the Ornstein–Uhlenbeck (OU) process as the reference dynamics. The OU process exhibits exponential mixing, which enables a more stable and analytically tractable setting for sample-based SB estimation. The paper establishs the first non-asymptotic generalization bound for estimating Schrödinger potentials $\\phi$ as the empirical risk minimizer and derives a convergence rate of $O(\\log^{3} (n) / n)$ for the KL divergence between the optimal coupling $\\pi^{*}$ and the empirical risk minimizer $\\hat{\\pi}$ (when the approximation error $\\Delta$ is small). Empirically, the paper extends the Light Schrödinger Bridge framework to an OU-based version (LightSB-OU), demonstrating improved stability and translation quality in synthetic, biological (single-cell RNA), and unpaired image translation tasks.", "questions": "- Beyond theoretical implications, are there practical improvements in unpaired translation performance (e.g., FID, MMD) that can be directly attributed to the tighter convergence rate?\n- The title emphasizes \"Unpaired Data Translation,\" but its primary contribution is theoretical. Perhaps a title that emphasizes the general results, such as \"Tight Bounds for Schrödinger Potential Estimation between Empirical Distributions\", would better reflect the scope of this paper? (Of course, this is left to the author's discretion.)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper provides a new tight finite-sample generalization bound for Schrödinger bridge (SB) estimation.\n- Replacing Brownian motion with an OU process is both theoretically novel and practically beneficial.\n- The proposed LightSB-OU algorithm effectively connects the theoretical framework to practical tasks. Notably, the OU process corresponds to the VP-SDE used in the diffusion model, while Brownian motion aligns with the VE-SDE. The experiments on both synthetic and real-world datasets demonstrate that the OU reference improves fidelity in unpaired data translation.\n- The manuscript is well organized and clearly written. The motivation, the implications of assumptions, and the meaning of theoretical results are well described.", "weaknesses": "- Although the derived convergence rate of $O( \\log^{3} (n) / n)$ is theoretically novel, there is no experimental study showing empirical convergence as a function of sample size. Including a convergence curve or ablation study on sample size would better demonstrate the practical significance of the bound.\n- While the bound is tight in a theoretical sense, it remains unclear how much this improved rate translates into practical accuracy gains over the previous $O(n^{-1/2})$ regime established by LightSB.\n- In the image-to-image translation experiments, only selected qualitative results are presented. Providing quantitative metrics, such as FID (for fidelity to target semantics) and LPIPS (for perceptual similarity between input and output), would significantly strengthen the empirical claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FPktdZBthd", "reviewer": "Reviewer_TRip", "summary": "In this paper, the authors introduce novel bounds for the estimation of the Kullback-Leibler divergence between a candidate distribution defined by its potential and the Schrodinger Bridge target. More precisely, they show that under strong assumptions on both the potential and using a Ornstein-Uhlenbeck reference dynamics they can provide a non-asymptotic upper bound on the empirical risk minimizer and the target Schrodinger Bridge. The bound is of order $O(1/n)$ where $n$ is the number of data points used in the approximation and $O(d)$ (up to logarithmic terms in both cases). Some of the technical lemmas seem to be inspired by [2]. In addition, of this main theoretical contribution, the authors consider the procedure of [1] but instead replace the Brownian reference motion by a Ornstein-Uhlenbeck process. In that case, they present experimental results on several benchmarks. \n\n[1] Korotin et al. (2024) -- Light Schrodinger bridge\n\n[2] Puchkin et al. (2025) -- Sample complexity of Schrodinger potential estimation", "questions": "* $\\hat{\\pi}$ in Page 3 hasn't been introduced yet. \n\n* See my question above on the time dependency. \n\n* Using a Ornstein-Uhlenbeck as a reference measure for Schrodinger Bridges is not new. For instance, it was used in [1]. In addition, it can be shown quite easily that using a Ornstein-Uhlenbeck as a reference path will lead to a quadratic cost function with a regularisation parameter that grows exponentially with the time of the process. It would be good if the authors could discuss this fact in the main paper. \n\n[1] Shi et al. -- Diffusion Schrödinger Bridge Matching", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The paper seems technical strong. In particular, there is an extensive discussion on the error bounds obtained in the literature and the ones obtained in this paper. In particular, the authors draw comparisons with [1]. \n\n* The theory on Schrodinger Bridge is still sparse and even though are extremely strong in the present work the authors derived a novel and interesting result. \n\n[1] Korotin et al. (2024) -- Light Schrodinger bridge", "weaknesses": "* The presentation of the paper is quite hard to follow. In particular, the introduction is extremely long and seems to merge both the contribution and the related works. In particular, while I understand that it is a choice of the authors, I found it quite hard to follow the related work section. In particular, I did not understand the discussion with [1,2]. Related to this issue, I found the related work to be quite poor. There is no mention on the impact of Schrodinger Bridge work and its application in Machine Learning. For example, Diffusion Schrodinger Bridge [3] but also Stochastic Interpolants [4] and Adversarial approaches [5] are competitive works which are not mentioned (some of them are compared with in the Numerical Experiments section but it is never mentioned why those works are relevant).\n\n* The assumptions are extremely strong. While I understand that the results obtained by the authors are new, the class of target distributions and reference potentials that is under consideration is extremely limited. This strongly limits the theoretical impact of the paper. \n\n* The methodology contribution is incremental. In the applications, the authors consider Light SB which indeed fits their framework. However, the only modification to this work is in the reference process considered. The results obtained are quite underwhelming (see Table 2 for instance where the methods is similar to Light SB in performance and not as good as OT-CFM [7]). I do appreciate however that the authors highlight the limitations of their approach \"We emphasize that these examples are not intended to suggest\nuniversal superiority but rather to showcase specific strengths of our method in certain cases.\"\n\n* In the theoretical results presented, even though this is discussed, I find the fact that the time must grow with the number of samples to be a bit concerning. Even though I agree that the doubly logarithmic dependence (which has a logarithmic effect on the regularisation due to the exponential convergence of the Ornstein-Uhlenbeck) can be mitigated, it is concerning that the time grows as the number of samples grows. I see this as one of  the main limitation of the paper. \n\n[1] Pooladian et al. (2024) -- Plug-in estimation of Schrodinger bridges\n\n[2] Tang et al. (2024) -- Simplified diffusion Schrodinger bridge\n\n[3] De Bortoli et al. (2021) -- Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling\n\n[4] Albergo et al. (2023) -- Stochastic Interpolants: A Unifying Framework for Flows and Diffusions\n\n[5] Gushchin et al. (2024) -- Adversarial Schrödinger Bridge Matching\n\n[6] Korotin et al. (2024) -- Light Schrodinger bridge\n\n[7] Tong et al. (2023) -- Improving and generalizing flow-based generative models with minibatch optimal transport", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NrGDzxo2u5", "reviewer": "Reviewer_Hswo", "summary": "This paper studies tight generalization bounds for the empirical risk minimizer within a class of Schrödinger Bridge (SB) potentials. The authors focus on a fundamental yet widely studied setting where the risk functional is the KL divergence and the underlying dynamics follow an Ornstein–Uhlenbeck (OU) process. Under a set of explicit assumptions (Assumptions 1–5), they establish a fast convergence rate of order $O(\\log^3 (n) / n$, which significantly improves upon prior results of order ($O(1\\sqrt{n}$). As a practical instantiation, the paper adapts the LightSB algorithm to the OU reference process and presents experiments on Gaussian mixture models (GMMs) and single-cell datasets.", "questions": "- Could the authors discuss the practical gap between the assumptions required by the theorem and the actual behavior of the LightSB implementation?\n\n- Could the authors provide an intuitive explanation for why the OU process yields a tighter bound? Specifically, since the required time horizon  $T$ is roughly proportional to $1/b$, a larger $b$ would make the terminal reference distribution more Gaussian-like, and the joint distribution of the reference dynamics nearly independent (memoryless), as discussed in [1]. Is this correct? Could author provide the better intuition?\n\n- Is it possible to extend the theoretical discussion to convergence guarantees for the gradient of the potential function? If so, are there existing works addressing this direction?\n\nReference\n\n[1] Domingo-Enrich et al., Adjoint Matching.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The primary contribution lies in deriving a notably tighter generalization bound by adopting the OU reference process. The improved rate represents a substantial theoretical advance over previous works.\n\n- The paper is well-structured and accessible, providing clear explanations of both the strengths and limitations of the proposed analysis. The authors carefully justify each assumption, discussing its practical implications and arguing convincingly that these assumptions are reasonable or attainable in practice.", "weaknesses": "While the theoretical contribution is valuable, my main concerns lie in the practical validation and experimental analysis.\n\n- This is my main concern. The empirical section does not quantitatively assess the gap between the empirical and ground-truth SB potentials with respect to the number of samples $n$. Such analysis would strengthen the connection between theory and practice. It would be particularly insightful to evaluate this gap in the case of Gaussian marginals, where the exact SB potential is analytically tractable. Furthermore, constraining the neural network parameters to satisfy the boundedness and growth conditions (related to constants $L$ and $M$) should be discussed or enforced. A direct comparison with the Brownian-motion-based LightSB (with fixed number of samples $n$) should be discussed. Finally, the influence of the time horizon on the convergence rate should be investigated.\n\n- In Table 2, the results focus solely on distributional alignment. It would be more comprehensive to also report transport cost metrics, as well as other measures such as the FID and transport cost in image-to-image translation tasks.\n\n- The parametrization of the exponential SB potential through a GMM may not be practical for high-dimensional or complex data. In many applications, algorithms estimate the gradient of the log potential (i.e., the control function) rather than the potential itself.\n\n\nOverall, I find the theoretical result strong and meaningful. With additional experiments addressing these points (especially the first bullet point), I would be inclined to raise my evaluation score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)