Optimal Transport for Deep Generative Models: State of the Art and Research Challenges

1 Motivation, Optimal Transport and Deep Generative Models

Optimal transport and deep generative models have attracted substantial attention in the machine learning and artificial intelligence community in recent years. Tremendous efforts have been made to leverage to combine the recent development of the optimal transport to formulating and learning deep generative models. The efforts have been resulting in a rich literature of related publications and methodologies in using optimal transport distance with deep generative models. Therefore, a comprehensive and systematic survey reviewing and categorizing existing approaches and methodologies, recognizing outstanding research questions, and discussing open challenges and future directions is imperative yet missing.

To this end, we would like to provide a systematic review of deep generative models using optimal transport. We categorize models and methods on the way optimal transport distance is used to formulate the problems. The goal is to help interested researchers address outstanding problems and more importantly, identify open research opportunities to use the tools from the established optimal transport theory in the deep generative model domain. To the best of our knowledge, this is the first comprehensive survey on deep generative models using optimal transport theory. The notable contributions of our survey can be summarized as follows: 1) summarizing and categorizing deep generative models based on their optimal transport distance formulation; 2) addressing the limitations of existing methods and suggest several open challenges and potential future research directions in the field of using optimal transport theory for deep generative models.

The rest of this paper is organized as follows. Sections 1.1 and 1.2 introduce the background and definitions of two main classes of deep generative models and optimal transport distances. Section 2 reviews optimal transport based deep generative models categorized by the formulation of optimal distance. Section 3 discusses the challenges and potential future research directions.

Notations. Let X be a compact metric space. We denote by P (X) the set of probability measures over X, each of which is referred to as a probability distribution. For any measurable function φ from X to R, and a metric c, Lipc (φ) denotes the Lipschitz constant of φ with respect to c and Lc ≜ {φ : Lipc (φ) ≤ 1}. We use capital letters to denote random variables, e.g. X, while bold lower case letters are used for vectors. Probability distribution of random variable X will be denoted as PX with the probability density as pX (x). When the context is clear the subscript of probability density will be removed, e.g. p (x).

1.1 Deep Generative Models

Deep generative model is a deep neural network based framework for estimating a probability distribution that is "close" to empirical data samples {xi}ni=1 which come from an unknown data distribution PX over the data space X. Model distribution Pθ over X is usually assumed as a latent variable model

where p (z) is some simple well-known distribution, e.g. normal or uniform distributions while the conditional distribution p (x | z) can be either deterministic (i.e. Dirac ) or stochastic. Since the model distribution Pθ is defined via deep neural networks, it may have no analytical form but can be easy to sample from. In the following sub-sections, we review two classes of deep generative models which are considered as dual (Generative Adversarial Network - GAN) and primal formulations (Variational Autoencoder - VAE) of the problem of minimizing the similarity between PX and Pθ.

Generative Adversarial Network (GAN)

Generative Adversarial Network (GAN) [Goodfellow et al., 2014] is a framework for estimating empirical data samples {xi}ni=1 by using two deep neural networks contesting with each other in a game. The difference between GAN and density estimation methods, i.e. maximum likelihood estimation (MLE), is that we do not need to define an explicit form of the estimator. In MLE, we usually define the model Pθ as an explicit parametric family of densities {pθ}θ∈Rd and try to solve the maximum likelihood problem maxθ∈Rd ∑ni=1 log pθ (xi) , which is equivalent to minimizing the Kullback-Leibler divergence between the model pθ and the empirical distribution of real data px ≈ 1n ∑ni=1 δxi , i.e. KL (PX ‖ Pθ) = ¯ px.

