Title: WASSERSTEIN PROXIMAL OF

Abstract: We introduce a new method for training GANs by applying the Wasserstein-2 metric proximal on the generators. This approach is based on the gradient operator induced by optimal transport theory, which connects the geometry of the sample space and the parameter space in implicit deep generative models. From this theory, we obtain an easy-to-implement regularizer for the parameter updates. Our experiments demonstrate that this method improves the speed and stability in training GANs in terms of wallclock time and Fréchet Inception Distance (FID) learning curves.

1. INTRODUCTION:

Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) are a powerful approach to learning generative models. Here, a discriminator tries to tell apart the data generated from a real source and the data generated by a generator, whereas the generator tries to fool the discriminator. This adversarial game is formulated as an optimization problem over an implicit generative model for the generator. An implicit generative model is a parametrized family of functions mapping a noise source to sample space. In trying to fool the discriminator, the generator should try to recreate the density distribution from the real source.

The problem of matching a target density can be formulated as the minimization of a discrepancy measure. The Kullback-Leibler (KL) divergence is known to be difficult when the distributions have a low dimensional support set, as is commonly the case in applications with structured data and high dimensional sample spaces. An alternative approach to define a discrepancy measure between densities is optimal transport, a.k.a. Wasserstein distance, or Earth Mover's distance. This has been used recently to define the loss function for learning generative models (Montavon et al., 2016; Frogner et al., 2015) . In particular, the Wasserstein GAN (Arjovsky et al., 2017) has attracted much interest in recent years.

Besides defining the loss function, optimal transport can also be used to introduce structures serving the optimization itself, in terms of the gradient operator. In full probability space, this is known as the Wasserstein steepest descent flow (Jordan et al., 1998; Otto, 2001) . In this paper we derive the Wasserstein steepest descent flow for deep generative models in GANs. We use the Wasserstein-2 metric function, which allows us to obtain a Riemannian structure and a corresponding natural (i.e., Riemannian) gradient. A well known example of a natural gradient is the Fisher-Rao natural gradient, which is induced by the KL divergence. In learning problems, one often finds that the natural gradients can offer advantages compared to the Euclidean gradient (Amari, 1998; 2016) . In GANs, because of the low dimensional support sets and the associated difficulties with the KL divergence, the Fisher-Rao natural gradient is problematic. Therefore, we propose to use the gradient operator induced by the Wasserstein-2 metric (Li & Montúfar, 2018a; b) .

We compute the proximal operator for the generators of GANs, where the regularization is the squared constrained Wasserstein-2 distance. In practice, the constrained distance can be approximated by a simple neural network. In implicit generative models, the constrained Wasserstein-2 metric exhibits a simple structure. We generalize the metric and introduce the relaxed proximal operator for generators, which allows us to further simplify the computation. The resulting relaxed proximal operator involves only the difference of outputs, so that the proximal computation has very simple parameter updates. The method can be easily implemented and used as a drop-in regularizer for the generator updates. This paper is organized as follows. In Section 2, we briefly introduce the Wasserstein natural gradient. A Wasserstein proximal method is introduced in Algorithm 1. In Section 3, we demonstrate the effectiveness of the proposed methods in experiments with various types of GANs. Section 4 reviews related work.

2. WASSERSTEIN PROXIMAL:

In this section, we briefly present optimal transport and its proximal operator on a parameter space. We then apply them to the optimization problems of GANs.

2.1. WASSERSTEIN NATURAL GRADIENT:

Optimal transportation defines a class of distance functions between probability densities. Given a pair ρ 0 , ρ 1 ∈ P p (R n ) of probability densities with finite p-th moment,

EQUATION

where the infimum is over all joint probability densities π(x, y) with marginals ρ 0 (x), ρ 1 (y). In the literature (see Villani, 2009) , W p is referred to as the Wasserstein-p distance. In this paper, we focus on the case p = 2, and further denote W 2 by W .

Following Benamou & Brenier (2000) , the Wasserstein-2 distance has a dynamical formulation as a trajectory transporting the initial density ρ 0 to the final density ρ 1 along a trajectory of minimal kinetic energy. The classic theory does not consider the setting where the density path is constrained to lie within a parametrized model. In the following we extend the classic theory to cover parameterized density models. Consider a parameterized probability ρ(θ, x), with parameter space Θ ⊂ R d . Suppose that ρ(θ, x) is locally injective as a mapping from Θ to P 2 (R n ). Then the Wasserstein-2 metric function constrained to the parameter space is given as follows (see Li & Montúfar, 2018a) .

Theorem 1 (Constrained Wasserstein-2 metric) The constrained Wasserstein-2 metric function d W : Θ × Θ → R + has the following formulation:

