Energy-guided Entropic Neural Optimal Transport
Authors: Petr Mokrov
Alexander Korotin
Alexander Kolesov
Nikita Gushchin
Evgeny Burnaev
Venue: ICLR 2024 poster (2024)
Keywords: energy-based model
generative model
optimal transport
entropic optimal transport
general optimal transport cost function
Primary area: generative models
TL;DR: We propose a novel energy-based method to compute entropic optimal transport with general cost functions.

Abstract: Energy-based models (EBMs) are known in the Machine Learning community for decades. Since the seminal works devoted to EBMs dating back to the noughties, there have been a lot of efficient methods which solve the generative modelling problem by means of energy potentials (unnormalized likelihood functions). In contrast, the realm of Optimal Transport (OT) and, in particular, neural OT solvers is much less explored and limited by few recent works (excluding WGAN-based approaches which utilize OT as a loss function and do not model OT maps themselves). In our work, we bridge the gap between EBMs and Entropy-regularized OT. We present a novel methodology which allows utilizing the recent developments and technical improvements of the former in order to enrich the latter. From the theoretical perspective, we prove generalization bounds for our technique. In practice, we validate its applicability in toy 2D and image domains. To showcase the scalability, we empower our method with a pre-trained StyleGAN and apply it to high-res AFHQ $512\times512$ unpaired I2I translation. For simplicity, we choose simple short- and long-run EBMs as a backbone of our Energy-guided Entropic OT approach, leaving the application of more sophisticated EBMs for future research. Our code is available at: https://github.com/PetrMokrov/Energy-guided-Entropic-OT