In GAN, we do not need to specify the density pθ which might not exist or be difficult to characterize in high dimensional data 1 . Instead, we can define a mechanism that can generate samples that are close to those generated from real data distribution PX (and approximated by empirical distribution ¯ PX). The generation process is defined through two steps: drawing a random variable Z from a fixed distribution PZ, e.g., uniform distribution in (0, 1) and mapping it to the same space of the real data samples, with a function Gθ : Z → X (usually a deep neural network) where θ is the parameter of the function. The induced distribution defined above is denoted as Pθ. We can vary θ to change the induced distribution and make it "close" to ¯ PX. To define closeness, some divergence between distributions is used. In vanilla GAN [Goodfellow et al., 2014], the authors used the Jensen-Shannon divergence between ¯ PX and Pθ

where Pm = Pθ+ ¯ PX 2 . A GAN model can be considered as a latent variable model in Eq 1 in which the conditional distribution P (x | z) is defined using a deterministic map Gθ : Z → X where Z is latent space, i.e., p (x | z) = 1x(G(z)) where 1x is the indicator function of x.

Variational Auto-Encoders

Variational Auto-encoders (VAE) is a sub-class of latent models in Eq. 1 in which the latent distributions are (standard) Gaussian distributions, i.e. p (z) = N (z | 0, I), and the conditional distribution p (x | z) = p (x | fθ (z)) is the likelihood with parameters defined using non-linear functions (typically a deep neural networks) fθ (z) parameterized by θ. For instance if the likelihood is a Gaussian distribution, then p (x | z) = N (µθ (z) , diag (σ2 θ (z))) where µθ (z) and σ2 θ (z) are non-linear functions with parameters θ. Unfortunately, marginal distribution p (x) is not tractable due to the non-linearity of deep neural networks used to define the conditional distributions. Variational inference therefore is used to approximate the posterior distribution PZ|X (and the marginal distribution PX). Typically, the variational conditional distribution qω (z | x) = N (µω (x) , Σω (x)) is used to approximate the conditional distribution where µω and Σω

1 Data may be in high dimensional space, but the support for density function may lie on low dimensional manifolds.

are two non-linear functions aka deep neural networks 2 with parameters ω. We can use the Kullback-Leibler divergence to measure the difference between empirical distribution PX and marginal distribution Pθ, i.e. infθ KL(PX, Pθ) which is equivalent to optimize the Evidence Lower BOund (ELBO) infθ,λ- E q(z|x) [lnp (x | z)] + E p(x) [KL (QZ|X , PZ )].

Applications of Deep Generative Models

Deep generative models have a wide range of applications real-world problems from image processing such as super resolution (generating images with higher pixel resolution) [Ledig et al., 2017; Vasu et al., 2018], image-to-image translation [Isola et al., 2017], photo inpainting [Pathak et al., 2016]; game simulation [Kim et al., 2020]; healthcare analytics [Frid-Adar et al., 2018; Yoon et al., 2019; Costa et al., 2017; Choi et al., 2017]; drug discovery [Jin et al., 2018; De Cao and Kipf, 2018]; finance [Takahashi et al., 2019]; bioinformatics [Marouf et al., 2020; Anand and Huang, 2018]; and e-commerce [Kumar et al., 2018]. Most of the applications are based on image-based generation tasks. The number of works that can generate data beyond images is still limited.

1.2 Optimal Transport Theory

In this section, we summarize optimal transport theory which provides the theoretical framework for defining problems of deep generative models

Primal Formulation

Given a continuous cost function c : X × X → R, and two probability distributions PX, PY ∈ P (X), the optimal transport distance of order p called Wasserstein-p distance is defined as the minimal cost for transporting from the density of px to that of py[Villani, 2008]

where Π is the set of joint distributions with marginal constraints of �X π (x, y) dy = px and �X π (x, y) dy = px .

Dual Formulation

Owing to the marginal constraint of the joint distribution, the optimization process in Eq. (3) is hard to employ. In practice, the celebrated dual from of the optimization in Eq. (3) for order p = 1 called Kantorovich-Rubinstein duality is often used [Villani, 2008]

This dual form leads to a simpler optimization in comparison with the primal form which is used to formulate the Wasserstein GAN models in section 2.1.

Relaxed Formulation

In practice, an entropic regularization version of Wasserstein distance in Eq. 3 which is faster to compute in discrete cases was proposed in [Cuturi, 2013] and called Sinkhorn distance.