d W (θ 0 , θ 1 ) 2 = inf 1 0 R n ∇Φ(t, x) 2 ρ(θ(t), x)dxdt : ∂ t ρ(θ(t), x) + ∇ • (ρ(θ(t), x)∇Φ(t, x)) = 0, θ(0) = θ 0 , θ(1) = θ 1 ,

where the infimum is among all feasible Borel potential functions Φ : [0, 1] × R n → R and continuous parameter paths θ : [0, 1] → R d . Here ∇• and ∇ are the divergence and gradient operators over R n .

We note that the constrained metric on parameter space can be different from the Wasserstein-2 distance on the full density set. The metric d W can be used to define a steepest descent optimization scheme. This can be formulated in two general ways.

One way is in terms of the corresponding Riemannian structure, i.e., an inner product between tangent vectors. A well known example is the Fisher natural gradient (Amari, 1998; 2016) . The constrained Wasserstein-2 metric allows us to obtain a Riemannian metric structure, from which we obtain the following constrained Wasserstein-2 gradient. We also call it Wasserstein natural gradient.

Theorem 2 (Wasserstein natural gradient) Given a loss function F : Θ → R, the Wasserstein gradient operator is given by

∇ W θ F (θ) = G(θ) -1 ∇ θ F (θ), where G(θ) = (G(θ) ij ) 1≤i,j≤d ∈ R d×d is given by G(θ) ij = R n ∇Φ i (x)∇Φ j (x)ρ(θ, x)dx.

Here for each:

i ∈ {1, • • • , d}, Φ i : R n → R is a solution (up to additive constants) of ∂ ∂θi ρ(θ, x) + ∇ • (ρ(θ, x)∇Φ i (x)) = 0.

Here ∇ W θ represents the natural gradient operator with respect to the constrained Wasserstein metric, ∇ θ represents the ordinary Euclidean gradient operator, and G is the matrix representing the Wasserstein Riemannian metric. The steepest descent flow is given by

EQUATION

The corresponding gradient descent iteration (forward Euler method) satisfies

θ k+1 = θ k -hG(θ k ) -1 ∇ θ F (θ k ),

where h > 0 is the step size. Often in practice, the computation of matrix G(θ) -1 is difficult.

The second way of obtaining a numerical scheme for equation 2 is in terms of the proximal operator. This is the backward Euler method, also named Jordan-Kinderlehrer-Otto (JKO) scheme (Jordan et al., 1998) , which is given by

EQUATION

)

Here, at each step, the distance of the parameter update acts as a regularization to the original loss function.

Computing d W is also often challenging. However, we can approximate the d W distance locally by a second order Taylor expansion. This approximation is particularly tractable within the parameterized setting that we discussed above.

This allows us to derive other first order schemes, such as the Semi-Backward Euler method:

Proposition 3 (Semi-Backward Euler method) The Semi-Backward Euler method for the gradient flow of loss function F : Θ → R is given by

θ k+1 = arg min θ∈Θ F (θ) + 1 h sup Φ R n Φ(x)(ρ(θ, x) -ρ(θ k , x)) - 1 2 (∇Φ(x)) 2 ρ(θ k , x)dx,

where the supremum is taken over Φ : R n → R with sufficient regularity for the integral to be well defined.

The Semi-Backward Euler method is often easier to approximate than the forward Euler method, because it does not require computing and inverting G(θ), and it is often simpler than the backward Euler method (JKO), because the constrained optimization over Φ is more tractable than the timedependent constraint involved in computing d W .

We implement the Semi-Backward Euler method in implicit generative model as follows. For each parameter θ ∈ R d , let the generator be given by g θ : R m → R n ; z → x = g(θ, z). This takes an input noise prior Z ∼ p(z) ∈ P 2 (R m ) to an output sample with density given by X = g(θ, Z) ∼ ρ(θ, x). Here R d is the parameter space, R m is the latent space, and R n is the sample space.

In this case, the update in Proposition 3 forms

θ k+1 = arg min θ∈Θ sup Φ F (θ) + 1 h E Z∼p(z) [Φ(g(θ, Z)) -Φ(g(θ k , Z)) - 1 2 ∇ x Φ(g(θ k , Z)) 2 ].

In practice, we apply a neural network to approximate variable Φ. See details in Appendix G.

2.2. REGULARIZATION ON GENERATORS:

In fact, the constrained Wasserstein-2 metric in implicit generative models allows for yet a simpler formulation. This reformulation allows us to define the relaxed Wasserstein metric, and further introduces a simple algorithm for proximal operator on generators.

Proposition 4 (Constrained Wasserstein-2 metric in implicit generative models)

