Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
Authors: Jaemoo Choi
Jaewoong Choi
Myungjoo Kang
Venue: NeurIPS 2023 poster (2023)
Keywords: Optimal Transport
Generative modeling
Generative adversarial network
TL;DR: We present a new framework on generative modeling using the semi-dual formulation of Unbalanced Optimal Transport.

Abstract: Optimal Transport (OT) problem investigates a transport map that bridges two distributions while minimizing a given cost function. In this regard, OT between tractable prior distribution and data has been utilized for generative modeling tasks. However, OT-based methods are susceptible to outliers and face optimization challenges during training. 
In this paper, we propose a novel generative model based on the semi-dual formulation of Unbalanced Optimal Transport (UOT). Unlike OT, UOT relaxes the hard constraint on distribution matching. This approach provides better robustness against outliers, stability during training, and faster convergence. We validate these properties empirically through experiments. Moreover, we study the theoretical upper-bound of divergence between distributions in UOT. Our model outperforms existing OT-based generative models, achieving FID scores of 2.97 on CIFAR-10 and 6.36 on CelebA-HQ-256. The code is available at \url{https://github.com/Jae-Moo/UOTM}.

Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
Jaeemo Choi* Seoul National University toony428snu. ac
Jaewoong Choi
Myungjoo Kang Seoul National University mkangesnu. ac.kr kr
Korea Institute for Advanced Study chwj 1bwj14750iias re.kr
Abstract
Optimal Transport (OT) problem investigates a transport map that bridges two distributions while minimizing given cost function. In this regard, OT between tractable prior distribution and data has been utilized for generative modeling tasks. However, OT-based methods are susceptible outliers and face optimization challenges during training. this paper, we propose novel generative model based the semi-dual formulation of Unbalanced Optimal Transport (UOT). Unlike OT, UOT relaxes the hard constraint on distribution matching. This approach provides better robustness against outliers, stability during training, and faster convergence. We validate these properties empirically through experiments. Moreover, we study the theoretical upper-hound of divergence between distributions in UOT. Our model outperforms existing OT-based generative models, achieving FID scores of 2.97 CIFAR-10 and 6.36 on CelehA-HQ-256 The code is available https://gi www./o///aeeeeeee -ao-MooeMMMne Moo/UOTM.
Introduction
Optimal Transport theory 77 explores the cost-optimal transport to transform one probability distribution into another. Since WGAN (6) OT theory has attracted significant attention in the field of generative modeling framework for addressing important challenges in this field. In particular, WGAN introduced the Wasserstein distance, an OT-based probability distance, as a loss function for optimizing generative models. WGAN measures the Wasserstein distance between the data distribution and the generated distribution, and minimizes this distance during training. The introduction of OT-based distance has improved the diversity (6 (2), convergence [63], and stability
[36, (22 of generative models, such as WGAN [6] and its variants (27 45 56). However, several works showed that minimizing the Wasserstein distance still faces computational challenges, and the
models tend to diverge without a strong regularization term, such as gradient penalty [51] 63
Recently, there has been surge of research directly modeling the optimal transport map between the input prior distribution and the real data distribution 2 3 49
61 82]. other words, the optimal transport serves the generative model itself. These approaches showed promising results that are comparable to WGAN models. However, the classical optimal transport-based approaches are known be highly sensitive outlier-like samples [8]. The existence of few outliers have a significant impact on the overall bT-ba distance and the corresponding transport map. This sensitivity can be problematic when dealing with large-scale real-world datasets where outliers and noises are inevitable.
To overcome these challenges, suggest new generative algorithm utilizing the semi-dual formu- lation of the Unbalanced Optimal Transport (UOT) problem [12] 44]. In this regard, we refer to our model the UOT-based generative model (UOTM). The UOT framework relaxes the hard constraints
Equal contribution. Correspondence to: Myungjoo Kang <ukangeenu.ac. ac. kr>
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Figure I: Generated Samples from UOTM trained on Left: CIFAR-10 and Right: CelebA-HQ
of marginal distribution in the OT framework by introducing soft entropic penalties. This soft con- straint provides additional robustness against outliers 66). Our experimental results demonstrate that UOTM exhibits such outlier robustness, as well as faster and more stable convergence compared existing OT-based models. Particularly, this better convergence property leads tighter matching of data distribution than the OT-based framework despite the soft constraint of UOT. Our UOTM achieves FID scores of 2.97 on CIFAR-10 10 and 6.36 on CelebA-HQ, outperforming existing OT-based adversarial methods by a significant margin and approaching state-of-the-arr performance. Further- more, this decent performance is maintained across various objective designs. Our contributions can be summarized as follows:
UOTM is the first generative model that utilizes the semi-dual form of UOT.
We analyze the theoretical upper-bound of divergence between marginal distributions in UOT, and validate our findings through empirical experiments.
We demonstrate that UOTM presents outlier robustness and fast and stable convergence.
To the best of our knowledge, UOTM is the first OT-based generative model that achieves near state-of-the performance on real- real-world image datasets.
2 Background
Notations Let X. y be two compact complete metric spaces, μ and v be probability distributions on X' and Y, respectively. We regard μ and the source and target distributions. In generative modeling tasks, μ and correspond to tractable noise and data distributions. For a measurable map T, Tafe represents the pushforward distribution of μ.. I[ v) denote the set of joint probability distributions on X' x y whose marginals are μ and v, respectively. M (X' x y) denote the set of joint positive measures defined on X x V. For convenience, for M+ X' )), let (o())) and (v) be the marginals with respect to A' and V) c(x, y) refers to the transport cost function defined on X' x y. Throughout this paper, we consider X = y C. IRd with the quadratic cost, c(x, y) = T|| - e12. where d indicates the dimension of data. Hcre, is given positive constant. For the precise notations and assumptions, see Appendix
Optimal Transport (OT) OT addresses the problem of searching for the most cost-minimizing way to transport source distribution je to target v, based on a given cost function c(:,)). )). In the beginning, Monge (53] formulated this problem with deterministic transport map. However, this formulation is non-convex and can be ill- po-ee depending on the choices of ju and v. To alleviate these problems, the Kantorovich OT problem (33] was introduced, which is convex relaxation of the Monge problem. Formally, the OT cost of the relaxed problem is given by as follows:
U c(x,W)dT(I
C(u,)) inf *ETT(L.)
where is a cost function, and æ coupling of p. and Unlike Monge problem, the minimizer " of Eq allways exists under some mild assumptions on (X,,), ()), v) and the cost ([77), Chapter 5). Then, the dual form of Eq given as:
U/ (u(z)dp())+ + / u(2)du(()
C(H,)=
(2)
((
2

---

where and are Lebesgue integrable with respect to measure μ and v, i... u E (') and E L (v).. For a particular case where e(z, y) is equal the distance function between a and y.-
then to and u is 1-Liipchitzz [77]. In such case, we call Eq a Kantorovich-RRuinnte For the general cost c(:,)) )) the Eq can be reformulated as follows ((77), Chapter 5):
duality.
Ux ('()ddu()) + I, o(v)do(s) EEL (v)
V( = sup
where the c-transform of is defined as r (x))) inf (c(2, v(y)). We call this formulation (Fq yEy
3 a semi-dual Cormuuuuttto formulation of OT.
type of optimal transport problem has emerged, which is called Unbalanced Optimal Transport (UOT) [2 44]. The Csiszär divergence Dy(μ(v) associated with f is generalization of f-divergence for the case where μ is not absolutely continuous with respect to y (See Appendix [A] for the precise definition). Here, the entropy function f: (0, x0) -> [0, oo] is assumed to be convex, lower semi-continuous, and non-negative. Note that the f-divergence families include a wide variety of divergences, such as Kullback-Leibles (KL) divergence and x² divergence. (See [71] for more examples of f-divergence and its corresponding generator f.) Formally, the UOT problem is formulated as follows:
Unbalanced Optimal Transport (UOT) Recently,
Ur cx, y)d v) + Dv, (Tolm) + (T,
inf Cub(1t,) v) = (XM+(Xx3)
(4)
The UOT formulation (Eq 4 has two key properties. First, UOT can handle the transportation of any positive measures by relaxing the marginal constraint I2
44 57]. allowing for greater flexibility in the transport problem. Second, UOT can address the sensitivity to outliers, which is a major limitation of OT. In standard OT, the marginal constraints require that even outlier samples are transported to the target distribution. This makes the OT objective (Eq to be significantly affected by a few outliers. This sensitivity of OT to outliers implies that the OT distance between two distributions can dominated by these outliers [8 666. On the other hand, UOT can exclude outliers from the consideration by flexibly shifting the marginals. Both properties are crucial characteristics of UOT. However, in this work, we investigate a generative model that transports probability measures. Because the total mass of probability measure is equal to 1, we focus the latter property.
23
3 Method
3.1 Dual and Semi-dual Formulation of UOT
Similar to OT, UOT also provides dual formulation 12 21 72
U/ -W((=((=))du(z)) S, -Wū((u()))ct(())
sup Cub(H,v) u(x)) Sc(()))
(5)
with u E C(A') U E C(Y) where C denotes set of continuous functions over its domain. Here, f denotes the convex conjugate of f.i.e. f"(y suPren ( (Ix, y) f())} for f ER [-00, oo).
Generalization of OT). Suppose V, and Vr, the convex indicator function of {1}. i.e. have zero values for x = and otherwise xo. Then, the objective in Eq becomes infinity if mo # l m A v, which means that UOT reduces into classical OT. Moreover, (x) = (x) = T. If we replace V and V3 with the identity function, Eq precisely recovers the dual form of OT (Eq 2.. Note that the hard constraint corresponds to
Remark 3.1 (UOT
T.
Now, we introduce the semi-dual formulation of UOT (72. For this formulation, we assume that yr and V are non-decreasing and differentiable functions (f" non-decreasing for the non-negative entropy function (65]):
Cut (1, v) = sup / M. (-u du(z) + / -V3(-o(v))ddb)) VEC
(6)
where (I)) denotes the c-transform of in Eg 3

---

Remark 3.2 (V' Candidate). Here, we clarify the feasible set of 4 in the semi-dual form of UOT. As aforementioned, the assumption for deriving iq that D is a non-decreasing, differentiable function. Recall that for any function f. its convex conjugate f is convex and lower semi-continuous [9] Also, for any convex and lower semi-continuous f. f a convex conjugate of f 7 Combining these two results, any non-decreasing, convex, and differentiable function can be candidate of V
3.2 Generative Modeling with the Semi-dual Form of UOT
In this section, we describe how we implement a generative model based on the semi-dual form (Eqq
6 of UOT. Following 18
28
41
54
61], we introduce T. to approximate v" as follows:
Te (x) arginf [c(2, y) vvy] nEY
v (z) c
cx, Tx)))) (T)(
(7)
Note that T. is measurable ((10). Prop 7.33). Then, the following objective J(v) can be derived from the equation inside supremum in Eq (6) and the right-hand side of Eq 7
J(o) = J Vi (- [[((] T,), - (T)()))) du(xx)
/ 1/ (-v(y)) dv(y).
(8)
In practice, there is closed-form expression of the optimal T. for each Hence, the optimization Tr for each v is required as in the parametrize to and T.- To with neural networks with parameter and A. Then, our learning objective Lr can be formulated as follows:
training of GAN [25]. In this work, we
LVg.To = inf J/ VI
(-inf inf [e (b(x)) us (Ta(2))) du(x) + / v (-va((v) (v)))))))())))) (9) To
(-
Finally, based on Eq 9 we propose a new training algorithm (Algorithm D. called the UOT-based generative model (UOTM). Similar to the training procedure of GANs, we alternately update the potential Us (lines 2-4) and generator Te (lines 5-7).
Algorithm 1 Training algorithm of UOTM
Require:
The source distribution j and the target distribution v. Non-decreasing, differentiable, convex function pair (1/ V2). Generator network Ta and the discriminator network ug. Total iteration number K.
1: for & = 0, 2,1,,..... K do
2: Sample batch X 14, Y v, NO..I).
3: C. [X] Σzex V (-c To(x, 2)) + us (To(x, 2))) + in Ever V5(-vs(1))
4: Update ø by using the loss L..
5: Sample a batch X u4 N
6: LT |X] EzEx (c (ax To =)) us (To(r,))))
7: Update @ by using the loss LT.
8: end for
we add stochasticity in our generator To by putting an auxiliary variable as an additional input. This allows to obtain the stochastic transport plan x (-1xx) for given input x E X, motivated by the Kantorovich relaxation
Following 26 42
33
The role of the auxiliary variable has been extensively and it has been shown to be useful in generative modeling.
discussed the literature [3] 26 42 80
Moreover, we incorporate R₁ regularization (60), Lreg = \||V
A (y))l for real data y E y, into the
objective (Eq which is a popular regularization method employed in various studies [3 51 80
3.3 Some Properties of UOTM
Divergence Upper-Bound of Marginals UOT relaxes the hard constraint of marginal distributions in OT into the soft constraint with the Csiszår divergence regularizer. Therefore, natural question is how much divergence is incurred in the marginal distributions because of this soft constraint in UOT.
The following theorem proves that the upper-bound of divergences between the marginals in UOT
is linearly proportional to in the cost function c(r, v). This result follows our intuition because down-scaling equivalent to the relative up- up-scaling of divergences De, De, Eq 4 (Notably, in our experiments, the optimization benefit outweighed the effect of soft constraint (Sec 5.2

---

Theorem 3.3. Suppose that and v are probability densities defined on X' and Y. Given the assumptions in Appendix A. suppose that p.v are absolutely continuous with respect to Lebesgue measure and N is continuously disferentiable. Assuming that the optimal potential v² = inf J (v) exists, e- is a solution of the following objective
J(0) = fxx
- (x()dd((())) + I, -o(y)d((u)
(10)
Note that the assumptions guarantee the existence of optimal transport map T between it and v. Furthermore, T* satisfies
where p((x) = V1'(-0" (x)) and i(y) = (22(-00(1))
(x) E aargnn wEV (c(xr,) v) - --" -u)l,
(11)
u-almost surely. In particular, Dv, (Tl (+lp) & (W2 v) where W2(je, D) is Wasserstein- 2 distance between μ and v.
Theorem 3.3 shows that T.. (Bq is a valid parametrization of the optimal transport map T that transports j to D. Moreover, if r is sufficiently small, then f and ú are close to and v, respectively. Therefore, we can infer that T. will transport μ to distribution that is similar to v.
Stable Convergence We discuss convergence properties of UOT objective and the corresponding OT objective J for the potential e through the theoretical findings of Gallouët et al. [21] Theorem 3.4 ((11). Under some mild assumptions in Appendix A the following holds:
1 Ep [IV (6-- + Ci [(v - / - + C2Ex [(v - -") (12)
(() - J(v))
2X
for some positive constant C1 and Cz. Furthermore, J(v) J(v') 2 21 E IIV (v v12)
Theorem 3.4 suggests that the UOT objective J gains stability over the OT objective J in two aspects. First, while J only bounds the gradient error ||V (vf r' J also bounds the function error (v-- = Second, J provides control over both the original function v and its e-transform of in L sense. We hypothesize this stable convergence property conveys practical benefits in neural network optimization. particular, UOTM attains better distribution matching despite its soft constraint (Sec
5.2 and faster convergence during training (Sec 5.3)
4 Related Work
Optimal Transport Optimal Transport (OT) problem addresses transport map between two
distributions that minimizes a specified cost function. This OT map has been extensively utilized in
various applications, such as generative modeling 2 45 61], point cloud approximation (50), and
domain adaptation (20). The significant interest in OT literature has resulted in the development of
diverse algorithms based on different formulations of OT problem, e.g., primary (Eq[), dual (Eq 2
and semi-dual forms (Eq 33. First, several works were proposed based on the primary form [4]] 81
These approaches typically involved multiple adversarial regularizers, resulting in a complex and challenging training process. Hence, these methods often exhibited sensitivity to hyperparameters. Second, the relationship between the OT map and the gradient of dual potential led to various dual
form based methods. In specific, when the cost function is quadratic, the OT map can be represented as the gradient of the dual potential 777. This correspondence motivated new methodology that parameterized the dual potentials to recover the OT map [2] [4]. Seguy et al. [64] introduced the entropic regularization to obtain the optimal dual potentials v and v, and extracted the OT map
from them via the barycentric projection. Some methods 40 49 explicitly utilized the convexity of potential by employing input-convex neural networks [T.
Recently, Korotin et al. [39] demonstrated that the semi-dual approaches [88 61 are the ones that best approximate the OT map among existing methods. These approaches [88
61 properly recovered OT maps and provided high performance in image generation and image translation tasks for large- scale datasets. Because our UOTM is also based on the semi-dual form, these methods show some connections to our work. If we let " and V in Eq 9 as identity functions, our Algorithm I Treduces to the training procedure of Fan et al. [8 (See Remark 3.2). For convenience, we denote Fan et al. [1] as a OT-based generative model (OTM). Moreover, note that Rout et al. [61] can be considered as a minor modification of parametrization from OTM. this paper, we denote Rout al. [6] as Optimal Transport Modeling (OTM*), and regard it as one of the main counterparts for comparison.
5

---

18
14
Cenerted dearty
Geraenc deaa
Density
Density
Target
13
Source
Source
(a) OTM
(b) UOTM
Figure 2: Outlier Robustness Test on Toy dataset with 1% outlier. For each subfigure, Left: Comparison of target density and generated density T and Right: Transport map of the trained model (zz (z)). While attempting to the outlier distribution, OTM generates undesired samples outside and learns the non-optimal transport map. In contrast, UTOM mainly generates in- distribution samples and achieves the optimal transport map. (For better visualization, the y-scale of the density plot is manually adjusted.)
7
6
7
9
1
9
6
7
9 0
5
7
8
9
4
6
3
9
3
3
1
3
9
D
6
6
8 0
3
6
5
4
2
0
1
/
3
4
4
9
6
Figure Outlier Robustness Test on Image dataset (CIFAR-I0+ 1% MNIST). Left: OTM exhibits
artifacts both in-distribution and outlier samples. Right: UOTM attains higher-fidelity samples while generating MNIST-like samples more sparingly, around 0.2%. FID scores of CIFAR- -0-liee samples are 13.82 for OTM and 4.56 for UOTM, proving that UOTM is more robust to outliers.
Unbalanced Optimal Transport The primal problem of UOT (Eq 4 relaxes hard marginal constraints of Kantorovich's 's problem [T through soft entropic penalties 12]. Most of the recent UOT approaches [1] [9
48 [58] estimate UOT potentials on discrete space by using the dual formulation of the problem. For example, Pham al. [58] extended the Sinkhorn method UOT, and Lübeek et al. [88 learned dual potentials through cyclic properties. To generalize UOT into the continuous case, Yang and Uhler [82] suggested the GAN framework for the primal formulation of unbalanced Monge OT. This method employed three neural networks, one for transport map, another for a discriminator, and the third for scaling factor. Balaji et al. [8] claimed that implementing GAN-like procedure using the dual form of UOT is hard to optimize and unstable. Instead, they proposed an alternative dual form of UOT, which resembles the dual form of OT with a Lagrangian regularizer. Nevertheless, Balaji al. g still requires three different neural networks as in [82] To the best of our knowledge, our work is the first generative model that leverages the semi- formulation of UOT. This formulation allows develop simple but novel adversarial training procedure that does not require the challenging optimization of three neural networks.
Experiments
In this section, we evaluate our model on the various datasets to answer the following questions:
1. Does UOTM offer more robustness outliers than OT-based model (OTM)? ($5.
2. Does the UOT map accurately match the source distribution to the target distribution? (5.2)
3. Does UOTM provide stable and fast convergence? ($5.3)
4. Does UOTM provide decent performance across various choices of V and V/ (5.44)
5. Does UOTM show reasonable performance even without regularization term Lreg ? 5.54
6. Does UOTM provide decent performance across various choices of 7 in c(t, y)? 5.4
Unless otherwise stated, we set V = 1 = 122 and D& as a KI. divergence in our UOTM model (Eq 9.. The source distribution μ is standard Gaussian distribution (00 I) with the same dimension as The target distribution v. In this case, the entropy function is given as follows:
x log - + 1, if > 0
T (x)
V (²(x)=c"
(13)
00,
if < 0²
In addition, when the generator Ta and potential uy are parameterized by the same neural networks as OTM* (61), we call them small model. When use the same network architecture in RGM

---

Table 1: Target Distribution Matching Test. UOTM achieves a better approximation target distribution v, Le., F4 v. indicates the results conducted by ourselves.
Table 2: Image Generation on CIFAR-10.
Class
Model
FID (
SNGAN+DGBIe
S
9.62
9.35
AutoGAN 24
12.4
8.60
Model Toy (r 0.1) Toy 0.02)
CIFAR-10
TransGAN 31
9.26
9.02
StyleGAN2 w/o ADA
35
8.32
9.18
Metric
PRu(Twplu) (t)
FID (1)
GAN
StyleGAN2 w/ ADA
35
2.92
9.83
DDGAN (T=1
16.68
OTM
0.05
0.05
7.68
DDGAN [0
9.63
Fixed-ul
0.02
0.004
7.53
RGM
2.47
9.68
UOTM
0.02
0.005
2.97
NCSN 68
25.3
8.87
DDPM
3.21
9,46
Table 3: Image Generation on CelebA-HQ.
Score SDE (VE)
669
2.20
9.89
Score SDE (VP)
69
2.41
9.68
Diffusion
Class
Model
FTD (t)
DDIM (50 steps)
67
4,67
8.78
CLD 15
2.25
Score SDE (VP)
69
7.23
Subspace Diffusion
2.17
9.94
Probability Flow
69
128.13
LSGM 74
2,10
9.87
LSGM 74
7.22
Diffusion
UDM
37
7.16
NVAE 73
23.5
7.18
DDGAN 80
7.64
Glow 38
48.9
3.92
RGM L3
7.15
VAE&EBM
PixelCNN 75
65.9
4,60
VAEBM 79
12.2
8.43
PGGAN 34
8.03
Recovery EBM
9,58
8.30
Adv. LAF 59
19.2
GAN
VQ-GAN 16
10.2
WGAN
55.20
DC-AE
55
15.8
WGAN-GP
39.40
6.49
StyleSwin
83
3.25
Robust-OT
g
21.57
OT-based
AE-OT-GAN 3
17.10
7.35
NVAE NCP-VAE
29.7
OTM (Small)
61
21.78
VAE
24.8
OTM (Large))
7.68
8.50
VAFRM 79
20.4
UOTM (Smally²
12.86
7.21
OTT-hased
UOTM
6.36
UOTM (Large)*
2.97+00007
9.68
13 refer them large model. Unless otherwise stated, consider the large model. For implementation details of experiments, please refer to Appendix B.2
5.1 Outlier Robustness of UOTM
One of the main features UOT is its robustness against outliers. In this subsection, we investigate how this robustness is reflected in generative modeling tasks by comparing our UOTM and OT-based model (OTM). For evaluation, we generated two datasets that have 1-2% outliers: Toy and Image datasets. The toy dataset is a mixture of samples from (11, 0.5²) (in-distribution) and A (-11, 0.5²) (outlier). Similarly, the image dataset a mixture of CIFAR- 10 (in-distribution) and MNIST
14 (outlier) in Balaji al. (8).
Figure 2 illustrates the learned transport map T and the probability density TH of generated distribution for OTM and UOTM models trained on the toy dataset. The density in Fig degmonstemsssssees that OTM attempts to both the in-distribution and outlier samples simultaneously. However, this attempt causes undesirable behavior of transporting density outside the target support. In other words, order address 1-2% of outlier samples. OTM generates additional failure samples that do not belong to the in-distribution or outlier distribution. On the other hand, UOTM focuses on matching in-distribution samples (b). Moreover, the transport maps show that OTM faces challenging optimization. The right-hand side of Figure 2a and 2b visualize the correspondence between the source domain samples z and the corresponding target domain samples T(x). Note that the optimal I should be monotone-increasing function. However, OTM failed to learn such transport map T. while UOTM succeeded.
Figure 3 shows the generated samples from OTM and UOTM models on the image dataset. A similar phenomenon is observed in the image data. Some of the generative images from OTM show some unintentional artifacts. Also, MNIST-like generated images display red or blue-tinted background. contrast, UOTM primarily generates in-distribution samples. In practice, UTOM generates MNIST- like data at very low probability of 0.2%. Notably, UOTM does not exhibit artifacts as in OTM. Furthermore, for quantitative comparison, we measured Fréchet Inception Distance (FID) score [29 by collecting CIFAR-10-BE samples. Then, compared this score with the model trained on clean
7

---

Source
Source
Source
(a) OTM
(b) Fixed-ji
(c) UOTM
Figure 4: Visualization of OT Map (ær, T())) trained on clean Toy data. The comparison suggests that Fixed-ji and UOTM models are closer the optimal transport map compared to OTM.
dataset. The presence of outliers affected OTM by increasing FID score over (7.68 13.82) while the increase was less than 2 in UOTM (2.97 - 4.56). In summary, the experiments demonstrate that UOTM more robust to outliers.
5.2 UOTM Generative Model
Target Distribution Matching As aforementioned in Sec 3.3, UOT allows some flexibility to the marginal constraints of OT. This means that the optimal generator Tÿ does not necessarily transport the source distribution precisely to the target distribution v, i.e., T ss v. However, the goal of the generative modeling task is to learn the target distribution v. In this regard, we assessed whether our UOTM could accurately match the target (data) distribution. Specifically, we measured KL divergence 778 between the generated distribution T# n and data distribution on the toy dataset (See the Appendix B.I for toy dataset details). Also, we employed the FID score for CIFAR- 10 dataset, because FID assesses the Wasserstein distance between the generated samples and training data in the feature space of the Inception network 770.
In this experiment, our UOTM model is compared with two other methods: OTM [8]] (Constraints in both marginals) and UOTM with fixed (Constraints only in the source distribution). We introduced Fixed-j variant because generative models usually sample directly from the input noise distribution μ.. This direct sampling implies the hard constraint on the source distribution. Note that this hard constraint corresponds setting (x) implementation details of UOTM with fixed μ..)
Eq 9 (Remark 3.1.
(See Appendix B.2 for
Table I shows the data distribution matching results. Interestingly, despite the soft constraint, our UOTM matches data distribution better than OTM both datasets. In the toy dataset, UOTM achieves similar KL divergence with and without fixed μ for each т, which is much smaller than OTM. This result can be interpreted through the theoretical properties in Sec Following Thm 3.3 both UOTM models exhibit smaller D (T for the smaller in the toy dataset. It is worth noting that, unlike UOTM, D (1)) does not change for OTM. This is because, in the standard OT problem, the optimal transport map π² does not depend while Fixed-p -u shows similar FID score to OTM, UOTM significantly outperforms both models. As discussed Thm 3.4 we interpret that relaxing both marginals improved the stability of the potential network v, which led to better convergence of the model on the more complex data. This convergence benefit can be observed in the learned transport map in Fig 4 Even on the toy dataset, the UTOM transport map presents better-organized correspondence between source and target samples than OTM and Fixed-j. In summary, our model is competitive in matching the target distribution for generative modeling tasks.
in Eq 1 Moreover, in CIFAR-10. 10,
Image Generation We evaluated our UOTM model on the two generative model benchmarks:
CIFAR-10 [4) (32 x 32) and CelebA-HO [46] (256 x 256). For quantitative comparison, we adopted and Inception Score (IS) 62 The qualitative performance of UOTM depicted in Fig As shown in Table 2 UOTM model demonstrates state-of-the off-the-rt art results among existing OT based methods, with an FID of 2.97 and of 9.68. Our UOTM outperforms the second- best performing OT-based model OTM(Large), which achieves an FID of 7.68, by a large margin. Moreover, UOTM(Small) model surpasses another UOT-based model with the same backbone network (Robust- OT), which achieves an FID of 21.57. Furthermore, our model achieves the state-of-the-art FID score 6.36 on CelebA-HQ (256x 256). To the best of our knowledge, UOTM is the first OT-based generative model that has shown comparable results with tfte-of-the-an models in various datasets.
FID
2

---

120
25
22.5
20.0
100
20
OTM
17.5
UOTM
UOTMIKL)
15.0
22
15
UOTM(z")
0 12.5
40
10
VOTM
20
2.5
200
400
600
800
1000
0.0
0.2
0.5
1.0
2.0
3
Epoch
xx10011
Figure FID Scores during Figure Ablation Study on Training on CIFAR-10.
Figure 7: Ablation Study on r
Regularizer Intensity X.
in c(x, v) = -| 1 v11..
5.3 Fast Convergence
The discussion in Thm suggests that UOTM offers some optimization benefits over OTM, such as faster and more stable convergence. In addition to the transport map visualization in the toy dataset 2 and 4, we investigate the faster convergence of UOTM on the image dataset. In CIFAR- 10, UOTM converges in 600 epochs. whereas OTM takes about 1000 epochs (Fig (5) To achieve the same performance as OTM, UOTM needs only 200 epochs of training, which about five times faster. In addition, compared to several approaches that train CIFAR-10 with NCSN+ (large model) 69). our model has the advantage in training speed; Score SDE 669 takes more than 70 hours for training 1IFAR-10, 48 hours for DDGAN [80]. and 35-40 hours for RGM on four Tesla V100 GPUs. OTM takes approximately 30-35 hours to converge, while our model only takes about 25 hours.
5.4 Ablation Studies
Generalized 1,, 2/2 We investigate the various choices of Csiszär divergence in UOT problem, i.e., Vi and V/2 in Eq 9 As dis- cussed in Remark 3.2 the necessary condition of feasible V is non-decreasing, convex, and differentiable functions. For instance, setting the Csiszår divergence Dy as f-divergence families such as KL divergence, x² divergence, and Pensen-Shanme isfies the aforementioned conditions. For practical optimization, we chose T (z) that gives finite values for all x, such as KL divergence (Eq 13) and x" divergence (Eq 14):
Table 4: Ablation Study Csiszår Divergence Da, (.I.).
(V1, V22
FID
(KL,KL)
2.97
divergence sat-
(x²,x3)
3.02
(KL., xP)
3.21
(x2. KL)
2.78
(Softplus, Saftplus)
3.17
(x 1)2, if 00 if <<<00
[- V''(x)
1 1
if x 2 2 if < -2
V(x) oc,
(14)
We assessed the performance of our UOTM models for the cases where D is KL divergence or x" divergence. Additionally, we tested our model when V = Softplus (Eq 37 which a direct parametrization of V²* Table 4 shows that our UOTM model achieves competitive performance across all five combinations of (Vi, V.)).
in Regularizer Lrey We conduct an ablation study to investigate the effect of the regularization term Ereg on our model's performance. The WGAN family is known to exhibit unstable training dynamics and to highly depend on the regularization term 6 27 51 [6].. Similarly. Figure 6 shows that OTM is also sensitive to the regularization term and fails to converge without (OTM shows an FID score of 152 without regularization. Hence, we excluded this result in Fig 6 for better visualization). In contrast, our model robust to changes in the regularization hyperparameter and produces reasonable performance even without the regularization.
in the Cost Function We performed an ablation study on 5 in the cost function c(z, y) T||æ - v112. In our model, the hyperparameter т controls the relative weight between the cost c(z, y) and the marginal matching Dv, , D Eq In Fig 7 our model maintains a decent performance of FID(< 5) for 7 (0.5, 1, x 10 However, the FID score sharply degrades at t = (0.1, 5) 10-3, Thm 3.3 suggests that smaller T leads to a tighter distribution matching. However, this cannot explain the degradation at = 0.1 10-3, We interpret this is because the cost function provides some regularization effect that prevents a mode collapse of the model (See the Appendix C.2 for detailed discussion).

---

Conclusion
In this paper, we proposed a generative model based on the semi-dual form Unbalanced Optimal Transport, called UOTM. Our experiments demonstrated that UOTM achieves better target distribution matching and faster convergence than the OT-based method. Moreover, UOTM outperformed existing OT-based generative model benchmarks, such CIFAR-10 and CelebA-HQ-256. The potential negative societal impact of our work is that the Generative Model often learns the dependence in the semantics of data, including any existing biases. Hence, deploying Generative Model in real world applications requires careful monitoring to prevent the amplification of existing societal biases present in the data. It important to carefully control the training data and modeling process of Generative Models to mitigate potential negative societal impacts.
Acknowledgements
This work was supported by KIAS Individual Grant [AP087501] via the Center for AI and Natural Sciences Korea Institute for Advanced Study, the NRF grant (2012RIA2C301083 and the MSIT/ITP 1171111709331 [2021-0-00077], [No. 2021-0-01343, Artificial Intelligence Graduate School Program(SNU))).
References
[1] Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Conference Machine Learning. pages 146-155, PMLR, 2017.
[2] Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. Ae-ot: new generative model hased on extended semi-discrete optimal transport. ICI. 2020, 2019.
[3] Dongsheng An, Yang Guo, Min Zhang, Xin Qi, Na Lei, and Xianfang Gu. Ae-ot-gan: Training gans from data specific latent distribution. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXVI 16, pages 548-564. Springer, 2020.
[4] Jyoti Aneja, Alex Schwing, Jan Kautz, and Arash Vahdat. A contrastive learning approach for training variational autoencoder priors. Advances in neural information processing systems, 34: 480 0000-49,.00 2021.
[5] Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining deep generative models via discriminator gradient flow. arXiv preprint urXiv:2012.0077 2020.
[6] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. International conference on machine learning. pages 214-223. PMIR, 2017.
[7] Shiri Avidan and Vitali Milman. The concept of duality in convex analysis, and the characterization of the legendre transform. Annals of mathematics, pages 661-66616666 2009.
[8] Yogesh Balaji. Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances in Neural Information Processing Systems, 33:12934 12944, 2020.
[9] Dimitri Bertsekas. Convex optimization theory, volume 1. Athena Scientific, 2009.
[10] Dimitri Bertsekas and Steven E Shreve. Stochastic optimal control: the discrete-time case, volume 5. Athena Scientific, 1996.
[11] Lactitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso, Unbalanced opti- mal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems, 34:23270-2322 2021.
[12] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François François-Xavier Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274 (11):300031223,, 201,
[13] Jaemoo Choi, Yesom Park, and Myungjoo Kang. Restoration based generative models. In Proceedings of the 40th International Conference on Machine Learning, volume 202. PMLR, 2023.
10

---

[14] Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012.
[15] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped langevin diffusion. The International Conference on Learning Representa- tions, 2022.
[16] Patrick Esser, Robin Rombach, and Bjom Ommer. Taming transformers for rightresolution image synthesis. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 12873-12883, 2021.
[17] Lawrence C Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
[18] Jiaojiao Fan, Shu Liu, Shaojun Ma, Yongxin Chen, and Hao-Min Zhou. Scalable computation of monge maps with general costs. In ICLR Workshop on Deep Generative Models for Highly Structured Data.
[19] Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In International Conference on Machine Learning, pages 3186-3197. PMLR, 2021.
[20] R Flamary, N Courty, D Tuia, and A Rakotomamonjy. Optimal transport for domain adaptation. JEEE Trans. Pattern Anal. Mach. Intell, 1, 2016.
[21] Thomas Gallouët, Roberta Ghezzi, and François- Xavier Vialard. Regularity theory and geometry of unbalanced optimal transport. arXiv preprint 1xXiv:2112.1105:: 2021.
[22] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energy- based models by diffusion recovery likelihood. Advances in neural information processing systems, 2021.
[23] Aude Genevay. Entropy-regularizzed optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUE), 2019.
[24] Xinyu Gong, Shiyu Chang, Yifan Jiang, and Zhangyang Wang. Autogan: Neural architecture search for generative adversarial networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 3224-3234, 2019.
[25] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-14) 2020.
[26] Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, (73(11))3377-3405 2017.
[27] Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
[28] Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. arXiv preprint arXiv:1904. 04546, 2019.
[29] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler. and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[30] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-688 6851, 2020.
[31] Yifan Jiang, Shiyu Chang, and Zhangyang Wang. Transgan: Two transformers can make one strong gan. arXiv preprint arXiv:2 102.07074, 1(3), 2021.
[32] Bowen Jing, Gabriele Corso, Renato Berlinghieri, and Tommi Jaakkola. Subspace diffusion generative models. arXiv preprint arXiv:2205.01 2022.
[33] Leonid Vitalevich Kantorovich. On problem of monge. Uspekhi Mat. Nauk, pages 225-226. 1948.
[34] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.10196.200 2017.
11

---

[35] Tero Karras, Miia Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in Neural Information Processing Systems, 13:12104-12114, 2020.
[36] Cheolhyeong Kim, Seungtae Park, and Hyung Ju Hwang. Local stability of wasserstein gans with abstract gradient penalty. IEEE Transactions on Neural Networks and Learning Systems, 33(9):4527-4537, 2021.
[37] Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and I1-Chul Moon. Score matching model for unbounded data score. arXiv preprint arXiv:2106.05522 2021.
[38] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. Advances in neural information processing systems, 31, 2018.
[39] Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein gans are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track.
[40] Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations.
2021.
[41] Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34: 14:14593-14605, 2021.
[42] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. arXiv preprini 2:XX:::200.122200 2022.
[43] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[44] Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and new hellinger-kantorovinnn kantorovich distance between positive measures. Inventiones mathematicae, 21 211(3):969-111 1117, 2018.
[45] Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein gan with quadratic transport cost. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4832-4841, 2019.
[46] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pages
3730-3738, 2015.
[47] Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large- scale optimal transport via adversarial training with cycle-consistency. arXiv preprint wwXiv:2033.06633 2020.
[48] Frederike Lübeck, Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semi- couplings. In NeurlPS 2022 Al for Science: Progress and Promises.
[49] Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning. pages 6672-6681. PMLR, 2020.
[50] Quentin Mérigot, Filippo Santambrogio. and Clément Sarrazin. Non-asymptotic convergence bounds for wasserstein approximation using point clouds. Advances in Neural Information Processing Systems, 34:1 14:122100-2888 12821, 2021.
[51] Lars Mescheder, Andreas Geiger, and Sebastian Nowozin. Which training methods for gans do actually converge? In International conference machine learning, pages 3481-3490. PMLR, 2018.
[52] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations.
[53] Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pages 666-704, 1781.
12

---

[54] Quan Hoang Nhan Dam, Trung Le, Tu Dinh Nguyen, Hung Bui, and Dinh Phung. Three- player wasserstein gan via amortised duality. In Proc. of the 28th Int. Joint Conf. on Artificial Intelligence (IJCAI), pages 1-11, 2019.
[55] Gaurav Parmar, Dacheng Li. Kwonjoon Lee, and Zhuowen Tu. Dual contradistinctive generative autoencoder. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 823-832, 2021.
[56] Henning Petzka, Asja Fischer, and Denis Lukovnikov. On the regularization of wasserstein gans. In International Conference on Learning Representations.
[57] Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Center for Research in Economics and Statistics Working Papers, (2017-86), 2017.
[58] Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In International Conference on Machine Learning, pages 7673-7682. PMLR, 2020.
[59] Stanislav Pidhorskyi, Donald A Adjeroh, and Gianfranco Doretto. Adversarial latent au- toencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14104-14113, 2020.
[60] Kevin Roth, Aurelien Lucchi, Sebastian Nowozin, and Thomas Hofmann. Stabilizing training of generative adversarial networks through regularization. Advances in neural information processing systems, 30, 2017.
[61] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference Learning Representations.
[62] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, Xi Chen, and Xi Chen. Improved techniques for training gans. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, 2016.
[63] Maziar Sanjabi, Jimmy Ba, Meisam Razaviyayn, and Jason D Lee. On the convergence and robustness of training gans with regularized optimal transport. Advances in Neural Information Processing Systems, 31, 2018.
[64] Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. In ICLR 2018- International Conference on Learning Representations, pages 2018.
[65] Thibault Séjouméé, Jean Feydy, François- François-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. arXiv e-prints, pages arXiv-1910, 2019.
[66] Thibault Séjourné, Gabriel Peyré, and François-Xaviea Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:221 2022.
[67] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models.
Advances in Neural Information Processing Systems, 2021.
[68] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural information Processing Systems, 32, 2019.
[69] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score- based generative modeling through stochastic differential equations. The International Conference on Learning Representations, 2021.
[70] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 2818-2828 2016. doi: 10.1 2010CV
[71] Dávid Terjék and Diego Sonzález-Sáncheer Optimal transport with f-divergence regularization and generalized sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics, pages 5135-5165 PMLR, 2022.
[72] Adrien Vacher and François- Xavverr Vialard. Stability of semi-dual unbalanced optimal transport: fast statistical rates and convergent algorithm. 2022.
[73] Arash Vahdat and Jan Kautz. Nvae: deep hierarchical variational autoencoder. Advances in Neural Information Processing Systems, 13:19667-19679, 2020.
13

---

[74]
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances Neural Information Processing Systems, 34:1 1287-11302, 2021.
[75]
Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pages 1747-1756. PMLR, 2016.
[76] Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
[77] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
[78] Qing Wang, Sanjeev R Kulkarni, and Sergio Verdú. Divergence estimation for multidimensional densities via k-nearest-neighhor distances. IEEE Transactions on Information Theory, 55(5): 2392-2405, 2009.
[79] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between variational autoencoders and energy- models. arXiv preprint urXiv:2010.0065: 2020.
[80] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprimt (Xxi:2112.078/// 2021.
[81] Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. In International Conference on Machine Learning, pages 6882-6892. PMLR. 2019.
[82] KD Yang and C Uhler. Scalable unbalanced optimal transport using generative adversarial networks. International Conference on Learning Representations, 2019.
[83] Bowen Zhang, Shuyang Gu, Bo Zhang, Jianmin Bao, Dong Chen, Fang Wen, Yong Wang, and Baining Guo. Styleswin: Transformer based gan for high-resolution image generation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 1 11304-11314, 2022.
14

---

A Proofs
Notations and Assumptions Let X' and y be compact complete metric spaces which are subsets of I", and μ., v/ be positive Radon measures of the mass 1. Let V 1 and 2 be a convex, differentiable. and V/ i (α) 42(x)) = oo
nonnegative function defined R. We assume that V (0) = (0) for x < 0. The convex conjugate of y is denoted by y We assume that y and V are differentiable and non-decreasing on its domain. Moreover, let c be a quadratic cost, c(x, y) = T vll2. where is positive constant.
First, for completeness, we present the derivation of the dual and semi-dual formulations of the UOT problem.
Theorem A.1
) The dual formulation of Unbalanced OT is given by
Ur - -V7( -W((-(()))dp(a) + I, -V((-u(v))duuu))
Cu) ((s)+u())(c((.))
sup
(15)
where u E C(2), E C()) with u(x) + mly) S cZ v). Note that strong duality holds.
Also, the semi-dual formulation of Unbalanced OT (4) is given by
Us -vi (-v (p))))(())+++))) + / - -Vi(-((y)ddeu))
Cuslu, u) = sup rEC
(16)
where (x) denotes the c-transform uf i.e., (x)) = inf (c(2, v) u(v)). NEY
Proof. The primal problem Eq can be rewritten as
inf J dxz, (Z.)))T(Z..)) J/ V/ TEM+(AX Jxx
mo(x)) dd(x))
du(x)) S,, T22
(i(U))
dv(y). (17)
dv(y))
We introduce slack variables and Q such that = no and Q = TT. By Lagrange multipliers where
and associated to the constraints, the Lagrangian form of Eq 17 is
C(T, P, == [ xxxy
clx, (c..))d(zz3)
(18)
+ f Vi
(P(x)) du(x)) + /, V2 du(())
(dQ(u)
du(u)) du(y)
(19)
+ / u(x) (dP(z) -/8 di di(iz, v)
v(v) (dQ(s) - 1 dz().)) (20)
The dual Lagrangian function is g(u, v) inf P.Q [(T,, P. P,Q,u, u. v). Then, g(u,v) v) reduces into three distinct minimization problems for each variables;
v) inf U u(x)) +
dP(xz)
du(x))
du(x))
| v(y)d(()) + Vo
(dQ(v)
inf
du(w)
(21)
(dv())
Q
/ (c(2,)) u(x) v(y)) do 9)
inf
xxx
Note that the first term of Eq 21 is
dP(x)) du(x)
[dP(x)
du(co) / 1i dp(x)))dp(x))
sup P
m(x)
dp(x))
(22)
Thus, Eq 21 can be rewritten as follows:
[ xxx
g(u, v)
inf
(c(x,v) - u(xx) - v(y)) do dzr (2r,w) w)
(23)
4 (-u(x)) du(x)) / vs (-v(z)) dz(x).
15

---

Suppose that c(x,y) v) - u(x)) - r(y) < 0 for some point (z. y). Then, by concentrating all mass of into the point, the value of the minimization problem of [2]ommsss -oo. Thus, avoid the trivial solution, we should set c(2; y) u(x) e(y) 2 0 almost everywhere. Within the constraint, the solution of the optimization problem of Eq 23 is c(2, y) = u(x) v(y) x-almost everywhere. Finally, we obtain the dual formulation of UOT:
Us -VI(- + J, V35 --3(
sup
(24)
M(r)+=())S((x..))
where c(z,y) = u(x) + o(y) m-almost everywhere (since vi and V2 is non-decreasing). In other words,
- ((w 1/x vi - + /, -
sup (L.))E((()
((
(25)
where is convex indicator function. Note that we assume i and 3 are convex, non-decreasing. and differentiable. Since 2 O, by letting terms in Eq 22 are finite values. Thus, we can apply
easily that all three theorem, which implies that the strong duality holds. Finally, combining the constraint u(x) + v(y) S c(z, y) with tightness condition, i.c. cxz, v) = u(x) + v(y) n-almost everywhere, we obtain the following semi-dual form:
-1 and v: 1,
1// -V (- inf [c(x, v) -(())))wv))))) (a()) + 1 - yEY
sup VEC
(26)
Now, we state precisely and prove Theorem 3.3
Lemma A.2. For any (u,v) E C(X) x C())) ler
V v((-u(=))dp(a) + - -Wi((-o(())d
T(u, v)
(27)
Then, under the assumptions in Appendix A the functional derivative of [u, v) is
fx (u(o))(((-w()))ddda))
I 5u(y) (- w(y) dy
(28)
where where (ōu, 86) E C(X) x C(V) denote the variations of (u, respectively.
Proof. For any (f.g) E C(X) x C()), let
T(f.9) [ - -V( -((-)(())dp(() + -V3( VV((g(())dv(v))
VV
Then, for arbitrarily given potentials (u, v) C(X) x C(V), and perturbation (17. () E C(X) x C()) which satisfies u(x) + n)(x) + v(y) + ((y) S c(r,y).
[(e) = I(+en,, u+e() = =/ no ((u + (n)()))dd(()== SwE' (u:
Note that L(e) is continuous function. Thus, the functional derivative of T(u, v), i.e. (()), given by calculus of variation (Chapter 8, Evans (7)
Jx bu(=))
I, (00())=3'(--(())) (ot))))))(+/oo////(((
(29)
where (du, 5v) E C(X)) x C(V) denote the variations of (u, w), respectively.
Theorem A.3. Suppose that μ and u are probability densities defined on X and y. Given the assumptions in Appendix A, suppose that u.v are absolutely continuous with respect to Lebesgue measure and N is continuously differentiable. Assuming that the optimal potential v² = inf (v) exists, V² is a solution of the following objective
J(v) [ - - (x) + (o())dd))y)
(30)
16

---

Note that the assumptions guarantee the existence of optimal transport map T benween and V. Furthermore, T satisfies
where (xxx) = (x) and i(y) =
T (x) arginl (c(ar v) v² (w)),
(31)
p.a.s..In particular, Dv, (7l Dv,(0l)) & TW2 v) where Wz(H, v) is a Wasserstein-2 dis- tance between μ and v.
Proof. By Lemma A derivative of Eq 27 is
J
I (oo)) be())
(32)
Note that such linearization must satisfy the inequality constraint uxx) + du(()) vgy) + do(y)
c(2, y) to give admissible variations. Then, the linearized dual form can be re- re-written as follows:
/ bu(z)dī(()) + I, bo(y)dū(y).
sup
(33)
(2.)))
where 5u(x)) + 5v(y) I () u(x) - v(v), p(xx = w (=(()(())) and ((y) = V (r((=(()))u(y) Thus, whenever (u* o²) is optimal, the linearized dual problem is non-positive due to the inequality constraint and tightness (i.e. ((zzy) ----- -u"(x) -² (y) = 0 r-almost everywhere),
which implies that Eq 33 achieves its optimum at (u* (u","")) Note that Eq 33 is a linearized dual problem of the following dual problem:
Us u(x)dp(w)) + /, u(y)
sup (())+
(34)
Hence, the UOT problem can be reduced into standard OT formulation with marginals ju and i. Moreover, the continuity of & gives that μA, v are absolutely continuous with respect to Lebesgue measure and have finite second moment (Note that X, X, are compact. Hence, T is bounded on v, and N are bounded). By Theorem 2. in Villani (76), there exists a unique measurable OT map T which solves Monge problem between and v. Furthermore, shown in Remark 5.13 in Villani et al. 777. T (I) E arginf wey [c(2, y) - v(y]].
Finally, if we constrain æ into [T,,, v), then the solution of Eq 4 is Cub = TW2 v). Thus, for the optimal x E MM where To = i and T 0. it is trivial that
T3 (2, v) 2 TW (7. 8) + Dv, (7lp) Dv, (olv), which proves the last statement of the theorem.
(35)
Theorem A.4 ((II). J is convex. Suppose v is a X-strongly convex, of and u are uniformly bounded on the support of p. and v, respectively, and Vi, V3 are strongly convex on every compact set. Then, it holds
1 1
Ji) - J(x')
R 2X
[IV (6
11 + Ci [(ue - u- + C2EN [(u v²)²], (36)
for some positive constant C1 and Cz. Furthermore, J(v) J(v') 2 E [ (112 11]
Proof. See
21
B Implementation Details
B.1
Data
Unless otherwise stated, the source distribution μ is a standard Gaussian distribution (0, with the same dimension the target distribution
Toy Data We generated 4000 samples for each Toy dataset: Toy dataset for Outlier Robustness
(Sec and Toy dataset for Target Distribution Matching (S2c In Fig of Sec 5.1 the target dataset consists of 99% of samples from \111, 0.5²) and 1% samples from ( 0.5% Tab 1
and 4f Sec 5.2 the source data is composed of 50% of samples from AV 0.5²) and 50% of
samples from A (1, 0.5). The target data is sampled from the mixture of two Gaussians: (-1100.5)
with probability of 1/3 and (2,, 0.5) with probability of 2/3.
17

---

CIFAR-10 We utilized all 50,000 samples. For each image æ, we applied random horizontal flip augmentation of probability 0.5 and transformed by 2x to scale the values within the [-1, 1] range. In the outlier experiment (Fig 3 of Sec 5.1). we added additional 500 MNIST samples to the clean CIFAR-10 dataset. When adding additional MNIST samples, each sample was resized to 3 32 x 32 by duplicating its channel to create 3-channeled image. Then, the samples were transformed by 2x 1. For the image generation task in See we followed the source distribution embedding of Rout et al. [6]] for the small model, i.e., the Gaussian distribution with the size of is bicubically upsampled to 3 x 32 x 32.
3 x
CelebA-HQ We used all 120,000 samples. For each image we resized it to 256 x 256 and applied random horizontal flip augmentation with a probability of 0.5. Then, we linearly transformed each image by 2r 1 to scale the values within [-1,
B.2 Implementation details
Toy data For all the Toy dataset experiments, we used the same generator and discriminator architectures. The dimension of the auxiliary variable is set one. For generator, passed through two fully connected (FC) layers with a hidden dimension of 128, resulting in 128-dimensional embedding. We also embedded data into the 128-dimensional vector by passing it through three- layered ResidualBlock [68]. Then, we summed up the two vectors and fed it to the final output module. The output module consists of two FC layers. For the discriminator, we used three layers of ResidualBlock and two FC layers (for the output module). The hidden dimension is 128. Note that the SiLU activation function is used. We used batch size of 256, learning rate of 10-1, and 2000 epochs. Like WGAN (6) we used the number of iterations of the discriminator per generator iteration as five. For OTM experiments, since OTM does not converge without regularization, we set Ri regularization to A = 0.01.
For the small model setting of UOTM, we employed the architecture of Balaji al. [8] Note that this is the same model architecture as in Rout et al. [6]]. We set a batch size of 128, 200 epochs, learning rate of 2 10-4 4 and 10-4 4 for the generator and discriminator, respectively. Adam optimizer with 31 = 0.5, 32 0.9 is employed. Moreover, we use R₁ regularization of X = 0.2. For the large model setting of UOTM, we followed the implementation of Choi al. [3] unless stated. We trained for 600 epochs with À = 0.2 for every iteration. Moreover, in OTM (Large) experiment, we trained the model for 100 epochs. The other hyperparameters are the same as UOTM. For the ablation studies in S.c we reported the best FII) score until 600 epoch for UOTM and 1000 epoch for OTM because the optimal epoch differs for different hyperparameters. Additionally, for the ablation study on in the cost function (Fig 7 we reported the best FID score among three Rı regularizer parameters: À = 0.2, 0.02, 1. To ensure the reliability of our experiments, we conducted three experiments with different random seeds and presented the mean and standard deviation in Tab 2 (UOTM (Large)).
CIFAR-10
CelebA-HQ (256 x 256) We set 10-5 and used R1 regularizer with À = 5 for every iteration. The model is trained for 450 epochs with an exponential moving average of 0.999. The other network architecture and optimizer settings are the same as Xiao et al. [80] and Choi et al. [3].
We used the implementation Wang et al. 78 for the KL divergence Tab [ of Sec 5.2 We set & = 2 (See Wang et al. (78). ). For the evaluation of image datasets, we used 50,000 generated samples to measure IS and FID scores.
Evaluation Metric
Ablation Study Csiszår Divergence To ensure self-comtainedness, include the definition of Softplus function employed in Table
Softplus(z) log(1 exp(x))
(37)
18

---

3.0
17.5
15.0
2.5
12.3
2.0
115
5.0
2.5
0.0
-3
-2
-1
2
-2
(a) KL
(b) X"
(c) Softplus
Figure 8: Various types of yr*
C
Additional Discussion on UOTM
C.1 Stability of UOTM
In this section, provide intuitive explanation for the stability and convergence of our model. Note that our model employs a non-decreasing, differentiable, and convex Vr* (Sec as depicted in Fig In contrast, OTM corresponds to the specific case of 1 (x)))) x. Now, recall our objective function for the potential Us and the generator Tg:
Lu
Σ
Vi (=c (o,t,, =)) us (To(I, +
Σ VEY
v2(-+s()))
(38)
X
Y
øEX
C
Σ FEX
To(x,) u (Tb(z,)))
(39)
X
Then, the gradient descent step of potential us for a single real data y with a learning rate y can be expressed as follows:
V VoLu = -V V =
=0 o+ ~V+ voooy))
(40)
Suppose that our potential network Us assigns a low potential for y. Since the objective of the potential network is to assign high values to the real data, this indicates that Us fails to allocate the appropriate potential to the real data point y. Then, because -v,(y) is large, (v)) in Eq 40 is large, as shown in Fig In other words, the potential network ug takes a stronger gradient step on the real data y where vo fails. Note that this does not happen for OTM because (p"(()) = for all æ. In this respect, UOTM enables adaptive updates for each sample, with weaker updates on well-perforning data and stronger updates on poor-performing data. We believe that UOTM achieves significant performance improvements over OTM by this property. Furthermore, UOTM attains higher stability because the smaller updates on well-behaved data prevent the blow-up model output.
C.2 Qualitative Comparison on Varying
In the ablation study r, UOTM exhibited a similar degradation in FID score(> 20) at 7 (
0.1,5. However, the generated samples are significantly different. Figure Dillustrates the generated samples from trained UOTM for the cases when is too small (r = 0.1) 1) and when 7 is too large t 5. When 7 is too small (Fig ((eft) the generated images show repetitive patterns. This lack of diversity suggests mode collapse. We consider that the cost function provides
regularization effect preventing mode collapse. For the source sample (noise) x, the cost function is defined as c(z, ((x)) π| T(x) Hence, minimizing this cost function encourages the generated images T(x) to disperse by aligning each image to the corresponding noise r. Therefore, increasing from 0.1 led to better FID results in 7
On the other hand, when T is too large Fig 9 (right), the generative samples tend to be noisy. Because the cost function is c(xr, T(x)) 71|2 [1 too large induces the generated image T(x) to be similar to the noise æ. This resulted the degradation of the FID score for = in Fig 7 As discussed in See 5.4 this phenomenon can also be understood through Thm 3.3
19

---

Figure 9: Generated samples from UOTM trained on CIFAR-10 with Left: = 001((110 3) and
Right: T 5(x10 3).
D Additional Results
Figure 10: Generated samples from UOTM trained CIFARIO for (V,, V22 (KI, KT)..
20

---

Figure 11: Generated samples from UOTM trained on CIFARIO for (V,, V22 - (x2 x²)x
Figure 12: Generated samples from UOTM trained on CIFAR10 for (V1, V22 = (KI, x2).
21

---

Coca
Figure 13: Generated samples from UOTM trained on CIFARIO for (V M2) - (x2,
Figure 14: Generated samples from UOTM trained on CIFARIO for (V1, V2) = (Softplus, Softllus..
22

---

Figure 15: Generated samples from UOTM trained on CelebA-HQ (256 256).
23


Reviews:
[{"review_id": "6CD7C2fR6F", "reviewer": "Reviewer_9t3v", "summary": "Standard optimal transport (OT) problem aims at comparing two probability distributions by finding an optimal coupling that achieves a minimum geometrical cost. A major bottleneck of standard OT is the equality of total transported mass between the underlying distributions.  This restraints based-OT data analysis on real machine learning pipelines among many other domains. To alleviate this problem, this constraint is relaxed and gives unbalanced optimal transport (UOT). This paper investigates based-UOT frameworks on generative modeling through semi-dual form of UOT. ", "questions": "### Minor typos/suggestions\n\n- L51: X and Y are compact complete spaces \n- L78: Csiszar divergences (add a reference) \n- L115: sem-dual --> semi-dual\n- L122: in Eq (9) $v$ should be indexed by $\\phi$\n", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and the approach is mostly well-presented. Addressing generative modeling through the UOT toolbox is novel up to my best knowledge. Numerical experiments in this paper show competitive performance of UOTM (UOT+generative modeling). A noted advantage of UOTM compared to OTM is its robustness to outliers and faster convergence. \n- Semi dual formulation of UOT is given with a theoretical guarantee of its stability compared to other OT objectives (see Theorem 3.4).\n- The code is attached in the supplementary and results are reproducible.", "weaknesses": "### Weakness/Questions\n- In the primal formulation of UOT (see Eq 4), I noticed that in the divergence terms $D_{¨\\Psi_1}$ and  $D_{¨\\Psi_2}$ there are not penalized by some positive tuning parameters as well considered in many UOT problem, I mean that one can expect the following formulation:\n$C_{ub}(\\mu, \\nu) = ... + \\lambda_1D_{¨\\Psi_1} + \\lambda_2D_{¨\\Psi_2}$ where $\\lambda_1, \\lambda_2$ control how much mass variations are penalized. \n- I guess the $\\tau$ hyperparameter has great importance in the given formulation of UOT via the cost $c(x, y) = \\frac \\tau2 \\||x - y\\||_2^2$:\n   - (i) In the experiments $\\tau$ should be as small as possible to get a good performance. Is the robustness of UOT is inherited from penalization of the cost or from properties of the divergence penalizations?\n   - (ii)  Let's consider the same cost function $c(x, y) = \\frac \\tau2 \\||x - y\\||_2^2$ with the dual formulation rather than the semi-dual. Can we expect that UOT with dual has an approximate performance like semi-dual UOT?\n- In L227 the mapping $T^\\star$ should be monotone increasing (as shown in Figure 2, b). Could you please explain this fact?\n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for spending time reading our manuscript carefully and providing thoughtful feedback. We hope our replies to be helpful in addressing the reviewer's concerns. \n\n$ $\n\n----\n> **Q1.** Introducing tuning parameters $\\lambda$ to the divergence terms, i.e., $\\lambda_{1} D_{\\Psi_{1}} +  \\lambda_{2} D_{\\Psi_{2}}$.\n\n**A.** From the primal form of UOT below, **adjusting parameters $\\lambda$ of the divergence terms is equivalent to adjusting the parameter $\\tau$ in the cost function $c(x, y) = \\frac{\\tau}{2} \\\\| x-y \\\\|_{2}^{2}$.** Specifically, multiplying a constant $k>0$ to $\\tau$ is equivalent to dividing both $\\lambda_{1}$ and $\\lambda_{2}$ by $k$. Instead of introducing tuning parameters to the divergence terms $D_{\\Psi_1}, D_{\\Psi_2}$, we controlled $c(x, y) = \\frac{\\tau}{2} \\\\| x-y \\\\|\\_{2}^{2}$ by the parameter $\\tau$. **The ablation study on $\\tau$ is presented in Fig 7**.\n$$\n    C\\_{ub}(\\mu, \\nu) := \\inf\\_{\\pi \\in \\mathcal{M}\\_{+}(\\mathcal{X}\\times\\mathcal{Y})} \\left[ \\int_{\\mathcal{X}\\times \\mathcal{Y}} c(x,y) d\\pi(x,y) + D_{\\Psi_1}(\\pi_0|\\mu) + D_{\\Psi_2}(\\pi_1|\\nu) \\right].\n$$\n\n$ $\n\n----\n> **Q2(1).** In the experiments, $\\tau$ should be as small as possible to get a good performance. Is the robustness of UOT inherited from cost or divergence penalization? \n\n**A.** **We consider that the robustness of UOT is inherited from the properties of divergence penalizations $D_{\\Psi_1}, D_{\\Psi_2}$.**\nIn Sec 5.1, we evaluated the outlier robustness of UOT by comparing it with OT-based Models (OTM). \nHere, note that **OTM has the cost penalization term only, not the divergence terms.**\nHence, if the robustness were due to the cost penalization, OTM would also exhibit robustness. \nHowever, the experimental results (Fig 2,3) contradicted this assumption by showing that OTM did not exhibit outlier robustness. \n\n$ $\n\n----\n> **Q2(2).**\nLet's consider the same cost function with the dual formulation rather than the semi-dual. Can we expect that UOT with dual has an approximate performance like semi-dual UOT?\n\n**A.** We appreciate the insightful comment. In our opinion, **the dual form of UOT would be a more challenging problem in terms of neural network training** (As a reminder, we included the equation below). If we adopt the dual form and parametrize two potentials $u, v$ by two neural networks, we would not possess a direct parametrization of transport map $T$ (by neural networks). Moreover, the training process would need to handle an inequality constraint $u(x)+v(y) \\leq c(x,y)$, which appears to be challenging,\n\n$$\nC\\_{ub}(\\mu, \\nu) = \\sup\\_{u(x)+v(y) \\leq c(x,y)} \\left[ \\int\\_{\\mathcal{X}} -\\Psi^{\\*}\\_{1} (-u(x)) d \\mu (x) + \\int\\_{\\mathcal{Y}} -\\Psi^{\\*}\\_{2} (-v(y)) d \\nu (y) \\right].\n$$\n\n$ $\n\n----\n> **Q3.**\nIn L227 the mapping $T^{\\star}$ should be monotone increasing (as shown in Figure 2, b). Could you please explain this fact?\n\n**A.** Thank you for the comment. This is because, **for the 1D distribution case, the optimal transport map $T^{\\star}$ takes the explicit form** for the quadratic cost function $c(x,y) = \\frac{\\tau}{2} \\\\| x-y \\\\|\\_{2}^{2}$. Intuitively, **$T^{\\star}$ maps the $i$-th largest source sample to $i$-th largest target sample.** Formally speaking, if we denote the CDFs (cumulative distribution function) of $\\mu$ and $\\nu$ as $F_{\\mu}$ and $F_{\\nu}$, $T^{\\star} = F_{\\nu}^{-1} \\circ F_{\\mu}$ ([78], Theorem 2.9). Therefore, $T^{\\star}$ should be monotone increasing. We agree with the reviewer that this statement requires further explanation. We would revise our manuscript accordingly.\n\n\n- Note that the optimal $T^{\\star}$ should be a monotone-increasing function. For the 1D distributions $\\mu$ and $\\nu$, the optimal $T^{\\star}$ is a map that transports the $k$-percentile source sample to $k$-percentile target sample ($0 \\leq k \\leq 1$) for $c(x,y) = \\frac{\\tau}{2} \\\\| x-y \\\\|_2^{2}$ ([78], Theorem 2.9).\n\n$ $\n\n----\n> **Minor typos/ Suggestions**\n\n**A.** Thank you for the careful advice. We would correct the manuscript accordingly.\n\n**References**\n\n[1] Santambrogio, Filippo. \"Optimal transport for applied mathematicians.\" Birkäuser, NY 55.58-63 (2015): 94."}, {"review_id": "9cJQjRaSO9", "reviewer": "Reviewer_TBHC", "summary": "In this paper, the authors propose a novel model (UOTM) that uses (as an objective function) a semi-dual form of the unbalanced optimal transport (UOT) problem. Since UOT relaxes the hard constraint of OT on distribution matching, they also provide the theoretical upper bound of divergences between marginals. The experiments conducted show that UOTM outperforms existing OT-based generative models (OTM) on CIFAR-10 and CelebA-HQ-256, and achieves comparable results to the other state-of-the-art baselines. Moreover, UOTM is more robust to outliers than other OTMs and provides stable and fast convergence.", "questions": "Minor comments:\n\nThe captions in Tab. 1, 2, 3, 4 and Fig. 5, 6, 7 could be more informative.\n\nSome figures (such as Fig. 5) are not referenced in the text.\n\nThe order of tables/figures needs to be changed (e.g., in the text, Table 1 is referenced after Table 3).\n\nl. 52: \"satisfies\" --> \"satisfy\".\n\nl. 67: \"couplings\" --> \"coupling\".\n\nEq. (4): ',' --> '.'.\n\nEq. (5): \"inf\" --> \"sup\" (already corrected in the supplement).\n\nl. 107: \".)\" --> \"):\".\n\nl. 115: \"sem-dual\" --> \"semi-dual\".\n\nl. 298: '.' --> ':'.", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed UOTM model has a solid theoretical background and appears original and significant.\n\n2. The experimental setup is appropriate, and the obtained results prove the superiority of the proposed solution concerning the state-of-the-art.\n\n3. UOTM is probably the first OT-based generative model that achieves state-of-the-art performance on real-world large-scale image datasets.", "weaknesses": "1. The presentation needs some improvement.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are deeply grateful to the reviewer for reading our paper and offering thoughtful feedback.\n\n$ $\n\n---\n\n> **Minor comments**\n\n**A.** Thank you for the valuable advice regarding the presentation of our work. Following the advice, we would revise our manuscript as follows:\n\n- The order and captions of tables/figures are revised.\n    - Table 1: Target Distribution Matching Test. UOTM achieves a better approximation of target distribution $\\nu$, i.e., $ T\\_{\\\\#} \\mu \\approx \\nu$.\n    - Table 2: Image Generation on CIFAR-10.\n    - Table 3: Image Generation on CelebA-HQ.\n    - Table 4: Ablation Study on Csiszar Divergence $D_{\\Psi_{i}}(\\cdot | \\cdot)$.\n    - Fig 5: FID Scores during Training on CIFAR-10.\n    - Fig 6: Ablation Study on Regularizer Intensity $\\lambda$.\n    - Fig 7: Ablation Study on $\\tau$ in $c(x,y)=\\frac{\\tau}{2} \\\\|x-y\\\\|_2^{2}$.\n    \n\n- We included a reference to Fig 5 in Section 5.3 (Line 283).\n    - In CIFAR-10, UOTM converges in 600 epochs, whereas OTM takes about 1000 epochs (Fig 5).\n\n$ $\n\n---\n\n> **Typos**\n\n**A.** We appreciate the reviewer's thoughtful comment. We would incorporate the recommended corrections into the manuscript."}, {"review_id": "PyAB3VcCs3", "reviewer": "Reviewer_6D69", "summary": "The paper derives a semi-dual formulation for the unbalanced optimal transport problem (using the known dual reformulation) and solves it with neural networks. The most notable contribution of the paper is that this formulation is applied to generative modeling and achieved nearly SOTA results on several standard datasets. It is also claimed that proposed method is more robust, stable and converges faster than existing OT methods.\n\nDespite the inspiring practical performance, it seems to me that the central theoretical result of the paper (which justifies the entire method) is incorrect. Overall, the paper lacks mathematical rigor and most derivations are done rather sparingly, raising the question whether they are actually correct and making it hard to check them. Plus there are some concerns regarding the experimental evaluations which I think should be more explicitly discussed.\n", "questions": "- What is the point of solving generative modeling with optimal transport maps?\n- Why is there no testing of U-OTM in the image generation with fixed first marginal (noise)? Can such an experiment be conducted? It would be interesting if there is an effect of relaxing the first marginal in your method, or it suffices to relax the second one only (for simplicity).\n- How did the authors perform comparison with OTM (see weaknesses section)?\n- It is known that the quality of the learned generative model significantly depends on the hyperparameters used. From remark 3.1 I see that the method OTM/OTM* is a particular case of the proposed unbalanced method with $phi^\\star(x)=psi^\\star(x)=x$. How did the authors perform the comparison with this method? Did the authors use the same code (e.g., your own) with the same hyperparameters (except for phi, psi)?\n- Is T_v in lines 116-117 a measurable function?\n- How do the authors measure KL divergence by using the samples (table 3)?\n- Could you please somehow (empirically) show that your method indeed learns the unbalanced OT map?\n- In lines 183-196, there is a discussion of some robust OT methods. Why is there no comparison with them? As far as I understand, this seems to be relevant (not absolutely sure).\n- In the proof of Theorem A.1 I do not understand why the strong duality holds (max dual=min primal)? Could you please detail this?\n", "limitations": "Yes\n\n**Update:** *reject* -> *borderline accept*", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Novel extension of dual (semi-dual) optimal transport methods to the unbalanced setup.\n- Good image generation results, code is available.\n- Faster/more stable convergence than prior OT methods; various studies present on the effects of different parameters (phi, psi, tau, etc.).\n", "weaknesses": "- Incorrect main theoretical results, lack of mathematical rigor in the presentation.\n- Insufficient clarity, vague experimental comparisons, limited details of main baselines.\n- One may argue that the algorithmical/theoretical part of the paper is incremental as c-transform based approaches are rather widespread nowadays.\n\nI discuss some weaknesses in detail below.\n\n*Wrong results.* Overall, the paper lacks mathematical rigor. This is a very big issue as it seems to me that informal statements of theorems (such as “under the suitable assumptions”, etc.) have led to incorrect theoretical results included in the paper. To be precise, I think the main theoretical result which forms the basis for the proposed algorithm is incorrect. I mean Theorem 3.3. One of the results claimed there is that $T_{v^*}$ is the optimal map between the rebalanced distributions $\\tilde{\\mu}$ and $\\tilde{\\nu}$.\n\nFirst, I do not understand here why the maximizer $v^*$ of the dual form actually exists (there is “sup” in all the dual forms, not “max”). This is not a trivial fact and should be proved if used. Second, I do believe that it is not true that $T_v^*$ is necessarily an optimal map even when $v^*$ exists. Let me try to provide a counterexample. Let $\\mu=\\delta_0$, $\\nu=\\delta_1$ and let $\\phi_1=\\phi_2$ for symmetry and consider cost $c(x,y)=|x-y|$. Due to symmetry, it is clear that $\\tilde{\\mu}=m\\cdot delta_0$ and $\\tilde{\\nu}=m\\cdot delta_1$ for some positive mass $m>0$. Hence optimal $v$ should be the optimal potential for the balanced transport problem between $\\delta_0$ and $\\delta_1$ (multiplied by mass $m$), here I follow eq. (28). The cost here is $1-0=1$ ($\\times m$), and I suppose that $v(x)=x$ is an optimal potential: $v^c=-v$ and $v^c(0)+v(1)=0+1=1=$cost ($\\times m$). At the same time, consider $T_v(x)=x+10000$. Then $|x-T_v(x)|-T_v(x)=-x=-v(x)=v^c(x)$, i.e., it is a minimizer in eq. (7). Clearly, $T_v$ maps $0$ to $10000$, so it is not an optimal transport map. Thus, I think the main theoretical result provided in the paper is not correct or may require additional assumptions.\n\nI tried to check the proof to see where there may be a mistake but encountered some difficulties. There the authors rather sparingly use the calculus of variations on the space of function without any formal definitions of what are these “derivatives”, whether they actually exist and skip some related questions. I believe this lack of details (and, to some extent, informalism) may lead to even more mistakes or gaps in their results; maybe this part should be checked by an expert in this particular type of mathematical apparatus. Even assuming this part is true, I am sufficiently sure that there is a mistake in lines (29)-(31). I do not understand where from $T_{v^*}$ appears there (this is not explained); explanations in line 568-569 seem to be wrong.\n\nOverall, this mistake which I pointed out seems to be very crucial for the paper: it turns out to be simply incorrect that the proposed method learns an OT map between distributions $\\tilde{\\mu}$ and $\\tilde{\\nu}$. Hence, it remains questionable what it actually learns and are there any guarantees.\n\nOne additional minor comment here: the function T is introduced in section 3.2 but there is no evidence, explanation or motivation explaining what this function is and how it is related to OT map (the answer only appears in the next section and seems not correct). The clarity of the exposition could be improved.\n\nTaking into account my comment about the incorrectness of the theoretical results, I also do not understand how Theorem 3.4 of [19] explains better stability or performance. To be precise, the theorem is entirely about the dual problem (c-transform-based, function $v$ only), while the authors solve the saddle point problem with $v$ and $T$. Obtaining the better value for the dual necessarily means a better dual variable. Yet for me it is not clear how this relates to the stability as the actual problem is min-max and involves two functions ($v$ and $T$, and we are finally interested in $T$, yes?). While I think there is some meaningful idea in the discussion in lines 148-157, it still seems a little bit speculative and should be made more convincing.\n\n*Limited details of comparisons.* The section with the comparisons is rather dense but still leaves some essential details in Appendix or even skips them. First, it was not evident for me how the authors applied their method to generative modeling (which transport cost has been used between the noise and real data). This (rather important) information seems to appear only in Appendix B.1 and, as far as I understand, the authors used the strategy similar to OTM* paper by Rout et. al. This could have been explained earlier. Second, it remains unexplained how the numbers in Table 1 are obtained (whether the authors reproduce all these methods or just took the metrics from the respective papers). Third, comparison with OT methods is somewhat vague. For example, the authors write that they compared with the OTM paper by Fan et. al. I have checked this paper and did not even find any experiments in image generation. Hence, I wonder how this comparison is done. Due to this, I can not judge how fair the comparisons are. I suggest the authors provide more details in the main text about all this.\n\n*Related work.* The semi-dual neural OT methods (a.k.a. c-transform-based) already existed for several years. The authors cite only rather recent papers in the field and do not mention the early papers where these approaches appeared, e.g., [1,2,3].\n\n[1] Henry-Labordere, P. (2019). (Martingale) Optimal Transport And Anomaly Detection With Neural Networks: A Primal-dual Algorithm. arXiv preprint arXiv:1904.04546.\n\n[2] Nhan Dam, Q. H., Le, T., Nguyen, T. D., Bui, H., & Phung, D. (2019). Three-player wasserstein gan via amortised duality. In Proc. of the 28th Int. Joint Conf. on Artificial Intelligence (IJCAI) (pp. 1-11).\n\n[3] Korotin, A., Li, L., Genevay, A., Solomon, J. M., Filippov, A., & Burnaev, E. (2021). Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34, 14593-14605.\n\n*Misprints.*\n- Line 99. Should be $R^d$ instead of $R$ (twice)?\n- Line 556. Should be $C(X)$, $C(Y)$ I guess?\n- Eq. (26). It seems like one has to assume that $\\Psi^*$ and $\\Phi^*$ are differentiable everywhere?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer providing valuable advice. We think answering the reviewer's comments has significantly improved our work. We hope our replies to be helpful in addressing the reviewer's concerns.\n\n### **1. Theoretical concerns**\n\n$ $\n\n### **1.1 Concerns already discussed in the manuscript**\n---\n> **W1.** The paper lacks mathematical rigor. There are informal statements of theorems such as “under the suitable assumptions”. \n\n**A1.** As mentioned in Line 51-52, **we presented the detailed assumptions of probabilistic space and measure in Appendix A (Line 523-525).** Specifically, $\\mathcal{X}$ and $\\mathcal{Y}$ are assumed to compact complete metric spaces which are subsets of $\\mathbb{R}^{d}$, and $\\mu,\\nu$ be positive Radon measures of mass 1 on $\\mathcal{X}$, $\\mathcal{Y}$, respectively. \n\n$ $\n\n---\n> **Limited details 1.** First, it was not evident to me how the authors applied their method to generative modeling (which transport cost has been used between the noise and real data). \n\n**A2.** As mentioned in Line 59, we adopt the **quadratic cost** $c(x,y)=\\frac{\\tau}{2}\\\\|x-y\\\\|_{2}^{2}$ throughout this paper. To enhance clarity of implementation, we included the **training algorithm of UOTM** in Line 125-126 of the manuscript. Moreover, we provided a detailed explanation of the connection between OTM* (Rout et al.) and our UOTM in Line 178-179.\n\n$ $\n\n---\n> **M3.** Eq 26, $\\Psi^*$ should be differential everywhere.\n\n**A3.** As mentioned in Line 106 and 524, we assumed $\\Psi^*$ s are differentiable.\n\n$ $\n\n### **1.2 Corrections to Theoretical Components**\n\n---\n> **W2.** There is no discussion on the existence of optimal $v$ and $T$. If exists, is $T_{v^*}$ an optimal map between the rebalanced distributions? The main theoretical result provided in the paper is not correct or may require additional assumptions.\n\n**A4.** We sincerely appreciate the reviewer for providing valuable feedbacks. We believe that these feedbacks have been a considerable help in revising our manuscript to provide a more mathematically rigorous statement. We added the additional assumptions on measures $\\mu, \\nu$ and $\\Psi^{*}$ to ensure the existence of solutions, and corrected Theorem 3.3 as follows (We removed Eq (29)-(31) accordingly):\n\n$ $\n\n> **Theorem 3.3** Suppose that $\\mu$ and $\\nu$ are probability densities defined on $\\mathcal{X}$ and $\\mathcal{Y}$. Given the assumptions in Appendix A, suppose that $\\mu, \\nu$ are absolutely continuous with respect to Lebesgue measure and $\\Psi^{\\*}$ is continuously differentiable. Assuming that the optimal potential $v^{\\star}=\\inf\\_{v\\in\\mathcal{C}}J(v)$ exists, $v^\\star$ is a solution of the following objective \n$$\n\\tilde{J}(v)= \\int_{\\mathcal{X}}{-v^c(x)d\\tilde{\\mu}(x)}+\\int_{\\mathcal{Y}}{-v(y)d\\tilde{\\nu}(y)},\n$$\n> where $\\tilde{\\mu}(x)={\\Psi_1^*}'(-{v^\\star}^c(x))\\mu(x)$ and $\\tilde{\\nu}(y)={\\Psi_2^*}'(-{v^\\star}(y))\\nu(y)$. Note that the assumptions guarantee the existence of optimal transport map $T^{\\star}$ between $\\tilde{\\mu}$ and $\\tilde{\\nu}$. Furthermore, $T^\\star$ satisfies\n$$\n    T^{\\star}\\in\\mathcal{T}:=\\\\{T_{v^\\star}:\\mathcal{X}\\rightarrow\\mathcal{Y}\\mid T_{v^\\star}(x)\\in\\arg\\min_{y\\in\\mathcal{Y}}\\left[c(x,y) - v^{\\star}(y)\\right]\\textit{ for measurable }T_{v^\\star}\\\\}.\\hspace{20pt} (a)\n$$\n> In particular, $D_{\\Psi_1}(\\tilde{\\mu}|\\mu)+D_{\\Psi_2}(\\tilde{\\nu}|\\nu)\\leq\\tau\\mathcal{W}_2^2(\\mu,\\nu)$ where $\\mathcal{W}_2(\\mu, \\nu)$ is a Wasserstein-2 distance between $\\mu$ and $\\nu$.\n\n$ $\n\nWe assume $\\mu,\\nu$ are absolutely continuous and $\\Psi^*$ is continuously differentiable. **These additional assumptions ensure the existence of the optimal potential $v^\\star$ and optimal transport map $T^\\star$ for the OT problem between $\\tilde{\\mu}$ and $\\tilde{\\nu}$.** The continuity of ${\\Psi^*}'$ gives that $\\tilde{\\mu},\\tilde{\\nu}$ are absolutely continuous with respect to Lebesgue measure and have a finite second moment (Note that $\\mathcal{X},\\mathcal{Y}$ are compact. Hence, ${\\Psi^*}'$ is bounded on $\\mathcal{X},\\mathcal{Y}$, and $\\mathcal{X},\\mathcal{Y}$ are bounded). These properties of $\\tilde{\\mu},\\tilde{\\nu}$ gives the existence of $v^\\star$ and $T^\\star$ [79].\n\nMoreover, thanks to the reviewer, through a rigorous investigation, we realized that $T_{v^\\star}$ might not serve as a transport map. Specifically, **the optimal transport map $T^\\star$ between $\\tilde{\\mu}$ and $\\tilde{\\nu}$ is contained in the set of introduced parametrization $\\mathcal{T}$ in Eq (a) ([71], Remark 5.13)**. However, we cannot ensure the uniqueness of the introduced parametrizations. The suggested counterexample $T_{v}(x)=x+10000$ is also a problem of this uniqueness. Both optimal transport $T^\\star (x)=x+1$ and $T_v (x)$ are included in the parametrization set, i.e., $T_{v},T^{\\star}\\in\\mathcal{T}$. **Nevertheless, note that the introduced assumption of absolute continuity excludes the suggested counterexample of $\\mu=\\delta_0,\\nu=\\delta_1$. Moreover, the assumed cost function $c(x,y)=\\|x-y\\|$ is different from our case.**\n\n$ $\n\n---\n> **W3.** It remains questionable if the proposed method learns OT map between $\\tilde{\\mu}$ and $\\tilde{\\nu}$.\n\n**A5.** In this part, we would like to advocate the $c$-transform-based parametrization (Eq (a)). **To the best of our knowledge, in the $c$-transform-based literature, various works [8, 37, 56, 76] employed this same parametrization of the OT maps $T$.** However, as discussed in [3, 4] and A4, the above parametrization does not provide a formal guarantee of convergence towards $T$.\n\nNevertheless, as reported in [56], the suggested parametrization converges to $T$ in practice. We agree with the reviewer that this part remains a gap in the literature, and further investigation is required.  However, our contribution is not limited to the theoretical parts. We kindly ask the reviewer to take into account the experimental contributions of our study as well.\n\n$ $\n\n---\n**Other concerns will be addressed in additional Official Comments.**"}, {"review_id": "3dX5uWVJjs", "reviewer": "Reviewer_Pz9z", "summary": "The authors propose to use semi-dual formulation of unbalanced optimal transport for generative modelling. Via the experiments on image datasets, the authors illustrate the advantages of the proposed method, namely outlier robustness, stability and fast convergence, and extensively compare with other methods.\n", "questions": "- I'm not very clear about the order of magnitude of $\\tau$ in the experiment. Is it $0.5, 1.2$, or $0.5 \\times 10^{-3}, 1.2\\times 10^{-3}$. The smaller it is, the more we are forced to respect the marginal constraint, i.e. approaching the balanced OT problem.\n\n- What is the definition of Softplus?\n\n- Regarding the theoritical justification of outlier robustness of UOT, the authors may consider to add reference [17]. I see that it is already cited, but not for that purpose.\n", "limitations": "The authors do not discuss the limitations, but do discuss the potential negative societal impact of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors have spent significant effort on the experiments, including extensive comparison with other competing methods, abalation study.\nThe literature review is also good. To this extent, I'm convinced that UOTM is indeed an efficient method for generative modelling under the presence of outliers.", "weaknesses": "However, this work is a direct implementation/application of [67] with no additional improvement or modification. The training procedure is also standard in GAN. For this reason, I would say the overall contribution is quite incremental, and this puts me on the borderline decision, where I'm fighting between the novelty/originality of the work and the effort that the authors spend on the experiments.\n\n- While the authors compare UOTM with various methods, I don't see the comparison with [76]. Is there any reason that [76] is not compared or considered in the experiments?\n\n- Given the fact that unbalancedness implies robustness, it's not new nor surprising to see UOTM outperforms OTM under the presence of noise. Probably it's worth discussing more about the comparison with other unbalanced-based methods (e.g. [8] or [17]), rather than balanced-based methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer for acknowledging that \"UOTM is indeed an efficient method for generative modeling under the presence of outliers\". Also, we agree with the reviewer that, in terms of method, our model is a direct parametrization of the semi-dual form of UOT [67]. Nevertheless, we would like to emphasize the contribution of our work.\n\n- Our work is the first attempt to introduce the semi-dual form of UOT into the generative modeling task. **The semi-dual form enables 2 network parametrizations of UOTM, while the other UOT-based approaches [1, 2] require 3 networks (Line 188-196).** We believe this factor simplifies the training process, contributing to the achievement of competitive results.\n\n- As the reviewer said, we demonstrated UOTM is an efficient method for generative modeling under the presence of outliers (Sec 5.1).\n\n- Moreover, we showed UOTM can be a competitive model for clean datasets as well. The soft constraint on distribution matching of UOT induces a concern about its performance on clean datasets. **To address this concern, we provided an upper-bound of the divergences between the marginals in UOT (Theorem 3.3). We evaluated UOTM on the image benchmark datasets, such as CIFAR-10 and CelebA-HQ (Sec 5.2).**\n\n- Furthermore, we provided **explanations for the observed experimental results where UOT-based model (UOTM) outperforms OT-based model (OTM) in terms of distribution matching, despite the soft constraint of UOT**. These explanations are in two folds:\n\n    - First, we suggested a theoretical explanation based on Theorem from [19]. UOT objective minimizes not only the gradient error $\\\\| \\nabla (v^{c}-v^{\\star c}) \\\\|$, but also the function errors $\\\\| v^{c}-v^{\\star c}\\\\|, \\\\| v-v^{\\star}\\\\|$ (Line 148-157). In contrast, OT objective only minimizes the gradient error $\\\\| \\nabla (v^{c}-v^{\\star c})\\\\|$. We hypothesize this property enables the neural network training easier for UOTM, leading to better distribution matching than OTM.\n\n    - Second, we provided some intuitive explanations from the engineering point of view in Appendix C.1. Compared to OTM, UOTM enables adaptive updates for each sample, with weaker updates on well-performing data and stronger updates on poor-performing data (Line 649-650). \n \n$ $\n\n------\n\n> **W1.** While the authors compare UOTM with various methods, I don't see the comparison with [76]. Is there any reason that [76] is not compared or considered in the experiments?\n\n**A1.** We conducted experiments of [76] on CIFAR-10. Despite extensive efforts in exploring various hyperparameters, [76] provided non-competitive results (with architecture \"Small\"), with FID scores around 80. Note that UOTM achieves a FID score of 12.86 on CIFAR-10 with the same Small backbone network.\n\n$ $\n\n------\n\n> **W2.** Given the fact that unbalancedness implies robustness, it's not new nor surprising to see UOTM outperforms OTM under the presence of noise. Probably it's worth discussing more about the comparison with other unbalanced-based methods (e.g. [8] or [17]), rather than balanced-based methods.\n\n**A2.** We appreciate the thoughtful advice. In the Outlier Robustness experiment (Sec 5.1), **our goal was to assess the robustness of the UOT-based model (UOTM) by comparing it with its OT-based counterpart (OTM).** Note that both UOTM and OTM are derived from the semi-dual form of each problem. Moreover, we aimed to provide a **better understanding of how the UOT-based model works** compared to the OT-based model. \n\nInstead, to demonstrate UOTM's superior performance among UOT-based models, **we compared UOTM with another UOT-based model (Robust-OT [8]) on image benchmarks such as CIFAR-10 (Table 1)**.\n\n$ $\n\n------\n\n> **Q1.**\nI'm not very clear about the order of magnitude of $\\tau$ in the experiment. Is it $0.5, 1, 2$ or $0.5 \\times 10^{-3}, 1 \\times 10^{-3}, 2 \\times 10^{-3}$. The smaller it is, the more we are forced to respect the marginal constraint, i.e. approaching the balanced OT problem.\n\n**A3.** We apologize for the confusion. In the ablation study on $\\tau$, $\\tau$ is chosen among $\\\\{0.1 \\times 10^{-3}, 0.5 \\times 10^{-3}, 1 \\times 10^{-3}, 2 \\times 10^{-3}, 5 \\times 10^{-3}\\\\}$. For clarity, we revised the description of the experimental setup in the manuscript as follows:\n\n- In Fig 7, our model maintains a decent performance of FID($\\leq 5$) for $\\tau = \\\\{0.5, 1, 2\\\\} \\times 10^{-3}$.\n\n$ $\n\n------\n\n> **Q2.**\nWhat is the definition of Softplus?\n\n**A4.** Thank you for the comment. To ensure self-containedness, we added the definition of Softplus in the appendix. The definition of Softplus is as follows:\n\n$$\n   \\operatorname{Softplus}(x) = \\log ( 1+ \\exp (x) )\n$$\n\n$ $\n\n------\n\n> **Q3.**\nRegarding the theoretical justification of the outlier robustness of UOT, the authors may consider adding a reference [17]. I see that it is already cited, but not for that purpose.\n\n**A5.** Thank you for the valuable advice. We added a reference [17] to support the outlier robustness of UOT as follows:\n\n- Line 87-88: Second, UOT can address the sensitivity to outliers, which is a major limitation of OT [17].\n- Line 212: One of the main features of UOT is its robustness against outliers [17].\n\n$ $\n\n**References**\n\n[8] Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. NeurIPS, 2020.\n\n[19] Thomas Gallouët, Roberta Ghezzi, and François-Xavier Vialard. Regularity theory and geometry of unbalanced optimal transport. arXiv preprint arXiv:2112.11056, 2021.\n\n[76] KD Yang and C Uhler. Scalable unbalanced optimal transport using generative adversarial networks. ICLR, 2019."}]

Author rebuttal:

Meta-review: This paper introduces a novel and interesting generalized formulation of the Optimal Transport problem that relies on a semi-dual form of the unbalanced OT problem. The paper was well received by the reviewers in general; they highlighted in particular the strength of the experimental results, the quality of writing, and the reproducibility of the results based on the provided codebase. 

Reviewer 6D69 had some concerns about the main theoretical result of the paper, unearthed through a very thorough and detailed read of the paper. Through a rich back-and-forth with this reviewer, the authors acknowledged the mistakes in the result but were able to identify the cause of the error and address it. The reviewer who raised this issue was mostly satisfied with the result. Thanks to reviewer 6D69, this paper has undergone thorough scrutiny, and assuming the authors will update their submission to fix the error and include a discussion summarizing the thread with the reviewer, this paper will come in a much stronger form, suitable for publication in the conference.
Decision: Accept (poster)