Title: Light Unbalanced Optimal Transport
Venue: NeurIPS 2024
Accepted: yes
Keywords: unbalanced optimal transport

light solver

entropy regularization

generative modeling
Abstract: While the continuous Entropic Optimal Transport (EOT) field has been actively developing in recent years, it became evident that the classic EOT problem is prone to different issues like the sensitivity to outliers and imbalance of classes in the source and target measures. This fact inspired the development of solvers that deal with the *unbalanced* EOT (UEOT) problem $-$ the generalization of EOT allowing for mitigating the mentioned issues by relaxing the marginal constraints. Surprisingly, it turns out that the existing solvers are either based on heuristic principles or heavy-weighted with complex optimization objectives involving several neural networks. We address this challenge and propose a novel theoretically-justified, lightweight, unbalanced EOT solver. Our advancement consists of developing a novel view on the optimization of the UEOT problem yielding tractable and a non-minimax optimization objective. We show that combined with a light parametrization recently proposed in the field our objective leads to a fast, simple, and effective solver which allows solving the continuous UEOT problem in minutes on CPU. We prove that our solver provides a universal approximation of UEOT solutions and obtain its generalization bounds. We give illustrative examples of the solver's performance.

Light Unbalanced Optimal Transport

Abstract

While the continuous Entropic Optimal Transport (EOT) field has been actively developing in recent years, it became evident that the classic EOT problem is prone to different issues like the sensitivity to outliers and imbalance of classes in the source and target measures. This fact inspired the development of solvers that deal with the unbalanced EOT (UEOT) problem -the generalization of EOT allowing for mitigating the mentioned issues by relaxing the marginal constraints. Surprisingly, it turns out that the existing solvers are either based on heuristic principles or heavyweighted with complex optimization objectives involving several neural networks. We address this challenge and propose a novel theoretically-justified, lightweight, unbalanced EOT solver. Our advancement consists of developing a novel view on the optimization of the UEOT problem yielding tractable and a non-minimax optimization objective. We show that combined with a light parametrization recently proposed in the field our objective leads to a fast, simple, and effective solver which allows solving the continuous UEOT problem in minutes on CPU. We prove that our solver provides a universal approximation of UEOT solutions and obtain its generalization bounds. We give illustrative examples of the solver's performance. The code is publicly available at https://github.com/milenagazdieva/LightUnbalancedOptimalTransport

1 Introduction

The computational optimal transport (OT) has proven to be a powerful tool for solving various popular tasks, e.g., image-to-image translation [68, 17, 50, 28], image generation [67, 13, 7] and biological data transfer [5, 42, 66]. Historically, the majority of early works in the field were built upon solving the OT problem between discrete probability measures [10, 53]. Only recently the advances in the field of generative models have led to explosive interest from the ML community in developing the continuous OT solvers, see [38] for a survey. The setup of this problem assumes that the learner needs to estimate the OT plan between continuous measures given only empirical samples of data from them. Due to convexity-related issues of OT problem [40], many works consider the EOT problem, i.e., use entropy regularizers which guarantee, e.g., the uniqueness of learned plans.

Meanwhile, researches attract attention to other shortcomings of the classic OT problem. It enforces hard constraints on the marginal measures and, thus, does not allow for mass variations. As a result, OT shows high sensitivity to an imbalance of classes and outliers in the source and target

measures [4] which are almost inevitable for large-scale datasets. To overcome these issues, it is common to consider extensions of the OT problem, e.g., unbalanced OT/EOT (UOT/UEOT) [8, 43]. The unbalanced OT/EOT formulations allow for variation of total mass by relaxing the marginal constraints through the use of divergences.

The scope of our paper is the continuous UEOT problem. It seems that in this field, a solver that is fast, light, and theoretically justified has not yet been developed. Indeed, many of the existing solvers follow a kind of heuristical principles and are based on the solutions of discrete OT. For example, [45] uses a regression to interpolate the discrete solutions, and [16, 36] build a flow matching upon them. Almost all of the other solvers [9, 70] employ several neural networks with many hyper-parameters and require time-consuming optimization procedures. We solve the aforementioned shortcomings by introducing a novel lightweight solver that can play the role of a simple baseline for unbalanced EOT.

Contributions. We develop a novel lightweight solver to estimate continuous unbalanced EOT couplings between probability measures ( M 4). Our solver has a non-minimax optimization objective and employs the Gaussian mixture parametrization for the UEOT plans. We provide the generalization bounds for our solver ( M 4.4) and experimentally test it on several tasks ( M 5.1, M 5.2).

Notations. We work in the Euclidian space ( R d , ∥ · ∥ ) . We use P 2 ,ac ( R d ) to denote the set of absolutely continuous (a.c.) Borel probability measures on R d with finite second moment and differential entropy. The set of nonnegative measures on R d with finite second moment is denoted as M 2 , + ( R d ) . We use C 2 ( R d ) to denote the space of all continuous functions ζ : R d → R for which ∃ a = a ( ζ ) , b = b ( ζ ) such that ∀ x ∈ R d : | ζ ( x ) | ≤ a + b ∥ x ∥ 2 . Its subspace of functions which are additionally bounded from above is denoted as C 2 ,b ( R d ) . For a.c. measure p ∈ P 2 ,ac ( R d ) (or M 2 , + ( R d ) ), we use p ( x ) to denote its density at a point x ∈ R d . For a given a.c. measure γ ∈ M 2 , + ( R d × R d ) , we denote its total mass by ∥ γ ∥ 1 def = ∫ R d × R d γ ( x, y ) dxdy . We use γ x ( x ) , γ y ( y ) to denote the marginals of γ ( x, y ) . They satisfy the equality ∥ γ x ∥ 1 = ∥ γ y ∥ 1 = ∥ γ ∥ 1 . We write γ ( ·| x ) to denote the conditional probability measure. Each such measure has a unit total mass. We use f to denote the Fenchel conjugate of a function f : f ( t ) def = sup u ∈ R { ut -f ( u ) } . We use I A to denote the convex indicator of a set A , i.e., I A ( x ) = 0 if x ∈ A ; I A ( x ) = + ∞ if x / ∈ A .

2 Background

Here we give an overview of the relevant entropic optimal transport (EOT) concepts. For additional details on balanced EOT, we refer to [10, 24, 53], unbalanced EOT - to [8, 43].

f -divergences for positive measures . For positive measures µ 1 , µ 2 ∈ M 2 , + ( R d ′ ) and a lower semi-continuous function f : R → R ∪ {∞} , the f -divergence between µ 1 , µ 2 is defined by

We consider f ( t ) which are convex, non-negative and attain zero uniquely when t = 1 . In this case, D f is a valid measure of dissimilarity between two positive measures (see Appendix C for details). This means that D f ( µ 1 ∥ µ 2 ) ≥ 0 and D f ( µ 1 ∥ µ 2 ) =0 if and only if µ 1 = µ 2 . In our paper, we also assume that all f that we consider satisfy the property that f is a non-decreasing function.

Kullback-Leibler divergence DKL [8, 62], is a particular case of such f -divergence for positive measures. It has a generator function f KL ( t ) def = t log t -t +1 . Its convex conjugate f KL ( t )=exp( t ) -1 . Another example is the χ 2 -divergence D χ 2 which is generated by f χ 2 ( t ) def = ( t -1) 2 when t ≥ 0 and ∞ otherwise. The convex conjugate of this function is f χ 2 ( t ) = -1 if t< -2 and 1 4 t 2 + t when t ≥ 2 .

Remark. By the definition, convex conjugates of f KL and f χ 2 divergences are proper, non-negative and non-decreasing. These properties are used in some of our theoretical results.

Entropy for positive measures. For µ ∈ M 2 , + ( R d ′ ) , its entropy [8] is given by

When µ ∈ P 2 ,ac ( R d ′ ) , i.e., ∥ µ ∥ 1 =1 , equation (1) is the usual differential entropy minus 1.

Classic EOT formulation (with the quadratic cost). Consider two probability measures p ∈ P 2 ,ac ( R d ) , q ∈P 2 ,ac ( R d ) . For ε > 0 , the EOT problem between p and q is

where Π( p, q ) is the set of probability measures π ∈P 2 ,ac ( R d × R d ) with marginals p , q (transport plans). Plan π ∗ attaining the minimum exists, it is unique and called the EOT plan.

Classic EOT imposes hard constraints on the marginals which leads to several issues, e.g., sensitivity to outliers [4], inability to handle potential measure shifts such as class imbalances in the measures p, q . The UEOT problem [70, 9] overcomes these issues by relaxing the marginal constraints [62].

Unbalanced EOT formulation (with the quadratic cost). Let D f 1 and D f 2 be two f -divergences over R d . For two probability measures p ∈ P 2 ,ac ( R d ) , q ∈ P 2 ,ac ( R d ) and ε > 0 , the unbalanced EOT problem between p and q consists of finding a minimizer of

see Fig. 1. Here the minimum is attained for a unique γ ∗ which is called the unbalanced optimal entropic (UEOT) plan. Typical choices of f i ( i ∈ [1 , 2] ) are f i ( t ) = τ i f KL ( t ) or f i ( t ) = τ i f χ 2 ( t ) ( τ i > 0 ) yielding the scaled DKL and D χ 2 , respectively. In this case, the bigger τ 1 ( τ 2 ) is, the more γ x ( γ y ) is penalized for not matching the corresponding marginal distribution p ( q ).

Remark 1. While the set M 2 , + ( R d × R d ) contains not only a.c. measures, infimum in problem (3) is attained for a.c. measure γ ∗ since -εH ( γ ∗ ) term turns to + ∞ otherwise. Thus, almost everywhere in the paper we are actually interested in the subset of a.c. measures in M 2 , + ( R d × R d ) .

̸

Remark 2. The balanced EOT problem (2) is a special case of (3). Indeed, let f 1 and f 2 be the convex indicators of { 1 } , i.e., f 1 = f 2 = I x =1 . Then the f -divergences D f 1 ( γ x ∥ p ) and D f 2 ( γ y ∥ q ) become infinity if p = γ x or q = γ y , and become zeros otherwise.

̸

Dual form of unbalanced EOT problem (3) is formulated as follows

It is known that there exist two measurable functions ϕ ∗ , ψ ∗ delivering maximum to (4). They have the following connection with the solution of the primal unbalanced problem (3):