2 called inference networks

The definition of Sinkhorn distance can be extended for continuous cases as follows

where KL(π | px ⊗py) ≜ � � ln π(x,y) px(x)py(y) dxdy is the relative entropy between the joint coupling and the marginals. However, the relaxed formulation of Sinkhorn distance in Eq. 5 is not a proper distance since Wϵ (PX, PX) ̸= 0 for ϵ > 0. Sinkhorn divergences [Ramdas et al., 2017] are introduced to resolve the drawback

which becomes to Wasserstein distance with ϵ → 0 and reaches MMD with ϵ → ∞ [Ramdas et al., 2017].

Approximated Formulation

Although Wasserstein distance on higher dimensions does not possess closed-form formulation in general, one can compute explicitly Wasserstein distance for one-dimensional measures [Villani, 2008]. This property motivates to approximate Wasserstein distance by projection probability distributions on high-dimensional to one-dimensional space [Bonneel et al., 2015]

where Sd = {θ ∈ Rd : ‖θ‖ = 1} is the unit d-dimensional sphere, and W (PXθ , PYθ ) is 1-dimensional Wasserstein distance between 1-dimensional projected measures with xθ = 〈x, θ〉 and yθ = 〈y, θ〉 for x, y ∈ X. The number of projections required to approximate sliced Wasserstein distance with a given accuracy are exponential proportion to data dimension [Kolouri et al., 2019a]. To mitigate the projection complexity, max-sliced Wasserstein distance was proposed [Deshpande et al., 2019]

A generalized version of sliced Wasserstein distance in Eq. 7 was introduced to replace the dot product in the projection with generalized Radon transform [Kolouri et al., 2019a] g : X × (Rd \{0}) → R , i.e. xθ = g (x, θ) where θ ∈ Ω ≜ Rd \ {0}

Similarly, maximum generalized sliced Wasserstein distances are defined in a similar fashion in Eq. 8. Recently, [Nguyen et al., 2021a] have introduced a slice-based variant called distributional sliced Wasserstein distance which that seeks for an optimal distribution over projections on the unit sphere.

set probabilitymeasures ρ on Sd satisfying E θ,θ′ ∼ρ C and not empty.

Wasserstein Distance for Structured Objects

Gromov-Wasserstein distance: let PX ∈ P (Rdx ) and PY ∈ P (Rdy ) be two probability measures on Euclidean spaces of different dimension dx and dy. Distance function cx : Rdx ×Rdx →R+ (resp. cy : Rdy ×Rdy →R+ ) defines the similarity between samples in PX (resp. PY ) . The GromovWasserstein distance of order p, GWp (PX, PY ), is defined as [Mémoli, 2011]

where Π is the set of joint distributions with marginal constraints defined similarly in Eq. (3).

2 Generative Models Using Optimal Transport

In this section we review and categorize deep generative models based on how optimal transport distance is used to formulate the objective functions. Figure 1 denotes the categorization of models based on the formulations of optimal transport described in Sec 1.2.

2.1 Dual Formulation

Wasserstein Generative Adversarial Network (WGAN)

One of the first works that had tried to use Wasserstein distance for learning deep generative models is the work by [Arjovsky et al., 2017]. Wasserstein distance was proposed to use for learning since it has the weak*-topology property that allows defining the distance between two distributions with non-overlapped supports and is expected to alleviate the mode-collapsing problem of GANs. In Wasserstein GAN, the authors suggested using the dual form Wasserstein distance of order 1 in Eq (4) instead of Jensen-Shannon divergence. The function φω (·) in Eq (4) is parameterized using a (deep) neural networks of parameters ω. However, one of the important constraints for φ is the Lipschitz condition, i.e. |φ(x) -φ(y)| ≤ c (x, y) 3 for all x, y ∈ X , which is not always satisfied by an arbitrary deep network. To endow the constraint, Arjovsky et al. suggest a heuristic to clip weights of network to some certain ranges, e.g. the clipped box of ω ∈ [-0.01, 0.01]. Despite the heuristically constraining the Lipschitz condition, Wasserstein GANs showed stability improvement in training processes. However, the weight clipping trick may lead to poor samples or fail to converge as observed in [Gulrajani et al., 2017]. There are the following works [Gulrajani et al., 2017; Petzka et al., 2018; Wei et al., 2018; Miyato et al., 2018] that solve the 1-Lipschitz constraint with principled approaches.