d W (θ 0 , θ 1 ) 2 = inf 1 0 E Z∼p(z) d dt g(θ(t), Z) 2 dt : d dt g(θ(t), Z) -∇ x Φ(t, g(θ(t), Z)) = 0, θ(0) = θ 0 , θ(1) = θ 1 ,

where the infimum is among all feasible Borel potential functions Φ :

[0, 1] × R n → R and continu- ous parameter paths θ : [0, 1] → R d .

Here the constrained Wasserstein metric requires that the derivative of the generator g w.r.t. θ ∈ R d be a gradient vector field of Φ w.r.t x ∈ R n . In other words, if we denote x(t) = g(θ(t), z), then

d dt x(t) = ∇ x Φ(t, x(t)). (Gradient constraint)

The gradient constraint is satisfied if the sample space is 1 dimensional, i.e., n = 1. In general, this is not true. Here Φ(t, x) is the other function depending on the parameter space Θ. Finding Φ involves computational difficulties. Fitting the gradient constraint is an open problem for the computations of Wasserstein proximal operator.

For simple computations, we withdraw the gradient constraint and consider a relaxed Wassersetin metric on parameter space:

d(θ 0 , θ 1 ) 2 = inf θ(t) 1 0 E Z∼p(z) d dt g(θ(t), Z) 2 dt : θ(0) = θ 0 , θ(1) = θ 1 .

We approximate the relaxed Wasserstein proximal operator based on the new metric d to obtain

EQUATION

)

where the infimum is among all feasible continuous parameter path θ :

[0, 1] → R d .

In fact, when the sample space is high dimensional, i.e., n > 1, the above update is not exactly the Wasserstein proximal. Instead, it simply regularizes the generator by the expectation of squared difference in sample space.

Algorithm 1 Relaxed Wasserstein Proximal, where F ω is a parameterized function to minimize Require: F ω , a parameterized function to minimize (e.g. Wasserstein-1 with a parameterized discriminator). g θ the generator. Require: h proximal step-size, B batch size. Require: Optimizer Fω and Optimizer g θ . Require: max iterations, and generator iterations for k = 0 to max iterations do Sample real data {x i } B i=1 and latent data

{z i } B i=1 ω k ← Optimizer Fω 1 B B i=1 F ω (g θ (z i )) for = 0 to generator iterations do Sample latent data {z i } B i=1 θ k ← Optimizer g θ 1 B B i=1 F ω (g θ (z i )) + 1 h g θ (z i ) -g θ k-1 (z i ) 2 end for end for

2.3. ILLUSTRATION OF WASSERSTEIN PROXIMAL:

We present a toy example to illustrate the effectiveness of Wasserstein proximal operator in GANs.

Consider a family of distribution with two weighted delta measures.

Let Θ = {θ = (a, b) : a < 0, b > 0}, and define ρ(θ, x) = αδ a (x) + (1 -α)δ b (x),

where α ∈ [0, 1] is a given ratio and δ a (x) is the delta function supported at point a. See Figure 2 .3.

In this model, for a loss function F : Θ → R, the proximal regularization is given as follows:

θ k+1 = arg min θ∈Θ F (θ) + 1 2h d(θ, θ k ) 2 ,

where θ = (a, b) and θ k = (a k , b k ). We check the following commonly used statistical distance (divergence) functions d between parameters θ and θ k . 1. Wasserstein-2 distance:

d W (θ, θ k ) 2 = α(a -a k ) 2 + (1 -α)(b -b k ) 2 ; 2. Euclidean distance: d E (θ, θ k ) 2 = (a -a k ) 2 + (b -b k ) 2 ;

3. Kullback-Leibler divergence:

d KL (ρ θ ρ θ k ) = R n ρ(θ, x) log ρ(θ, x) ρ(θ k , x) dx = ∞; 4. L 2 -distance: d L 2 (ρ θ , ρ θ k ) 2 = R n |ρ(θ, x) -ρ(θ k , x)| 2 dx = ∞.

Here the KL divergence and L 2 -distance cannot measure the difference of probability models. The Wasserstein-2 and Euclidean distances still work in these cases. In addition, the Euclidean distance d E does not depend on the structure of model ρ(θ, x), while the constrained Wasserstein-2 metric d W does.

Proposition 5 Given θ * = (a * , b * ) ∈ Θ, consider the Wasserstein-1 metric as the loss function, i.e.,

F W1 (θ) = W 1 (ρ θ , ρ θ * ) = α|a -a * | + (1 -α)|b -b * |. Denote θ k+1 W = arg min θ F W1 (θ) + 1 2h d W (θ, θ k ) 2 , and θ k+1 E = arg min θ F W1 (θ) + 1 2h d E (θ, θ k ) 2 . For each stepsize h > 0, then F W1 (θ k+1 E ) ≥ F W1 (θ k+1 W ).