Published as a conference paper at ICLR 2024
ENERGY-GUIDED ENTROPIC NEURAL OPTIMAL TRANSPORT
Petr Mokrov' Alexander Korotin Alexander Kolesov', Nikita Gushchin', Evgeny Burnaey
Skolkovo Institute of Science and Technology, Moscow, Russia Artificial Intelligence Research Institute, Moscow, Russia (petr. www.rovoovvvvooooo a. korotin (awiinnnsskkeeeeee @skoltech. ru
ABSTRACT
Energy-based models (EBMs) are known in the Machine Learning community for decades. Since the seminal works devoted to EBMs dating back to the noughties, there have been a lot of efficient methods which solve the genera- tive modelling problem by means of energy potentials (unnormalized likelihood functions). In contrast, the realm of Optimal Transport (OT) and, in particular, neural OT solvers is much less explored and limited by few recent works (ex- cluding WGAN- based approaches which utilize OT as a loss function and do not model maps themselves).
our work, we bridge the gap between EBMs and Engropy-regularizeed OT. We present novel methodology which allows uti- lizing the recent developments and technical improvements of the former in order enrich the latter. From the theoretical perspective, prove generalization bounds for our technique. In practice, we validate its applicability in toy 2D and image domains. To showcase the scalability, pre-trained StyleGAN and apply it to high-res AFHQ 512 x 512 unpaired 121 translation. For simplicity, choose simple short- and long-run EBMs a backbone of our Energy-guided Entropic OT approach, leaving the application of more sophisticated EBMs for future research. Our code is available at: https; //github. com/ Pet rMokrev
empower method with
Figure 1: AFHQ 512 x 512 Car+Dog unpaired translation by our Energy-guided EOT solver applied in the latent space of StyleGAN2-ADA. Our approach does not need data2latent encoding. Left: source samples; right: translated samples.
INTRODUCTION
The computational Optimal Transport (OT) field is an emergent and fruitful area the Machine Learning research which finds its applications in generative modelling (Arjovsky et al. 2017 Gul-
rajani et al. 2017: Deshpande et al., 2018), domain adaptation (Nguyen et al. 2021 Shen et al. 2018 Wang et al. 2022), unpaired image-to- to-igage translation (Xie al. [2099] Hu al.) )..) datasets
manipulation (Alvarez- & Fusi 2020), population dynamics (Ma et al. 2021 Wang et al. 2018).
gradient flows modelling (Alvarez-Melis et al.
2022
Mokrov et al. 2021 barycenter estimation

---

Published as a conference paper at ICLR 2024
(Korotin et al. 2022a; Fan et al. 2021). The majority of the applications listed above utilize OT as a loss function, e.g., have WGAN-like objectives which compare the generated (fake) and true data distributions. However, for some practical use cases, e.g., unpaired image-to-image translation (Korotin et al. 2023b), it is worth modelling the OT maps or plans by themselves.
The existing approaches which recover OT plans based various theoretically-adviii advised techniques.
Some of them (Makkuva et al.) 2020 Korotin et al.] (021a)) utilize the specific form of the cost
function, e.g., squared Euclidean distance. The others (ie et al. 2019 Lu et al. (2020) modify GAN
objectives with additional OT regularizer, which results in biased OT solvers (Gazdieva et al. 2022
Thm. 1). The works (Fan al., 2023) Korotin et al. 2023b; Rout al. 2022) take advantage of dual OT problem formulation. They are capable of tackling unbiased large-scale continuous OT with general cost functions but may yield fake solutions (Korotin al.
2023a). To overcome this issue, (Korotin et al.] 2023a) propose to use strictly convex regularizers which guarantee the uniqueness of recovered OT plans. And one popular choice which has been extensively studied both discrete (Cuturi 2013) and continuous (Genevay et al. 2016 Clason et al. 2021 settings is the Entropy. The well -studidd methodological choices for modelling Entropy-
OT (EOT) include (a) stochastic dual maximization approach which prescribes alternating optimization of dual potentials (Seguy et al. 2018| Daniels et al.] (202T) and (b) dynamic setup having connection to Schrödinger
bridge problem (Bortoli et al. 2021] Gushchin al. 2023 Chen et al. 2022). In contrast to the methods presented in the literature, we come up with an approach for solving EOT built upon EBMs.
Contributions. We propose novel energy-based view on the EOT problem.
1. We take advantage of weak dual formulation for the EOT problem and distinguish the EBM-related nature of dual potential which originates due to this formulation (...))
2. We propose paarretccall-grounnde yet tosy--t--mpleeenttt modifications to the standard EBMs training procedure which makes them capable of recovering the EOT plans ($3.2).
3. We establish generalization bounds for the EOT plans learned via our proposed method
3.3
4. We showcase our algorithm's 's performance on low- and moderate-dimennsional toy setups and large-scale 512x x 512 images transfer tasks solved with help of a pre-trained StyleGAN ($5)
Notations. Throughout the paper, X' and y are compact subsets of the Euclidean space, i.e., X C R D.
and y C Pv The continuous functions on X' are denoted as C(X). In turn, P(X') are the sets of Borel probability distributions on X. Given distributions P E P(x') and Q E (()), [I(,, Q) designates the set of couplings between the distributions IP and Q. i.e. probability distributions on product space X x y with the first and second marginals given by P and Q, respectively. We use IIP)) to denote the set of probability distributions on X' x y with the first marginal given by P. The
d?(xr) absolutely continuous probability distributions on X' are Pac (X). For P E (X)(X)) we use du and dQ(s) to denote the corresponding probability density functions. Given distributions f and defined dy
on Z, μ < p means that p is absolutely continuous with respect to p.
BACKGROUND
2.1 OPTIMAL TRANSPORT
The generic theory behind OT could be found in (Villani al. 2009 Santambrogio 2015). For the specific details on BOT, see (Genevay et al. 2016 Genevay 2019).
Let IP P(X) and Q E P()).. The primal OT problem due Kantorovich (Villani et al. 2009 is:
O Q def
inf TETT(P..)
/ c(1, y)dT(a
In the equation above, e: X x y + knowledge of how data from the source and target distribution should be aligned. Typically, the cost function c(æ, y) is chosen to be Euclidean norm || - v112 yielding the I-Wasserstain distance (W1) or halved squared Euclidean norm 2 1| vll2 yielding the square of 2- Wasssersteinn distance (WP)..
a continuous cost function which reflects a practitioner's
The distributions x* E [IP, Q) which minimize objective T are called the Optimal Transport plans. Problem (1) may have several OT plans (Peyré et al.) 2019 Remark 2.3) and in order to impose the uniqueness and obtain a more tractable optimization problem, a common trick is to regularize (T) with strictly convex (w.r.tt distribution π) functionals R. P(X' x )) + R.
2

---

Published as a conference paper at ICLR 2024
Entropy-regularizzed Optimal Transport. In our work, utilize the popular Entropic regularization
(Cuturi 2013) which has found its applications in various works (Solomon et al. 2015 Schiebinger
et 2019 Rukhaia 202T). This is mainly because of amenable sample complexity (Genevay 2201 2019 83) and tractable dual representation of the Fegropy-regulariize OT problem which can be leveraged by, e-g., Sinkhorn's algorithm (Cuturi [2013] Vargas al., (021). Besides, the EOT objective is
known to be strictly convex (Genevay et al. 2016) thanks to the strict convexity of Entropy H and
KL divergence (Santambrogio) 2015) Nutz 2021 Nishiyama 2020) appearing in EOT formulations.
Let e > 0. The EOT problem can be formulated in the following ways:
! (P. EOT BOT (TP,
+eKL (T X Q) eH(T), ([
(2)
def
min (EITP xxy
c(a, v)d
(3)
(4)
These formulations are equivalent when and Q are absolutely continuous w.r.t. the corresponding standard Lebesgue measures since KL (TT|F x
H (=()==)))d H (Q) (1)
((Q) + HPP). In other words, the equations
and (4) are the same up to additive constants.
In the remaining paper, we will primarily work with the EOT formulation and, henceforth, we will additionally assume P (x), Q (>) when necessary.
Let x- E [1( Q) be the solution of EOT problem. The measure disintegration theorem yields: drr
dr (y/2z)d<"""" (x) dr (y|==)P(P))
Distributions क. (-(x) will play an important role in our analysis. fact, they constitute the only ingredient needed to (stochastically) transform source point æ E X' to target samples /1. V2. E Y w.r.t. HOT plan. We say that distributions
are the optimal conditional plans.
EOT problem weak OT problem (Gozlan et al.) 2017] [Backhoff- Veraguas et al.] 2019). Given a weak transport cost C X P()') which penalizes the displacement of point E (+|xx) E P())), the weak OT problem is given by
weak OT (WOT) problem. EOT problem (4) be understood the so-called
X into distribution
wOTc
def
=E / C(I,, (+1x)) da dr(x).
inf
(5)
=dP(w)
EOT formulation (4) is particular case of weak OT problem 5 for weak transport cost:
= S,, c(x, v)dx(vl>) (=((-))))
CEOT (+1xx)) (1))))
(6)
Note that weak cost C is strictly convex and lower semicontinuous, as is the case for CEoT, the solution for (5) exists and unique (Backhoff-
al. 2019).
Weak OT dual formulation of the EOT problem. Similar to the case of classical Kantorovich
OT (1). the weak OT problem permits the dual representation. Let C()). Following (Backhoff)
Veraguas et al. 2019 Eq. (1.3)) one introduces weak C-transform f : X + R by
(x)
def inf
{C(E,))
f(y)dply)
MEP())
For our particular case of EOT-advised weak OT cost (6). equation reads
Oaam (x) = min
{ c(a, v)dn(v) - cH() - J,. f()ddp()
def min MEP())
(8)
pEP())
Note that the existence and uniqueness of the minimizer for g follows from Weierstrass theorem (Santambrogio) 2015) Box 1.1.) along with lower semicontinuity and strict convexity of Ga.f in The dual weak functional Pe C()) RR for primal WOT problem (5)
F (f) J S { + /, f(o)dQ(v.
Thanks to the compactness of X and V, there is the strong duality (Gozlan et al. 2017 Thm. 9.5):
( min Gz. (2)))))))))))))))>))))))))))))() / [(w))Q()
EOT c.e (P, ! sup
m H.E)()))
sup F (f). SECO")
feC(y)
We say that g the weak dual objective. will play an important role in our further analysis.
3

---

Published as a conference paper at ICLR 2024
2.2
ENERGY- BASED MODELS
The EBMs are a fundamental class of deep Generative Modelling techniques (LeCun et al., 2006 Salakhutdinov et al. 2007) which parameterize distributions of interest μ P()) by means of the
distribution density:
du(y) exp exp(-F(y))
(10)
dy
In the equation above E V + R the Energy function (negative unnormalized log-likelhhoo), and Z = Jy (xy the normalization constant, known as the partition function.
Let μ. E P()) be true data distribution which is accessible by samples and po(y), 6 E be parametric family of distributions approximated using, e.g., a deep Neural Network Ea,, which imitates the Energy function in (T).. In EBMs framework, one tries bring the parametric distribution μo to the reference one μe by optimizing the KL divergence between them. The minimization of KL (иll)e) done via gradient descent by utilizing the well-known gradient (Xie et al. 2016):
a
I,
a Foly)dp(y) 09
J
KL. (Mllwo) 08
(yl dua(y)- 88
(11)
The expectations on the right-hand side of (11) are estimated by Marlo- which requires samples from ft and μag. While the former are given, the latter are usually obtained via Unadjusted Langevin Algorithm (ULA) (Roberts & Tweedie, 1996). It iterates the discretized Langevin dynamics
Y44 Yi 7 Ô By EYt + VnEi: & 100,1,,
(12)
starting from simple prior distribution Yo no. for L steps, with a small discretization step > 0. In practice, there have been developed lot of methods, which improve circumvent the procedure above by informative initialization (Hinton) 2002: Du & Mordatch] 2019). more sophisticated MCMC
approaches (Lawson al. 2019; Qiu et al. 2020 Nijkamp al. 2022). regularizations (Du et al.
2021; Kumar et al. 2019)). explicit auxiliary generators (Xie et al. 2018 Yin et al. [0022 Haan et al.
2019 Gao al.) 2020p. The application of these EBM improvements for the HOT problem is
fruitful avenue for future work. For more in-depth discussion of the methods for training EBMs,
see a recent survey (Song & Kingma| 202I).
3 TAKING UP EOT PROBLEM WITH EBMS
In this section, we connect EBMs and the EOT problem and exhibit our proposed methodology. At first, we present some theoretical results which characterize weak dual objective (9) and its optimizers (...). Secondly, we develop the optimization procedure (33.2) and corresponding algorithm capable of implicitly recovering EOT plans. Thirdly, we establish generalization bounds for our proposed
method ($3.)) All proofs are situated in Appendix B
3.1 ENERGY-GUIDED REFORMULATION OF WEAK DUAL EOT
We start our analysis by taking a close look at objective (9) The following proposition characterizes
the inner min,,, optimization problem arising in (9).
Theorem (Optimizer of weak tron-trannform). Let f C(V) and x X. Then inner weak dual objective min,eep(y) G (H) (8) permits the unique minimizer μt. which is given by
dp((v) def
1
[[() - c(x,y)
exp
(13)
dy
Z(f
where Z(f,z) x) def Jy exp [(v))c(a.)) dy.
By substituting minimizer
13
obtain the close form for the weak CEOr transffrmmf:m:
(CEO (x) = GG.r = - log (ffx)) = - log (S, exp
f(y) (f())=
c(x, (()))c(x,)) y)
dy).
(14)
8
The equation I4 resembles (c, E) trannsooom Genevay
2019 Eq. 4.15) appearing in standard semi-dual EOT formulation (Genevay 2019 $4.3). For completeness, we shortly introduce
4

---

Published as a conference paper at ICLR 2024
the dual EOT and semi-dual EOT problems in Appendix relegating readers Genevvy 2019 for a more thorough introduction. In short, the particular form of weak dual EOT objective. which differs from semi-dual EOT objective, and allows to utilize EBMs, as show in 3.2 Thanks to 14 objective
( permits the reformulation:
sup FEC()) {
J log 2(, + (wnuuw)) (15)
BOT o. (I,, Q) sup Cr (f) fEC()
For given f E C()), consider the distribution dr/ (x, y) def duf (P)) We prove, that the optimization of weak dual objective (15) brings n- closer to the optimal plan '²"
Theorem (Bound on the quality of the plan recovered from the dual variable). For brevity, define the optimal value of ( by C'Coo def EOT (I,, Q). For every I C()) it holds that
Fe (f) (1)= KL KL (T)((()))) dp(x) = =KL (=*1(===)
(16)
From Theorem it follows that given an approximate maximizer f of dual objective (13), one immediately obtains distribution x¹ which is close to the optimal plan π². Actually, n' is formed by conditional distributions y (Theorem whose energy functions are given by f (adjusted with transport cost c). Below we show that f in (15) can be optimized analogously to EBMs as well.
3.2 OPTIMIZATION PROCEDURE
Following the standard machine learning practices, we parameterize functions f C()) as neural networks fo with parameters 8 E θ and derive the loss function corresponding to (15) by:
/ log x))P(x)) = fo(v)dQ(v)
I(O)
(17)
The conventional way to optimize loss functions such as 17 is the stochastic gradient ascent. In the following result, derive the gradient of L(0) 0.
Theorem 3 (Gradient of the weak dual loss L(O)). It holds true that:
L00 - --)
a (olw) duch (v)dP(x)+ / 2 30 (o()) /
(18)
Formula (88) resembles gradient of Energy-based loss, formula (TT). This allows us to look at EOT problem from the perspectives of EBMs. In order to emphasize the novelty of our approach, and, simultaneously, establish the deep connection between the optimization of weak dual objective in form 15 and EBMs, below characterize the similarities and differences between standard EBMs optimization procedure and our proposed HOT-encouraged gradient ascent following (L(
Differences. In contrast to the of EBMs, potential fe. optimized by means of loss function L, does not represent an energy function by itself. However, the tandem of cost function e and So helps recover the Energy functions of conditional distributions u!o:
(v) clz,0) fo(y) Mi
E
Therefore, one can sample from distributions udf following ULA 12 or using more advanced
MCMC approaches (Girolami & Calderhead 2011: Hoffman et al. 2014 Samsonov et al. (202)). In
practice, when estimating (18), we need samples (æ1, v1 ), (12, y/2 )...........................)).. NI yN) from distribution
/ (x, y) del dul0 (x)). They could be derived through the simple two-stage procedure:
1. Sample I1 II... IN ~ IP i..., derive random batch from the source dataset.
2. Sample yn |2) ~ unf ... yN N / e.g. performing Langevin steps 12
Similarities. Besides a slightly more complicated two-stage procedure for sampling from generative distribution So the gradient ascent optimization with (18) is similar to the gradient descent with
5

---

Published as a conference paper at ICLR 2024
[1). This allows practitioner to adopt the existing practically efficient architectures of EBMs, e.g..
Du Mordatch| 2019; Du et al. 2021 Gao et al. 2021: Zhao et al.) 2021). in order to solve EOT.
findings and detail optimization procedure in the Algorithm [ The procedure is basic, i.c. for the sake of simplicity, we specifically remove all technical tricks which typically used when optimizing EBMs (persistent replay buffers (Tieleman, 20088, temperature adjusting, etc.). Particular implementation details are given in the experiments section ($5).
Algorithm. We summarize
We want to underline that theoretical and practical setup allows performing paarretically-rroundd truly conditional data generation by means of EBMs, which unlocks the data-to-data translation applications for the EBM community. Existing approaches leveraging such applications with Energy- inspired methodology lack theoretical interpretability, see discussions in $4.
Algorithm 1:
Entropic Optimal Transport via Energy- -Basgd Modelling
Input Source and target distributions P' and Q, accessible by samples;
Entropy regularization coefficient & >> 0, cost function c(z, y) R= R8 D2 + R; number of Langevin steps K > 0, Langevin discretization step size 1) 0: basic noise std do > 0; potential network fo R., R. batch size N 0. Output:traines potential network fø- recovering optimal conditional EOT plans
for i 1, do
Derive batches { n=1 X P, PunK H Q of sizes N;
Sample basic noise (0) ~ NNO,, oo) of size N;
for k 1,2,.... K do
Sample 7(k) = (2/h1 N
(2) m==1 where 2
NO, 1):
Obtain y(k) {yA j==l with Langevin step:
(k) i
(k) Yn
(k-1) + stoppgraad
Ø By [fo(y) c( vn v)1 ( + vizi. (k)
Y/
7+-
Σ fo
Un
(K)
Σ fo (vm) V..YY
aL.
Perform gradient step over by using :0:
3.3 GENERALIZATION BOUNDS FOR LEARNED ENTROPIC PLANS
Below, we sort out the question of how far a recovered plan is from the true optimal plan nT"
In practice, the source and target distributions and Yu = {ym 1 M m=1 ~ Q. i.e., finite datasets. Besides, the available potentials f come from restricted functional class F C()), e.g., f are neural networks. Therefore, in practice, actually optimize the following empirical counterparl of the weak dual objective (15)
given by empirical samples XN {I J P
N
M
For Coor
def max JEF
Σ
Σ f(mm) log ((f + M m==1
max JEF
8 N
n=1
and recover f def arg max JEF
Fi Chor (f) J.A. A question arises: how close is म to the OT plan n"?
Our Theorem 4 below characterizes the error between .ं and nt arising due to approximation (J is restricted), and estimation (finite samples of P. Q are available) errors. To bound the estimation error, we employ the well-known Rademacher complexity (Shalev-Shwartz & Ben-David, 20144) 526).
Theorem 4 (Finite sample learning guarantees). Denote the functional class of weak Cror-transforms corresponding to F by FCam l-e log Z(f,-) f F). Let Rx(F, Q) and RM (F IP) denote the Rademacher complexities of functional classes and F w.r.. Q and IP for sample sizes N and M, respectively. Then the following upper bound on the error between x and π5 holds:
Estimation evoor
Approsimation error
E KL ( (T") (T*11²)] 55-
(J IP) (
4RM [T, + 1
max P (f) (19) FEF

---

Published as a conference paper at ICLR 2024
where the expectation is taken over random realizations of datasets XN ~ P.P.Y YM Q of N, M.
We note that there exist many statistical bounds for EOT (eenvvay] 2019 Genevay et al., 2019
Rigollet & Stromme 2022, Mena & Niles- Weed 2019 Luise et al. 2019 del Barrio et al. 2023 yet they are mostly about sample complexity of EOT. L... estimation of the OT cost value, or accuracy of the estimated barycentric projection - Jy y dn" (yyl:) in the non-parametric setup. In contrast to these works, our result is about the estimation of the entire OT plan x² in the parametric setup. Our Theorem could be improved by deriving explicit numerical bounds. This can be done by analyzing particular NNs architectures, similar to (Klusowski & Barron 2018 Sreekumar al. 2021). We leave the corresponding analysis to follow-up research.
4 RELATED WORKS
In this section, we look over existing works which are the most relevant to our proposed method. We divide our survey into two main parts. Firstly, we discuss the EBM approaches which tackle similar practical problem setups. Secondly, we perform an overview of solvers dealing with Entropy- regularized OT. The discussion of general-purpose OT solvers is available in Appendix A.2
4.1 ENERGY-BASE MODELS FOR UNPATRED DATA-TO-DATA TRANSLATION
Given a source and target domains X' and V, accessible by samples, the problem of unpaired data- to-data translation (Zhu et al., 2017) to transform a point x E X' from the source domain to corresponding points vi 322 x's content. In order to solve this problem. (Zhao & Chen 202T Zhao al., 202T) propose to utilize a pretrained EBM of the target distribution Q. initialized by source samples F ~ P. In spite of plausibly-looking practical results, the theoretical properties of this approach remain unclear. Furthermore, being passed through MCMC, the obtained samples may lose the conditioning on the source samples. contrast, our proposed approach free from the aforementioned problems and can be tuned to reach the desired tradeoff between the conditioning power and data variability. The authors of (Xie et al, (02T) propose to cooperatively train CycleGAN and EBMs to solve unpaired 12I problems. However, in their framework, EBMs just help to stabilize the training of 12I maps and can not be considered as primal problem solvers.
C y from the target domain while "preserving" some notion of
4.2 ENTROPY-REGULARIE
OT
To the best of our knowledge, all continuous EOT solvers are based either on KL-guided formulation
(Genevay et al. 2016, Seguy et al. 2018 Daniels al., 2021) unconditional entropic one
33
with its connection to the Schrödinger bridge problem (Finlay et al.) 2020 Bortoli et al. 2021 T
Gushchin et al. 2023 Chen al. 2022 Shi al. 2023). Our approach seems to be the first which
lakes advantage of conditional entropic formulation (4). Methods (Genevay et al. 2016 Seguy et al.
2018 Daniels et al.
202T exploit dual form of 2). see (Genevay) 2019 Eq. 4.2), which is an unconstrained optimi iatioi problem w.r.t. a couple of dual potentials (ur, v). However, (Genevay
al. 2016 | et al., 2018) do not provide direct way for sampling from optimal conditional
plans π [v since it requires the knowledge of target distribution Q. In order to leverage this issue, (Daniels et al. 202T) proposes to employ a separate acore-based model approximating Q. At the inference stage (Daniels al. 2021) utilizes MCMC sampling, which makes this work to be the closest to ours. The detailed comparison is given below:
1.
The authors of (Daniels ct al. 2021 optimize dual potentials (u, v) following the dual form of () This procedure is unstable for small it requires the exponentiation of large numbers which are of order e!_ At the same time, a "small E" regime is practically important for some downstream applications when one needs close-to- dodie---ettteeiiittiee plan between X' and y domains. On the contrary, our Energy- based approach does not require exponent computation and can be adapted for small by proper adjusting of ULA (12) parameters (step size, number of steps, etc.).
2. In (Daniels et al. 2022), it is mandatory to have three models, including third-party casre-based model. Our algorithm results in a single potential fo capturing all the information about the OT conditional plans and only optionally may be combined with extra generative model (85.3).
The alternative EOT solvers (Finlay et al. 2020: Bortoli al., 2021 Gushchin et al. 2023 Chen ct al. 2022] Vargas et al.] |021]|Shi et al., 2023) are based on the connection between primal EOT
(3) and the Schrödinger bridge problem. The majority of these works model the EOT plan
7

---

Published as a conference paper at ICLR 2024
time-dependennt stochastic process with learnable drift and diffusion terms, starting from P the initial time and approaching Q at the final time. This requires essource-consumen techniques to solve stochastic differential equations. Morcover, the aforementioned methods work primarily with the quadratic cost and hardly could he accommodated for a more general case.
5 EXPERIMENTAL ILLUSTRATIONS
In what follows, we demonstrate the performance of our method on toy 2D scenario, Gaussian-to- Gaussian and nigh-dimensional AFHQ Cat/Wild-+Dog image transformation problems solved using the latent space of pretrained StyleGAA2-ADA, (Karras et al. 2020). In the first two experiments the cost function is chosen be squared halved 2a norm: c(z, y) 2 || yIβ3. while in the latter
case, it is more tricky and involves StyleGAN generator. An additional experiment with Colored MNIST images translation setup is considered in Appendix C.I
Our code is written in PyTorch and publicly available at https://githiin /ttp://ggttuu.... com/ rot htwp://nerrgggggeeeeeeeee
ropicii The actual neural network architectures as well as practical training setups are disclosed in the corresponding subsections of Appendix D
5.1
TOY 2D
We apply our method for 2D Baausiann-SSiisvene case and demonstrate the qualitative results on Figure for Entropy regularization coefficients 0.1, 0.001. Figure 26 shows that our method succeeds in transforming source distribution P to target distribution Q for both Entropy regularization coefficients. In order to ensure that our approach learns optimal conditional plans x (y) (r) well, and correctly solves EOT problem, we provide Figures 2c and 2d On these images, pick several points x E X and demonstrate samples from the conditional plans x( [æ), obtained either by our method (TT (11(-1>) wf*) or by discrete EOT solver (Flamary et al. 2021). In contrast to approach, the samples generated by the discrete EOT solver come solely from the training dataset. Yet these samples could be considered fine approximation of ground truth in 2D.
Source samples y Cordit. camples [1
Source samples x Condit, comples y Rioo
Scunce datriuulonn
Fiiled distrisution
Source samples Cordit samples y-
Source samples Londi: samples y- Reoe
Target ditribution y
Fited distrioution
-1 (d) Discrete conditional plans (Dor(-(a): 10 (up), 5 10 (down)
(a) Input and target distributions P and Q
(b) Our fitted distributi- ons; = 10-1 (up), 10 (down).
(c) Our fitted conditional plans T(+2z); E = 10-1 (up), 10 (down)
Figure Performance of Energy-guided EOT on Gaussian -wisssol 2D setup.
5.2 GAUSSIAN-TO-GAUSSSIA
transformation tasks in various dimensions (D2 D2 2, 16, 64, 128), for which the exact optimal EOT plans are analytically known (Janati) et al. 2020 We choose & = 0.1. 1, 10 and compare the performance of our approach with those,
Here we validate our method
described in $.2 We report the BW2 -VVP metric, see Appendix D.2 for the explanation, between the learned and optimal n* plans in Table [ As
see, method manages to recover the optimal plan rather well compared to baselines. Technical peculiarities are disclosed Appendix [
8

---

Published as a conference paper at ICLR 2024
Method
2
16
64
128
2
16
64
128
16
64
128
Ours
0.01
0.2
0.36
0.97
0.03
0.1
0.26
0.31
0.11
0.09
0.18
0.29
Gushchin
2023
0.08
0.19
0.04
0.12
[Borloh el at
2021
1.97
4.22
3.44
5.87
0.88
1.29
2.32
2.43
m/a
m/a
n/a
n/a
Chen
[Air]
1.7
2.64
n/a
n/a
Chnh et mi
212
Joint)
0.45
4.8
5.6
5.28
0.07
0.21
0.34
0.58
0.88
2.12
2.77
3.45
Wargas et al
2021
2.96
2.94
4.06
4.0
0.3
0.9
1.34
1.8
0.1
0.21
0.72
1.14
DDmels
2021
n/a
1.36
4.62
5.33
3.22
5.57
3.13
4,98
Seguy 0 a.
2018
11.3
28.6
39.4
57.4
6.77
14.6
25.6
47.1
10.2
14.6
28.9
40.1
(a) E
0.1
(b) E
(c) E = 10
Table 1: Performance (BW3-UVP)) of tneegy-ggidded EOT (ours) and baselines on
Gaussian->Gaustnn tasks for dimensions D = 2, 16, 64, 128 and reg. coefficients 0.1, 1, 10.
5.3 HIGH-DIMENSIONAL UNPAIRED IMAGE-TO-IMAGE TRANSLATION
In this subsection, we deal with the large-scale unpaired I2I setup, As with many other works in EBMs, e.g.. (Zhao & Chen 2021 Tiwary et al. 2022], we consider learning in the latent space. We pick pre-trained StyleGANV2-Ade (Kar- ras et al. 2020) generator G for Dogs AFHQ 512 512 dataset and consider Cat-+Dog (and Wild->Dog) unpaired translation. As TP, we use the dataset of images of cats (or wild); as Q. we use AO0 1512), i.c., the latent distribution of the StyleGAN. We use our method with learn the EOT between P and Q with cost ] || image generated from the latent code z. Note that our method trains only one MLP network fo acting on the latent space, which is then used for inference (combined with G). Moreover, our approach does not need a generative model of the source distribution P. and does not need encoder (data2latent) networks. The gualitative results are provided in Figures and 3 Our method allows us translate the images from one domain to the other while maintaining the similarity with the input image. For more examples and qualitative comparisons, compare our approach with popular unpaired I2I models ILVR (Choi et al.) [202T). SDEdit (Meng al., 2022), EGSDE (Zhao al. 2022). CycleGAN (Zhu al.. 2017). MUNIT (Huang al. 2018) and StarGANv2 (Choi et al. (020)) the obtained FID metrics are reported in Table As we can see,
Figure 3: AFII 512 512 Wild~>Dog unpaired 121 by our method the latent space of StyleGAN2-ADA. Left: source; right: translated.
1 to
G(=)ll i.e., e² between the input image and the
Appendix C.2 For the quantitative analysis,
our approach achieves comparable-to to-S---T----TAAA quality.
Method
Ours
ILVR
SDEdit
EGSDE
CycleGAN
MUNIT
StarGAN v2
Cat + Dog FID
56.6
74.37
74.17
51.04
85.9
104.4
54.88
Wild + Dog FID
65.8
75.33
68.51
50.43
Table Baselines FID for Cat -> Dog and Wild -> Dog-
6 DISCUSSION
Our work paves a principled connection between EBMs and EOT. The latter an emergent problem in generative modelling, with potential applications like unpaired data-to-data translation (Korotin lt al. 20236). Our proposed FEBM-based learning method for EOT is theoretically grounded and we provide proof-of-concept experiments. We believe that our work will inspire future studies that will further empower EOT with recent EBMs capable of efficiently sorting out truly large-scale setups
(Du al. 2021 Gao al. 2021 Zhao et al. (02TT.
The limitations of our method roughly match those of basic EBMs. Namely, our method requires using MCMC methods for training and inference. This may be time-consuming. For the extended discussion of limitations, see Appendix
The broader impact of our work is the same as that of any generative modelling research. Generative models may be used for rendering, image editing, design, computer graphics, etc. and simplify the existing digital content creation pipelines. At the same time, it should be taken into account that the rapid development of generative models may also unexpectedly affect some jobs in the industry.
'II scores of the baselines are taken from (Zhao et al. (022). In order estimate FID, we downscale the generated images 256 256 for fair comparison with the alternative methods.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGEMENTS
The work was supported by the Analytical center under the RF Government (subsidy agreement 000000173002211500000 Grant No. 70-2021-00145 02. 11.201....
REFERENCES
David Alvarez-Melis and Nicolo Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems,
2020.
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. /tt.ttss:://eeeeeeiiiie.eeee htt.://eeee//oormme-oo-o/oooO/
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev. Neural op- timal transport with general cost functionals. The Twelfth International Conference on Learning Representations, 2024. URL
net/ forum 2ooouniiddgli22
Julio Backhoff- Backhoff-Veraguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):203, 2019.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL httptttt/ppeeeeeiiieeeee http://www.....eoooroomoeeo/eooa
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Represen- tations, 2022. httthtttt:///onntt
net htt.:/oormmnddiiiiiiita
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Condi- tioning method for denoising diffusion probabilistic models. in 2021 1 ieee. In CVF international conference on computer vision (ICCV), pp. 14347-14356, 2021.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pallern recognition, pp. $188-8197, 2020.
Christian Clason, Dirk A Lorenz, Hinrich Mahler, and Benedikt Wirth. Entropic regularization of continuous optimal transport problems. Journal of Mathematical Analysis and Applications, 494 (1):124432, 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 34:12955-12965, 2021.
Eustasio del Barrio, Alberto González Sanz, Jean-Michel Loubes, and Jonathan Niles Weed. An improved central limit theorem and fast convergence rates for entropic transportation costs. SIAM Journal on Mathematics of Data Science, 5(3):639-669, 2023.
Ishan Deshpande, Ziyu Zhang, and Alexander G Schwing. Generative modeling using the sliced wasserstein distance. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3483-3491, 2018.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
10

---

Published as a conference paper at ICLR 2024
Yilun Du, Shuang Li, B. Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. In Proceedings of the 38th International Conference on Machine Learning (ICML-21,, 2021.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen Chen. Scalable computations of wasserstein barycenter via imput convex neural networks. Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1571-1581. PMLR, 18-24 Jul 2021. URL htttps://ppocceeddingg
mlr. wttppesss/v133/aa222222 l
Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL Featured Certification.
hittps://ppenrevee net/ wttp:///oooommmmm21n00oeemm
Chris Finlay, Augusto Gerolin, Adam M Oberman, and Aram-Alexandre Pooladian. Learning normalizing flows from entropy- -anrtp---katoovicchhh potentials. arXiv preprint arXiv:2006.06033 2020.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. URL http://tt:///mt /mp:: ww.. www.oo/ggpppp/////e.eo 2200///2222244521 ht.ml
Ruiqi Gao, Erik Nijkamp, Diederik P Kingma, Zhen Xu, Andrew M Dai, and Ying Nian Wu. Flow contrastive estimation of energy-based models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7518-7528, 2020.
Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energy-based models by diffusion recovery likelihood. In International Conference on Learning Representations, 2021. htatt::://peereeeeew.eeeee/ee/
Milena Gazdieva, Litu Rout, Alexander Korotin, Andrey Kravchenko, Alexander Filippov, and Evgeny Burnaev. An optimal transport perspective on unpaired image super-resolution. arXiv preprint arXiv:2202.01 2022.
Aude Genevay. Entropy-regularized optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUB), 2019.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. Advances in neural information processing systems, 29, 2016.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pp. 1574-1583, PMLR, 2019.
Mark Girolami and Ben Calderhead. Riemann manifold langevin and hamiltonian monte carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 73(2): 123-214, 2011.
Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273(11):3327-4405 2017.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnev.. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023.
11

---

Published as a conference paper at ICLR 2024
Tian Han, Erik Nijkamp, Xiaolin Fang, Mitch Hill, Song-Chun Chun Zhu, and Ying Nian Wu. Divergence triangle for joint training of generator model, energy-based model, and inferential model. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 8670-8679, 2019.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 1778)1177--1800, 2002.
Matthew D Hoffman, Andrew Gelman, et al. The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15()):1933-1633 2014.
Weining Hu, Meng Li, and Xiaomeng Ju. Improved cyclegan for image-to-image translation.
Xun Huang, Ming- Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to-image translation. In Proceedings of the European conference on computer vision (ECCV), pp. 172-189, 2018.
Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33:10468-10433 2020.
Tero Karras, Miik Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in neural information processing systems, 33: 11044144-,, 2020.
Jason M Klusowski and Andrew R Barron. Approximation by combinations of relu and squared relu ridge functions with Ed and e controls. IEEE Transactions on Information Theory, 64(12): 7649-7656, 2018.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations. 2021a. URLphttss:////peeeeeee
net http://wwww.ee//ouutmmmmammm httt://ooormmmlddebooeoeo
EXsa
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021b. URL /ttppss//ope...re...ee. net. ret.:/orrumi--33PAS555
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Advances in Neural Information Processing Systems, 2022a. URL https GiEnzxTnaMi
Kyunghyun Cho Advances //openreview.oo http::/eee/tfoooeeeoeeoo
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein GANs are not optimal transport? In Thirt Conference on Neural Information Processing
Systems Datasets and Benchmarks Track, 2022b. URL httlhtttp:///o/mm hta.t::://opennee.eww.eee.e wet. formm?
1d=vtsapi-dGic
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport.
In The Eleventh International Conference on Learning Representations, 2023a. URL https
//openreview. net. mnee/forumidddiieeoeeeeeeet Zuc. Uma
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In
The Eleventh International Conference on Learning Representations, 2023b. URL https
//openreview. wet.iiewwwwettee.... toteet/fooommmmddBBBR....
Rithesh Kumar, Sherjil Ozair, Anirudh Goyal, Aaron Courville, and Yoshua Bengio. Maximum
entropy generators for energy-basee models. arXiv preprint #rXiv:/901.085000 2019.
John Lawson, George Tucker. Bo Dai. and Rajesh Ranganath. Energy-inspired models: Learning with sampler-induced distributions. Advances Neural Information Processing Systems, 32, 2019.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. tutorial on based learning. Predicting structured data, 10)), 2006.
12

---

Published as a conference paper at ICLR 2024
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations. 2023. bttps:::::
Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large-scale optimal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.06633 2020.
Giulia Luise, Saverio Salzo, Massimiliano Pontil, and Carlo Ciliberto. Sinkhorn barycenters with free support via frank-wolfe algorithm. Advances in neural information processing systems, 32, 2019.
Shaojun Ma, Shu Liu, Hongyuan Zha, and Haomin Zhou. Learning stochastic behaviour from aggregate data. In International Conference on Machine Learning, pp. 7258-7267. PMLR, 2021.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference Machine Learning, pp. 6672-6681. PMLR, 2020.
Simone Di Marino and Augusto Gerolin. An optimal transport approach for the schrödinger bridge problem and convergence of sinkhorn algorithm. Journal of Scientific Computing, 85(2):27, 2020.
Gonzalo Mena and Jonathan Niles Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. Advances in Neural Information Processing Systems, 32, 2019.
Chenlin Meng. Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2022. URL http::/s////peeeeeeee.....
wwwp:/wwwwwwwwwweeeeeeeeeeeeeeee formm?
id=aBsCjcPu_tt
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 15256, 2021.
Tuan Nguyen, Trung Le, He Zhao, Quan Hung Tran, Truyen Nguyen, and Dinh Phung. Most: Multi-source domain adaptation via optimal transport for student-teacher learning, In Uncertainty in Artificial Intelligence, pp. 225-235. PMLR, 2021.
Erik Nijkamp. Mitch Hill, Tian Han, Song-Chun Zhu, and Ying Nian Wu. On the anatomy of meme-based maximum likelihood learning of energy- based models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5272-5280, 2020.
Erik Nijkamp. Ruiqi Gao, Pavel Sountsov, Srinivas Vasudevan, Bo Pang, Song-Chun Zhu, and
Ying Nian Wu. MCMC should mix: Learning energy-based model with neural transport latent space MCMC. In International Conference on Learning Representations, 2022. URL https
//openreview net. Btte/fformmmId=CC930aaaa
Tomohiro Nishiyama. Convex optimization on functionals of probability densities. arXiv preprint
arXiv:2002.00 208.040002.06888, 2020.
Marcel Nutz. Introduction entropic optimal transport. Lecture notes, Columbia University, 2021.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning, 11(5-6)3355-67, 2019.
Yixuan Qiu, Lingsong Zhang, and Xiao Wang. Unbiased contrastive divergence algorithm for training energy- latent variable models. In International Conference on Learning Representations, 2020.
Philippe Rigollet and Austin J Stromme. On the sample complexity of entropic optimal transport. arXiv preprint arXiv:2206. 13472, 2022.
Gareth Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli, pp. 341-363, 1996.
13

---

Published as a conference paper at ICLR 2024
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022. URL [RRl/ttts///opeeeeee// htte:/Eorunnna11ss12222
Giorgi Rukhaia. A FreeForm Optics Application of Entropic Optimal Transport. PhD thesis, Université Paris sciences et lettres, 2021.
Ruslan Salakhutdinov, Andriy Mnih, and Geoffrey Hinton. Restricted boltzmann machines for collaborative filtering. In Proceedings of the 24th international conference on Machine learming, pp. 791-798, 2007.
Sergey Samsonov, Evgeny Lagutin, Marylou Gabrié, Alain Durmus, Alexey Naumov, and Eric Moulines. Local-global MCMC kernels: the best of both worlds. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho Advances in Neural Information Processing Systems, 2022. http:/ttss///ppeeeeeeeee
net www.eeeeetttorreeeeee ottp://oorrmmm1t//baaaaaac3e
Filippo Santambrogio. Optimal transport for applied mathematicianss. Birkäuser, NY, 55(58-63):94, 2015.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal- analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):9289943, 2019.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Confer- Learning Representations, 2018. URL CRLhttt::/// /ttp////pp///eieeee/e//e /ttpppp////eeerrreeeee net wwwp:/rutttoume..eem forum?ide B121plbRM
Shai Shwler-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
Jian Shen, Yanru Qu, Weinan Zhang, and Yong Yu. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
forum? Lttp://rmm-2000000H00T
Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1-11, 2015.
Yang Song and Diederik P Kingma. How to train your emergy-based models. arXiv preprint
000.022001.03288, 2021.
Sreejith Sreekumar, Zhengxin Zhang, and Ziv Goldfeld. Non-asymptotic performance guarantees for neural estimation of f-divergences. In International Conference on Artificial Intelligence and Statistics, pp. 3322-3330. PMLR, 2021.
Amirhossein Taghvaci and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application improved training for gans. arXiv preprint arXiv:1902.07197. 2019.
Tijmen Ticleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference on Machine learning, pp. 1064-1071, 2008.
Piyush Tiwary, Kinjawl Bhattacharyya, and Prathosh AP. Boundary preserving twin energy-based- models for image to image translation. 2022.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood, Entmpy, 23(9): 134, 2021.
14

---

Published as a conference paper at ICLR 2024
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Rui Wang, Ruiyi Zhang, and Ricardo Henao. Wasserstein uncertainty estimation for adversarial domain matching. Frontiers in big Data, 5, 2022.
Yisen Wang, Bo Dai, Lingkai Kong, Sarah Monazam Erfani, James Bailey, and Hongyuan Zha. Learning deep hidden nonlinear dynamics from aggregate data, arXiv preprint arXiv: 1807.08237 2018.
Jianwen Xie, Yang Lu, Song-Chum Zhu, and Yingnian Wu. A theory of generative convnet. In International Conference on Machine Learning, pp. 3635-2644, PMLR, 2016.
Jianwen Xie, Yang Lu, Ruiqi Gao, Song-Chun Zhu, and Ying Nian Wu. Cooperative training of descriptor and generator networks. IEEE transactions on pattern analysis and machine intelligence, 42(1):27-45, 2018.
Jianwen Xie, Zilong Zheng, Xiaolin Fang, Song- Song-Chun Zhu, and Ying Nian Wu. Learning cycle- consistent cooperative networks via alternating meme teaching for unsupervised cross-domain translation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10430-10440, 2021.
Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. In International Conference Machine Learning. pp. 6882-6892. PMLR, 2019.
Xuwang Yin, Shiying Li, and Gustavo K Rohde, Learning energy-based models with adversarial training. In Computer Rision-ECCy 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings. Part V, pp. 209-226. Springer, 2022.
Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations. Advances in Neural Information Processing Systems, 35:3609-3622 -3233, 2022.
Yang Zhao and Changyou Chen. Unpaired image-to-image translation via latent energy transport. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16418-16427, 2021.
Yang Zhao, Jianwen Xie, and Ping Li. Learning energy-based generative models via coarse-to- fine expanding and sampling. In International Conference on Learning Representations, 2021.
Jun Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. Proceedings of the IEEE international conference computer vision, pp. 2223-2232. 2017.
15

---

Published as a conference paper at ICLR 2024
A
EXTENDED BACKGROUND AND RELATED WORKS
A.1
QUAL/SEMI-DUAL EOT PROBLEMS AND THEIR RELATION TO WEAK DUAL EOT PROBLEM
Dual formulation of the EOT problem. Primal EOT problem (2) has the dual reformulation Genevay et al. 2016). Let E C(2) and v E C()). Define the dual functional by
(u(u,)) v) def u(z)dP(x)+ ((y)d((w)= - exp
(((x) + v(v) c(x.))
d d x (z, v).
XxY
(20)
Then the strong duality holds (Genevay et al. 2016 Proposition 2.1), i.c.
BOT 0.5
sup Fe(u, v).
(21)
(EC(X).EEC())
The sup here may not be attained in C()), yet it is common to relax the formulation and consider E _ (IP) and E _ (Q) instead (Marino & Gerolin 2020). In this case, the sup becomes max (Genevay 2019. Theorem 7). The potentials u², , which constitute a solution of the relaxed (2T) are called the (Entropic) Kantorovich potentials. The optimal transport plan n" which solves the primal problem 2 could be recovered from pair of Kantorovich potentials (u", ",") as
11" (x) + v²()) c(x,y) dT(())dQ(v)
d (x, v) exp
(22)
From the practical viewpoint, the dual objective (1 is unconstrained maximization problem which can be solved by conventional optimization procedures. The existing methods based on (21 as well as their limitations and drawbacks are discussed in the related works section, $4.2
Semi-dual formulation of the EOT problem. Objective (2T) is a convex optimization problem. By fixing function C()) and applying the first-order optimality conditions for the marginal optimization problem suPued(()) F' (u, v). one can recover the solution of
Uhe def
arg max Fe v) mEC(X)
(23)
the closed-form:
(0* log (// exp
e(v)) - c(x,,)) aQ(w)
(24)
Function of is called (c, -ttranssmoo (Genevay| 2019). By substituting the argument C(X) of the max with BhE in equation (2T) and performing several simplifications, one can recover the objective
sup / ohe (x)dP(()) I (())d(().
BOT Q)
(25)
vEC())
This one is called the semi-dual formulation of EOT problem (Genevay 2019
$4.3). It is not so popular as the classical dual problem (2T) since the estimation of oh is triviaivial, yet it has a direct relation formulation (.) which forms the basis of our proposed method. In order comprehend this relation, below we compare (c, -ttraassoooo and min, G2, (J4), given by (14).
Correspondence between (semi-) dual EOT and weak dual EOT. As we already pointed out in the main part of the manuscript, equation (14) which then appears in weak dual objective (15) looks
similar to (c, --ttansfform 24 The difference is the integration measure, which is Q in the case of (4 and the standard Lebesgue one in our case [4)..
From the theoretical point of view, such dissimilarity not significant. That why semi-dual (25) and weak dual ( optimization problems are expected to share such properties as convergence, existence of optimizers and so on. In particular, there relation between potentials u, v that appear in (semi-) dual EOT problems (21 25), and our optimized potential f from (9).
Let Q E (a), and Eq y IR be the energy function of Q, i.e.,
dQ(s) dy
exp (-Eq(v)).
Consider the parameterization of potentials by means of f as follows:
v(y) +- f(v) (o(())
(26)
16

---

Published as a conference paper at ICLR 2024
Then,
v(y) - cx
5 (x)) - log
exp
dQ(x)
S exp
[(y) - c( v)
dQ(u) dy dy
- log
exp (1)
8
f(y) - cxt, v)
- log
exp
exp exp(Eg(u)) exp (- (B)( dy
+ Const Q)
800 log Z(, (onst(Q)
(27)
By substituting 26
27 in 25
we obtain that semi-dual EOT objective
25
recovers our weak dual
objective 15) up to reparameterization 26
{ U0e (0)()))(()) + ((0))Q(1)) vEC(V)
EOT (P, Q) 6,6
sup
fEC())) {S = Log + f(oaaau }
Consty (Q) sup
(cH((Q)
Strictly speaking, the optimization class =f different from C(Y) that appears in (15). However, under mild assumptions on Eg the corresponding optimization problems are similar. One important consequence of the observed equivalence is the existence optimal potential f (not necessary to be continuous) which solves weak dual objective p It can be expressed through optimal Kantorovich potential o* by f =v² -EE-.
C(D)) in the equation above is
From the practical point of view, the difference between (4) and 24 is much more influential.
Actually, it is the particular form of internal weak dual problem solution (4) that allows us utilize
EBMs, see $32
A.2 DISCRETE AND CONTINUOUS OT SOLVERS REVIEW
Discrete OT. The discrete OT (DOT) is the specific domain in OT research area, which deals with distributions supported on finite discrete sets. There have been developed various methods for solving DOT problems (Peyré et al. 2019), the most efficient of which deals with discrete EOT (Cuturi 2013). In spite of good theoretically -ereticcallygrouuneed convergence guarantees, it is hard to adopt the DOT solvers for ditti---isti sampling and mapping, which limits their applicability in some real-world scenarios.
Continuous OT. In the continuous setting, the source and target distributions become accessible only by samples from (limited) budgets. In this case, OT plans are typically parameterized with neural networks and optimized with the help of SGD-like methods by deriving random batches from the datasets. The approaches dealing with such practical setup are called continuous OT solvers.
There exists lot of continuous OT solvers (Makkuva al. 2020 Korotin al. 202 la Fan al. 2023 Xie et al. 2019 Rout et al.) 2022 Gazdieva et al. 2022 Korotin et al. 2022b) Taghvaei & Jalai 2019 Liu al. 2023). However, the majority of these methods model OT as a deterministic map which for each input point assigns a single data point rather than distribution π (z) Only limited number of approaches are capable of solving OT problems which require stochastic mapping, and, therefore, potentially applicable for our EOT case. Here we exclude methods designed specifically for EOT and cover them in the further narrative.
The recent line of works (Korotin et al.) 2023b Asadulaev et al.) 2024) considers dual formulation of weak OT problem (Gozlan et al./ 2017) and comes up with max min objective for various weak
(Korotin et al. 2023b a) and even general (Asadulaev et al. 2024) cost functionals. However, their proposed methodology requires the estimation of weak cost by samples, which complicates its application for EOT. An alternative concept (Xie et al. (019)) works with primal OT formulation (I) and lifts boundary source and target distribution constrants by WGAN losses. It also utilizes sample estimation of corresponding functionals and can not be directly adapted for EOT setup.
17

---

Published as a conference paper at ICLR 2024
B PROOFS
B.1 PROOF OF THEOREM
I
Proof. In what follows, analyze the optimizers of the objective min G (1) introduced in
8 Let μ E P())). We have:
G / c(r, v)dp(v) + / log
dulul du(v)
-11 s(v)dp(y)
dy
dx, v) - (()) + +oog du(v) du(()
dy
/ ( log x)) + log Z( - ( - cx, 4)
du(y) dy
+ log
du(y)
does not depend
= lug duá(u) m1g
- log Z(f: Z(f.z)) + 2( ( = -e log Z(f,2) x) + =KL (μllи²)
duf(v) dy
du(y) du(v)
log dy
(28)
The last equality holds true thanks to the fact that mn E Pac())) and Vy E y duf(w) dy
0. This
leads to the conclusion that the absolute continuity of μ. (μ < X, where À the Lebesgue measure on )) is equivalent the absolute continuity of w.r.t. lol then the last equality in the derivations above reads as +00 = + From (88 we conclude that 4d arg min Gx. (p).
(c.). In particular, if $ Pa
PEP())
B.2 PROOF OF THEOREM 2
Proof. Recall that we denote the optimal value of weak dual objective BOT,, (P, Q) given by formula (4), thanks to the strong duality, i.c.,
9 by Feeer equals to
clx, v) (z, v) f H H(T (1T()))))PP((
(x,
(29)
OBOT
XxX
For potential f E C()), our Theorem
yields:
Furm (f) / 92, g (1)dd(())))))+/ f(y)d((v) Eq. 14
& log (() / f(v)d((v) (30)
In what follows, we combine 29 and 30
Fw.
Frror (f)
f c(2,y)dr y)d (x,))= / H( ((="(olz))d(())+) / log z) dl((x)
S/1 f()) dQ(y)
=dr (a))
=d (v)
Jrn dr) (x. - ((*x))))) + / log ())
f(v) - c(x.v)) y) dx dx'(x, +
-E xxX
log Z(f,z) x) dn (I.0) H(TT
xxy
oxp ((o)))
exp( ((v) ((0)-ote
-o{ / log Z(f,x)
f(y) c(2,))
dm (x: v) (o.,.) - H(* ((T(()))))P())
exp
xxX
E
Anlivi
(d((()) d (x, v) - / (yl=)))()))))
--
log
Xxy
dy
18

---

Published as a conference paper at ICLR 2024
[
(d) dr (v(z) di (x) +e- dy
/x log
(dT"(())) dr (v(x)))((x)
- x 1y log
x
dy
=dP(x)
1 { log
(dx" (u)xz) (dm"(u)lz)
(udl())) ) dr (vlla) farp(ao) ) dy
log
dy
/ KL (+x())=)) dlP(.))
In the last transition of the derivations above, we use the equality
J log
dx (u)(x)
(u)( dr d (ylx) KI. (+(x)))l)
log
dy
dy
holds true due the same reasons, as explained in the proof of Theorem T after formula (8
We are left to show that fxx KL (IT (122) diP(x) = KL (**1x//) Since dr (z. dr (v)x))P(x) and dx (z, v) dud (1)))( we derive:
/ KL ( (+xx)l144 dlp()) = J / log x
dd (vlx)
dr (wll)))dP)()) =
1 log
dd (v)x))P(())
dr (ul=))P(())
dut (w)
dr (xx y)
Ixxy
log
dr/ (x,
dr (I, y) = KL KL (='(l=/)
which completes the proof.
B.3
PROOF OF THEOREM 3
Proof. The direct derivations for 18 read as follows:
0 (0)
- / a log x(o. + /, 50 (o()))(()))
a
00
00
/
0 exp
[So(y) c(r,y) dy dP(x)) + / 00 fo())
x
z( x)
00 y
-E =/1 I(ferx) /
So(u)
fo(y) - c(x.))) y)
dy dd(x)) + / a
36
00 fo(v)d(()))
exp
/ y 00 fo(w)
fo(y) c(x y)
dy dl(xx) + / a
00 fe (o
exp
2(Jo,
=dw!? (v)
/
dud (v)dP(x)+ + / a (o
30 fo(y)
which finishes the proof.
B.4 PROOF OF THEOREM 4
To begin with, for the ease of reading and comprehension, we recall the statistical learning setup from 333 In short, X and YM are empirical samples from IP and Q, respectively, F C(V) is function class which we are looking for a potential f. which optimizes the empirical weak dual objective
N
M Σ m=1
Fag ( = - N
Σ
log Z(J xn) + M
f(ym),
7=
19

---

Published as a conference paper at ICLR 2024
Fu (f). Additionally, we introduce f def arg max fEF Feor (f). ie., /= arg maxx Copr optimizes weak dual objective (9) in the function class under consideration. Following statistical generalization practices, our analysis utilizes Rademacher complexity. We recall the definition below.
Definition (Rademacher complexity R. (N (1), Let F be function class and p be distribution. The Rademacher complexity of F with respect to μ of sample size N is
def
sup Σ ((xn)on } SEF n=1
RN 11)
where \N μ are mutualll independent, fontn=1 are mutually independent Rademacher random variables, ie.,, Prob( = =1] Prob (0n -1}) 0.5, and the expectation taken with respect to both {an] and {on)n=)
L.
Now we are ready to verify our Theorem
Proof. Our Theorem yields that
eKL TT
Can
Fur
Below we upper-bound the right-hand side of the equation above.
Coor (f)
Ceoo
Fw Cron (5²)
F. (S²) -Fay (J) + Coor (f)
F Cum (J)
5
F
FFeoor (1=))
+
(31)
Fuoo
(f¹)
Feeor (f)
(32)
Te (
Fry Cran (f)
(33)
Analysis of (31). Equation (31) relates approximation error and depends on the richness of class F. The detailed investigation of how could the approximation error be treated in the context of our Energy-guided EOT setup and, more generally, EBMs is an interesting avenue for future work.
Analysis of 32). Similar to (Taghvaei & Jalali, 2019 Theorem 3.4), we estimate 32 using the
Rademacher complexity bounds. First, we need the following technical Lemma.
Lemma 1. For each particular samples XN, YM, there exists f E such that: Fe (ft) Fu (f)
FC (f) 1 Fin (f)
Proof. Let's consider Por (1+)
Fra Ceor
(f). There are two possibilities.
1. F'um (1²) > Fru (f).
Since vf E F: Fu
Fron (f) 2 Cran (f), and, in particular, Fy (f) 2 Fu
Cror (f=), it holds:
F (5*) Fuu (Π) Firm (f*) Fu ( S
F (fz)- Fry (5²) S
Fnn Cum (1F)
Fru
Chor )
i.c., set f ft f5.
2. Pu (f) > Foy (1*).
Cao
Similar to the previous case, VS E Fu (f²) 2 F ( = F (1²) 2 F ( Analogous derivations read as:
From ( Fre (
Frw ( Cooor (1+)
Cea
Foy (f) F (f) S F (f) - Fru (f),
i.e., set f and finish the proof of the Lemma.
20

---

Published as a conference paper at ICLR 2024
Now we continue the proof of our Theorem. For particular samples XN and YM. consider f from Lemma We derive:
F (f*) - FCtor (f) S | Ceor (f)
F (f)
1/1 -e log x)) dPx))
-E log Z(f. and
f(ym)
F(y)dQ(()) ,
N
M
{/1 = log Z(f, 2) dF(x))
log Z(f, Inn)
f [ymm)
F())Q(()
Σ
M
1/1 e log Z(F, x)) dP(x))
- log x()
vn))
+ f F()DRO
M
flym)
N
M
N
CEOT (In)
/ f(y)d(())
M
[(m)
sup fef
Gror (x))d((x)
+ suup FEF
N
M
==1
Σ h(I,))
/ f(y)d(()) fer
(ym)
sup REFOOOT
h(x)dF(x)
sup
X
N
M
Recall that F.Cror
(fCror F)
1-e
log ((1
F). Thanks well-known Lemma 26.2), it holds:
Rademacher bound
Shalev-Shwart & Ben-David 2014
h(()))))
E
sup her
h())dP(a) /EF y
Σ
S 2R (FCem 1 IP),
N
n=1
M
f((m))
E
sup
f(y)d(())
Σ
S 2R (F,0),
M
where the expectations are taken with respect to samples NN and YM- Combining the results above, conclude:
Free CEOT
(f
Fu (f] s 2R (FCun P) + 2R (F, (Q)).
(34)
Analysis of 33 Similar to the previous case, obtain the inequality: Fg
(
Fg (f) S
sup beF
1 (())dP())
Σ h(zn)
N
+ sup feF
/
((v)d(())
Σ
M
{1
N
M
m2
Therefore,
E
Fru
Fro
2R (JrhoO PP) + (F,
(35)
Cum
By gathering equations
31
34
35
we prove the theorem:
E [LL
T || E Fw,
Funt
Ceor
17
(p))
+E - Fren (5))
Fru (/) _ea -1E B Coon (f) Fo
Fuy
5
F
Fu (f+)
+ E [ (FCOm , + 4R as (F, U?)
C EXTENDED EXPERIMENTS
C.I COLORED MNIST
In this subsection, we consider Colored MNIST (Gushchin et 2023 55.3). Following Gushchin
et al. 2023), we set source and target distributions P and C to be colored handwritten images of digits "2" and "3" accordingly. The entropic regularization coefficients are in range = 0.01, 0.1,
21

---

Published as a conference paper at ICLR 2024
The qualitative results of learning model directly in the data space are presented in Figure 4 As we see, our learned EOT plans successfully preserve color and geometry of the transformed images. Gen- slightly noised since we add noise to target images when training for stability. For quantitative analysis and comparison with competitive methods, we borrow the results on the ColoredMNIST transfer problem for (Bortoli et 202T). (Daiell et al. (021), (Gushchin 2023) from (Gushchin et al., 2023). Additionally, we run the code for the recent (Shi al. on our own. The methods generally work for different due to their principles, and we choose
0.01
0.1
LPIPS (VGG)
0.043
0.063
0.11
Table 3: LPIPS variance of Our method for ColoredMNIST "2"-"3" transfer.
erated images (Figures 4b
4c 4dl
al.
2023
as an admissible entropic regularization power for all methods except (Daniels et al. 2021 which struggles for small e, see the discussion in & 4.2 For it, we choose & = 25. The obtained FID metrics are reported in Table 4. For the qualitative performance baselines, see (Gushchin
8
al. 2023 Fig. 2). Besides, we provide Table 3 with LPIPS metric to show that the diversity of our method increases with &.
Method
Ours
Bortoli et al.
2021
(Shi ct al.
2023
Gushchin et al.
2023
Daniels et al 2021
FID
109
93
91.3
6.3
14.7
25)
Table 4: Baselines FID for ColoredMNIST "2" "3" transfer,
We honestly state that the FID of our approach is not good. One reason is that the default Langevin dynamic produces slightly noisy samples. FID is known to terribly react to any noise. Secondly, we emphasize that we adapt the simplest long-run EBMs with persistent replay buffer (Nijkamp et al.. 2020) for the experiment, see Appendix D.4 the details. We leave the applications of modern EBMs which can generate sharp data (Du & Mordatch
2019, Du al. 2021 for future research.
2
3
3
3
3 3
3
3
3
3
3
3
3
33
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
7
3
2
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
(a) x P
(b) y ~ *(1x),)= 0.01
(c) y #(|=). 8
0.1
(d) y #(-1=)
Figure 4: Quantitative performance of Energy-guided EOT on Colored MNIST.
C.2 EXTENDED HIGHI- DIGUIDDMENSIONAL UNPAIRED IMAGE-TO-IMAGE TRANSLATION
In this section, we provide additional quantitative results and comparisons for our considered high- dimensional 12I setup. Table 5 we show uncurated samples from our approach learned on 512 512 AFHQ Cat->Do and Wild-Dog image transfer problems. To compare our visual results with alternatives, we demonstrate the pictures generated by (Zhao et al. 2022) and (Daniels et al. 202l solvers, see Figure 6 The former demonstrates SOTA results, see Table but has relation to OT. The latter is the closest approaah to ours. For (Daniels et al. 2021 } we trained their algorithm in the same setup
used, with the latent space of the StylaGAN and transport cost c(z, y) 1 || G(=)Il3, see $ 5.3 We found that their method works only for & = 10000 yielding unconditional generation. It concordance with our findings about the approach, see discussion in and the insights from the original work, see (Daniels et al. 202T $5.1).
D EXPERIMENTAL DETAILS
General Details. For the first two experiments, we take the advantage of replay buffer B constructed as described in (Du & Mordatch 2019). When training, the ULA algorithm is initialized by samples from B with probability p 0.95 and from Gaussian noise with probability 1 0.05. For the last two image-data experiments, we also use a similar replay buffer but with p = 1, i.e., we do not update B with short-run samples.
22

---

Published as a conference paper at ICLR 2024
(a) 512 512 AFHQ Cag
(b) 512 x 512 AFHQ Wild-Dog
Figure 5: Uncurated Image-to- mmagett-Immage translation by method in the latent space of StyleGAN.
(a) 256 256 samples obtained with Zhao al. (0022)
(b) 512 x 512 samples obtained with
[Daniels al. [202I), 10000
Figure 6: Image-to-Image Cat->Dog translation by alternative methods
D.1 Toy 2D DETAILS
We parameterize the potential fo as MLP with two hidden layers and 0.2) the activation function. Each hidden layer has 256 neurons. The hyperparameters of Algorithm are as follows: K = 100, og = 1, = 1024, see the meaning of each particular variable in the algorithm listing. The Langevin discretization steps 0.05 for & = 0.1 and 0.005 for g = 0.001. The reported numbers are chosen for reasons of the training stability.
Computation complexity. The experiment was conducted on a single GTX 1080 Ti and took approximately two hours for each entropy regularization parameter value.
23

---

Published as a conference paper at ICLR 2024
D.2 GAUSSIAN-TO GAUSSSIA DETAILS
For the source and target distributions, we choose P (0, Ex) and Q - (O, Ey) with Ex and Ey chosen random. For the reproducibility, these parameters are provided in the code.
The details of baseline methods (Table D are given in Appendix E For each particular parameter E, D. our trained potential fe is given by MLP with three hidden layers, 512 neurons each. and ReLU activation. The same architectural setup is chosen for [Gushchin et. al.] (Gushchin
al. K
2023) for fair comparison. The hyperparameters of Algorithm I are the same for each E, D: 100,00 do = I.N = 1024, 0.1. At the inference stage, we run ULA steps for K 700 iterations. The only parameter which we choose to be dependent on E, D is the learning rate. affects the stability of the training. The particular learning rates which are used in our experiments More specific training peculiarities could be found in our code.
are given in Table
Bures- Wasserstein UVP metric. The BW2-UVP metric (Korotin et al. 2021b Eq. 18) is the Wasserstein-2 2 distance between distributions Ti and IT2 that are coarsened to Gaussians and nor- malised by the variance of distribution n2:
UVP( der
W (N En, (. "xz)). Var(T))
100%
In our experiment, T2 the optimal plan n² which is known to be Gaussian, and my is the learned plan nt whose mean and covariance are estimated by samples.
Computation complexity. Each experiment with particular e, D takes approximately 12 hours on a single GTX 1080 Tí.
D
2
16
64
128
1
5
4. 2.
2. 4-
1.
4.
2.
2
5.
5
10
10
10-7
10
10-6
10-5
10-5
10-3
10-3
10"
10-3
10-5
Table Learning rates for Gaussian-to- to--ausian experiment; 0.1, 1, 10 and D 2, 16, 64, 128.
D.3 HIGHI-DIMENSIONAL UNPAIRED IMAGE-TO-IMAGE TRANSLATION DETAILS
General details. In this experiment, we learn EOT between a source distribution of images P E P(BS5512X312) and Q = (O, 1s12) E ((2512), which the latent distribution of the pretrained StyleGAN G. We use non-euclidean cost c(2,
_ || G(=) 1/ Below we describe the primary
idea behind this choice. Consider the pushforward distribution Qumbi def G₁Q. In case, it is the parametric distribution of AFHQ Dogs. Thanks to our specific cost function, the learned optimal conditional plans n (-|2r) between IP and Q help to approximate (given is sufficiently small) the standard Euclidean Optimal Transport between and Qumbi which is the motivating problem of several OT researches (Makkuva al. 2020 Korotin et al. 2021a). The corresponding (stochastic)
mapping given by pushforward distributions Gyæ/ (:(x).. practice, we sample from xf (-Iar) using our -usi-gu MCMC and then pass the obtained samples through G. Note that our setup seems to be the first theoretically-advises attempt to leverage W OT between 512 x 512 images.
Technical details. The AFHQ dataset is taken from the StarGAN v2 CCoi et al. 2020) github:
The dataset includes three groups of high-quality 512 x 512 images: Dogs, Cats and Wilds (wildlife animals). The latter two groups are used as the source distributions P. The pretrained (on AFHQ Dogs) StyleGAN2-ADA Karras et al. (2020) is taken from the official Py Torch implementation:
ht htt:::/ggituuubu com/ NWwals//aaaie.a WWlabs/stylegaan2 ada -pytorch
As a potential fe which operates in the 512-dimensional latent space of the StyleGAN model, we choose fulle-conneeted MLP with ReLU activations and three hidden layers with 1024, 512 and 256 neurons, accordingly. The training hyperparameters are: = 100, /n = 0.008, on = 1.0, N = 128. For both Cat and Wild->Dog experiments, we train model for 11 epochs with learning
24

---

Published as a conference paper at ICLR 2024
rate 10-4 which starts monotonically decreasing to 10-5 after the fifth epoch. At inference, initialize the sampling procedure (in the latent space) with standard Normal noise. Then we repeat Langevin steps for Kisi = 1000 iterations with / 0.008. After that, additional Keeline 1000 steps are performed with /'m = 0.0005. The obtained latent codes are then passed through the StyleGAN generator, yielding the images from the AFHQ Dog dataset.
Computation complexity. The training takes approximately a day on A100 GPUs. The inference (as described above) takes about 20 minutes per batch on single A100 GPU.
D.4 COLORED MNIST DETAILS
For generating Colored MNIST dataset, make of the code, provided by the authors of
(Gushchin et al. 2023). The dataset consists of colored images of digit "2" (== 7K) and colored
images of digit "3" (Rs 7K). The images scaled resolution 32 32.
To solve the problem in view, we adapt the base EBM code from Nijkamp et al. 2020
We do nothing but embed the cost function gradient's computation when performing Langevin steps, leaving all the remaining technical details unchanged. In particular, we utilize simple CNNs with LeakyReLUUnnggat. seeeggttiv__eeee seeattiieelleeeeeeee_eeeee__eeee 0.05) activations as our learned potential fo- We pick batch size N = 256 and initialize the persistent replay buffer at random using Uniform| - 1. 1]3x32x32 distribution. We use Adam optimizer with learning rate gradually decreasing from The reported images # to approximately 7000 training iterations.
10-5.
For each entropic coefficient e = 0.01,0.1, 1, we run experiments with the parameters given in T For training stability, we add Gaussian noise (0, 9 77) to target samples when computing loss estimate Ż in Algorithm
0.01
0.1
K E
[500, 1000, 2000)
{500, 1000, 2000}
{500, 1000, 2000) (0.2, 0.3)
/0 (..1.
(U.1, 0.3)
Table 6: Training parameters for Colored MNIST; e = 0.01, 0.1,
At the inference stage, we initialize the MCMC chains with source data samples. The ULA steps repeated for Kist 2000 iterations with the same Langevin discretization step size rj as the one used at the training stage. For each E, the reported images 4 are picked for those parameters set K, which we found to be the best in terms of qualitative performance.
For LPIPS calculation, we use the official code:
https: http://t.b..////ccaaaaaaaeeeeeee
where we pick VGG backbone for calculating lpips features. calculate the resulting metric, we sample 18 target images from / (-Ix) for each test source image v. For every pair of these 18 images, we compute LPIPS and report the average value (along the generated target images and source ones).
Computation complexity. It takes approximately one day on one V100 GPU to complete an image data experiment for each set of parameters.
E DETAILS OF THE BASELINE METHODS
In this section, we discuss details the baseline methods with which we compare our method on the Caausian-to-Gauusian transformation problem.
Daniels et.al. Daniels et al. 202 We use the code from the authors' repository
https:
for their evaluation in the Gaussian case. We employ their configuration www.:main/cca/eccm/ceaaaaaaaaeae p.y..
25

---

Published as a conference paper at ICLR 2024
(eguy et.al.] (Seguy al. 2018 We the part of the code of SCONES conrespond- ing to learning dual OT potentials blob/ h//main/// http//miiaaaatppp////// py and the barycentric projection blob/main/ jtriin//ppooo....o py in the Gaussian case with configuration
[Chen et.al.] (Joint) (Chen et al. 2022
We utilize the official code from
https: httpps://githhub..cmm/hh htt.:/ccoommhhhii-BPPea-
with their configuration blob/ma i the checkerboard-to-nnn toy experiment, changing the number of steps of dynamics from 100 to 200 steps. Since their hyper-parameters. are developed for their 2-dimensional experiments, we increase the number of iterations for dimensions 16, 64 and 128 15 000.
g. py for
We also take the code from the We base our configuration on the authors' one (blob/me Allobmmain, htt//maain//ee configs seaffiggs/effulue ddfflltt __s___moo________ to__ttt__ppp___ www_spppisoooo_se .y) for the moon-to- spiral experiment. As earlier, we increase the number of steps of dynamies up to 200. Also, we change the number of training epochs for dimensions 16, 64 and 128 to 2,4 and 8 correspondingly.
[Chen et.al.] (Att) (Chen et al. 2022
same repository
above.
[e Bortoli et.al.] Bortoli et al. 2021 We utilize the official code from
bridge
with their configuration blob/ http/maann/cooooo hatp///oo///aaa//eeeeeeeeesessss /ata/saaaeeee/eeeeeaaaaaaaaa 22. yaml for toy problems. We increase the amount of steps of dynamics to 200 and the number of steps of IPF procedure for dimensions 16, 64 and 128 to 30, 40 and 60, respectively.
VVarga et.al.] Vargas et al.
2021
We use the official code from
with
hyper-parameters (asa//d_exxatee red_exxmmpees.. haapplsss..pyyll.. pyynb.
Toy We set the number of steps to 200. As earlier, we in-
from
crease the number of steps of IPF procedure for dimensions 16, 64 and 128 to 1000, 3500 and 5000, respectively.
|Vargas et.al.] Vargas et al. 202
We tested the official code from
https: /tte:://giiien weuihhhubb...
Instead of Gaussian processes, we used neural network as for [NOT]. We use N 200 discretization steps as for other SB solvers, 5000 IPF iterations, and 512 samples from distributions Po and P1 each of them. We use the Adam optimizer with = 10-4 for optimization.
GGushhhin et.al.] (Gushchin et al. 2023) We use the code provided by the authors. In our experi- ments, we use exactly the same hyperparameters for this setup as the authors (Gushchin et al.) 2023 Appendix B), except the number of discretization steps N, which we set to 200 as well as for other Schrödinger Bridge based methods.
F EXTENDED DISCUSSION OF LIMITATIONS
In general, the main limitation of our approach is the usage of MCMC. This procedure is time- consuming and requires adjusting several hyperparameters. Moreover, in practice. it may not always converge the desired distribution which introduces additional biases. The other details of launching our proposed algorithm arise due to its connection to EBM's learning procedure. It is known that EBMs for generative modelling could be trained by two different optimization regimes: short-run (non-convergent) training and long-run (convergent) training (Nijkamp et al.) 2020). In the first regime, the learned potential does not necessarily represent the energy function of the learned distribution. Because of this. the short-run mode may not always be adapted for Energy- guided EOT, since it seems crucial for fe to represent the true component of the conditional Energy potentials E (v) c(z.))- fo(g) In particular, for our Colored MNIST experiment, we found " the short-run regime be unstable and utilize exclusively long-run mode. At the same time, for
26

---

Published as a conference paper at ICLR 2024
Toy 2D and Gaussian-te Crassianntt-Gaaste Gaussian experiments well as for latent-space diggs-iieeiiional 12I setup, non- converrnn training was successful.
27


Reviews:
[{"review_id": "EIIIuasW95", "reviewer": "Reviewer_mVby", "summary": "The paper proposes a novel method for computing entropic optimal transport problem (EOT), utilizing techniques from energy-based models (EBM). Specifically, the paper uses the weak dual form to reformulate the EOT problem into an optimization task over space of functions, where the objective takes form similar to the EBM, i.e. in exponential form. The paper then proceeds to parametrize the function space by neural networks, and applies the algorithm to various tasks. Theoretical results are provided regarding why the proposed method approximates the optimal EOT coupling, and how the estimation and approximation contributes to generalization. Experiments are implemented, where various technical subtleties are also addressed.", "questions": "Please see above (section Weaknesses) for details. An additional question: is it possible to give a characterization of $\\epsilon$ dependence of computation/estimation/approximation? The reviewer understands that this is additional work, so simple answers such as exponential/polynomial dependence would be also good.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is overall well written and presented, and the the ideas are original to the knowledge of the reviewer. The discussion of all results seem plenty and extensive. Some strengths:\n1. The paper points out that the semi dual form of EOT works well for approximation of optimal coupling. The reviewer finds it interesting, as in semi dual form, essentially only one of the two equations of the Schrödinger system is satisfied, thus the other marginal usually lacks control. However, as shown in the theorem 2, even if only one potential is parametrized, constructing a joint distribution by taking conditionals to be normalized exponential models gives clean approximation of the optimal coupling, with bounds on approximation gap. As the construction naturally extrapolates, the proposed method not only computes EOT, but also helps generative sampling.\n2. The proposed usage of EBM seems principled, as EOT semi dual form admits exponential form, which enables application of well-studied sampling methods, and the corresponding gradients all have simple feasible forms.\n3. The experiments seem plenty and sufficient, and various training techniques for EBM are also discussed.", "weaknesses": "Some weaknesses:\n1. One major concern is the bound in Theorem 4, where a classical bound of error illustrating the balance between approximation and estimation is provided, using Rademacher complexity and optimality gap. However, it is unclear what is expected as the overall rate from this bound, as the choice of parametric class remains heuristic. This can be important, as generative tasks usually operates in high dimensions, and the dependence in dimensionality seems crucial to justify the applicability. There are plenty of works, for example, please see [1,2] for approximation and estimation using neural networks. Clarifying the proper choice of the parametric class and giving an explicit balancing characterization would give a clearer picture.\n2. The paper seems to focus more on applicability of EBM, though this class of methods usually requires sampling, which creates major computational burden and additional estimating error. A full computational and space-time complexity analysis seems needed, as even regardless of the NN optimization, the construction of the loss function requires significant computation to obtain an oracle. Furthermore, it would be interesting to also see how sampling error enters the bound in Theorem 4.\n\n[1] Klusowski, Jason M., and Andrew R. Barron. \"Approximation by combinations of ReLU and squared ReLU ridge functions with $\\ell^ 1$ and $\\ell^ 0$ controls.\" IEEE Transactions on Information Theory 64.12 (2018): 7649-7656.\n\n[2] Sreekumar, Sreejith, Zhengxin Zhang, and Ziv Goldfeld. \"Non-asymptotic performance guarantees for neural estimation of f-divergences.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2021.", "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": "lZBlYNkbEH", "reviewer": "Reviewer_7M7Q", "summary": "The paper reformulates the weak dual of the energy-optimal transport problem to identify an expectation inside the loss function that results in their loss function being trained in a similar way to energy-based models. This allows energy-based training procedures to be used to solve optimal transport problems.", "questions": "Is there a relationship between $(c, \\epsilon)$-transform and the Cole–Hopf transformation in PDEs ? They do seem very similar. It would also be very interesting to apply this to more than just continuous distribution but discrete distribution like languages so doing task like language translation would be interesting to see.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The arguments in the paper are clear and straightforward. The paper is well structured with the contributions highlighted clearly. The background is well-presented and I don't see typos. Figure 2 is well-made and shows the efficacy of their method.", "weaknesses": "The biggest weakness of their proposed method uses energy-based training which involves MCMC. I am unclear if this is ideal as MCMC can be tricky. It would be interesting to see if the unpaired image-to-image task can be done with other OT methods to better see how useful this particular formulation and method is.", "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": "87IDWFVND8", "reviewer": "Reviewer_2Kak", "summary": "This work bridges the gap between Energy-Based Models (EBMs) and Entropy-regularized Optimal Transport (EOT). In particular, it demonstrates that solving EOT is, to some extent, equivalent to EBM training. Thus, a novel methodology is introduced that leverages recent EBM developments to enhance the EOT solver. The approach is theoretically underpinned by generalization bounds and validated through practical applications in 2D and image domains.", "questions": "- It appears that equation 18 is akin to the maximum likelihood training of the EBM defined in equation 13. Could you offer a more intuitive explanation of why maximizing the likelihood of equation 13 is valid and how it equates to entropy-regularized optimal transport?\n- Can equation 17 be optimised using alternative EBM training techniques, such as noise contrastive estimate, and score matching?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This research uncovers the connection between energy-based models and entropy-regularized optimal transport, opening up new applications for EBMs, including unpaired data-to-data translation.", "weaknesses": "My primary concern centres around the scalability of the proposed approach. The training process hinges on simulating MCMC, which poses significant challenges when dealing with high-dimensional datasets. While promising results have been demonstrated in experiments on high-dimensional unpaired image-to-image translation, it is worth noting that this approach couples with a pretrained GAN model and conducts training in latent spaces. I hold reservations about its direct applicability to image spaces.", "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)