Title: Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
Venue: NeurIPS 2023
Accepted: yes
Keywords: Optimal Transport

Generative modeling

Generative adversarial network
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

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 https://github.com/Jae-Moo/UOTM .

1 Introduction

Optimal Transport theory [57, 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 as a 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, 27], convergence [63], and stability [36, 52] 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 a surge of research on directly modeling the optimal transport map between the input prior distribution and the real data distribution [2, 3, 49, 61, 82]. In other words, the optimal transport serves as the generative model itself. These approaches showed promising results that are comparable to WGAN models. However, the classical optimal transport-based approaches are known to be highly sensitive to outlier-like samples [8]. The existence of a few outliers can have a significant impact on the overall OT-based 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, we suggest a new generative algorithm utilizing the semi-dual formulation of the Unbalanced Optimal Transport (UOT) problem [12, 44]. In this regard, we refer to our model as the UOT-based generative model (UOTM) . The UOT framework relaxes the hard constraints

10

of marginal distribution in the OT framework by introducing soft entropic penalties. This soft constraint provides additional robustness against outliers [8, 66]. Our experimental results demonstrate that UOTM exhibits such outlier robustness, as well as faster and more stable convergence compared to existing OT-based models. Particularly, this better convergence property leads to a 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 and 6.36 on CelebA-HQ, outperforming existing OT-based adversarial methods by a significant margin and approaching state-of-the-art performance. Furthermore, 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-art performance on real-world image datasets.

2 Background

Notations Let X , Y be two compact complete metric spaces, µ and ν be probability distributions on X and Y , respectively. We regard µ and ν as the source and target distributions. In generative modeling tasks, µ and ν correspond to tractable noise and data distributions . For a measurable map T , T # µ represents the pushforward distribution of µ . Π( µ, ν ) denote the set of joint probability distributions on X ×Y whose marginals are µ and ν , respectively. M + ( X ×Y ) denote the set of joint positive measures defined on X ×Y . For convenience, for π ∈ M + ( X ×Y ) , let π 0 ( x ) and π 1 ( y ) be the marginals with respect to X and Y . c ( x, y ) refers to the transport cost function defined on X ×Y . Throughout this paper, we consider X = Y ⊂ R d with the quadratic cost, c ( x, y ) = τ ∥ x -y ∥ 2 2 , where d indicates the dimension of data. Here, τ is a given positive constant. For the precise notations and assumptions, see Appendix A.

Optimal Transport (OT) OT addresses the problem of searching for the most cost-minimizing way to transport source distribution µ to target ν , based on a given cost function c ( · , · ) . In the beginning, Monge [53] formulated this problem with a deterministic transport map. However, this formulation is non-convex and can be ill-posed depending on the choices of µ and ν . To alleviate these problems, the Kantorovich OT problem [33] was introduced, which is a convex relaxation of the Monge problem. Formally, the OT cost of the relaxed problem is given by as follows:

where c is a cost function, and π is a coupling of µ and ν . Unlike Monge problem, the minimizer π ⋆ of Eq 1 always exists under some mild assumptions on ( X , µ ) , ( Y , ν ) and the cost c ([77], Chapter 5). Then, the dual form of Eq 1 is given as:

where u and v are Lebesgue integrable with respect to measure µ and ν , i.e., u ∈ L 1 ( µ ) and v ∈ L 1 ( ν ) . For a particular case where c ( x, y ) is equal to the distance function between x and y , then u = -v and u is 1-Lipschitz [77]. In such case, we call Eq 2 a Kantorovich-Rubinstein duality . For the general cost c ( · , · ) , the Eq 2 can be reformulated as follows ([77], Chapter 5):

where the c -transform of v is defined as v c ( x ) = inf y ∈Y ( c ( x, y ) -v ( y )) . We call this formulation (Eq 3) a semi-dual formulation of OT .

Unbalanced Optimal Transport (UOT) Recently, a new type of optimal transport problem has emerged, which is called Unbalanced Optimal Transport (UOT) [12, 44]. The Csiszàr divergence D f ( µ | ν ) associated with f is a generalization of f -divergence for the case where µ is not absolutely continuous with respect to ν (See Appendix A for the precise definition). Here, the entropy function f : [0 , ∞ ) → [0 , ∞ ] 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-Leibler (KL) divergence and χ 2 divergence. (See [71] for more examples of f -divergence and its corresponding generator f .) Formally, the UOT problem is formulated as follows:

The UOT formulation (Eq 4) has two key properties. First, UOT can handle the transportation of any positive measures by relaxing the marginal constraint [12, 23, 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 1) to be significantly affected by a few outliers. This sensitivity of OT to outliers implies that the OT distance between two distributions can be dominated by these outliers [8, 66]. 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 on the latter property.

3 Method

3.1 Dual and Semi-dual Formulation of UOT

Similar to OT, UOT also provides a dual formulation [12, 21, 72]:

with u ∈ C ( X ) , v ∈ C ( Y ) where C denotes a set of continuous functions over its domain. Here, f ∗ denotes the convex conjugate of f , i.e., f ∗ ( y ) = sup x ∈ R {⟨ x, y ⟩ -f ( x ) } for f : R → [ -∞ , ∞ ] .