On each step of the update, the solution obtained by Wasserstein proximal decreases the objective function further than the one by Euclidean proximal. Here the proof is based on a simple fact of the shrinkage operator, see details in Appendix B.

This example introduces a case that Wasserstein-2 proximal works better than Euclidean proximal for the Wasserstein-1 loss function.

3. EXPERIMENTS ON GANS:

Here we present numerical experiments using the Relaxed Wasserstein Proximal (RWP) algorithm and the Semi-Backward Euler (SBE) method in order to perform Wasserstein gradient-descent on various GANs. We find that the Relaxed Wasserstein Proximal provides both better speed (measured by wallclock) and stability in training GANs.

3.1. RESULTS OF RELAXED WASSERSTEIN PROXIMAL:

The Relaxed Wasserstein Proximal (RWP) algorithm is intended to be an easy-to-implement, drop-in replacement to improve speed and convergence of GAN training. It does this by applying regularization on the generator during training. This is novel as most GAN training focuses on regularizing the discriminator, e.g. with a gradient penalty (Gulrajani et al., 2017b; Petzka et al., 2017; Kodali et al., 2018; Adler & Lunz, 2018; Miyato et al., 2018) , and there has been limited exploration in regularizing the generator (Chen et al., 2016) . Specifically, we modify the update rule for the generator by:

• Update for number of iterations before updating the discriminator:

θ ← Optimizer θ Original loss + 1 2h g θ -g θ k-1 2

So two hyperparameters are introduced: the proximal step-size h, and the number of iterations . In some GANs, one may update the discriminator a number of times and then update the generator a number of times, and then repeat; we will call one loop of this update an outer-iteration.

A more detailed description of the algorithm is given in Appendix D.

We test the Relaxed Wassersteing Proximal regularization on three GAN types:

• Standard GANs (Goodfellow et al., 2014) ,

• WGAN-GP (Gulrajani et al., 2017a) , and • DRAGAN (Kodali et al., 2018) .

We use the CIFAR-10 dataset (Krizhevsky, 2009) , and the aligned and cropped CelebA dataset (Liu et al., 2015) . And we utilize the DCGAN (Radford et al., 2015) architecture for the discriminator and generator. To measure the quality of generated samples, we employ the Fréchet Inception Distance (FID) (Heusel et al., 2017) both to measure performance and to measure convergence of GAN training (lower FID is better); we used 10,000 generated images to measure the FID. For CIFAR-10, we measure the FID every 1000 outer-iterations, and for CelebA we measure the FID every 10,000 outer-iterations.

Our particular hyperparameter choices for training are given in Appendix C. Note that since we intend RWP to be a drop-in regularization, the non-RWP hyperparameters (i.e. not h nor ) are chosen to work well before applying RWP.

To summarize our results, the Relaxed Wasserstein Proximal regularization improves both the speed (wallclock) and stability of convergence. It is a tricky to compare the result of using RWP, as it performs multiple generator iterations. We thus align the comparison according to wallclock time (this procedure was also used by Heusel et al., 2017) . In Figure 2 we see that our regularization improves convergence speed (measured in wallclock time), and also obtains a lower FID for all GAN types. In particular, in DRAGAN we see a 20% improvement in sample quality according to the FID.

The same results are also found for the CelebA dataset, shown in Figure 3 . We note that multiple generator iterations will sometimes prevent Standard GANs on CelebA from learning initially, at which point we restart the algorithm, and once it starts learning then the run is successful. This is practically very easy to detect and provides minimal trouble, so Figure 3 focuses on successful runs. We predict this defect will be rectified with a more stable loss function, such as WGAN-GP, or with different h's and 's.

We also examine the effect of multiple generator updates compared to discriminator updates. More specifically, in RWP since we update the generator multiple times before updating the discriminator, then it is worth examining the effect of not using the regularization. We see in Figure 4 that even using the most stable GAN type out of the three -WGAN-GP -if we omit regularization then the FID has high variance and even tends to rise in the end. But with RWP, the FID converges with more stability and achieves a lower FID.

Samples from the models are provided in Appendix E. We also performed latent space walks (Radford et al., 2015) to show RWP regularization does not cause the GAN to memorize. For details see Appendix F.

Figure 2 : The effect of using RWP regularization, on the CIFAR-10 dataset. The experiments are averaged over 5 runs. The bold lines are the average, and the enveloping lines are the minimum and maximum. From the three graphs, we see that using the easy-to-implement RWP regularization improves speed as measured by wallclock time, and it also is able to achieve a lower FID. Here we see RWP regularization improves the speed (via wallclock time), and achieves a lower FID. We note multiple generator iterations might cause initial learning to fail, but once it starts then it remains successful. This is practically easy to detect, so we show successful runs.