Remark. For some of our results, we will need to restrict potentials ( ϕ, ψ ) in problem (4) to the space C 2 ,b ( R d ) ×C 2 ,b ( R d ) . Since established variants of dual forms [8] typically correspond to other functional spaces, we derive and theoretically justify a variant of the dual problem (4) in Appendix A.3. Note that it may hold that optimal potentials ψ ∗ , ϕ ∗ / ∈ C 2 ,b ( R D ) , i.e., the supremum is not achieved within our considered spaces. This is not principle for our subsequent derivations.

Computational UEOT setup. Analytical solution for the unbalanced EOT problem is, in general, not known. 1 Moreover, in real-world setups where unbalanced EOT is applicable, the measures p , q are typically not available explicitly but only through their empirical samples (datasets).

We assume that data measures p, q ∈ P 2 ,ac ( R d ) are unknown and accessible only by a limited number of i.i.d. empirical samples { x 0 , ..., x N }∼ p , { y 0 , ..., y M }∼ q . We aim to approximate the optimal UEOT plan γ ∗ solving (3) between the entire measures p, q . The recovered plans should allow the out-of-sample estimation, i.e., generation of samples from γ ∗ ( ·| x new ) where x new is a new test point (not necessarily present in the train data). Optionally, one may require the ability to sample from γ ∗ x .

The described setup is typically called the continuous EOT and should not be twisted up with the discrete EOT [53, 10]. There the aim is to recover the (unbalanced) EOT plan between the empirical measures ˆ p = 1 N ∑ N i =1 δ x i , ˆ q = 1 M ∑ M j =1 δ y j and out-of-sample estimations are typically not needed.

3 Related Work

Nowadays, the sphere of continuous OT/EOT solvers is actively developing. Some of the early works related to this topic utilize OT cost as the loss function [27, 25, 2]. These approaches are not relevant to us as they do not learn OT/EOT maps (or plans). We refer to [39] for a detailed review.

At the same time, there exist a large amount of works within the discrete OT/EOT setup [10, 15, 69, 51], see [53] for a survey. We again emphasize that solvers of this type are not relevant to us as they construct discrete matching between the given (train) samples and typically do not provide a generalization to the new unseen (test) data. Only recently ML community started developing out-ofsample estimation procedures based on discrete/batched OT. For example, [19, 56, 32, 14, 48, 57] mostly develop such estimators using the barycentric projections of the discrete EOT plans. Though these procedures have nice theoretical properties, their scalability remains unclear.

Balanced OT/EOT solvers. There exists a vast amount of neural solvers for continuous OT problem. Most of them learn the OT maps (or plans) via solving saddle point optimization problems [3, 18, 37, 22, 60, 50]. Though the recent works [28, 61, 11, 41, 30] tackle the EOT problem (2), they consider its balanced version. Hence they are not relevant to us. Among these works, only [41, 30] evade non-trivial training/inference procedures and are ideologically the closest to ours. The difference between them consists of the particular loss function used. In fact, our paper proposes the solver which subsumes these solvers and generalizes them for the unbalanced case. The derivation of our solver is non-trivial and requires solid mathematical apparatus, see M 4.

Unbalanced OT/EOT solvers. A vast majority of early works in this field tackle the discrete UOT/UEOT setup [6, 20, 54] but the principles behind their construction are not easy to generalize to the continuous setting. Thus, many of the recent papers that tackle the continuous unbalanced OT/EOT setup employ discrete solutions in the construction of their solvers. For example, [45] regress neural network on top of scaling factors obtained using the discrete UEOT while simultaneously learning the continuous OT maps using an ICNN method [47]. In [16] and [36], the authors implement Flow Matching [44, FM] and conditional FM on top of the discrete UEOT plans, respectively. The algorithm of the latter consists of regressing neural networks on top of scaling factors and simultaneously learning a conditional vector field to transport the mass between re-balanced measures. Despite the promising practical performance of these solvers, it is still unclear to what extent such approximations of UEOT plans are theoretically justified.

The recent papers [70, 9] are more related to our study as they do not rely on discrete OT approximations of the transport plan. However, they have non-trivial minimax optimization objectives solved using complex GAN-style procedures. Thus, these GANs often lean on heavy neural parametrization, may incur instabilities during training, and require careful hyperparameter selection [46].

For completeness, we also mention other papers which are only slightly relevant to us. For example, [23] considers incomplete OT which relaxes only one of the OT marginal constraints and is less general than the unbalanced OT. Other works [12, 4] incorporate unbalanced OT into the training objective of GANs aimed at generating samples from noise.

In contrast to the listed works, our paper proposes a theoretically justified and lightweight solver to the UEOT problem, see Table 1 for the detailed comparison of solvers.

4 Light Unbalanced OT Solver

In this section, we derive the optimization objective ( M 4.1) of our U-LightOT solver, present practical aspects of training and inference procedures ( M 4.2) and derive the generalization bounds for our solver ( M 4.4). We provide the proofs of all theoretical results in Appendix A.

4.1 Derivation of the Optimization Objective

Following the learning setup described above, we aim to get a parametric approximation γ θ,w of the UEOT plan γ ∗ . Here θ, ω are the model parameters to learn, and it will be clear later why we split them into two groups. To recover γ θ,ω ≈ γ ∗ , our aim is to learn θ, ω by directly minimizing the DKL divergence between γ θ,w and γ ∗ :

The main difficulty of this optimizing objective (6) is obvious: the UEOT plan γ ∗ is unknown . Fortunately, below we show that one still can optimize (6) without knowing γ ∗ .

Recall that the optimal UEOT plan γ ∗ has the form (5). We first make some changes of the variables. Specifically, we define v ∗ ( y ) def = exp { 2 ψ ∗ ( y ) -∥ y ∥ 2 2 ε } . Formula (5) now reads as

Since the conditional plan has the unit mass, we may write

where c v ∗ ( x ) def = ∫ R d exp { ⟨ x,y ⟩ ε } v ∗ ( y ) dy is the normalization costant ensuring that ∫ R d γ ∗ ( y | x ) dy = 1 .

Consider the decomposition γ ∗ ( x, y ) = γ ∗ x ( x ) γ ∗ ( y | x ) . It shows that to obtain parametrization for the entire plan γ ∗ ( x, y ) , it is sufficient to consider parametrizations for its left marginal γ ∗ x and the conditional measure γ ∗ ( y | x ) . Meanwhile, equation (8) shows that conditional measures γ ∗ ( ·| x ) are entirely described by the variable v ∗ . We use these observations to parametrize γ θ,w . We set

where u ω and v θ parametrize marginal measure γ ∗ x and the variable v ∗ , respectively. In turn, the constant c θ ( x ) def = ∫ R d exp { ⟨ x,y ⟩ ε } v θ ( y ) dy is the parametrization of c v ∗ ( x ) . Next, we demonstrate our main result which shows that the optimization of (6) can be done without the access to γ ∗ .

Theorem 4.1 (Tractable form of DKL minimization) . Assume that γ ∗ is parametrized using (9) . Then the following bound holds: ε DKL ( γ ∗ ∥ γ θ,w ) ≤L ( θ, w ) -L ∗ , where

and constant ( -L ∗ ) is the optimal value of the dual form (4) . The bound is tight in the sence that it turns to 0 = 0 when v θ ( y )=exp { 2 ψ ∗ ( y ) -∥ y ∥ 2 / 2 ε } and u ω ( x )= c θ ( x ) exp { 2 ϕ ∗ ( x ) -∥ x ∥ 2 / 2 ε } .

In fact, (10) is the dual form (4) but with potentials ( ϕ, ψ ) expressed through u ω , v θ (and c θ ):

Our result can be interpreted as the bound on the quality of approximate solution γ θ,ω (3) recovered from the approximate solution to the dual problem (4). It can be directly proved using the original ( ϕ, ψ ) notation of the dual problem, but we use ( u ω , v θ ) instead as with this change of variables the form of the recovered plan γ θ,ω is more interpretable ( u w defines the first marginal, v θ - conditionals).

Instead of optimizing (6) to get γ θ,ω , we may optimize the upper bound (10) which is more tractable. Indeed, (10) is a sum of the expectations w.r.t. the probability measures p, q . We can obtain Monte-Carlo estimation of (10) from random samples and optimize it with stochastic gradient descent procedure w.r.t. ( θ, ω ) . The main challenge here is the computation of the variable c θ and term ∥ u ω ∥ 1 . Below we propose a smart parametrization by which both variables can be derived analytically.

4.2 Parametrization and the Optimization Procedure

Parametrization. Recall that u ω parametrizes the density of the marginal γ ∗ x which is unnormalized. Setting x = 0 in equation (7), we get γ ∗ ( y | x = 0) ∝ v ∗ ( y ) which means that v ∗ also corresponds to an unnormalized density of a measure. These motivate us to use the unnormalized Gaussian mixture parametrization for the potential v θ ( y ) and measure u ω ( x ) :

Here θ def = { α k , r k , S k } K k =1 , w def = { β l , µ l , Σ l } L l =1 with α k , β l ≥ 0 , r k , µ l ∈ R d and 0 ≺ S k , Σ l ∈ R d × d . The covariance matrices are scaled by ε just for convenience.

For this type of parametrization, it holds that ∥ u ω ∥ 1 = ∑ L l =1 β l . Moreover, there exist closed-from expressions for the normalization constant c θ ( x ) and conditional plan γ θ ( y | x ) , see [41, Proposition def

def

x

T

S

k

x

+2

2

3.2]. Specifically, define

r

k

(

x

)

=

r

k

+

S

k

x

and

˜

α

k

(

x

)

=

α

k

exp

{

ε

}

. It holds that

Using this result and (11), we get the expression for γ θ,w :

Training. We recall that the measures p, q are accessible only by a number of empirical samples (see the learning setup in M 2). Thus, given samples { x 1 , ..., x N } and { y 1 , ..., y M } , we optimize the empirical analog of (10):

using minibatch gradient descent procedure w.r.t. parameters ( θ, w ) . In the parametrization of v θ and u ω (11), we utilize the diagonal matrices S k , Σ l . This allows decreasing the number of learnable parameters in θ and speeding up the computation of inverse matrices S -1 k , Σ -1 l .

Inference. Sampling from the conditional and marginal measures γ θ ( y | x ) ≈ γ ∗ ( y | x ) , u ω ≈ γ ∗ x is easy and lightweight since they are explicitly parametrized as Gaussian mixtures, see (12), (11).

4.3 Connection to Related Prior Works

