Potential Flow Generator with $L_2$ Optimal Transport Regularity for Generative Models

Introduction

Many of the generative models, for example, generative adversarial networks (GANs) (Goodfellow et al., 2014; Arjovsky et al., 2017; Salimans et al., 2018) and flow-based models including normalizing flows (Rezende and Mohamed, 2015; Kingma and Dhariwal, 2018; Chen et al., 2018), aim to find a generator that could map the input distribution to the target distribution.

In many cases, especially when the input distribution is purely noises, the specific maps between input and output are of little importance as long as the generated distributions match the target ones. However, in other cases like image-to-image translations, where both input and target distributions are distributions of images, the generators are required to have additional regularity such that the input individuals are mapped to the “corresponding” outputs in some sense. If paired input-output samples are provided, $L_{p}$ penalty could be hybridized into generators loss functions to encourage the output individuals fit the ground truth (Isola et al., 2017). For the cases without paired data, a popular approach is to introduce another generator and encourage the two generators to be the inverse maps of each other, as in CycleGANs (Zhu et al., 2017), DualGANs (Yi et al., 2017) and DiscoGANs (Kim et al., 2017), etc. However, such a pair of generators is not unique and lacks clear mathematical interpretation about its effectiveness.

In this paper we introduce a special generator, i.e., the potential flow generator, with $L_{2}$ optimal transport regularity. It is not only a map from the input distribution to the target distribution, but also the optimal transport map with squared $L_{2}$ distance as transport cost. In Fig. 1 we provide a schematic comparison between generators with and without optimal transport regularity. Note that the generator with optimal transport regularity has the characteristic of “proximity”in that the inputs tend to be mapped to nearby outputs. As we will show later, this “proximity” characteristic of $L_{2}$ optimal transport regularity could be utilized in image translation tasks. Compared with other approaches like CycleGANs, the $L_{2}$ optimal transport regularity has a much clearer mathematical interpretation.

There have been other approaches to learn the optimal transport map in generative models. For example, Seguy et al. (2017) proposed to first learn the regularized optimal transport plan and then the optimal transport map, based on the dual form of regularized optimal transport problem. Also, Yang and Uhler (2018) proposed to learn the unbalanced optimal transport plan in an adversarial way derived from a convex conjugate representation of divergences. In the W2GAN model proposed by Leygonie et al. (2019), the discriminator’s objective is the 2-Wasserstein metric so that the generator is the $L_{2}$ optimal transport map. All the above approaches need to introduce, and are limited to, specific loss functions to train the generators. Our proposed potential flow generator takes a different approach in that with up to a slight augmentation of the original generator loss functions, our potential flow generator could be integrated into to a wide range of generative models with various generator loss functions, including different versions of GANs and flow-based models. This simple modification makes our method easy to adopt on various tasks considering the existing rich literature and the future developments of generative models.

Our main contributions are the following:

We propose a potential flow generator with $L_{2}$ optimal transport regularity for a wide range of generative models. Two different versions are proposed: the discrete version and the continuous version.

We compare the two versions and verify the correctness and robustness of the continuous potential flow generator by comparing it with analytical $L_{2}$ optimal transport maps. We also show the effectiveness of our methods in different generative models for more complicated distributions.

We show the effectiveness of our generator with $L_{2}$ optimal transport regularity in image translation tasks with unpaired training data.

In Section 2 we give a formal definition of optimal transport map and the motivation to apply $L_{2}$ optimal transport regularity to generators. In Section 3 we give a detailed formulation of potential flow generator and the augmentation to the original loss functions. Results are then provided in Section 4. We include a discussion and conclusion in Section 5.

Generative Models and Optimal Transport Map

First, we introduce the concept of push forward, which will be used extensively in the paper.