Figure 4 : An experiment demonstrating the effect of performing 10 generator iterations per outer-iteration with and without RWP, where an outer-iteration is a single loop of: a number of discriminator iterations, then a number of generator iterations. This experiment goes to 1,000,000 outer-iterations to show long-term behavior. With RWP regularization we obtain convergence, as well as lower FID. Without RWP, the training is highly variable and the FID is even on a rising trend in the end. Under review as a conference paper at ICLR 2019

3.2. RESULTS OF SEMI-BACKWARD EULER METHOD:

The training of Semi-Backward Euler (SBE) is a more complicated. Here we attempt to approximate three functions: the usual discriminator and generator, and the potential function Φ p . The algorithm and particular hyperparameter settings are presented in the appendix in Section G. We present our attempts at optimizing over the three networks in Figure 5 . Since both the standard WGAN-GP and the SBE on WGAN-GP had the same generator iterations, then we align according to this. As we see, the Semi-Backward Euler method is comparable to norm WGAN-GP. We leave deeper investigation of the Semi-Backward Euler method for future work.

4. RELATED WORKS:

In the literature, many different aspects of optimal transport have been applied into machine learning and GANs.

1. Loss function. Many studies apply the Wasserstein distance as the loss function. There are mainly two reasons for using the Wasserstein loss function (Frogner et al., 2015; Montavon et al., 2016) . On the one hand, the Wasserstein distance is a statistical distance depending on the metric of the sample space. So it introduces a statistical estimator, named the the minimal Wasserstein estimator (Bassetti et al., 2006) , depending on the geometry of the data. On the other hand, the Wasserstein distance is useful for comparing probability distributions supported on lower dimensional sets. This is often intractable for other divergence functions. In GANs, these properties have been leveraged in Wasserstein GAN (Arjovsky et al., 2017) . In this case, the loss function is chosen as the Wasserstein-1 distance function. In its computations, the discriminator, also called the Kantorovich dual variable, needs to satisfy the 1-Lipschitz condition. Many studies work on the regularization of the discriminator in order to satisfy this condition (Gulrajani et al., 2017b; Petzka et al., 2017) .

2. Gradient flows in full probability set. The Wasserstein-2 metric provides a metric tensor structure (Lott, 2007; Otto, 2001; Li, 2018) , under which the probability space forms an infinite dimensional Riemannian manifold, named the density manifold (Lafferty, 1988) . The gradient flow in the density manifold links with many transport-related partial differential equations (Villani, 2009; Nelson, 1985) . A famous example is that the Fokker-Planck equation, the probability transition equation of Langevin dynamics, is the gradient flow of the KL divergence function. In this perspective, two angles have been developed in the learning communities. Firstly, many groups try to leverage the gradient flow structure in probability space supported on the parameter space. They study the stochastic gradient descent by the transition equation in the probability over parameters (Mei et al., 2018) . Secondly, many nonparametric models have been studied, such as the Stein gradient descent method (Liu, 2017) . It can be viewed as the generalization of Wasserstein gradient flow. In addition, Frogner & Poggio (2018) consider an approximate inference method for computing Wasserstein gradient flow in full probability set. Here an approximation towards Kantorovich dual variables is introduced.

3. Gradient flow constrained on parameter space. The Wasserstein structure can also be constrained on parameter space. Carlen & Gangbo (2003) studied the constrained Wasserstein gradient with fixed mean and variance. Here the density subset is still infinite dimensional. Many approaches also focus on Gaussian families or elliptical distributions (Takatsu, 2011) . The Wasserstein gradient flow in Gaussian family has been studied by Malagò et al. (2018) .

Compared to previous works, our approach applies the Wasserstein gradient to work on general implicit generative models.

5. DISCUSSION:

In this work, we apply the constrained Wasserstein gradient and its relaxations on implicit generative models. Whereas much work has focused on regularizing the discriminator, in this work we focus on regularizing the generator. For Wasserstein GAN (with gradient penalty), we compute the Wasserstein-2 gradient flow of Wasserstein-1 distance on parameter space. Experimentally, the proposed method allows us to obtain a better minimizer in the sense of FID, with faster convergence speeds in wall-clock time.

In particular, we denote g θ (σ, σ) = σ T G(θ)σ, where G(θ) = (G(θ) ij ) 1≤i,j≤d ∈ R d×d is the associated metric tensor defined in Theorem 2.

Here we assume that G(θ) is smooth and positive definite, so that (Θ, g θ ) forms a smooth Riemannian manifold. In this case, Theorem 2 studies the constrained Wassertein gradient operator in parameter space.