̸

Remark 3.1 ( UOT as a Generalization of OT ) . Suppose Ψ 1 and Ψ 2 are the convex indicator function of { 1 } , i.e. have zero values for x = 1 and otherwise ∞ . Then, the objective in Eq 4 becomes infinity if π 0 = µ or π 1 = ν , which means that UOT reduces into classical OT. Moreover, Ψ ∗ 1 ( x ) = Ψ ∗ 2 ( x ) = x . If we replace Ψ ∗ 1 and Ψ ∗ 2 with the identity function, Eq 5 precisely recovers the dual form of OT (Eq 2). Note that the hard constraint corresponds to Ψ ∗ ( x ) = x .

̸

Now, we introduce the semi-dual formulation of UOT [72]. For this formulation, we assume that Ψ ∗ 1 and Ψ ∗ 2 are non-decreasing and differentiable functions ( f ∗ is non-decreasing for the non-negative entropy function f [65]):

where v c ( x ) denotes the c -transform of v as in Eq 3.

Remark 3.2 ( Ψ ∗ Candidate ) . Here, we clarify the feasible set of Ψ ∗ in the semi-dual form of UOT. As aforementioned, the assumption for deriving Eq 6 is that Ψ ∗ 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 is a convex conjugate of f ∗ [7]. Combining these two results, any non-decreasing, convex, and differentiable function can be a candidate of Ψ ∗ .

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 (Eq 6) of UOT. Following [18, 28, 41, 54, 61], we introduce T v to approximate v c as follows:

Note that T v 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:

In practice, there is no closed-form expression of the optimal T v for each v . Hence, the optimization T v for each v is required as in the generator-discriminator training of GAN [25]. In this work, we parametrize v = v ϕ and T v = T θ with neural networks with parameter ϕ and θ . Then, our learning objective L v ϕ ,T θ can be formulated as follows:

Finally, based on Eq 9, we propose a new training algorithm (Algorithm 1), called the UOT-based generative model (UOTM) . Similar to the training procedure of GANs, we alternately update the potential v ϕ (lines 2-4) and generator T θ (lines 5-7).

Algorithm 1 Training algorithm of UOTM

Require: The source distribution µ and the target distribution ν . Non-decreasing, differentiable, convex function pair (Ψ ∗ 1 , Ψ ∗ 2 ) . Generator network T θ and the discriminator network v ϕ . Total iteration number K .

1: for k = 0 , 1 , 2 , . . . , K do

2: Sample a batch X ∼ µ , Y ∼ ν , z ∼ N ( 0 , I ) .

3: L v = 1 | X | ∑ x ∈ X Ψ ∗ 1 ( -c ( x, T θ ( x, z )) + v ϕ ( T θ ( x, z ))) + 1 | Y | ∑ y ∈ Y Ψ ∗ 2 ( -v ϕ ( y ))

4: Update ϕ by using the loss L v .

5: Sample a batch X ∼ µ , z ∼ N ( 0 , I ) .

6: L T = 1 | X | ∑ x ∈ X ( c ( x, T θ ( x, z )) -v ϕ ( T θ ( x, z ))) .

7: Update θ by using the loss L T .

8: end for