The idea of using Gaussian Mixture parametrization for dual potentials in EOT-related tasks first appeared in the EOT/SB benchmark [29]. There it was used to obtain the benchmark pairs of probability measures with the known EOT solution between them. In [41], the authors utilized this type of parametrization to obtain a light solver (LightSB) for the balanced EOT.

Our solver for unbalanced EOT (10) subsumes their solver for balanced EOT as well as one problem subsumes the other for the special case of divergences, see the remark in M 2. Let f 1 , f 2 be convex indicators of { 1 } . Then f 1 ( t ) = t and f 2 ( t ) = t and objective (10) becomes

r

T

k

x

Here line (15) depends exclusively on θ and exactly coincides with the LightSB's objective, see [41, Proposition 8]. At the same time, line (16) depends only on ω , and its minimum is attained when u ω = p . Thus, this part is not actually needed in the balanced case, see [41, Appendix C].

4.4 Generalization Bounds and Universal Approximation Property

Theoretically, to recover the UEOT plan γ ∗ , one needs to solve the problem L ( θ, ω ) → min θ,ω as stated in our Theorem 4.1. In practice, the measures p and q are accessible via empirical samples X def = { x 1 , ..., x N } and Y def = { y 1 , ..., y M } , thus, one needs to optimize the empirical counterpart ̂ L ( θ, ω ) of L ( θ, ω ) , see (14). Besides, the available potentials u ω , v θ over which one optimizes the objective come from the restricted class of functions. Specifically, we consider unnormalized Gaussian mixtures u ω , v θ with K and L components respectively. Thus, one may naturally wonder: how close is the recovered plan γ ̂ θ, ̂ ω to the UEOT plan γ ∗ given that ( ̂ θ, ̂ ω ) = arg min θ,ω ̂ L ( θ, ω ) ? To address this question, we study the generalization error E DKL ( γ ∗ ∥ γ ̂ θ, ̂ ω ) , i.e., the expectation of DKL between γ ∗ and γ taken w.r.t. the random realization of the train datasets X ∼ p , Y ∼ q .

̂

̂ θ, ω

Let ( θ, ω ) = arg min ( θ,ω ) ∈ Θ × Ω L ( θ, ω ) denote the best approximators of L ( θ, ω ) in the admissible class. From Theorem 4.1 it follows that that E DKL ( γ ∗ ∥ γ ̂ θ, ̂ ω ) can be upper bounded using E ( L ( ̂ θ, ̂ ω ) -L ∗ ) . The latter can be decomposed into the estimation and approximation errors:

We establish the quantitative bound for the estimation error in the proposition below.

Proposition 4.2 (Bound for the estimation error) . Let p, q be compactly supported and assume that f 1 , f 2 are Lipshitz. Assume that the considered parametric classes Θ , Ω ( ∋ ( θ, ω )) consist of unnormalized Gaussian mixtures with K and L components respectively with bounded means ∥ r k ∥ , ∥ µ l ∥ ≤ R (for some R > 0 ), covariances sI ⪯ S k , Σ l ⪯ SI (for some 0 < s ≤ S ) and weights a ≤ α k , β l ≤ A (for some 0 < a ≤ A ). Then the following holds:

where O ( · ) hides the constants depending only on K,L,R,s,S,a,A,p,q,ε but not on sizes M,N .

This proposition allows us to conclude that the estimation error converges to zero when N and M tend to infinity at the usual parametric rate. It remains to clarify the question: can we make the approximation error arbitrarily small ? We answer this question positively in our Theorem below.

Theorem 4.3 (Gaussian mixture parameterization for the variables provides the universal approximation of UEOT plans) . Let p and q be compactly supported and assume that f 1 , f 2 are Lipshitz. Then for all δ > 0 there exist Gaussian mixtures v θ , u ω (11) with scalar covariances S k = λ k I d ≻ 0 , Σ l = ζ l I d ≻ 0 of their components that satisfy DKL ( γ ∗ ∥ γ θ,ω ) ≤ ε -1 ( L ( θ, ω ) -L ∗ ) < δ.

In summary , results of this section show that one can make the generalization error arbitrarily small given a sufficiently large amount of samples and components in the Gaussian parametrization. It means that theoretically our solver can recover the UEOT plan arbitrarily well.

5 Experiments

In this section, we test our U-LightOT solver on several setups from the related works. The code is written using PyTorch framework and is publicly available at https://github.com/milenagazdieva/LightUnbalancedOptimalTransport .

The experiments are issued in the convenient form of *.ipynb notebooks. Each experiment requires several minutes of training on CPU with 4 cores. Technical training details are given in Appendix B.

5.1 Example with the Mixture of Gaussians

We provide an illustrative 'Gaussians Mixture' example in 2D to demonstrate the ability of our solver to deal with the imbalance of classes in the source and target measures. We follow the experimental

(a) Input, target measures

(b) U-LightOT,

τ

= 10

2

(c) U-LightOT,

τ

= 10

1

(d) U-LightOT,

τ

= 10

0

x~p

y~q

x~pw(-)

y~Vθ.w(·1x)

setup proposed in [16, Figure 2] and define the probability measures p , q as follows (Fig. 2a):

We test our U-LightOT solver with scaled DKL divergences, i.e., f 1 ( t ) , f 2 ( t ) are defined by τ · f KL ( t ) = τ ( t log( t ) -t + 1) where τ > 0 . We provide the learned plans for τ ∈ [1 , 10 1 , 10 2 ] . The results in Fig. 2 show that parameter τ can be used to control the level of unbalancedness of the learned plans. For τ = 1 , our U-LightOT solver truly learns the UEOT plans, see Fig. 2d. When τ increases, the solver fails to transport the mass from the input Gaussians to the closest target ones. Actually, for τ = 10 2 , our solutions are similar to the solutions of [41, LightSB] solver which approximates balanced EOT plans between the measures. Hereafter, we treat τ as the unbalancedness parameter. In Appendix C, we test the performance of our solver with D χ 2 divergence.

Remark. Here we conduct all experiments with the entropy regularization parameter ε = 0 . 05 . The parameter ε is responsible for the stochasticity of the learned transport γ θ ( ·| x ) . Since we are mostly interested in the correct transport of the mass (controlled by f 1 , f 2 ) rather than the stochasticity, we do not pay much attention to ε throughout the paper.

5.2 Unpaired Image-to-Image Translation

Another popular testbed which is usually considered in OT/EOT papers is the unpaired image-toimage translation [71] task. Since our solver uses the parametrization based on Gaussian mixture, it may be hard to apply U-LightOT for learning translation directly in the image space. Fortunately, nowadays it is common to use autoencoders [59] for more efficient translation. We follow the setup of [41, Section 5.4] and use pre-trained ALAE autoencoder [55] for 1024 × 1024 FFHQ dataset [34] of human faces. We consider different subsets of FFHQ dataset ( Adult , Young , Man , Woman ) and all variants of translations between them: Adult ↔ Young and Woman ↔ Man .

The main challenge of the described translations is the imbalance of classes in the images from source and target subsets, see Table 2. Let us consider in more detail Adult → Young translation. In the FFHQ dataset, the amount of adult men significantly outnumbers the adult women , while the amount of young men is smaller than that of young women . Thus, balanced OT/EOT solvers are

expected to translate some of adult man representatives to young woman ones. At the same time, solvers based on unbalanced transport are supposed to alleviate this issue.

Baselines. We perform a comparison with the recent procedure [16, UOT-FM] which considers roughly the same setup and demonstrates good performance. This method interpolates the results of unbalanced discrete OT to the continuous setup using the flow matching [44]. For completeness, we include the comparison with LightSB and balanced optimal transport-based flow matching (OTFM) [44] to demonstrate the issues of the balanced solvers. We also consider neural network based solvers relying on the adversarial learning such as UOT-SD [9] and UOT-GAN[70].

Young→Adult

Adult→Young

Man→Woman

Woman→Man

0

100

90

80

=10

7=20

入=0.005

T =50

入=0.05

7 = 102

7

=250

入

=103

=0.5

0

=10

入=0.005

=20

=50

入=0.05

O

QT=10²

7=103

(keep)

T =10

=20

入=0.005

=50

入=0.05

7=102

O

T =10

=20

入=0.005

7=50

入= 0.05

7=102

OT=103

70

。

U-LightOT (ours)

U-LightOT (ours)

入=0.5

U-LightOT (ours)

U-LightOT (ours)

A

60

.

.

UOT-FM

UOT-SD

.

.

UOT-FM

UOT-SD

AcC.

.

.

UOT-FM

UOT-SD

.

.

UOT-FM

UOT-SD

50

.

UOT-GAN

.

UOT-GAN

UOT-GAN

UOT-GAN

.

OT-FM

OT-FM

OT-FM

OT-FM

40

LightSB

LightSB

LightSB

LightSB

40

60

80

100

60

70

80

90

100

60

70

80

90

100

40

60

80

100

Acc.(target)

Acc.(target)

Acc.(target)

Acc.(target)

Metrics. We aim to assess the ability of the solvers to perform the translation of latent codes keeping the class of the input images unchanged, e.g., keeping the gender in Adult → Young translation. Thus, we train a 99% MLP classifier for gender using the latent codes of images. We use it to compute the accuracy of preserving the gender during the translation. We denote this number by accuracy (keep) .

Meanwhile, it is also important to ensure that the generated images belong to the distribution of the target images rather than the source ones, e.g., belong to the distribution of young people in our example. To monitor this property, we use another 99% MLP classifier to identify whether each of the generated images belong to the target domain or to the source one. Then we calculate the fraction of the generated images belonging to the target domain which we denote as accuracy (target).

Results. Unbalanced OT/EOT approaches are equipped with some kind of unbalancedness parameter (like τ in our solver) which influences the methods' results: with the increase of unbalancedness, accuracy of keeping the class increases but the accuracy of mapping to the target decreases. The latter is because of the relaxation of the marginal constraints in UEOT. For a fair comparison, we aim to compute the trade-offs between the keep/target accuracies for different values of the unbalancedness. We do this for our solver and UOT-FM. We found that GAN-based unbalanced approaches show unstable behaviour, so we report UOT-SD and UOT-GAN's results only for one parameter value.