Lipschitz Constraint Variants

[Gulrajani et al., 2017, Corollary 1] have proved that the optimal critic function, φ∗, has gradients norm 1 almost everywhere. Hence, they proposed to train Wasserstein GAN with gradient penalty instead of weight clipping which results with a more stable and improved performance called WGAN gradient penalty (WGAN-GP) RGP = E P ¯X [(‖∇ωφ(x)‖2 -1) 2 ] where ¯X = tX + (t - 1)Y for t ∼ Uni (0, 1) and X ∼ PX, Y ∼ Pθ which respectively are real and generated samples. This is motivated by the fact that the optimal critic function φ∗ contains straight lines with gradient norm 1 connecting coupled points from PX, and Pθ(cf. [Gulrajani et al., 2017, Proposition 1]). [Petzka et al., 2018] then proposed an alternative to gradient penalty called Lipschitz penalty which replace a l2 norm of gradient in GP to l1, RLP = E P ¯X [(max {0, ‖∇wf (x)‖-1}) 2 ]. Gradient and Lipschitz penalty regularizers only enforce Lipschitz constraint on a local data domain (not ∀x, y ∈ X), therefore [Wei et al., 2018] further improved the WGAN-GP by incorporating a consistency term to the objective function of WGAN-GP to enforce the (global) Lipschitz constraints of Wasserstein GAN. They defined a practical form of term CT as follows RCT = E x∼PX [max (0, c (x1, x2) + 0.1 · d (x1-, x2-) - K)], where x1, x2 and x1-, x2- are perturbed embedded real data via critic φ, i.e. x1,2 = φdrop (x) where φdrop is a hidden layers dropout of φ 4 . Two vectors x1-and x2- are outputs of the last layer of φdrop corresponding to x1,x2. Spectral normalization technique proposed by [Miyato et al., 2018] allows to set the upper bound of the Lipschitz constant of φ. In this setting, the weight matrix in each layer is normalized by the spectral norm of that matrix which is equivalent to the largest

singular value. A recent work by [Avraham et al., 2019] has introduced an additional low dimensional representation (latent space) in parallel with original data and used Wasserstein distances on both latent space and original space 5 .

3 In [Arjovsky et al., 2017], authors concretely chose the cost function c (x, y) = ‖x - y‖ which allows weight clipping inducing Lipschitz condition.

4 Note that since dropout is stochastic,φdrop produces two different networks for x1and x2.

Another line of work that aims to remove the Lipschitz condition in the dual form of Wasserstein distance is to reformulate it. The work of [Liu et al., 2018] redefined the dual form in Eq. (3) as two-step of optimization: solving linear programming to approximate discretized critic function T and then regress the critic φ to fit T. Recent work of [Dam et al., 2019] introduced a new function called mover to get rid of the Lipschitz constraint but a new function to optimize. Their new objective function becomes a min-max-min loss.

Customized Transport Loss Wasserstein GAN