Following [26, 42], we add stochasticity in our generator T θ by putting an auxiliary variable z as an additional input. This allows us to obtain the stochastic transport plan π ( ·| x ) for given input x ∈ X , motivated by the Kantorovich relaxation [33]. The role of the auxiliary variable has been extensively discussed in the literature [13, 26, 42, 80], and it has been shown to be useful in generative modeling. Moreover, we incorporate R 1 regularization [60], L reg = λ ∥∇ x v ϕ ( y ) ∥ 2 2 for real data y ∈ Y , into the objective (Eq 9), which is a popular regularization method employed in various studies [13, 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, a 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 ( x, y ) . This result follows our intuition because down-scaling τ is equivalent to the relative up-scaling of divergences D Ψ 1 , D Ψ 2 in Eq 4. (Notably, in our experiments, the optimization benefit outweighed the effect of soft constraint (Sec 5.2).)

Theorem 3.3. Suppose that µ and ν are probability densities defined on X and Y . Given the assumptions in Appendix A, suppose that µ, ν are absolutely continuous with respect to Lebesgue measure and Ψ ∗ is continuously differentiable. Assuming that the optimal potential v ⋆ = inf v ∈C J ( v ) exists, v ⋆ is a solution of the following objective

where ˜ µ ( x ) = Ψ ∗ 1 ′ ( -v ⋆c ( x )) µ ( x ) and ˜ ν ( y ) = Ψ ∗ 2 ′ ( -v ⋆ ( y )) ν ( y ) . Note that the assumptions guarantee the existence of optimal transport map T ⋆ between ˜ µ and ˜ ν . Furthermore, T ⋆ satisfies

µ -almost surely. In particular, D Ψ 1 (˜ µ | µ )+ D Ψ 2 (˜ ν | ν ) ≤ τ W 2 2 ( µ, ν ) where W 2 ( µ, ν ) is a Wasserstein2 distance between µ and ν .

Theorem 3.3 shows that T v ⋆ (Eq 7) is a valid parametrization of the optimal transport map T ⋆ that transports ˜ µ to ˜ ν . Moreover, if τ is sufficiently small, then ˜ µ and ˜ ν are close to µ and ν , respectively. Therefore, we can infer that T v ⋆ will transport µ to a distribution that is similar to ν .

Stable Convergence We discuss convergence properties of UOT objective J and the corresponding OT objective ˜ J for the potential v through the theoretical findings of Gallouët et al. [21].

Theorem 3.4 ([21]) . Under some mild assumptions in Appendix A, the following holds:

for some positive constant C 1 and C 2 . Furthermore, ˜ J ( v ) -˜ J ( v ⋆ ) ≥ 1 2 λ E ˜ µ [ ∥∇ ( v c -v ⋆c ) ∥ 2 2 ] .

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 c -v ⋆c ) ∥ 2 2 , J also bounds the function error ( v c -v ⋆c ) 2 . Second, J provides control over both the original function v and its c -transform v c in L 2 sense. We hypothesize this stable convergence property conveys practical benefits in neural network optimization. In 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 a 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 1), dual (Eq 2), and semi-dual forms (Eq 3). First, several works were proposed based on the primary form [47, 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 [77]. This correspondence motivated a new methodology that parameterized the dual potentials to recover the OT map [2, 64]. Seguy et al. [64] introduced the entropic regularization to obtain the optimal dual potentials u 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 [1].

Recently, Korotin et al. [39] demonstrated that the semi-dual approaches [18, 61] are the ones that best approximate the OT map among existing methods. These approaches [18, 61] properly recovered OT maps and provided high performance in image generation and image translation tasks for largescale datasets. Because our UOTM is also based on the semi-dual form, these methods show some connections to our work. If we let Ψ ∗ 1 and Ψ ∗ 2 in Eq 9 as identity functions, our Algorithm 1 reduces to the training procedure of Fan et al. [18] (See Remark 3.2). For convenience, we denote Fan et al. [18] 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. In this paper, we denote Rout et al. [61] as Optimal Transport Modeling (OTM*) , and regard it as one of the main counterparts for comparison.

2.5

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

x

0.0

0.2

0.5

0.8

1.0

1.2

1.5

1.8

Density

Target density

Generated density

3

2

1

0

1

2

3

Source

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

Target

(a) OTM

1.0

0.5

0.0

0.5

1.0

1.5

x

0.0

0.2

0.5

0.8

1.0

1.2

1.5

1.8

Density

Target density

Generated density

3

2

1

0

1

2

3

4

Source

0.6

0.8

1.0

1.2

1.4

Target

(b) UOTM

6

7

996

99(acl

49679

3

1933

3

190

3

34490

Unbalanced Optimal Transport The primal problem of UOT (Eq 4) relaxes hard marginal constraints of Kantorovich's problem (Eq 1) through soft entropic penalties [12]. Most of the recent UOT approaches [11, 19, 48, 58] estimate UOT potentials on discrete space by using the dual formulation of the problem. For example, Pham et al. [58] extended the Sinkhorn method to UOT, and Lübeck et al. [48] 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 a transport map, another for a discriminator, and the third for a scaling factor. Balaji et al. [8] claimed that implementing a 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 et al. [8] 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-dual formulation of UOT. This formulation allows us to develop a simple but novel adversarial training procedure that does not require the challenging optimization of three neural networks.

5 Experiments

In this section, we evaluate our model on the various datasets to answer the following questions:

Does UOTM offer more robustness to outliers than OT-based model (OTM)? (§5.1)

Does the UOT map accurately match the source distribution to the target distribution? (§5.2)

Does UOTM provide stable and fast convergence? (§5.3)

Does UOTM provide decent performance across various choices of Ψ ∗ 1 and Ψ ∗ 2 ? (§5.4)

Does UOTM show reasonable performance even without a regularization term L reg ? (§5.4)

Does UOTM provide decent performance across various choices of τ in c ( x, y ) ? (§5.4)

Unless otherwise stated, we set Ψ := Ψ 1 = Ψ 2 and D Ψ as a KL divergence in our UOTM model (Eq 9). The source distribution µ is a standard Gaussian distribution N ( 0 , I ) with the same dimension as the target distribution ν . In this case, the entropy function Ψ is given as follows:

In addition, when the generator T θ and potential v ϕ are parameterized by the same neural networks as OTM* [61], we call them a small model. When we use the same network architecture in RGM

[13], we refer to them as large model. Unless otherwise stated, we 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 of 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 N (1 , 0 . 5 2 ) (in-distribution) and N ( -1 , 0 . 5 2 ) (outlier). Similarly, the image dataset is a mixture of CIFAR-10 [43] (in-distribution) and MNIST [14] (outlier) as in Balaji et al. [8].

Figure 2 illustrates the learned transport map T and the probability density T # µ of generated distribution for OTM and UOTM models trained on the toy dataset. The density in Fig 2a demonstrates that OTM attempts to fit both the in-distribution and outlier samples simultaneously. However, this attempt causes undesirable behavior of transporting density outside the target support. In other words, in order to 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 the in-distribution samples (Fig 2b). 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 x and the corresponding target domain samples T ( x ) . Note that the optimal T ⋆ should be a monotone-increasing function. However, OTM failed to learn such a 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 a red or blue-tinted background. In contrast, UOTM primarily generates in-distribution samples. In practice, UTOM generates MNISTlike data at a 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-like samples. Then, we compared this score with the model trained on a clean

2

1

0

1

2

Source

2

1

0

1

2

3

4

Target

(a) OTM

3

2

1

0

1

2

Source

3

2

1

0

1

2

3

Target

(b) Fixed-

µ

2

1

0

1

2

Source

2

1

0

1

2

3

Target

(c) UOTM

dataset. The presence of outliers affected OTM by increasing FID score over 6 (7 . 68 → 13 . 82) while the increase was less than 2 in UOTM (2 . 97 → 4 . 56) . In summary, the experiments demonstrate that UOTM is more robust to outliers.

5.2 UOTMas 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 ν , i.e., T # µ ≈ ν . However, the goal of the generative modeling task is to learn the target distribution ν . In this regard, we assessed whether our UOTM could accurately match the target (data) distribution. Specifically, we measured KL divergence [78] between the generated distribution T # µ and data distribution ν on the toy dataset (See the Appendix B.1 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 [70].

In this experiment, our UOTM model is compared with two other methods: OTM [18] (Constraints in both marginals) and UOTM with fixed µ (Constraints only in the source distribution). We introduced Fixedµ 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 to setting Ψ ∗ 1 ( x ) = x in Eq 9 (Remark 3.1). (See Appendix B.2 for implementation details of UOTM with fixed µ .)

Table 1 shows the data distribution matching results. Interestingly, despite the soft constraint, our UOTMmatches data distribution better than OTM in 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 3.3. Following Thm 3.3, both UOTM models exhibit smaller D KL ( T # µ | ν ) for the smaller τ in the toy dataset. It is worth noting that, unlike UOTM, D KL ( T # µ | ν ) does not change for OTM. This is because, in the standard OT problem, the optimal transport map π ⋆ does not depend on τ in Eq 1. Moreover, in CIFAR-10, while Fixedµ shows a similar FID score to OTM, UOTM significantly outperforms both models. As discussed in 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 a better-organized correspondence between source and target samples than OTM and Fixedµ . In summary, our model is competitive in matching the target distribution for generative modeling tasks.

Image Generation We evaluated our UOTM model on the two generative model benchmarks: CIFAR-10 [43] ( 32 × 32 ) and CelebA-HQ [46] ( 256 × 256 ). For quantitative comparison, we adopted FID [29] and Inception Score (IS) [62]. The qualitative performance of UOTM is depicted in Fig 1. As shown in Table 2, UOTM model demonstrates state-of-the-art results among existing OT-based methods, with an FID of 2.97 and IS 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 (RobustOT), which achieves an FID of 21.57. Furthermore, our model achieves the state-of-the-art FID score of 6.36 on CelebA-HQ (256 × 256). To the best of our knowledge, UOTM is the first OT-based generative model that has shown comparable results with state-of-the-art models in various datasets.

0

200

400

600

800

1000

Epoch

0

20

40

60

80

100

120

FID

OTM

UOTM

5.3 Fast Convergence

The discussion in Thm 3.4 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 (Fig 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 is 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 [69] takes more than 70 hours for training CIFAR-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 We investigate the various choices of Csiszàr divergence in UOT problem, i.e., Ψ ∗ 1 and Ψ ∗ 2 in Eq 9. As discussed in Remark 3.2, the necessary condition of feasible Ψ ∗ is non-decreasing, convex, and differentiable functions. For instance, setting the Csiszàr divergence D Ψ as f -divergence families such as KL divergence, χ 2 divergence, and Jensen-Shannon divergence satisfies the aforementioned conditions. For practical optimization, we chose Ψ ∗ ( x ) that gives finite values for all x , such as KL divergence (Eq 13) and χ 2 divergence (Eq 14):

We assessed the performance of our UOTM models for the cases where D Ψ is KL divergence or χ 2 divergence. Additionally, we tested our model when Ψ ∗ = Softplus (Eq 37), which is a direct parametrization of Ψ ∗ . Table 4 shows that our UOTM model achieves competitive performance across all five combinations of (Ψ ∗ 1 , Ψ ∗ 2 ) .

λ in Regularizer L reg We conduct an ablation study to investigate the effect of the regularization term L reg 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, 56]. Similarly, Figure 6 shows that OTM is also sensitive to the regularization term and fails to converge without it (OTM shows an FID score of 152 without regularization. Hence, we excluded this result in Fig 6 for better visualization). In contrast, our model is 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 τ in the cost function c ( x, y ) = τ ∥ x -y ∥ 2 2 . In our model, the hyperparameter τ controls the relative weight between the cost c ( x, y ) and the marginal matching D Ψ 1 , D Ψ 2 in Eq 4. In Fig 7, our model maintains a decent performance of FID( ≤ 5 ) for τ = { 0 . 5 , 1 , 2 } × 10 -3 . However, the FID score sharply degrades at τ = { 0 . 1 , 5 } × 10 -3 . Thm 3.3 suggests that a smaller τ 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 a detailed discussion).

0.0 0.2

0.5

1.0

2.0

5

10

15

20

25

FID

OTM

UOTM(KL)

UOTM(

2

)

0

1

2

3

4

5

(×10

3

)

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

22.5

FID

UOTM

6 Conclusion

In this paper, we proposed a generative model based on the semi-dual form of 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 as 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 a Generative Model in real-world applications requires careful monitoring to prevent the amplification of existing societal biases present in the data. It is important to carefully control the training data and modeling process of Generative Models to mitigate potential negative societal impacts.

References

Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning , pages 146-155. PMLR, 2017.

Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. Ae-ot: A new generative model based on extended semi-discrete optimal transport. ICLR 2020 , 2019.

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.

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-493, 2021.

Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining deep generative models via discriminator gradient flow. arXiv preprint arXiv:2012.00780 , 2020.

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

Shiri Artstein-Avidan and Vitali Milman. The concept of duality in convex analysis, and the characterization of the legendre transform. Annals of mathematics , pages 661-674, 2009.

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.

Dimitri Bertsekas. Convex optimization theory , volume 1. Athena Scientific, 2009.

Dimitri Bertsekas and Steven E Shreve. Stochastic optimal control: the discrete-time case , volume 5. Athena Scientific, 1996.

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

Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis , 274 (11):3090-3123, 2018.

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.

Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine , 29(6):141-142, 2012.

Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped langevin diffusion. The International Conference on Learning Representations , 2022.

Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12873-12883, 2021.

Lawrence C Evans. Partial differential equations , volume 19. American Mathematical Society, 2022.

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 .

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.

R Flamary, N Courty, D Tuia, and A Rakotomamonjy. Optimal transport for domain adaptation. IEEE Trans. Pattern Anal. Mach. Intell , 1, 2016.

Thomas Gallouët, Roberta Ghezzi, and François-Xavier Vialard. Regularity theory and geometry of unbalanced optimal transport. arXiv preprint arXiv:2112.11056 , 2021.

Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energybased models by diffusion recovery likelihood. Advances in neural information processing systems , 2021.

Aude Genevay. Entropy-regularized optimal transport for machine learning . PhD thesis, Paris Sciences et Lettres (ComUE), 2019.

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.

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-144, 2020.

Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis , 273(11):3327-3405, 2017.

Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of wasserstein gans. Advances in neural information processing systems , 30, 2017.

Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. arXiv preprint arXiv:1904.04546 , 2019.

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.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Yifan Jiang, Shiyu Chang, and Zhangyang Wang. Transgan: Two transformers can make one strong gan. arXiv preprint arXiv:2102.07074 , 1(3), 2021.

Bowen Jing, Gabriele Corso, Renato Berlinghieri, and Tommi Jaakkola. Subspace diffusion generative models. arXiv preprint arXiv:2205.01490 , 2022.

Leonid Vitalevich Kantorovich. On a problem of monge. Uspekhi Mat. Nauk , pages 225-226, 1948.

Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.10196 , 2017.

Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in Neural Information Processing Systems , 33:12104-12114, 2020.

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.

Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and Il-Chul Moon. Score matching model for unbounded data score. arXiv preprint arXiv:2106.05527 , 2021.

Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems , 31, 2018.

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 .

Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations , 2021.

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:14593-14605, 2021.

Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. arXiv preprint arXiv:2201.12220 , 2022.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new hellinger-kantorovich distance between positive measures. Inventiones mathematicae , 211(3):969-1117, 2018.

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.

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.

Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Largescale optimal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.06635 , 2020.

Frederike Lübeck, Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semicouplings. In NeurIPS 2022 AI for Science: Progress and Promises .

Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning , pages 6672-6681. PMLR, 2020.

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:12810-12821, 2021.

Lars Mescheder, Andreas Geiger, and Sebastian Nowozin. Which training methods for gans do actually converge? In International conference on machine learning , pages 3481-3490. PMLR, 2018.

Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations .

Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci. , pages 666-704, 1781.

Quan Hoang Nhan Dam, Trung Le, Tu Dinh Nguyen, Hung Bui, and Dinh Phung. Threeplayer wasserstein gan via amortised duality. In Proc. of the 28th Int. Joint Conf. on Artificial Intelligence (IJCAI) , pages 1-11, 2019.

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.

Henning Petzka, Asja Fischer, and Denis Lukovnikov. On the regularization of wasserstein gans. In International Conference on Learning Representations .

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Center for Research in Economics and Statistics Working Papers , (2017-86), 2017.

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.

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

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.

Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations .

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, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 29, 2016.

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.

Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. In ICLR 2018International Conference on Learning Representations , pages 1-15, 2018.

Thibault Séjourné, Jean Feydy, François-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. arXiv e-prints , pages arXiv-1910, 2019.

Thibault Séjourné, Gabriel Peyré, and François-Xavier Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:2211.08775 , 2022.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. Advances in Neural Information Processing Systems , 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

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.

Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 2818-2826, 2016. doi: 10.1109/CVPR.2016.308.

Dávid Terjék and Diego González-Sánchez. Optimal transport with f -divergence regularization and generalized sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics , pages 5135-5165. PMLR, 2022.

Adrien Vacher and François-Xavier Vialard. Stability of semi-dual unbalanced optimal transport: fast statistical rates and convergent algorithm. 2022.

Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in Neural Information Processing Systems , 33:19667-19679, 2020.

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems , 34:11287-11302, 2021.

Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning , pages 1747-1756. PMLR, 2016.

Cédric Villani. Topics in optimal transportation , volume 58. American Mathematical Soc., 2021.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

Qing Wang, Sanjeev R Kulkarni, and Sergio Verdú. Divergence estimation for multidimensional densities via k -nearest-neighbor distances. IEEE Transactions on Information Theory , 55(5): 2392-2405, 2009.

Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between variational autoencoders and energy-based models. arXiv preprint arXiv:2010.00654 , 2020.

Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:2112.07804 , 2021.

Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. In International Conference on Machine Learning , pages 6882-6892. PMLR, 2019.

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

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11304-11314, 2022.

A Proofs

Notations and Assumptions Let X and Y be compact complete metric spaces which are subsets of R d , and µ , ν be positive Radon measures of the mass 1. Let Ψ 1 and Ψ 2 be a convex, differentiable, nonnegative function defined on R . We assume that Ψ 1 (0) = Ψ 2 (0) = 1 , and Ψ 1 ( x ) = Ψ 2 ( x ) = ∞ for x < 0 . The convex conjugate of Ψ is denoted by Ψ ∗ . We assume that Ψ ∗ 1 and Ψ ∗ 2 are differentiable and non-decreasing on its domain. Moreover, let c be a quadratic cost, c ( x, y ) = τ ∥ x -y ∥ 2 2 , 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 ([12, 21, 72]) . The dual formulation of Unbalanced OT (4) is given by

where u ∈ C ( X ) , v ∈ C ( Y ) with u ( x ) + v ( y ) ≤ c ( x, y ) . Note that strong duality holds.

Also, the semi-dual formulation of Unbalanced OT (4) is given by

where v c ( x ) denotes the c -transform of v , i.e., v c ( x ) = inf y ∈Y ( c ( x, y ) -v ( y )) .

Proof. The primal problem Eq 4 can be rewritten as

We introduce slack variables P and Q such that P = π 0 and Q = π 1 . By Lagrange multipliers where u and v associated to the constraints, the Lagrangian form of Eq 17 is

The dual Lagrangian function is g ( u, v ) = inf π,P,Q L ( π, P, Q, u, v ) . Then, g ( u, v ) reduces into three distinct minimization problems for each variables;

Note that the first term of Eq 21 is

Thus, Eq 21 can be rewritten as follows:

Suppose that c ( x, y ) -u ( x ) -v ( y ) < 0 for some point ( x, y ) . Then, by concentrating all mass of π into the point, the value of the minimization problem of Eq 23 becomes -∞ . Thus, to avoid the trivial solution, we should set c ( x, y ) -u ( x ) -v ( y ) ≥ 0 almost everywhere. Within the constraint, the solution of the optimization problem of Eq 23 is c ( x, y ) = u ( x ) + v ( y ) π -almost everywhere. Finally, we obtain the dual formulation of UOT:

where c ( x, y ) = u ( x ) + v ( y ) π -almost everywhere (since Ψ ∗ 1 and Ψ ∗ 2 is non-decreasing). In other words,

where ı is a convex indicator function. Note that we assume Ψ ∗ 1 and Ψ ∗ 2 are convex, non-decreasing, and differentiable. Since c ≥ 0 , by letting u ≡ -1 and v ≡ -1 , we can easily see that all three terms in Eq 25 are finite values. Thus, we can apply Fenchel-Rockafellar's theorem, which implies that the strong duality holds. Finally, combining the constraint u ( x ) + v ( y ) ≤ c ( x, y ) with tightness condition, i.e. c ( x, y ) = u ( x ) + v ( y ) π -almost everywhere, we obtain the following semi-dual form:

Now, we state precisely and prove Theorem 3.3.

Lemma A.2. For any ( u, v ) ∈ C ( X ) ×C ( Y ) , let

Then, under the assumptions in Appendix A, the functional derivative of I ( u, v ) is

where where ( δu, δv ) ∈ C ( X ) ×C ( Y ) denote the variations of ( u, v ) , respectively.

Proof. For any ( f, g ) ∈ C ( X ) ×C ( Y ) , let

Then, for arbitrarily given potentials ( u, v ) ∈ C ( X ) ×C ( Y ) , and perturbation ( η, ζ ) ∈ C ( X ) ×C ( Y ) which satisfies u ( x ) + η ( x ) + v ( y ) + ζ ( y ) ≤ c ( x, y ) ,

Note that L ( ϵ ) is a continuous function. Thus, the functional derivative of I ( u, v ) , i.e. L (0) , is given by calculus of variation (Chapter 8, Evans [17])

where ( δu, δv ) ∈ C ( X ) ×C ( Y ) denote the variations of ( u, v ) , respectively.

Theorem A.3. Suppose that µ and ν are probability densities defined on X and Y . Given the assumptions in Appendix A, suppose that µ, ν are absolutely continuous with respect to Lebesgue measure and Ψ ∗ is continuously differentiable. Assuming that the optimal potential v ⋆ = inf v ∈C J ( v ) exists, v ⋆ is a solution of the following objective

where ˜ µ ( x ) = Ψ ∗ 1 ′ ( -v ⋆c ( x )) µ ( x ) and ˜ ν ( y ) = Ψ ∗ 2 ′ ( -v ⋆ ( y )) ν ( y ) . Note that the assumptions guarantee the existence of optimal transport map T ⋆ between ˜ µ and ˜ ν . Furthermore, T ⋆ satisfies

µ -a.s..In particular, D Ψ 1 (˜ µ | µ ) + D Ψ 2 (˜ ν | ν ) ≤ τ W 2 2 ( µ, ν ) where W 2 ( µ, ν ) is a Wasserstein-2 distance between µ and ν .

Proof. By Lemma A, derivative of Eq 27 is

Note that such linearization must satisfy the inequality constraint u ( x ) + δu ( x ) + v ( y ) + δv ( y ) ≤ c ( x, y ) to give admissible variations. Then, the linearized dual form can be re-written as follows:

where δu ( x ) + δv ( y ) ≤ c ( x, y ) -u ( x ) -v ( y ) , ˜ µ ( x ) = Ψ ∗ 1 ′ ( -u ( x )) µ ( x ) , and ˜ ν ( y ) = Ψ ∗ 2 ′ ( -v ( y )) ν ( y ) . Thus, whenever ( u ⋆ , v ⋆ ) is optimal, the linearized dual problem is non-positive due to the inequality constraint and tightness (i.e. c ( x, y ) -u ⋆ ( x ) -v ⋆ ( y ) = 0 π -almost everywhere), which implies that Eq 33 achieves its optimum at ( u ⋆ , v ⋆ ) . Note that Eq 33 is a linearized dual problem of the following dual problem:

Hence, the UOT problem can be reduced into standard OT formulation with marginals ˜ µ and ˜ ν . Moreover, the continuity of Ψ ∗′ gives that ˜ µ, ˜ ν are absolutely continuous with respect to Lebesgue measure and have a finite second moment (Note that X , Y are compact. Hence, Ψ ∗′ is bounded on X , Y , and X , Y are bounded). By Theorem 2.12 in Villani [76], there exists a unique measurable OT map T ⋆ which solves Monge problem between ˜ µ and ˜ ν . Furthermore, as shown in Remark 5.13 in Villani et al. [77], T ⋆ ( x ) ∈ arginf y ∈Y [ c ( x, y ) -v ( y )] .

Finally, if we constrain π into Π( µ, ν ) , then the solution of Eq 4 is C ub = τ W 2 2 ( µ, ν ) . Thus, for the optimal π ∈ M + where π 0 = ˜ µ and π 1 = ˜ ν , it is trivial that

which proves the last statement of the theorem.

Theorem A.4 ([21]) . J is convex. Suppose v is a λ -strongly convex, v c and v are uniformly bounded on the support of µ and ν , respectively, and Ψ ∗ 1 , Ψ ∗ 2 are strongly convex on every compact set. Then, it holds

for some positive constant C 1 and C 2 . Furthermore, ˜ J ( v ) -˜ J ( v ⋆ ) ≥ 1 2 λ E ˜ µ [ ∥∇ ( v c -v ⋆c ) ∥ 2 2 ] .

B Implementation Details

B.1 Data

Unless otherwise stated, the source distribution µ is a standard Gaussian distribution N ( 0 , I ) with the same dimension as the target distribution ν .

Toy Data We generated 4000 samples for each Toy dataset: Toy dataset for Outlier Robustness (Sec 5.1) and Toy dataset for Target Distribution Matching (Sec 5.2). In Fig 2 of Sec 5.1, the target dataset consists of 99% of samples from N (1 , 0 . 5 2 ) and 1% of samples from N ( -1 , 0 . 5 2 ) . In Tab 1 and Fig 4 of Sec 5.2, the source data is composed of 50% of samples from N ( -1 , 0 . 5 2 ) and 50% of samples from N (1 , 0 . 5) . The target data is sampled from the mixture of two Gaussians: N ( -1 , 0 . 5) with a probability of 1 / 3 and N (2 , 0 . 5) with a probability of 2 / 3 .

CIFAR-10 We utilized all 50,000 samples. For each image x , we applied random horizontal flip augmentation of probability 0.5 and transformed it by 2 x -1 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 × 32 by duplicating its channel to create a 3-channeled image. Then, the samples were transformed by 2 x -1 . For the image generation task in Sec 5.2, we followed the source distribution embedding of Rout et al. [61] for the small model, i.e., the Gaussian distribution µ with the size of 3 × 8 × 8 is bicubically upsampled to 3 × 32 × 32 .

CelebA-HQ We used all 120,000 samples. For each image x , we resized it to 256 × 256 and applied random horizontal flip augmentation with a probability of 0.5. Then, we linearly transformed each image by 2 x -1 to scale the values within [ -1 , 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 z is set to one. For a generator, we passed z through two fully connected (FC) layers with a hidden dimension of 128, resulting in 128-dimensional embedding. We also embedded data x into the 128-dimensional vector by passing it through threelayered 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 a batch size of 256, a learning rate of 10 -4 , 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 R 1 regularization to λ = 0 . 01 .

CIFAR-10 For the small model setting of UOTM, we employed the architecture of Balaji et al. [8]. Note that this is the same model architecture as in Rout et al. [61]. We set a batch size of 128, 200 epochs, a learning rate of 2 × 10 -4 and 10 -4 for the generator and discriminator, respectively. Adam optimizer with β 1 = 0 . 5 , β 2 = 0 . 9 is employed. Moreover, we use R 1 regularization of λ = 0 . 2 . For the large model setting of UOTM, we followed the implementation of Choi et al. [13] unless stated. We trained for 600 epochs with λ = 0 . 2 for every iteration. Moreover, in OTM (Large) experiment, we trained the model for 1000 epochs. The other hyperparameters are the same as UOTM. Tt For the ablation studies in Sec 5.4, we reported the best FID 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 1 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)).

CelebA-HQ (256 × 256) We set τ = 10 -5 and used R 1 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. [13].

Evaluation Metric We used the implementation of Wang et al. [78] for the KL divergence in Tab 1 of Sec 5.2. We set k = 2 (See Wang et al. [78]). For the evaluation of image datasets, we used 50,000 generated samples to measure IS and FID scores.

Ablation Study on Csiszàr Divergence To ensure self-containedness, we include the definition of Softplus function employed in Table 4.

3

2

1

0

1

2

3

x

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

y

(a) KL

3

2

1

0

1

2

3

x

1

0

1

2

3

4

5

y

(b)

χ

2

3

2

1

0

1

2

3

x

0.0

0.5

1.0

1.5

2.0

2.5

3.0

y

(c) Softplus

C Additional Discussion on UOTM

C.1 Stability of UOTM

In this section, we provide an intuitive explanation for the stability and convergence of our model. Note that our model employs a non-decreasing, differentiable, and convex Ψ ∗ (Sec 3) as depicted in Fig 8. In contrast, OTM corresponds to the specific case of Ψ ∗ ( x ) = x . Now, recall our objective function for the potential v ϕ and the generator T θ :

Then, the gradient descent step of potential v ϕ for a single real data y with a learning rate of γ can be expressed as follows:

Suppose that our potential network v ϕ 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 v ϕ fails to allocate the appropriate potential to the real data point y . Then, because -v ϕ ( y ) is large, Ψ ∗′ ( -v ϕ ( y )) in Eq 40 is large, as shown in Fig 8. In other words, the potential network v ϕ takes a stronger gradient step on the real data y where v ϕ fails. Note that this does not happen for OTM because Ψ ∗′ ( x ) = 1 for all x . In this respect, UOTM enables adaptive updates for each sample, with weaker updates on well-performing 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 of the model output.

C.2 Qualitative Comparison on Varying τ

In the ablation study on τ , UOTM exhibited a similar degradation in FID score( ≥ 20 ) at τ ( × 10 -3 ) = 0 . 1 , 5 . However, the generated samples are significantly different. Figure 9 illustrates the generated samples from trained UOTM for the cases when τ is too small ( τ = 0 . 1 ) and when τ is too large τ = 5 . When τ is too small (Fig 9 (left)), the generated images show repetitive patterns. This lack of diversity suggests mode collapse. We consider that the cost function c provides some regularization effect preventing mode collapse. For the source sample (noise) x , the cost function is defined as c ( x, T ( x )) = τ ∥ x -T ( x ) ∥ 2 2 . Hence, minimizing this cost function encourages the generated images T ( x ) to disperse by aligning each image to the corresponding noise x . Therefore, increasing τ from 0.1 led to better FID results in Fig 7.

On the other hand, when τ is too large Fig 9 (right), the generative samples tend to be noisy. Because the cost function is c ( x, T ( x )) = τ ∥ x -T ( x ) ∥ 2 2 , too large τ induces the generated image T ( x ) to be similar to the noise x . This resulted in the degradation of the FID score for τ = 5 in Fig 7. As discussed in Sec 5.4, this phenomenon can also be understood through Thm 3.3.

D Additional Results

100.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nHowever, 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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe authors do not discuss the limitations, but do discuss the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nYes\n\n**Update:** *reject* -> *borderline accept*", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

comments: 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.

STRENGTHS:
1. The proposed UOTM model has a solid theoretical background and appears original and significant.

2. The experimental setup is appropriate, and the obtained results prove the superiority of the proposed solution concerning the state-of-the-art.

3. 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.

QUESTIONS:
Minor comments:

The captions in Tab. 1, 2, 3, 4 and Fig. 5, 6, 7 could be more informative.

Some figures (such as Fig. 5) are not referenced in the text.

The order of tables/figures needs to be changed (e.g., in the text, Table 1 is referenced after Table 3).

l. 52: "satisfies" --> "satisfy".

l. 67: "couplings" --> "coupling".

Eq. (4): ',' --> '.'.

Eq. (5): "inf" --> "sup" (already corrected in the supplement).

l. 107: ".)" --> "):".

l. 115: "sem-dual" --> "semi-dual".

l. 298: '.' --> ':'.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nStandard 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n### 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\nQUESTIONS:\n### 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\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}