For convenience, we visualize the accuracies' pairs for our solver and its competitors in Fig. 3. The evaluation shows that our U-LightOT method can effectively solve translation tasks in high dimensions ( d = 512 ) and outperforms its alternatives in dealing with class imbalance issues. Namely, we see that our solver allows for achieving the best accuracy of keeping the attributes of the input images than other methods while it provides good accuracy of mapping to the target class. While balanced methods and some of the unbalanced ones provide really high accuracy of mapping to the target, their corresponding accuracies of keeping the attributes are worse than ours meaning that we are better in dealing with class imbalance issues. As expected, for large parameter τ , the results of our U-LightOT solver coincide with those for LightSB which is a balanced solver. Meanwhile, our solver has the lowest wall-clock running time among the existing unbalanced solvers, see Table 3 for comparison. We demonstrate qualitative results of our solver and baselines in Fig. 4. The choice of unbalancedness parameter for visualization of our method and UOT-FM is detailed in Appendix B.

We present the results of the quantitative comparison in the form of tables in Appendix D.1. In Appendix C, we perform the ablation study of our solver focusing on the selection of parameters τ , ε and number of Gaussian mixtures' components.

6 Discussion

Potential impact. Our light and unbalanced solver has a lot of advantages in comparison with the other existing UEOT solvers. First, it does not require complex max-min optimization. Second, it provides the closed form of the conditional measures γ θ ( y | x ) ≈ γ ∗ ( y | x ) of the UEOT plan.

ndul

UOT-SD

UOT-GAN

OT-FM

LightSB

UOT-FM

andu

as-1on

UOT-GAN

OT-FM

LightSB

UOT-FM

(c) Man → Woman

andu

aS-1on

UOT-GAN

OT-FM

LightSB

Input

aS-1on

UOT-GAN

OT-FM

LightSB

sno

(d) Woman → Man

Moreover, it allows for sampling both from the conditional measure γ θ ( y | x ) and marginal measure u ω ( x ) ≈ γ ∗ x ( x ) . Besides, the decisive superiority of our lightweight and unbalanced solver is its simplicity and convenience of use. Indeed, it has a straightforward and non-minimax optimization objective and avoids heavy neural parametrization. As a result, our lightweight and unbalanced solver converges in minutes on CPU. We expect that these advantages could boost the usage of our solver as a standard and easy baseline for UEOT task with applications in different spheres.

The limitations and broader impact of our solver are discussed in Appendix E.

ACKNOWLEDGEMENTS. The work of Skoltech was supported by the Analytical center under the RF Government (subsidy agreement 000000D730321P5Q0002, Grant No. 70-2021-00145 02.11.2021). We thank Kirill Sokolov, Mikhail Persiianov and Petr Mokrov for providing valuable feedback and suggestions for improving the proofs and clarity of our paper.

References

R. Agrawal and T. Horel. Optimal bounds between f-divergences and integral probability metrics. Journal of Machine Learning Research , 22(128):1-59, 2021.

M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning , pages 214-223. PMLR, 2017.

A. Asadulaev, A. Korotin, V. Egiazarian, and E. Burnaev. Neural optimal transport with general cost functionals. In The Twelfth International Conference on Learning Representations , 2024.

Y. Balaji, R. Chellappa, and S. Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances in Neural Information Processing Systems , 33: 12934-12944, 2020.

C. Bunne, S. G. Stark, G. Gut, J. S. Del Castillo, M. Levesque, K.-V. Lehmann, L. Pelkmans, A. Krause, and G. Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature Methods , 20(11):1759-1768, 2023.

L. Chapel, R. Flamary, H. Wu, C. Févotte, and G. Gasso. Unbalanced optimal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems , 34:23270-23282, 2021.

T. Chen, G.-H. Liu, and E. Theodorou. Likelihood training of schrödinger bridge using forwardbackward sdes theory. In International Conference on Learning Representations , 2021.

L. Chizat. Unbalanced optimal transport: Models, numerical methods, applications . PhD thesis, Université Paris sciences et lettres, 2017.

J. Choi, J. Choi, and M. Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. In Advances in Neural Information Processing Systems , volume 36, 2024.

M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems , pages 2292-2300, 2013.

M. Daniels, T. Maunu, and P. Hand. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems , 34:12955-12965, 2021.

Q. Dao, B. Ta, T. Pham, and A. Tran. Robust diffusion gan using semi-unbalanced optimal transport. arXiv preprint arXiv:2311.17101 , 2023.

V. De Bortoli, J. Thornton, J. Heng, and A. Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

N. Deb, P. Ghosal, and B. Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances in Neural Information Processing Systems , 34: 29736-29753, 2021.

P. Dvurechenskii, D. Dvinskikh, A. Gasnikov, C. Uribe, and A. Nedich. Decentralize and randomize: Faster algorithm for Wasserstein barycenters. In Advances in Neural Information Processing Systems , pages 10760-10770, 2018.

L. Eyring, D. Klein, T. Uscidda, G. Palla, N. Kilbertus, Z. Akata, and F. Theis. Unbalancedness in neural monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations , 2024.

J. Fan, A. Taghvaei, and Y. Chen. Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462 , 2020.

J. Fan, S. Liu, S. Ma, H.-M. Zhou, and Y. Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. URL https://openreview.net/forum?id=2mZSlQscj3 . Featured Certification.

K. Fatras, Y. Zine, R. Flamary, R. Gribonval, and N. Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. In AISTATS 2020-23nd International Conference on Artificial Intelligence and Statistics , volume 108, pages 1-20, 2020.

K. Fatras, T. Séjourné, R. Flamary, and N. Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In International Conference on Machine Learning , pages 3186-3197. PMLR, 2021.

R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, et al. Pot: Python optimal transport. The Journal of Machine Learning Research , 22(1):3571-3578, 2021.

M. Gazdieva, L. Rout, A. Korotin, A. Kravchenko, A. Filippov, and E. Burnaev. An optimal transport perspective on unpaired image super-resolution. arXiv preprint arXiv:2202.01116 , 2022.

M. Gazdieva, A. Korotin, D. Selikhanovych, and E. Burnaev. Extremal domain translation with neural optimal transport. In Advances in Neural Information Processing Systems , volume 36, 2023.

A. Genevay. Entropy-Regularized Optimal Transport for Machine Learning . Theses, PSL University, Mar. 2019. URL https://theses.hal.science/tel-02319318 .

A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics , pages 1608-1617. PMLR, 2018.

N. Gozlan, C. Roberto, P.-M. Samson, and P. Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis , 273(11):3327-3405, 2017.

I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville. Improved training of Wasserstein GANs. In Advances in Neural Information Processing Systems , pages 5767-5777, 2017.

N. Gushchin, A. Kolesov, A. Korotin, D. Vetrov, and E. Burnaev. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems , 2023.

N. Gushchin, A. Kolesov, P. Mokrov, P. Karpikova, A. Spiridonov, E. Burnaev, and A. Korotin. Building the bridge of schr \ " odinger: A continuous entropic optimal transport benchmark. arXiv preprint arXiv:2306.10161 , 2023.

N. Gushchin, S. Kholkin, E. Burnaev, and A. Korotin. Light and optimal schr \ "odinger bridge matching. In arXiv preprint arXiv:2402.03207 , 2024.

G. H. Hardy, J. E. Littlewood, and G. Pólya. Inequalities . Cambridge university press, 1952.

J.-C. Hütter and P. Rigollet. Minimax estimation of smooth optimal transport maps. 2021.

H. Janati, B. Muzellec, G. Peyré, and M. Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems , 33: 10468-10479, 2020.

T. Karras, S. Laine, and T. Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4401-4410, 2019.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

D. Klein, T. Uscidda, F. Theis, and M. Cuturi. Generative entropic neural optimal transport to map within and across spaces. arXiv preprint arXiv:2310.09254 , 2023.

A. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems , 34:14593-14605, 2021.

A. Korotin, L. Li, J. Solomon, and E. Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=3tFAs5E-Pe .

A. Korotin, A. Kolesov, and E. Burnaev. Kantorovich strikes back! wasserstein gans are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track , 2022.

A. Korotin, D. Selikhanovych, and E. Burnaev. Kernel neural optimal transport. In International Conference on Learning Representations , 2023. URL https://openreview.net/forum? id=Zuc_MHtUma4 .

A. Korotin, N. Gushchin, and E. Burnaev. Light schr \ " odinger bridge. In The Twelfth International Conference on Learning Representations , 2024.

T. Koshizuka and I. Sato. Neural lagrangian schr \ " { o } dinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations , 2022.

M. Liero, A. Mielke, and G. Savaré. Optimal entropy-transport problems and a new hellingerkantorovich distance between positive measures. Inventiones mathematicae , 211(3):969-1117, 2018.

Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

F. Lübeck, C. Bunne, G. Gut, J. S. del Castillo, L. Pelkmans, and D. Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semi-couplings. arXiv preprint arXiv:2209.15621 , 2022.

M. Lucic, K. Kurach, M. Michalski, S. Gelly, and O. Bousquet. Are GANs created equal? a large-scale study. In Advances in neural information processing systems , pages 700-709, 2018.

A. Makkuva, A. Taghvaei, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning , pages 6672-6681. PMLR, 2020.

T. Manole, S. Balakrishnan, J. Niles-Weed, and L. Wasserman. Plugin estimation of smooth optimal transport maps. arXiv preprint arXiv:2107.12364 , 2021.

M. Mohri, A. Rostamizadeh, and A. Talwalkar. Foundations of machine learning . MIT press, 2018.

P. Mokrov, A. Korotin, and E. Burnaev. Energy-guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations , 2024.

Q. M. Nguyen, H. H. Nguyen, Y. Zhou, and L. M. Nguyen. On unbalanced optimal transport: Gradient methods, sparsity and approximation error. arXiv preprint arXiv:2202.03618 , 2022.

T. T. Nguyen, H. D. Nguyen, F. Chamroukhi, and G. J. McLachlan. Approximation by finite mixtures of continuous density functions that vanish at infinity. Cogent Mathematics & Statistics , 7(1):1750861, 2020.

G. Peyré, M. Cuturi, et al. Computational optimal transport. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

K. Pham, K. Le, N. Ho, T. Pham, and H. Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In International Conference on Machine Learning , pages 7673-7682. PMLR, 2020.

S. Pidhorskyi, D. A. Adjeroh, and G. Doretto. Adversarial latent autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14104-14113, 2020.

A.-A. Pooladian and J. Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109.12004 , 2021.

P. Rigollet and A. J. Stromme. On the sample complexity of entropic optimal transport. arXiv preprint arXiv:2206.13472 , 2022.

R. Rockafellar. Duality and stability in extremum problems involving convex functions. Pacific Journal of Mathematics , 21(1):167-187, 1967.