As Wasserstein distance is defined based on cost function c (·, ·) which is reflected in the Lipschitz condition in the dual formulation. WGAN-GP used the l2 norm of the gradient in the loss function which corresponds to the l2 norm used for cost function c. Similarly, WGAN-LP implies the l1 norm used. The underlying cost function reflects the geometry of the space generated data lying on. The work of [Adler and Lunz, 2018] has generalized the cost function of norms on Euclidean space to norms on Banach spaces such as Sobolev norms and Lp norms. These extensions allow practitioners more ranges of the cost functions to choose from to emphasize features they wish for generated data. The quadratic cost function was proposed to use in [Liu et al., 2019] in which they used two-step optimization in [Liu et al., 2018] for learning. A recent work by [Korotin et al., 2021] also used the quadratic cost in their Wasserstein distance and used the input convex neural networks (ICNN) [Amos et al., 2017] for approximating the distance. As Wasserstein distance is known as an effective distance in differentiating images [Rubner et al., 2000], it was used as a transport loss for Wasserstein GAN models [Dukler et al., 2019]. Though three-player WGAN [Dam et al., 2019] was developed to overcome the Lipschitz constraint, their model is not limited to any specific transport loss as long as it is lower semicontinuous and bounded [Villani, 2008].

2.2 Primal Formulation

Wasserstein Auto-Encoder (WAE)

[Tolstikhin et al., 2018] have introduced to use Wasserstein distance to formulate auto-encoder instead of KL divergence used in VAE in Sec 1.1. In the paper, they consider a deterministic conditional distribution pθ (x | z) via a map Gθ : Z → X. It turns out that the Wasserstein distance of order 1 between the empirical data distribution PX and the generated data distribution Pθ under cost function c can be represented as W1 (PX, Pθ) = infQ: QZ=PZ E QZ|X, PX [c (x, Gθ (z))]

where q (z) = � q (z | x) p (x) dx is the marginal distribution of Z. In practical, they relax the constraint of QZ = PZ by using Lagrangian multiplier and define the WAE objective

5 In fact, they used Monge formulation [Villani, 2003, p.4] on latent space, a strengthened version of Wasserstein distance in which the constraint of one-to-one mapping between PX and Pθ is endowed.)

where qω (z | x) is an encoder parameterized by ω, and DZ is an arbitrary divergence between QZ and PZ. Authors suggested using variants of generative models such as GAN [Goodfellow et al., 2014] or the maximum mean discrepancy (MMD) [Gretton et al., 2012; Li et al., 2015] as divergence DZ. Recent works have been exploring to use difference divergences for Dz, for instance, [Zhang et al., 2019] approximate QZ as a Gaussian and use the prior PZ as a Gaussian distribution which leads a analytical form of Wasserstein distance between two Gaussians for Dz. [Kolouri et al., 2019b] used sliced Wasserstein distance in Eq. (7) to approximate the difference between QZ and PZ. Sinkhorn divergence is also used to characterize distinction in [Patrini et al., 2020].

2.3 Relaxed and Approximated Formulation

Instead of using the original version of Wasserstein (in both dual and primal forms), there is a line of work that dedicated to using variants of sliced Wasserstein distance in Sec 1.2 for learning deep generative models These works used the primal form of sliced Wasserstein distance as signals to update the generators. However, since data lies on high dimensional spaces, the number of projections for a good approximation is huge. [Deshpande et al., 2018; 2019; Kolouri et al., 2019a] have suggested to use a learnable function, which is learned to most discriminate two distributions, to map data before projecting to compute the (max/generalized) sliced Wasserstein distance. In [Wu et al., 2019], authors used sliced Wasserstein distance in both primal and dual form to learn Wasserstein GAN and Wasserstein Autoencoder.

As Sinkhorn divergence in Eq. 6 is computed using the Sinkhorn algorithm which is amenable to automatic differentiation framework, it was used to learn deep generative models [Frogner et al., 2015]. However, when dealing with highdimensional data, choosing cost function c is a critical task owing to the curse of dimensionality. In [Genevay et al., 2018], authors proposed to learn cost function c via a deep neural network φ to map data from original space onto Euclidean space Rd which has a similar role as critics in WGAN, i.e. cω (x, y) = (‖φω (x) -φω (y)‖). Typically, embedding function from data space X to lower dimension is learned with a maximization problem in a similar fashion with MMD models [Li et al., 2017] which leads to the optimization problem minθ maxω Sϵ (PX, Pθ). One of the benefits of learning deep generative models with Sinkhorn divergences is that its natural gradient can be approximated accurately with low complexity [Shen et al., 2020]. This will help the training process converge faster.