Given two Polish space $X$ and $Y$, $B(X)$ and $B(Y)$ the Borel $\sigma$-algebra on $X$ and $Y$, and $\mathcal{P}(X),\mathcal{P}(Y)$ the set of probability measures on $B(X)$ and $B(Y)$. Let $f:X\rightarrow Y$ be a Borel map, and $\mu\in\mathcal{P}(X)$. We define $f_{\#}\mu\in\mathcal{P}(Y)$, the push forward of $\mu$ through $f$, by

With the concept of push forward, we can formulate the goal of GANs and and flow-based models as to train the generator $G$ such that $G_{\#}\mu$ is equal to or at least close to $\nu$ in some sense, where $\mu$ and $\nu$ are the input and target distribution, respectively. Usually, the loss functions for training the generators are metrics of closeness that vary for different models. For example, in continuous normalizing flows (Chen et al., 2018), such metric of closeness is $\text{KL}(G_{\#}\mu||\nu)$ or $\text{KL}(\nu||G_{\#}\mu)$. In Wasserstein GANs (WGANs) (Arjovsky et al., 2017), the metric of closeness is the Wasserstein-1 distance between $G_{\#}\mu$ and $\nu$, which is estimated in a variational form with discriminator neural networks. As a result, the generator and discriminator neural networks are trained in an adversarial way:

where $D$ is the discriminator neural network and the Lipschitz constraint could be imposed via the gradient penalty (Gulrajani et al., 2017), spectral normalization (Miyato et al., 2018), etc.

Now we can introduce the concept of Monge’s optimal transport problem and the optimal transport map as following:

and the minimizer $f^{*}\in T$ is called the optimal transport map.

In general, Monge’s problem could be ill-posed in that $T$ could be empty set or there is no minimizer in $T$. Moreover, the optimal transport map could be non-unique, for example, if transport cost $c(x,y)=\|x-y\|_{1}$. However, for the special case we consider, there exists a unique solution to Monge’s problem.

Informally speaking, with $L_{2}$ transport cost the optimal transport map has the characteristic of “proximity”, i.e. the inputs tend to be mapped to nearby outputs. In image translation tasks, such “proximity” characteristic would be helpful if we could properly embed the images into Euclidean space such that our preferred input-output pairs are close to each other.

Apart from image translations, the optimal transport problem with $L_{2}$ transport cost is important in many other aspects. For example, it is closely related to gradient flow (Ambrosio et al., 2008), Fokker-Planck equations (Santambrogio, 2017), porous medium flow (Otto, 1997), etc.

Potential Flow Generator

One of the most important results for the case we are studying is that the optimal map is of the form $f(x)=\nabla\Phi(x)$, where the so-called Brenier potential $\Phi(x)$ is convex and lower semi-continuous. We refer the readers to Gangbo and McCann (1996) and McCann and Guillen (2011) for the details. Note that the convexity constraint for $\Phi$ is essential for optimality. Here is an example of non-convex $\Phi$ leading to sub-optimal transport map: let $\mu$ and $\nu$ both be uniform distribution on $$, then the identical map which has zero cost is the optimal map, while we can construct a non-convex$\Phi(x)=x-\frac{1}{2}x^{2}$which leads to the sub-optimal transport map$f(x)=1-x$.

In the context of generative models, it looks plausible to represent the Brenier potential $\Phi(x)$ with a neural network and induce the generator from the gradient. However, the constraint of convexity is not trivial for neural networks. Therefore, we turned to another characterization of optimal transport map, i.e., the potential flow formulation which was firstly proposed by Benamou and Brenier (2000): suppose both $\mu$ and $\nu$ admit probability density $\rho_{\mu}$ and $\rho_{\nu}$, and consider all smooth enough density field $\rho(t,x)$ and velocity field $v(t,x)$, where $t\in[0,T]$, subject to the continuity equation as well as initial and final conditions

The above equation actually says that such velocity field will induce a transport map: we can construct an ordinary differential equation (ODE)

and the map between initial point to final point gives the transport map from $\mu$ to $\nu$.

As is proposed by Benamou and Brenier (2000), for the transport cost function $c(x,y)=||x-y||^{2}$, the minimal transport cost is equal to the infimum of

among all $(\rho,v)$ satisfying equation (4). The optimality condition is given by

In other words, the optimal velocity field is actually induced from a flow with time-dependent potential $\phi(t,x)$.

Note that different from the Brenier potential formulation, there is no explicit requirement for convexity in the potential flow formulation. This helps us to apply this formulation to design our generators with $L_{2}$ optimal transport regularity.

2 Potential Flow Generator

The potential $\phi(t,x)$ is the key function to estimate, since the velocity field could be obtained by taking the gradient of the potential and consequently the transport map could be obtained from the ODE Eqn. 5. There are two strategies to use neural networks to represent $\phi$.

One can take advantage of the fact that the time-dependent potential field $\phi$ is actually uniquely determined by its initial condition from Eqn. 7, and use a neural network to represent the initial condition of $\phi$, i.e. $\phi(0,x)$, while approximating $\phi(t,x)$ via time discretization schemes.

Alternatively, one can use a neural network to represent $\phi(t,x)$ directly and later apply the PDE regularity for $\phi(t,x)$ in Eqn. 7.

We name the generators defined in the above two approaches as discrete potential flow generator and continuous potential flow generator, respectively, and give a detailed formulation as follows.

Finally, we can build the transport map from Eqs. 5:

The discrete potential flow generator has built-in optimal transport regularity since the optimal condition (Eqn. 7) is encoded in the time discretization (Eqn. 8). However, such discretization also introduces nested gradients, which dramatically increases computational cost when the number of total steps $n$ is increased. In our examples, we set $T=1$ and $n=4$.

2.2 Continuous Potential Flow Generator

With the velocity field we could estimate the transport map by solving the ODE (Eqn. 5) using any numerical ODE solver. As an example, we can use the first-order forward Eular scheme, i.e.

In the continuous potential flow generator, increasing the number of total steps would not introduce high order differentiations, therefore we could have very small time steps for a better precision of ODE solver. In practice, we set $T=1$, $n=10$ in image translation tasks and $n=100$ elsewhere for continuous potential flow generators. Different from the discrete potential flow generator, the optimal condition (Eqn. 7) is not encoded in the continuous potential flow generator, therefore we need to penalize Eqn. 7 in the loss function, as we will discuss in the next subsection.

One may come up with another strategy of imposing the $L_{2}$ optimal transport regularity: to use a vanilla generator, which is a neural network directly mapping from inputs to outputs, and penalize the $L_{2}$ transport cost, i.e., the loss function is

where $L_{original}$ is the original loss function for the generator, and $\alpha$ is the weight for the transport penalty. We emphasize that such strategy is much inferior to penalizing the PDE (Eqn. 7) in the continuous potential flow generator. When training the vanilla generator with $L_{2}$ transport penalty, no matter how we weight the $L_{2}$ transport cost penalty, in principle we always have to make a trade off between “matching the generated distribution with the target one” and “reducing the transport cost” since there is always a conflict between them, and consequently $G_{\#}\mu$ will be biased towards $\mu$. On the other hand, there is no conflict between matching the distributions and penalizing Eqn. 7 in the continuous potential flow generator. As a consequence, the continuous potential flow generator is robust with respect to different weights for the PDE penalty. We will show this in Section 4.

3 Training the Potential Flow Generator

While the optimal condition (Eqn. 7) has been considered in the above two generators, the constraints of initial and final conditions are so far neglected. Actually, the constraint of initial and final conditions gives the principle to train the neural network: we need to tune the parameter $\theta$ in the neural network so that $G_{\#}\mu$ matches $\nu$. This could be done by GANs or flow-based models.

In the cases where the available data are samples from $\mu$ and $\nu$, we could apply GAN models to train the potential flow generator. In particular, we replace the generator in GANs with the potential flow generator $G$, feed samples from $\mu$ into $G$, and view the outputs as “fake data” of $\nu$. Both the discrete and continuous potential flow generator could be applied to such cases.

For the discrete potential flow generator, since the optimal transport regularity is already built in, the loss for training $G$ is simply the GAN loss for the generator, i.e.

where $L_{GAN}$ actually depends on the specific version of GANs. For example, if we use WGANs with gradient pernalty, then

where $D$ is the discriminator neural network.

For the continuous potential flow generator, as mentioned above, we also need to penalize the PDE (Eqn. 7) for the optimal transport regularity. Inspired by the applications of neural networks in solving forward and backward problems of PDEs, as in PINNs (Raissi et al., 2017a, b), we penalize the squared residual of the PDE on the so-called “residual" points. In particular, the loss for continuous potential flow generator would be

where $\{(t_{i},x_{i})\}_{i=1}^{N}$ are the residual points for the estimating the residual of PDE (Eqn. 7), and $\lambda$ is the weight for the PDE penalty. While there could be other strategies to select the residual points, here we set them as the points on “trajectories” of input samples, i.e.

3.2 Loss in flow-based models

where as in the GAN model, $\{(t_{i},x_{i})\}_{i=1}^{N}$ are the residual points for estimating the residual of PDE (Eqn. 7), and $\lambda$ is the weight for the PDE penalty.

To estimate $\log p_{G_{\#}\mu}(y)$, we have the ODE that connects the probability density at inputs and outputs of the generator:

with initial condition $U(0)$ as $y=G(x)$ and $U(T)$ as the corresponding $x=G^{-1}(y)$. Note that with the maximum likelihood training, the density of $\mu$ could be unnormalized, since multiplications with $p_{\mu}$ would merely lead to a constant difference in the loss function.

In practice, we also need to discretize Eqns. 19 and 20 properly when calculating $\log p_{G_{\#}\mu}(y)$. For example, we use the first-order Euler scheme with time step $\Delta t=0.01$ for total time interval $T=1.0$. Note that the discrete potential flow cannot be trivially applied in flow-based models since we found that the time step size is too large to calculate the density accurately.

Results

In this section we show the numerical results of potential flow generators in different generative models. In specific, in Section 4.1 we show the results of potential flow generators applied on 2D problems and image translation tasks with GAN models, while in Section 4.2 we show the results in a flow-based model. The numerical experiments are based on Python package TensorFlow(Abadi et al., 2015).

In this subsection, we apply the potential flow generator to several two dimensional problems where samples from $\mu$ and $\nu$ are available.

We first study the following two problems where we know analytical solutions for the optimal transport maps.

In the first problem we assume that both $\mu$ and $\nu$ are Gaussian distributions. Suppose $\mu=\mathcal{N}(m_{1},\Sigma_{1})$, $\nu=\mathcal{N}(m_{2},\Sigma_{2})$, then (Gelbrich, 1990) the minimum transport cost between $\mu$ and $\nu$ will be

which is also known as the squared Wasserstein-2 distance between two Gaussian distributions. In particular, we consider the case

In this case we can check that $f((x,y))=(2x,0.5y)$ is the optimal transport map.

In the second problem we assume that $\mu$ and $\nu$ are concentrated on concentric rings. In polar coordinates, suppose that $\mu$ is the distribution with radius $r$ uniformly distributed on $[0.5,1]$, while angular $\theta$ uniformly distributed on $[0,2\pi]$, and suppose $\nu$ is the distribution with radius $r$ uniformly distributed on $[2,2.5]$, while angular $\theta$ uniformly distributed on $[0,2\pi]$. In this case we can check that $f((r,\theta))=(r+1.5,\theta)$ in polar coordinates is the optimal transport map.

In both cases, we have 40000 samples from each of $\mu$ and $\nu$ as our training data. Samples from $\mu$ and $\nu$ as well as the optimal transport map in both cases are illustrated in Fig. 2(a) and Fig. 3(a).

As for the GAN loss for generators we use the sliced Wasserstein distanceStrictly speaking, there is no “adversarial” training when we use sliced Wasserstein loss since the distance is estimated explicitly rather than represented by an other neural network. However, the idea of computing the distance between fake data and real data coincide with other GANs, especially WGANs. Therefore, in this paper we view sliced Wasserstein distance as a special version of GAN loss., due to its relatively low computational cost, robustness, and clear mathematical interpretation in low dimensional problems Deshpande et al. (2018). In particular, we estimate the sliced Wasserstein distances between samples of $G_{\#}\mu$ and $\nu$ using 1000 random projection directions. The batch size is set as 1000.

For all the cases we run the code three times using different random seed. In Fig. 2 we illustrate the transport maps in one of the three runs for each case. A more systematic and quantitative comparison is provided in Table 1 and Table 2, where the best results are marked as bold.

As we already mentioned, the $L_{2}$ transport penalty would make $G_{\#}\mu$ biased towards $\mu$ to reduce the transport cost from $\mu$ to $G_{\#}\mu$. This is clearly shown in both problems where the penalty weight is $\alpha=0.1$. Actually, we observed more significant biases with larger penalty weights. For the cases with smaller penalty weights $\alpha=0.01,0.001$, in some of the runs, while $G_{\#}\mu$ are close to $\nu$, the estimated transport maps are far from the optimal ones, which shows that the $L_{2}$ transport penalty cannot provide sufficient regularity if the penalty weight is too small. These numerical results are consistent with our earlier discussion about the intrinsic limitation of the $L_{2}$ transport penalty.

On the other hand, the potential flow generators give better matching between $G_{\#}\mu$ and $\nu$, as well as smaller errors between the estimated transport maps and the analytical optimal transport maps. Notably, in both problems the continuous potential flow generators give good results with a wide range of PDE penalty weights ranging from $0.1$ to $10$, which shows the superiority of PDE penalty in the continuous potential flow generators compared with the transport penalty in vanilla generators. We also report that while in the first problem the discrete potential flow generator achieves a comparable result with the continuous potential flow generators, in the second problem we encountered “NAN" problems during training the discrete potential flow generator in some of the runs. This indicates that the discrete potential flow generator is not as robust as the continuous one, which could be attributed to the high order differentiations in the discrete potential flow generators.

Apart from the previous two problems, we also applied the continuous potential flow generators to another three pairs of $(\mu,\nu)$ distributions, which are illustrated in Fig. 4. Here we use WGAN-GP instead of SWG since we found that it provides better GAN loss for complicated distributions. The PDE penalty weight are set as $1.0$.

From Fig. 4 we can see the match between $G_{\#}\mu$ and $\nu$ in each of the problems, which shows the effectiveness of the potential flow generator. We could also see that samples of $\mu$ tend to be mapped to nearby positions, which demonstrates the characteristics of “proximity” in the potential flow generator maps due to the $L_{2}$ optimal transport regularity.

1.2 Image Translations: the MNIST and CelebA dataset

In this section, we aim to show the capability of potential flow generator in dealing with high dimensional problems, and also to show the potential of optimal transport map in tasks of image translations using unpaired training data. In particular, we apply the continuous potential flow generator, due to its robustness compared with the discrete potential flow generator, on the following two tasks:

Translation between the MNIST images (LeCun et al., 2010). We divide the MNIST training dataset into two clusters: (a) images of digits 0 to 4, and (b) images of digits 5 to 9. We view the two clusters of images as samples of $\mu$ and $\nu$, respectively, i.e., we want to find the optimal transport map from images of digits 0 to 4 to images of digits 5 to 9.

Translation between the CelebA images (Liu et al., 2015). We randomly pick 60000 images from the CelebA training dataset and divide them into two clusters: (a) no-smiling face images, i.e., images with attribute “smiling” labeled as false, and (b) smiling face images, i.e., images with attribute “smiling” labeled as true. We crop the images so that only faces remain on the images, and view the two clusters as samples of $\mu$ and $\nu$, respectively, i.e., we want to find to the optimal transport map from no-smiling face images to smiling face images.

Note that in both tasks there are no paired images in training data, and in the first task we do not distinguish digits among $\mu$ samples and $\nu$ samples.

For the first problem on the MNIST dataset, we embed the images into 100-dimensional Euclidean space, i.e. the total components in PCA is 100. We randomly pick images of digits from 0 to 4 from the test set and show the corresponding inputs and outputs in Fig. 5. Although the digits are mixed in the training dataset, we can see that after training, the potential flow generator tends to translate images of digit 0 to digit 6, digit 1 to digit 7, digit 3 to digit 5 or 8, and digit 4 to digit 9. This is consistent with our previous discussion about the characteristics of “proximity” in that the input digits and output digits “look similar”, and the corresponding embedding vectors should be close in the $L_{2}$ distance.

For the second problem on the CelebA dataset, we embed the images into a 700-dimensional Euclidean space, i.e. the total components in PCA is 700. We randomly pick no-smiling images from the training and testing datasets, and show the corresponding inputs and outputs in Fig. 6. We can see that for most of the images, the potential flow generator could translate the no-smiling faces to smiling faces while keeping the other attributes. We can also see the failure for other images, especially for side faces, which could be partially attributed to the fact that these images are outliers in PCA. The reconstructed output images are kind of blurred, since it is difficult to learn the high order modes of PCA.

2 Potential Flow Generator in flow-based models

Conclusions

In this paper we propose potential flow generators with $L_{2}$ optimal transport regularity. In particular, we propose two versions: the discrete one and the continuous one. Both potential flow generators could be integrated in GAN models with various generator loss functions, while the latter one could also be integrated into flow-based models. For the discrete potential flow generator, the $L_{2}$ optimal transport is directly encoded in, while for the continuous potential flow generator we only need a slight augmentation to the original generator loss functions to impose the $L_{2}$ optimal transport regularity.

We first show the correctness of potential flow generators in estimating $L_{2}$ optimal transport map by comparing with analytical reference solutions. We report that the continuous potential flow generator outperforms the discrete one in robustness. We also illustrated the characteristic of “proximity” for potential flow generator due to $L_{2}$ optimal transport regularity, and consequently the effectiveness of the potential flow generator in image translation tasks using unpaired training data from the MNIST dataset and the CelebA dataset.

Apart from image-to-image translations, it is also possible to apply the potential flow generator to other translation tasks, text-to-text translations for example, if the translation objects could be properly embedded into Euclidean space. It is also interesting to study the application of potential flow generator in other generative models apart from the ones in this paper. Moreover, while in this paper we use PCA for image embedding, a possible improvement is to integrate potential flow generator with other embedding techniques like autoencoders and graph embedding methods. We leave these possible improvements to future work.

Acknowledgement

This work is supported by the ARO MURI project (W911NF-15-1-0562), the DOE PhILMs project (DE-SC0019453), and the DARPA AIRA project (HR00111990025).

References