R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022.

L. Rout, A. Korotin, and E. Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations , 2022.

V. Seguy, B. B. Damodaran, R. Flamary, N. Courty, A. Rolet, and M. Blondel. Large scale optimal transport and mapping estimation. In International Conference on Learning Representations , 2018.

T. Séjourné, G. Peyré, and F.-X. Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:2211.08775 , 2022.

S. Shalev-Shwartz and S. Ben-David. Understanding machine learning: From theory to algorithms . Cambridge university press, 2014.

A. Taghvaei and A. Jalali. 2-Wasserstein approximation via restricted convex potentials with application to improved training for GANs. arXiv preprint arXiv:1902.07197 , 2019.

A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning , pages 9526-9536. PMLR, 2020.

F. Vargas, P. Thodoroff, A. Lamacraft, and N. Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy , 23(9):1134, 2021.

G. Wang, Y. Jiao, Q. Xu, Y. Wang, and C. Yang. Deep generative learning via schrödinger bridge. In International Conference on Machine Learning , pages 10794-10804. PMLR, 2021.

Y. Xie, M. Chen, H. Jiang, T. Zhao, and H. Zha. On scalable and efficient computation of large scale optimal transport. In International Conference on Machine Learning , pages 6882-6892. PMLR, 2019.

Y. Xie, Y. Luo, and X. Huo. An accelerated stochastic algorithm for solving the optimal transport problem. arXiv preprint arXiv:2203.00813 , 2022.

K. D. Yang and C. Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations , 2018.

J.-Y. Zhu, R. Zhang, D. Pathak, T. Darrell, A. A. Efros, O. Wang, and E. Shechtman. Toward multimodal image-to-image translation. Advances in neural information processing systems , 30, 2017.

A Proofs

A.1 Proof of Theorem 4.1

Proof. For convenience, we split this proof into several steps.

Step 1. To begin with, we recall that the primal UEOT problem (3) attains a minimizer γ ∗ . Then by substituting γ ∗ directly in (3), we get

Here in line (18), we use equation (5) which establishes the connection between γ ∗ and the optimal potentials ϕ ∗ , ψ ∗ delivering maximum to the dual UEOT problem (4). Substituting these optimal potentials to the dual problem, we get

Note that (19) equals to (20) due to the equivalence between primal and dual UEOT problems, see M 2. We will use this fact in further derivations.

For every function f , its convex conjugate

Then for arbitrary x ∈ R d , taking f = f 1 , u = γ ∗ x ( x ) p ( x ) , t = ϕ ∗ ( x ) , we get

Similarly, for arbitrary y ∈ R d

Since (19) = (20), inequalities (22), (23), (24) actually turn to equalities.

Step 2. Since we consider convex, lower semi-continuous functions f , it holds that f = f . Then, substituting f in formula (21), we derive

Thus, for arbitrary x ∈ R d , taking f = f 1 , t = -γ ∗ x ( x ) p ( x ) , we get that -f 1 ( γ ∗ x ( x ) p ( x ) ) ≤ u γ ∗ x ( x ) p ( x ) + f 1 ( -u ) ∀ u ∈ R . At the same time, from step 1 , we know that -f 1 ( γ ∗ x ( x ) p ( x ) ) = γ ∗ x ( x ) p ( x ) ϕ ∗ ( x ) + f 1 ( -ϕ ∗ ( x )) , see (22). Then

Similarly, for arbitrary y ∈ R d

Step 3. Now we are ready to prove the main result. We note that:

While the ground-truth UEOT plan γ ∗ ( x, y ) and the optimal dual variables ϕ ∗ ( x ) , ψ ∗ ( y ) are connected via equation (5), our parametrized plan γ θ,ω ( x, y ) can be expressed using ϕ θ,ω ( x ) , ψ θ ( y ) as γ θ,ω ( x, y ) = exp { ε -1 ( ϕ θ,ω ( x ) + ψ θ ( y ) -∥ x -y ∥ 2 / 2 ) } , see (9). Then

A.2 Proof of Proposition 4.2

We begin with proving the auxiliary theoretical results (Propositions A.1, A.2) which are needed to prove the main proposition of this section.

Proposition A.1 (Rademacher bound on the estimation error) . It holds that

where F 1 = { f 1 ( -ϕ θ,ω ) | ( θ, ω ) ∈ Θ × Ω } , F 2 = { f 2 ( -ψ θ ) | θ ∈ Θ } for ϕ θ,ω ( x ) = ε log u ω ( x ) c θ ( x ) ∥ x ∥ 2 2 , ψ θ ( y ) = ε log v θ ( y ) + ∥ y ∥ 2 2 , and R N ( F 1 , p ) , R M ( F 2 , q ) denote the Rademacher complexity [63, M 26] of the functional classes U , V w.r.t. to the sample sizes N , M of distributions p , q .

Proof of Proposition A.1. The derivation of this fact is absolutely analogous to [41, Proposition B.1], [50, Theorem 4] or [64, Theorem 3.4].

Proposition A.2 (Bound on the Rademacher complexity of the considered classes) . Let 0 < a ≤ A , let 0 < u ≤ U , let 0 < w ≤ W and V > 0 . Consider the class of functions

Following [41, Proposition B.2], we call the functions of the class V as constrained log-sum-exp quadratic functions. We assume that f 1 , f 2 are Lipshitz functions and measures p , q are compactly supported with the supports lying in a zero-centered ball of a radius R > 0 . Then

where the constants C 0 , C 1 do not depend on sizes N , M of the empirical samples from p, q .

Proof of Proposition A.2. Thanks to [41, Proposition B.2], the Rademacher complexities of constrained log-sum-exp quadratic functions x ↦→ log u ω ( x ) , x ↦→ log c θ ( x ) and y ↦→ log v θ ( y ) are known to be bounded by O ( 1 √ N ) or O ( 1 √ M ) respectively. According to the definition of Rademacher complexity, for single quadratic functions x ↦→ x T x 2 ( y ↦→ y T y 2 ) it is just equal to zero. Then, using the well-known scaling and additivity properties of the Rademacher complexity [63], we get that x ↦→-ε log u ω ( x ) + ε log c θ ( x ) -∥ x ∥ 2 2 and y ↦→-ε log v θ ( y ) -∥ y ∥ 2 2 are bounded by O ( 1 √ N ) and O ( 1 √ M ) respectively. The remaining step is to recall that f 1 ( x ) and f 2 ( y ) are Lipschitz. Therefore, according to Talagrand's contraction principle [49], the Rademaher complexities of F 1 and F 2 are also bounded by O ( 1 √ N ) and O ( 1 √ M ) , respectively.

Proof of Proposition 4.2. The proof directly follows from Propositions A.1 and A.2.

A.3 Proof of Theorem 4.3

To begin with, we provide a quick reminder about the Fenchel-Rockafellar theorem which is needed to derive the dual form of problem (3).

Theorem A.3 (Fenchel-Rockafellar [58]) . Let ( E,E ′ ) and ( F, F ′ ) be two couples of topologically paired spaces. Let A : E ↦→ F be a continuous linear operator and A : F ′ ↦→ E ′ be its adjoint. Let

f and g be lower semi-continuous and proper convex functions defined on E and F respectively. If there exists x ∈ dom f s.t. g is continuous at Ax , then

and the min is attained. Moreover, if there exists a maximizer x ∈ E then there exists y ∈ F ′ satisfying Ax ∈ ∂g ( y ) and Ay ∈ ∂f ( -x ) .

Wenote that in the below theorem C 2 ( R d ) does not denote the space of twice differentiable continuous functions. The exact definition of this space and C 2 ,b ( R d ) is given in notations part.

Theorem A.4 (Dual form of problem (3)) . The primal UEOT problem (3) has the dual counterpart (4) where the potentials ( ϕ, ψ ) belong to the space C 2 ,b ( R d ) × C 2 ,b ( R d ) . The minimum of (3) is attained for a unique γ ∗ ∈ M 2 , + ( R d × R d ) .

Proof. We recall that in the primal form, the minimization is performed over functions γ belonging to M 2 , + ( R d × R d ) . In this proof, we suppose that this space is endowed with a coarsest topology σ ( M 2 , + ( R d × R d )) which makes continuous the linear functionals γ ↦→ ∫ ζdγ, ∀ ζ ∈ C 2 ( R d × R d ) . Then the topological space ( M 2 , + ( R d × R d ) , σ ( M 2 , + ( R d × R d )) ) has a topological dual ( M 2 , + ( R d × R d ) , σ ( M 2 , + ( R d × R d )) ) ′ which, actually, is (linear) isomorphic to the space C 2 ( R d × R d ) , see [26, Lemma 9.9]. This fact opens an opportunity to apply the well-celebrated FenchelRockafellar theorem. For this purpose, we will consider the following spaces: E def = C 2 ( R d ) ×C 2 ( R d ) , F def = C 2 ( R d × R d ) and their duals E ′ def = M 2 , + ( R d ) ×M 2 , + ( R d ) and F ′ def = M 2 , + ( R d × R d ) .

Step 1. Recall that the convex conjugate of any function g : M 2 , + ( R d × R d ) → R ∪{ + ∞} is defined for each ζ ∈ ( M 2 , + ( R d × R d )) ′ ∼ = C 2 ( R d × R d ) as g ( ζ ) = sup γ ∈M 2 , + ( R d × R d ) {⟨ γ, ζ ⟩ -g ( γ ) } . For the convenience of further derivations, we introduce additional functionals corresponding to the summands in the primal UEOT problem (3):

For our purposes, we need to calculate the convex conjugates of these functionals. Fortunately, convex conjugates of f -divergences F 1 ( γ x ) and F 2 ( γ y ) are well-known, see [1, Proposition 23], and equal to

To proceed, we calculate the convex conjugate of P ( γ ) :

Here in the transition from (31) to (32), we keep in mind our prior calculations of DKL in (28). Recall that DKL is non-negative and attains zero at the unique point

Thus, we get

Step 2. Now we are ready to apply the Fenchel-Rockafellar theorem in our case. To begin with, we show that this theorem is applicable to problem (3), i.e., that the functions under consideration satisfy the necessary conditions. Indeed, it is known that the convex conjugate of any functional (e.g., F 1 ( · ) , F 2 ( · ) , P ( · ) ) is lower semi-continuous and convex . Besides, the listed functionals are proper convex 2 . Indeed, the properness of F 1 ( · ) and F 2 ( · ) follows from the fact that f -divergences are known to be lower-semicontinuous and proper themselves, while properness of P ( · ) is evident from (32).