3 Challenges and Discussions

Since optimal transport in machine learning is a fastdeveloping and promising area, there are potential open questions and challenges in applying optimal transport for deep generative models particularly. In this section we would like to highlight a number of open challenges for future research of using optimal transport for learning deep generative models research in both practical and theoretical aspects.

Cost function. Wasserstein distance has the benefit of taking into account the geometry of data space of distributions which is reflected through the cost function c. Therefore choosing the appropriate cost function for each data type is one of the key factors in designing successful learning models. As described in Sec 2, the popular cost functions used in major classes of models are lp (p = 1, 2) norms which imply that data lie on Euclidean spaces. Some others used norms on Banach spaces or Wasserstein distance. However, which cost function is suitable for use with a given data type/dataset is still an open question. Some works do not specify the cost function but design deep neural network to learn to map from data spaces to Euclidean spaces and use norms on them. These methodologies are merely heuristic and do not possess any theoretical guarantee.

Approximation formulation comparison. There are different forms to approximate Wasserstein distance however it is not clear which approximation is tighter than the others. For instance, all formulations are approximated using the mini-batch approximation of intractable expectation computation. The approximation in the dual formulation of WGAN also comes from deep neural network approximation of the critic, and soft constraint of 1-Lipschitz condition for the critics. From the experimental perspective, it is imperative to have benchmarks on real-world datasets for comparison between variants of Wasserstein distance. From the theoretical perspective, there are bound between Wasserstein distances and sliced variants [Bonneel et al., 2015] but how good a mini-batch approximation of these formulations in terms of asymptotic and gradient properties is raised as potential challenges.

Better Wasserstein approximation and extensions. Beyond four approximated formulations which have presented in Sec 1.2, the community is actively proposing some other approximation methods which are probably amenable for learning with deep generative models. For instance, recent work on minibatch Wasserstein [Fatras et al., 2020; Nguyen et al., 2021b] with the theoretical bound on the approximation of intractable Wasserstein distance has been proposed. Applying these new approximations for learning deep generative models is a fruitful direction. There are several extensions of optimal transport to more general settings: multiple measures called multi-marginal OT, unnormalized ones called unbalanced OT [Chizat et al., 2018]. Investigations on these extensions are still an open direction with few recently published works [Balaji et al., 2020; Cao et al., 2019].

Beyond Euclidean data generation. Many data in realworld applications have an underlying structure beyond Euclidean data such as social networks, molecular graph (graph data), financial data, sentences (sequence data). Most of the existing works of deep generative models focused and demonstrated on synthetic and image data. Data generation using deep neural networks on graph or sequence data is an open problem with few works published [Yoon et al., 2019; Golany et al., 2020; De Cao and Kipf, 2018; Jin et al., 2018; Liao et al., 2019]. Optimal transport based deep generative models can provide a flexible framework for learning these complex structured data. If one can define a suitable cost function c for a corresponding data type, models in Sec 2 can be used to generate the desired data with some considerations. For instance, we can use (soft) dynamic time wrapping (DWT) [Staib et al., 2017; Blondel et al., 2020] for realvalued time series, and Gromov-Wasserstein distance in Eq. (9) for graphs. When the dual form models such as WGAN are used, the challenge is to enforce the Lipschitz constraint under the new cost function 6 . Gradient penalty or Lipschitz penalty in WGAN is not valid as they have been defined on lp norms. If we use primal formmodels such as WAE, the new cost function can be applied directly. However, the challenge may come from the complexity of computing these metrics which can be a burden when training with a large mini-batch or dataset.

4 Conclusion

Our work provides a systematic review of using optimal transport for deep generative models. We present a taxonomy of deep generative models based on variants of optimal transport distance including primal, dual, relaxed, and approximated formulations, followed by a detailed introduction, comparison, and discussion connections and differences of them. We also discuss the current challenges of the current models in terms of cost function awareness and comparison of different approximations of optimal transport distance. We finally describe visions and directions of data generation with non-Euclidean data.