B PROOFS OF WASSERSTEIN NATURAL GRADIENT:

Proof of Theorem 1 The distance d W can be written into the action function in Wasserstein statistical manifold. In other words, consider

d W (θ 0 , θ 1 ) 2 = inf 1 0 θ(t) T G(θ(t)) θ(t) : θ(0) = θ 0 , θ(1) = θ 1

where the infimum is taken over θ(t) ∈ C 1 ([0, 1], Θ). Following the definition of metric tensor in definition 6, we have

θ(t) T G(θ(t)) θ(t) = R n (∇Φ(t, x)) 2 ρ(θ(t), x)dx, with Φ(t, x) satisfying ∂ t ρ(θ(t), x) = ∇ θ ρ(θ(t), x) θ(t) = -∇(ρ(θ(t), x)∇Φ(t, x)).

We finish the proof.

Proof of Theorem 2:

The gradient operator on a Riemannian manifold (Θ, g θ ) is defined as follows.

For any σ ∈ T θ Θ, then the Riemannian gradient

∇ W θ F (θ) ∈ T θ Θ satisfies g θ (σ, ∇ W θ F (θ)) = (∇ θ F (θ), σ). In other words, θT G(θ)∇ W θ F (θ) = ∇ θ F (θ) T σ. Since θ ⊂ R d and G(θ) is positive definite, then ∇ W θ F (θ) = G(θ) -1 ∇ θ F (θ).

Proof of Proposition 3. We next present the derivation of the proposed semi-backward method.

EQUATION

and

EQUATION

Proof of Claim. We next prove the claim. Denote the geodesic path θ

* (t), t ∈ [0, 1], with θ * (0) = θ, θ * (1) = θ k , s.t. d W (θ, θ k ) 2 = 1 0 ( d dt θ * (t)) T G(θ * (t)) d dt θ * (t)dt.

We reparameterize the time of θ * (t) into the time interval [0, h]. Denote τ = ht and θ(τ

) = θ * (ht). Thus θ(τ ) = θ k + θ-θ k h τ + O(τ 2 ) and d dτ θ(τ ) = θ-θ k h + O(τ ), d W (θ, θ k ) 2 =h h 0 d dτ θ(τ ) T G(θ(τ )) d dτ θ(τ )dτ =h h 0 ( θ -θ k h + O(h)) T G(θ k + O(h))( θ -θ k h + O(h))dτ =(θ -θ k ) T G(θ k )(θ -θ k ) + O(h 2 ),

which proves equation 6.

We next prove equation 7. On the L.H.S. of equation 7,

∇ θ ρ(θ k , x)(θ -θ k ) = ρ(θ, x) -ρ(θ k , x) + O(h).

From the definition of G(θ),

1 2 (θ -θ k ) T G(θ k )(θ -θ k ) = 1 2 R n (∇Φ(x)) 2 ρ(θ k , x)dx, where -∇ • (ρ(θ k , x)∇Φ(x)) = ∇ θ ρ(θ k , x)(θ -θ k ) = ρ(θ k , x) + O(h).

On the R.H.S. of equation 7, the maximizer

EQUATION

Applying equation 8 into the R.H.S. of equation 7, we have