Now we consider the linear operator A : C 2 ( R d ) × C 2 ( R d ) ↦→ C 2 ( R d × R d ) which is defined as A ( ϕ, ψ ) : ( x, y ) ↦→ ϕ ( x ) + ψ ( y ) . It is continuous, and its adjoint is defined on M 2 , + ( R d × R d ) as A ( γ ) = ( γ x , γ y ) . Indeed, ⟨ A ( γ ) , ( u, v ) ⟩ = ⟨ γ, A ( u, v ) ⟩ = ∫ R d ∫ R d γ ( x, y )( u ( x ) + v ( y )) dxdy = ∫ R d γ x ( x ) u ( x ) dx + ∫ R d γ y ( y ) v ( y ) dy . Thus, the strong duality and the existence of minimizer for (3) follows from the Fenchel-Rockafellar theorem which states that problems

and

are equal. The uniqueness of the minimizer for (3) comes from the strict convexity of P ( · ) (which holds thanks to the entropy term). Note that the conjugate of the sum of F 1 and F 2 is equal to the sum of their conjugates since they are defined for separate non-intersecting groups of parameters.

Next we prove that the supremum can be restricted to C 2 ,b ( R d × R d ) . Here we use ' ∧ ' to denote the operation of taking minimum between the function f : R d → R and real value k : ( f ∧ k )( x ) def = min( f ( x ) , k ) . Then analogously to [26, Theorem 9.6], we get:

Since the another inequality is obvious, the two quantities are equal which completes the proof.

Proof of Theorem 4.3. Our aim is to prove that for all δ > 0 there exist unnormalized Gaussian mixtures u ω and v θ s.t. L ( θ, ω ) -L ∗ < δε . We define

Then from (4) and Theorem (A.4), it follows that L ∗ = inf ( ϕ,ψ ) ∈C 2 ,b ( R d ) ×C 2 ,b ( R d ) J ( ϕ, ψ ) . Finally, using the definition of the infimum, we get that for all δ ′ > 0 there exist some functions ( ̂ ϕ, ̂ ψ ) ∈ C 2 ,b ( R d ) ×C 2 ,b ( R d ) such that J ( ̂ ϕ, ̂ ψ ) < L ∗ + δ ′ . For further derivations, we set δ ′ def = δε 2 and pick the corresponding ( ̂ ϕ, ̂ ψ ) .

Step 1. We start with the derivation of some inequalities useful for future steps. Since ( ϕ, ψ ) ∈ C 2 ,b ( R d ) ×C 2 ,b ( R d ) , they have upper bounds ̂ a and ̂ b such that for all x, y ∈ R d : ̂ ϕ ( x ) ≤ ̂ a and ̂ ψ ( y ) ≤ ̂ b respectively. We recall that by the assumption of the theorem, measures p and q are compactly supported. Thus, there exist balls centered at x = 0 and y = 0 and having some radius R > 0 which contain the supports of p and q respectively. Then we define

We get that

Importantly, for all x and y within the supports of p and q it holds that ˜ ϕ ( x ) = ̂ ϕ ( x ) and ˜ ψ ( y ) = ̂ ψ ( y ) , respectively. Then

Combining (38) and (39), we get that J ( ˜ ϕ, ˜ ψ ) ≤ J ( ̂ ϕ, ̂ ψ ) < L ∗ + δ .

Before moving on, we note that functions exp { ˜ ϕ ( x ) / ε } and exp { ˜ ψ ( y ) / ε } are continuous and nonnegative. Therefore, since measures p and q are compactly supported, there exist some constants e min , h min > 0 such that exp { ˜ ϕ ( x ) / ε } > e min and exp { ˜ ψ ( y ) / ε } > h min for all x and y from the supports of measures p and q respectively. We keep these constants for future steps.

Step 2. This step of our proof is similar to [41, Theorem 3.4]. We get that

From this we can deduce that y ↦→ exp( ˜ ψ ( y ) / ε ) is a normalizable density since it is bounded by the unnormalized Gaussian density. Moreover, we see that it vanishes at the infinity. Thus, using the result [52, Theorem 5a], we get that for all δ ′′ > 0 there exists an unnormalized Gaussian mixture v ˜ θ = v ˜ θ ( y ) such that

Following the mentioned theorem, we can set all the covariances in v ˜ θ to be scalar, i.e., define v ˜ θ ( y ) = ∑ K k =1 ˜ α k N ( y | ˜ r k , ε ˜ λ k I d ) for some K and ˜ α k ∈ R + , ˜ r k ∈ R d , ˜ λ k ∈ R + ( k ∈ { 1 , . . . , K } ). For our future needs, we set

For simplicity, we consider the other mixture v θ ( y ) def = v ˜ θ ( y ) exp( -∥ y ∥ 2 2 ε ) which is again unnormalized and has scalar covariances, see the proof of [41, Theorem 3.4] for explanation. We denote the weights, means, and covariances of this mixture by α k ∈ R + , r k ∈ R d and λ k ∈ R + , respectively. We derive that

Then

From this, we see that exp ( ˜ ϕ ( x ) / ε ) c θ ( x ) tends to zero while x approaches infinity. It means that x ↦→ exp ( ˜ ϕ ( x ) / ε ) c θ ( x ) corresponds to the normalizable density. Using [52, Theorem 5a], we get that for all δ ′′ > 0 there exists an unnormalized Gaussian mixture u ˜ ω such that

Analogously with v ˜ θ , we can set all the covariances in u ˜ ω to be scalar, i.e., define u ˜ ω = ∑ L l =1 ˜ β l N ( x | ˜ µ l , ε ˜ ζ l I d ) for some L , ˜ µ l ∈ R d , ˜ ζ l ∈ R + ( l ∈ { 1 , ..., L } ) . Moreover, we consider u ω ( x ) = u ˜ ω ( x ) exp {-∥ x ∥ 2 2 ε } which is again an unnormalized density with scalar covariances. We denote the weights, means, covariances of this mixture by β l ∈ R + , µ l ∈ R d , ζ l ∈ R + respectively. Next we recall the equation (42) and get that

Step 4. Nowweturn to other expressions. Using the property that a function t ↦→ log t is 1 t min -Lipshitz on interval [ t min , + ∞ ) we get

Step 3. At this point, we will show that for every δ ′′ > 0 , there exists an unnormalized Gaussian mixture u ˜ ω which is δ ′′ -close to exp { ˜ ϕ ( x ) / ε } c θ ( x ) . Using the closed-form expression for c θ ( x ) from [41, Proposition 3.2], we get that

Similarly, we get that

We use this inequality, monotonicity of logarithm function, and (41), to derive

Recall that f 1 and f 2 are non-decreasing functions. Moreover, they are Lipshitz with the constants L 1 , L 2 respectively. Thus, we get

Integrating these inequalities over all x and y in supports of p and q respectively, we get

Finally, combining (44) and (52), we get

which completes the proof.

Remark. In fact, the assumption of the Lishitzness of f 1 ( x ) , f 2 ( x ) can be omitted. Indeed, under the " everything is compact " assumptions of Proposition 4.2 and Theorem 4.3, inputs to f 1 ( · ) , f 2 ( · ) also always belong to certain compact sets. The convex functions are known to be Lipshitz on compact subsets of R , see [31, Chapter 3, M 18], and we actually do not need the Lipschitzness on the entire R .

B Experiments Details

B.1 General Details

To minimize our objective (14), we parametrize α k , r k , S k of v θ and β l , µ l , Σ l of u ω in (11). Here we follow the standard practices in deep learning and parametrize logarithms log α k , log β l instead of directly parameterizing α k , β l . In turn, variables r k , µ l are parametrized directly as multidimensional vectors. We consider diagonal matrices S k , Σ l and parametrize them via their diagonal values log( S k ) i,i and log(Σ l ) i,i respectively. We initialize the parameters following the scheme in [41]. In all our experiments, we use the Adam optimizer.

B.2 Details of the Experiment with Gaussian Mixtures

We use K = L = 5 , ε = 0 . 05 , lr = 3 e -4 and batchsize 128 . We do 2 · 10 4 gradient steps. For the LightSB algorithm, we use the parameters presented by the authors in the official repository.

B.3 Details of the Image Translation Experiment

We use the code and decoder model from

We download the data and neural network extracted attributes for the FFHQ dataset from

In the Adult class we include the images with the attribute Age ≥ 44 ; in the Young class - with the Age ∈ [16 , 44] . Weexcluded the images with faces of children to increase the accuracy of classification per gender attribute. For the experiments with our solver, we use weighted DKL divergence with parameters τ specified in Appendix C, and set K = L = 10 , ε = 0 . 05 , lr = 1 , and batch size to 128 . We do 5 · 10 3 gradient steps using Adam optimizer [35] and MultiStepLR scheduler with parameter γ = 0 . 1 and milestones = [500 , 1000] . For testing [41, LightSB] solver, we use the official code (see the link above) and instructions provided by the authors.

Baselines. For the OT-FM and UOT-FM methods, we parameterize the vector field ( v t,θ ) t ∈ [0 , 1] for mass transport using a 2-layer feed-forward network with 512 hidden neurons and ReLU activation. An additional sinusoidal embedding[65] was applied for the parameter t . The learning rate for the Adam optimizer was set to 1e-4. To obtain an optimal transport plan π ∗ ( x, y ) discrete OT solvers from the POT [21] package were used. These methods are built on the solutions (plans π ∗ ( x, y ) ) of discrete OT problems, to obtain them we use the POT [21] package. Especially for the UOT-FM, we use the ot.unbalanced.sinkhorn with the regularization equal to 0 . 05 . We set the number of training and inference time steps equal to 100 . To obtain results of UOT-FM for Fig. 3, we run this method for 3K epochs with parameter reg _ m ∈ [5 e -4 , 5 e -3 , 5 e -2 , 0 . 5 , 1 , 10 , 10 2 , 10 3 , 10 4 , 10 5 , 10 6 ] and reported the mean values of final metrics for 3 independent launches with different seeds. In Tables 20, 21, 22, for each translation we report the results for one chosen parameter specified in

(a) U-LightOT, τ = 10

(b) U-LightOT,

τ

= 5

(c) U-LightOT,

τ

= 2

(d) U-LightOT, τ = 1

Appendix D.1. We use the corresponding checkpoints of UOT-FM to visualize its performance in Fig. 4. For [9, UOT-SD], [70, UOT-GAN] we use the official code provided by the authors, see the links:

https://github.com/Jae-Moo/Unbalanced-Optimal-Transport-Generative-Model https://github.com/uhlerlab/unbalanced_ot

While both UOT-SB, UOT-GAN methods were not previously applied to the FFHQ dataset, we set up a grid search for the parameters and followed the instructions provided by the authors for parameter settings. Both for UOT-SB, UOT-GAN, we used a 3-layer neural network with 512 hidden neurons, and ReLU activation was used for the generator networks and the potential and discriminator, respectively. Adam optimizer [35] with lr = 10 -5 and lr = 10 -4 was used to train the networks in UOT-SB and UOT-GAN, respectively. We train the methods for 10K iterations and set a batch size to 128. For UOT-SD, we used DKL divergence with their unbalancedness parameter τ =0 . 002 . For other parameters, we used the default values provided by the authors for CIFAR-10 generation tasks. For all baseline models which use entropy regularization, we set ε = 0 . 05 .

C Additional Discussion & Experiments with f -divergences

Details about f -divergences between positive measures. In the classic form, f -divergences are defined as measures of dissimilarity between two probability measures. This definition should be revised when dealing with measures of arbitrary masses. In the paragraph below we show that if the function f is convex, non-negative, and attains zero uniquely at point { 1 } then D f ( µ 1 ∥ µ 2 ) is a valid measure of dissimilarity between two positive measures.

Let µ 1 , µ 2 ∈ M 2 , + ( R d ′ ) be two positive measures. The f -divergence satisfies D f ( µ 1 ∥ µ 2 ) ≥ 0 which is obvious from the non-negativity of f . From the definition of D f ( µ 1 ∥ µ 2 ) and the fact that function f attains zero uniquely at a point { 1 } , we obtain that D f ( µ 1 ∥ µ 2 ) = 0 if and only if µ 1 ( x ) = µ 2 ( x ) holds µ 2 -everywhere. Actually, µ 1 ( x ) = µ 2 ( x ) should hold for all x as µ 1 must be absolutely continuous w.r.t. µ 2 (otherwise D f ( µ 1 ∥ µ 2 ) is assumed to be equal + ∞ ).

The usage of D χ 2 divergence. We tested the performance of our solver with scaled DKL divergences in the main text, see M 5.1, M 5.2. For completeness, here we evaluate our solver with D χ 2 divergence in Gaussian Mixture experiment. We use the same experimental setup as in M 5.1 and present the qualitative results in Fig. 5. Interestingly, the solver's results differ from those which we obtain for DKL divergence. For D χ 2 divergence, supports of learned plans' marginals constitute only parts of source and target measures' supports when τ = 1 . The issue disappears with a slight increase of τ , i.e., for τ = 2 . At the same time, a further increase of τ is useless, since the learned plans fail to deal with class imbalance issue. Thus, parameter τ should be adjusted heuristically. In the case of DKL divergence, supports coincide for all τ , see Fig. 2. This motivates us to use DKL divergences in our main experiments.

Parameter τ in Gaussian Mixture experiment. An ablation study on unbalancedness parameter τ in Gaussian Mixture experiment is conducted in M 5.1 and above in this section. For completeness, we also perform the quantitative assessment of our solver with DKL divergence for different unbalancedness parameters τ . We compute the normalized OT cost ( E x ∼ p E y ∼ γ ( y | x ) ( x -y ) 2 2 ) between the source and generated distributions, and the Wasserstein distance between the generated and target

distributions (computed by a discrete OT solver). For completeness, we additionally calculate the metrics for the balanced [41, LightSB] approach. The results are presented in Table 4.

Results. Recall that the unbalanced nature of our solver leads to two important properties. Firstly, our solver better preserves the properties of the input objects than the balanced approaches -indeed, it allows for preserving object attributes (classes) even in the case of class imbalance. Secondly, due to the relaxed boundary condition for the target distribution, the distribution generated by our solver is naturally less similar to the target distribution than for balanced methods.

The above intuitive reasoning is confirmed by the metrics we obtained. Indeed, as the τ parameter increases, when our method becomes more and more similar to balanced approaches, the normalized OT cost between the source and generated distributions increases, and the Wasserstein distance between learned and target distributions decreases. LightSB [1] baseline, which is a purely balanced approach, shows the best quality in terms of Wasserstein distance and the worst in terms of OT cost.

Parameters τ , ε in image experiments. The effect of entropy regularization parameter ε is well studied, see, e.g., [28, 41]. Namely, increasing the parameter ε stimulates the conditional distributions γ θ ( y | x ) to become more dispersed. Still, below we provide an additional quantitative analysis of its influence on the learned translation. Besides, we address the question how does the parameter τ influence the performance of our solver in image translation experiments? To address this question, we learn the translations Young → Adult , Man → Woman on FFHQ dataset varying the parameters τ , ε , see M 5.2 for the experimental setup details. We test our solver with scaled DKL divergence training it for 5K iterations. Other hyperparameters are in Appendix B. In Tables 5, 8, 11, 14, we report the accuracy of keeping the attributes of the source images (e.g., gender in Young → Adult translation). In Tables 6, 9, 12, 15, we report the accuracy of mapping to the correct target class (e.g., adult people in Young → Adult translation). In Tables 7, 10, 13, 16, we report FD metrics which is defined as Frechet distance between means and covariances of the learned and the target measures. For convenience, we additionally illustrate the results of ablation studies on Fig. 6.

Results show that increase of ε negatively influences both accuracy of keeping the attributes of the source images and FD of generated latent codes which is caused by an increased dispersity of γ θ ( y | x ) .

Interestingly, accuracy of mapping to the correct target class does not have an evident dynamics w.r.t. ε . At the same time, when τ increases , the learned plans provide worse accuracy for keeping the input latents' class but better FD of generated latent codes and accuracy of mapping to the target

Young→Adult

Adult→Young

Man→Woman

Woman→Man

100

D

90

eep)

80

(keep)

(keep)

(keep)

70

60

=0.01

Acc.

=0.01

AcC.

=0.01

Acc.

=0.01

A

。

=0.05

=0.05

=0.05

=0.05

50

。

=0.1

=0.1

=0.1

=0.1

=0.5

=0.5

=0.5

.

=0.5

40

=1

.

=1

=1

.

=1

40

60

80

100

60

70

80

90

100

60

70

80

90

100

40

60

80

100

Acc. (target)

Acc. (target)

Acc.(target)

Acc.(target)

class. It is an expected behavior since for bigger τ , the constraints on the marginals of the learned plans become more strict. That is, we enforce the marginals of the learned plans to be closer to source and target measures which allows learning more accurate mappings to target measure but does not allow keeping the source classes in the case of imbalance issues. Interestingly, in Adult → Young translation, FD of learned latents and accuracy of mapping to the target do not change much for τ ≥ 10 2 while accuracy of keeping the attributes exhibits a significant drop between τ = 10 2 and τ = 10 3 . Thus, we can treat τ = 10 2 as optimal since it provides the best trade-off between the quality of learned translations and their ability to keep the features of input latents. In the case of Young → Adult translation, the values of accuracy and FD exhibit significant differences for considered τ . Thus, we may consider a more detailed scale and choose τ = 2 . 5 · 10 2 as the optimal one.

It is important to note that our method offers a flexible way to select a domain translation configuration that allows for better preserving the properties of the original objects or generating a distribution closer to the target one. The final optimal configuration selection remains at the discretion of the user. The highlighted values in the Tables are used for comparison with other approaches in M D.1.

Remark. FD should be treated as a relative measure of similarity between learned and target measures. The results obtained by balanced solver [41, LightSB] (equivalent to ours for big τ ) are considered as a gold standard.

Number of Gaussian components in potentials. For completeness, we perform an ablation study of our U-LightOT solver with different number on Gaussian components ( K , L ) in potentials v θ and u ω , respectively. We run the solver in Young → Adult translation with 5K steps, ε = 0 . 05 and set τ = 250 . The quantitative results (accuracy of keeping the class, accuracy of mapping to the target, FD of generated latents vs target latents) are presented in the Tables below.

The results show that in the considered task, our solver provides good performance even for small number of Gaussian components. This can be explained by the smoothness of the latent representations of data in ALAE autoencoder.

D Additional Experimental Results

D.1 Quantitative comparison with other methods in Image-to-Image translation experiment

In this section, we provide additional results of quantitative comparison of our U-LightOT solver and other unbalanced and balanced OT/EOT solvers in image translation experiment. Tables 20, 21, 22, provide values used for plotting Fig. 3 in the main text. The unbalancedness parameters used for our U-LightOT solver and [16, UOT-FM] are specified in th Tables below. For obtaining the result of [9, UOT-SD], we use their unbalancedness parameter τ = 0 . 002 . For other details on methods' parameters used to obtain the results below, see Appendix B. Note that we do not include FID metric for assessing the quality of generated images since we found that it is not a representative metric for assessing the performance of models performing the translation of ALAE latent codes.

D.2 Outlier robustness property of U-LightOT solver

To show the outlier robustness property of our solver, we conduct the experiment on Gaussian Mixtures with added outliers and visualize the results in Fig. 7. The setup of the experiment, in general, follows the Gaussian mixtures experiment setup described in section 5.2 of our paper. The difference consists in outliers (small gaussians) added to the input and output measures.

(a) Measures

p, q

(b) LightSB

(c) U-LightOT (ours)

The results show that our U-LightOT solver successfully eliminates the outliers and manages to simultaneously handle the class imbalance issue. At the same time, the balanced LightSB [41] solver fails to deal with either of these problems.

E Limitations and Broader Impact

Limitations. One limitation of our solver is the usage of the Gaussian Mixture parametrization which might restrict the scalability of our solver. This points to the necessity for developing ways to optimize objective (4) with more general parametrization, e.g., neural networks.

Broader impact. Our work aims to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: For each contribution listed in abstract and introduction we provide the links to the sections about them.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

3. Theory Assumptions and Proofs

Justification: We discuss the limitations of our solver in M 6.

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide the proofs and the assumptions in Appendix A.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide a full list of the experimental details in Appendix B. We use publicly available datasets. The code for the experiments is provided in supplementary material.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code for our solver is included in supplementary material and will be made public after acceptance of our paper. We use publicly available datasets.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide the experimental details in Appendix B.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: The focus of the papers is mostly theoretical. The experiments are provided just for illustration of the derived theoretical results.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The computational resources are discussed in

9. Code Of Ethics

M 5.

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conforms with NeurIPS Code of Ethics. We discuss the societal impact of our paper in M 6.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the societal impact of our paper in M 6.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The research does not release data or models that have a high risk for misuse and does not need safeguards.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cited the creators all of the used assets.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: The code is included in supplementary material. The license for the code will be provided after the paper acceptance.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not include crowdsourcing experiments or research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not include crowdsourcing experiments or research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes U-LightOT, a lightweight solver for the Unbalanced Entropic Optimal Transport (UEOT) problem. This method uses a Gaussian Mixture approximation for the potential $v_{\\theta}(y)$ and measure $u_{w}(x)$. This paper proves that under this approximation, the KL divergence to the ground truth UEOT plan has a tractable form. U-LightOT is evaluated on Gaussian Mixture and Unpaired Image-to-Image Translation tasks.\n\nSTRENGTHS:\n1. This paper provides a theoretical analysis of the generalization bounds and the universal approximation property for the Gaussian mixture parametrization.\n2. The proposed method is a lightweight solver for the UEOT problem, which requires several minutes of CPU training for the experiment in Sec 5.\n3. This paper is easy to follow.\n\nWEAKNESSES:\n1. The optimization objective and Gaussian Mixture approximation in Sec 4 are similar to [1].\n\n2. While this paper provides the universal approximation property for Gaussian Mixure approximation, I have concerns about whether this Gaussian Mixture parametrization can achieve decent results for more complex distributions, such as in generative modeling within the data space on CIFAR-10. \n\n[1] Korotin, Alexander, Nikita Gushchin, and Evgeny Burnaev. \"Light schr\\\" odinger bridge.\" ICLR 2024.\n\nQUESTIONS:\n1. In the Unpaired Image-to-Image Translation task, Table 2 only presents the accuracy of keeping the attributes of the source images. However, since the goal of this task is semantic translation, the accuracy of the target semantic is also required. For example, in the Young-to-Adult task, the accuracy of whether the generated image is indeed an adult image. Could you provide this target semantic accuracy results?\n2. In the Appendix, Tables 5 and 7 show the Frechet distance (FD) between the learned and target measures. I believe this FD metric evaluates whether the semantic translation is successful, at the marginal level. Generally, increasing τ\\tauτ decreases (improves) FD metrics in Table 5 and 7. Could you clarify how the optimal $\\tau$ is selected? I am curious because when $\\tau$ is overly large, U-LightOT achieves worse accuracy compared to other models in Table 2.\n3. Could you present the FD results (Tables 5 and 7) for the other models?\n4. For the optimal $\\tau=500$ in the Man-to-Woman task, Table 6 shows an accuracy of 83.85 at best, while Table shows an accuracy of 92.85. Could you clarify which result is correct?\n\nLIMITATIONS:\nThe authors addressed the limitations and broader impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe proposal focuses on developing a fast solver for the unbalanced entropy-regularized optimal (EOT) transport between continuous Radon measures. The authors utilize the dual formulation of unbalanced EOT and use the relationship between the optimal potentials (i.e., the dual variables) and the primal transport plan. They then consider a parameterization of the transport plan and plan to minimize the KL divergence between this parameterized transport plan and the optimal one. Given that the optimal transport plan is unknown, the authors first use the relationship between the primal and dual solutions to reparameterize the dual variables and then derive a tight upper bound for the KL between the optimal plan and the parameterized one, which they propose to minimize. To deal with the normalization terms in their upper bound, the authors use a similar framework to that of Gushchin et al. [29] and assume the reparameterized dual variables are unnormalized Gaussian mixtures; this assumption enables analytic solutions to the otherwise difficult-to-calculate terms in the upper bound. Lastly, the authors provide a generalization error bound for their proposed framework. The paper provides two small-scale numerical examples to demonstrate their solver's efficiency: 1) two-dimensional Gaussian mixtures and 2) unpaired-image-to-image translation in the embedding space of an autoencoder, specifically ALAE, on an unbalanced subset of FFHQ dataset for Adult, Young, Man, Woman face.\n\nSTRENGTHS:\n+ The paper is very well written and straightforward to follow.\n+ The clever parameterizations used in this paper (while they appear in some prior work), provide a unique approach for solving the UEOT problem between continuous measures.\n+ The provided generalization error bounds (while straightforward to derive), are important and certainly add value to the paper.\n+ The method is easy to implement and fast to train. Quick convergence on the CPU is a notable achievement unlocked by this work.\n\nWEAKNESSES:\n- One major weakness is that the paper does not discuss how $K$ and $L$, i.e., the number of Gaussians in the mixtures, affect the results. The generalization error bound mentions that $K$ and $L$ will appear as constants in the error bound, but the practical implications of the choice of $K$ and $L$ are missing from the paper. \n\n- Experiments are relatively modest: 2 experiments in low dimensions. It would be beneficial to have an experiment on robustness to outliers, as this is included in the main claim.\n\n- The paper claims to have a fast solver but lacks a detailed speed comparison for either experiment. It would be great to have a wall-clock comparison of competing methods.\n\n- The Gaussian mixture assumption limits the method's applicability to only low-dimensional problems, and it is not clear whether this limitation can be overcome.\n\nQUESTIONS:\n* How does the performance change as a function of $K$ (assuming $L=K$)?\n\nLIMITATIONS:\nLimitations are provided in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

comments: SUMMARY:
This work focuses on the largely computationally intractable efforts in unbalanced OT dual form where neural networks are used as a proxy (used as potentials) in order to approximate Wasserstein distances. In this work, the authors set out to significantly reduce this optimization procedure by decomposing the join optimal solution into conditionals which allows for both easier inference and a reduction in the number of parameters required. Experimental results are then carried out to show the success of this method beyond the improved efficiency.

STRENGTHS:
[+] The reduction in efficiency seems to be quite strong and effective way to reduce the overall parameters required to approximate OT distances
[+] The theoretical results are sound and well motivated.
[+] A generalization bound is also presented, attesting to the soundness one achieves with this light variation.

WEAKNESSES:
[-] Appears to be specific only to the case of having KL divergence penalties for the mass constraints.
[-] Paper can appear a bit difficult and dense to read.

QUESTIONS:
(1) The reduction you get appears to have (perhaps even superficially) some relationship to the way WAE decomposes the coupling into conditionals. More coincidentally, WAE also uses conditional Gaussians to parametrize the encoder distribution, although for different purposes. Do you have any comments if there is any deeper link here?

(2) Do you have any intuition if one were to use other penalties beyond KL to enforce the mass constraint?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel approach to solving the continuous Unbalanced Entropic Optimal Transport (UEOT) problem. The authors introduce a lightweight, theoretically-justified solver that addresses the challenges of sensitivity to outliers and class imbalance in traditional Entropic Optimal Transport (EOT). The proposed method features a non-minimax optimization objective and employs Gaussian mixture parametrization for UEOT plans, resulting in a fast, simple, and effective solver. The authors provide theoretical guarantees for their solver's performance and apply it to simulated and image data.\n\nSTRENGTHS:\n- The paper is well-written and easy to follow\n- The paper describes well related literature and clearly motivates the approach / why there is a need for this solver\n- The paper introduces a novel way to solve UEOT problems using Gaussian mixtures, even if the approach was previously used for balanced EOT problems as mentioned by the authors.\n- The authors thoroughly study generalization bounds.\n- The authors consider a wide range of competing methods.\n\nWEAKNESSES:\n- While the authors provide generalisation bounds, it would be helpful to assess the performance of the method on the UEOT plan between Gaussian distributions, see Janati et al. ,2020\n- As mentioned by the authors, the Gaussian mixture approach is likely to work only in low dimensions. It would be interesting to see when it fails, e.g. using the benchmark above.\n\nQUESTIONS:\n- The authors state that for OT-FM and UOT-FM in the FFHQ dataset, they use a 2-layer feed-forward network with 512 hidden neurons and ReLU activation. Where does this parameterization come from? It seems to be relatively small for a flow matching architecture on images, and does not seem to be the architecture used in the original papers.\n- Why are FID scores not reported for the image translation tasks?\n\nLIMITATIONS:\nThe authors have considered the limitations and potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a lightweight solver for Unbalanced Entropic Optimal Transport (UEOT) that does not rely on neural network parametrization. Instead, the authors parameterize the potential functions of UEOT using Gaussian Mixture Models (GMM). This parametrization enables the derivation of a tractable joint coupling. By incorporating the parameterized potential into the dual objective, the authors achieve a simple and tractable loss function. Additionally, the paper provides a universal approximation result for GMM parametrization. Experiments are conducted on toy data (GMM) and image-to-image (I2I) translation.\n\nSTRENGTHS:\n- The paper proposes a simple and fast UEOT algorithm.\n- The paper justifies the GMM parametrization by presenting generalization bounds.\n- The paper demonstrates applicability to large-scale tasks such as I2I translation when combined with an autoencoder (AE).\n- The paper is well-written, clear, and easy to follow.\n\nWEAKNESSES:\n- The method of parameterizing the potential function using GMMs was already proposed in LightSB [1]. The only change here is the switch to a UOT objective, making the methodological contribution minimal. Aside from the universal approximation result, the theoretical contributions are also limited.\n\n- The experiments are not comprehensive. First, the experiments are conducted only on face-related data. More diverse datasets should be included. Second, the fairness of the comparisons is questionable. In the I2I experiments, the authors use the ALAE autoencoder, while some of comparison methods are implemented directly in the image space. All of the comparisons should be implemented in the latent space for fairness. Third, since U-LightOT are implemented on a latent space that captures attributes well, the attribute accuracy is expected to be high. Other than accuracy, more general metrics such as c-FID or FID should be used for comparison. Fourth, there is a lack of ablation studies on the number of Gaussian modals $N$ and $M$. This is very important and expected to be sensitive hyperparameter, thus, I believe authors should provide ablation studies on this parameter. Overall, the practical utility of the approach is questionable.\n\n[1] Light Schrodinger Bridge, ICLR, 2024.\n\nQUESTIONS:\n- How does the performance change when parameterizing very high-dimensional and multi-modal GMM data with fewer or more $N,M$? \n- In the I2I experiments, how does the number of modes in the GMM affect performance?\n- In toy data experiments, does U-LightOT has lower transport plan costs and smaller Wasserstein distances between the target and generated distributions compared to other comparisons?\n\nLIMITATIONS:\nDiscussed in Weakness section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}