R n Φ * (x)(ρ(θ, x) -ρ(θ k , x)) - 1 2 ∇Φ * (x) 2 ρ(θ k , x)dx = R n Φ * (x)[-∇ • (ρ(θ k , x)∇Φ * (x)] - 1 2 ∇Φ * (x)ρ(θ k , x)dx = R n ∇Φ * (x) 2 ρ(θ k , x) - 1 2 ∇Φ * (x) 2 ρ(θ k , x)dx = 1 2 R n ∇Φ * (x) 2 ρ(θ , x)dx.

Comparing the L.H.S. and R.H.S. of equation 7, we prove the claim.

From the claim,

θ k+1 = arg min θ∈Θ F (θ) + 1 h d W (θ, θ k ) 2 2 = arg min θ∈Θ F (θ) + 1 2h (θ k -θ) T G(θ k )(θ k -θ) + O(h) = arg min θ∈Θ F (θ) + 1 h sup Φ R n Φ(x)(ρ(θ, x) -ρ(θ k , x)) - 1 2 ∇Φ(x) 2 ρ(θ k , x)dx + O(h).

Thus we derive a consistent numerical method in time, known as the Semi-backward method:

θ k+1 = θ k -hG(θ k ) -1 ∇ θ F (θ k+1 ).

Proof of Proposition 4. This result is proven in Li & Osher (2018) . We present it here for the completion of paper. The implicit model is given by the following push-forward relation. Denote

EQUATION

Given the gradient constraint

d dt g(θ(t), z) = ∇Φ(t, g(θ(t), z)),

we shall show that the probability density transition equation of g(θ(t), z) satisfies the constrained continuity equation

EQUATION

and

EQUATION

On the one hand, consider

EQUATION

where the second equality holds from the push forward relation in equation 9.

On the other hand, consider

EQUATION

where ∇, ∇• are gradient and divergence operators w.r.t. x ∈ R n . The second to last equality holds from the push forward relation equation 9, and the last equality holds using the integration by parts w.r.t. x. Since equation 12 equals equation 13 for any f ∈ C ∞ c (R n ), we prove equation 10. In addition, by the definition of the push forward operator equation 9, we have

E Z∼p(z) d dt g(θ(t), Z) 2 = R n ∇Φ(t, g(θ(t), z)) 2 p(z)dz = R n ∇Φ(t, x) 2 ρ(θ(t), x)dx.

Thus we prove equation 11.

Proof of Proposition 5. This example allows us to compute the proximal operator explicitly. On the one hand, we compute the Wasserstein proximal operator explicitly:

θ k+1 W = (a k+1 W , b k+1 W ) = arg min θ F W1 (θ) + 1 2h d W (θ, θ k ) 2 = arg min (a,b) α|a -a * | + (1 -α)|b -b * | + 1 2h (α|a -a k | 2 + (1 -α)|b -b k |). I.e., a W k+1 = arg min a |a -a * | + 1 2h |a -a k | 2 , b W k+1 = arg min b |b -b * | + 1 2h |b -b k | 2 .

Here

a k+1 W = shrink a * (a k , h) =    a k -h if a k > a * + h; a k + h if a k < a * -h; a * otherwise.

Similarly, b k+1 W = shrink b * (b k , h). On the other hand, we calculate the Euclidean proximal operator explicitly:

θ k+1 E = (a k+1 E , b k+1 E ) = arg min θ F W1 (θ) + 1 2h d E (θ, θ k ) 2 = arg min (a,b) α|a -a * | + (1 -α)|b -b * | + 1 2h (|a -a k | 2 + |b -b k | 2 ).

I.e.,

a k+1 E = arg min a α|a -a * | + 1 2h |a -a k | 2 , b k+1 E = arg min b (1 -α)|b -b * | + 1 2h |b -b k | 2 .

Here

a k+1 E = shrink a * (a k , αh) =    a k -αh if a k > a * + αh; a k + αh a k < a * -αh; a * otherwise. Similarly, b k+1 E = shrink b * (b k , (1 -α)h).

Here we only need to check that for all possible cases,

F W1 (θ k+1 E ) > F W1 (θ k+1 W ). If a k > a * + h and b k > b * + h, then F W1 (θ k+1 W ) =α[(a k -a * -h) + h 2 ] + (1 -α)[(b k -b * -h) + h 2 ] =α(a k -a * ) + (1 -α)(b k -b * ) - h 2 ,

and

F W1 (θ k+1 E ) =α[(a k -a * -αh)] + (αh) 2 2h + (1 -α)[(b k -b * -αh)] + (1 -α) 2 h 2 2h =α(a k -a * ) + (1 -α)(b k -b * ) - h 2 [α 2 + (1 -α) 2 ]. Since α ∈ [0, 1], then α 2 + (1 -α) 2 ≤ [α + (1 -α)] 2 = 1, then F W1 (θ k+1 W ) ≤ F W1 (θ k+1 E ).

In other cases, the proof follows similarly. We finish the proof.

C HYPERPARAMETERS FOR RELAXED WASSERSTEIN PROXIMAL EXPERIMENTS:

The following hyperparameter settings for the Relaxed Wasserstein Proximal experiments in Section 3.1 are:

• A batch size of 64 for all experiments.

• For CIFAR-10 with WGAN-GP: The Adam optimizer with learning rate 0.0001, β 1 = 0.5, and β 2 = 0.9 for both the generator and discriminator. We used a latent space dimension of 128, h = 0.1, and = 10 generator iterations. • For CIFAR-10 with Standard and DRAGAN: The Adam optimizer with learning rate 0.0002, β 1 = 0.1, and β 2 = 0.999 for both the generator and discriminator. We used a latent space dimension of 100, h = 0.2, and = 5 generator iterations. • For aligned and cropped CelebA with Standard: The Adam optimizer with learning rate 0.0002, β 1 = 0.5, and β 2 = 0.999 for both the generator and discriminator. We used a latent space dimension of 100, h = 0.2, and = 5 generator iterations.

D A PRACTICAL DESCRIPTION OF THE RELAXED WASSERSTEIN PROXIMAL:

As mentioned in Section 3.1, the Relaxed Wasserstein Proximal is meant to be an easy-to-implement, drop-in regularization. For instructional purposes, we take a specific example to showcase the algorithm: Relaxed Wasserstein Proximal on Standard GANs (with non-saturating gradient for the generator):

• Given:

-A generator g θ , and discriminator D ω , -The distance function F ω (g θ ) = E x∼real [log(D ω (x))] -E z∼N (0,1) [log(1 -D ω (g θ (z))], F LATENT SPACE WALK Radford et al. (2015) suggest that walking in the latent space could detect whether a generator was memorizing. We see in Figure 8 and Figure 9 that we have smooth transitions, so this is not the case for GANs with RWP regularization.

Figure 8 : A latent space walk for a with RWP regularization on Standard GANs, on CelebA.

As we have smooth transitions, this shows the generator is not overfitting. The latent space walk is done by interpolating between 4 points in the latent space.

Figure 9 : A latent space walk for a network with RWP regularization on WGAN-GP, on CIFAR-10.

As we have smooth transitions, this shows the generator is not overfitting. The latent space walk is done by interpolating between 4 points in the latent space.

G ALGORITHM AND PARTICULAR HYPERPARAMETERS FOR THE SEMI-BACKWARD EULER METHOD:

The specific hyperparameter settings used for the Semi-Backward Euler (SBE) on WGAN-GP, trained on CIFAR-10, are:

• A batch size of 64.

• The DCGAN architecture for the discriminator and generator. A one-hidden-layer fullyconnected network (a.k.a. MLP) for the potential Φ p . We also used layer-normalization (Lei Ba et al. ( 2016)) for each layer. • We used the Adam optimizer with learning rate 0.0002, β 1 = 0.1, and β 2 = 0.999 for both the generator, discriminator, and potential Φ p . We used a latent space dimension of 100, and h = 0.2. • Every outer-iteration loop, we updated the discriminator 5 times (as suggested in WGAN-GP), the generator once, and the potential 5 times. Note an outer-iteration is defined as one loop of: updating the discriminator a number of times, updating the potential a number of times, and updating the generator a number of times.

:



Figures and Tables:

Figure 1: Illustration of the example from Section 2.3. The Wasserstein proximal penalizes parameter steps in proportion to the mass being transported, which results in updates pointing towards the minimum of the loss function. The Euclidean proximal penalizes all parameters equally, which results in updates naively orthogonal to the level sets of the loss function.

Figure 3: The effect of Relaxed Wasserstein Proximal (RWP) regularization on Standard GANs, on the CelebA dataset.The experiment was averaged over 5 runs. The bold lines are the average, and the enveloping lines are the minium and maximum. Here we see RWP regularization improves the speed (via wallclock time), and achieves a lower FID. We note multiple generator iterations might cause initial learning to fail, but once it starts then it remains successful. This is practically easy to detect, so we show successful runs.

Figure 5: The effect of the Semi-Backward Euler (SBE) method, on the CIFAR-10 dataset. As we observe, the training is comparable to the standard way of training using the WGAN-GP loss.The experiment was averaged over 5 runs. The bold lines is the average, and the enveloping lines are the minimum and maximum.

Figure 6: A sample of images generated by RWP regularization on Standard GANs, on CelebA.

Figure 7: A sample of images generated by RWP regularization on WGAN-GP, on CIFAR-10.

Peer review (iclr2019):
 The paper intends to utilize natural gradient induced by Wasserstein-2 distance to train the generator in GAN. Starting from the dynamical formulation of optimal transport, the authors propose the Wasserstein proximal operator as a regularization, which is simple in form and fast to compute. The proximal operator is added to training the generator, unlike most other regularizations that focus on the discriminator. This is an interesting direction. The motivation is clear but by so many steps of approximation and relaxation, the authors didn’t address what is the final regularization actually corresponding to? Personally I am not convinced that theoretically the proposed training method is better than the standard SGD. The illustration example in the paper is not very helpful as it didn’t show how the proposed proximal operator works. The proximal operator serves as a regularization and it introduces some error, I would like to know how does this carry over to the whole training procedure. In GAN, the optimal discriminator depends on the current generator. Many approaches to GAN training (i.e. WGAN-GP) advocates to update the generator once in every “outer-iteration”. I am not sure how the proposed approach fit in those training schemes. In the simulation, the difference is not very significant, especially in FID vs iteration number. This could be due to parameter tuning in standard WGAN-GP. I encourage more simulation studies and take more GAN structures into consideration. Lastly, the stability mentioned in the paper lacks a formal definition. Is it the variance of the curves? Is it how robust the model is against outer iterations?