Year: 2018

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.

text: 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 Fréchet Inception Distance (FID) learning curves.

heading: 1 INTRODUCTION
text: 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 Wasserstein2 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 & Montú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.

heading: 2 WASSERSTEIN PROXIMAL
text: 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.

heading: 2.1 WASSERSTEIN NATURAL GRADIENT
text: Optimal transportation defines a class of distance functions between probability densities. Given a pair ρ0, ρ1 ∈ Pp(Rn) of probability densities with finite p-th moment,
Wp(ρ0, ρ1) p = inf ∫ Rn×Rn ‖x− y‖pπ(x, y)dxdy, (1)
where the infimum is over all joint probability densities π(x, y) with marginals ρ0(x), ρ1(y). In the literature (see Villani, 2009), Wp is referred to as the Wasserstein-p distance. In this paper, we focus on the case p = 2, and further denote W2 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 Θ ⊂ Rd. Suppose that ρ(θ, x) is locally injective as a mapping from Θ to P2(Rn). Then the Wasserstein-2 metric function constrained to the parameter space is given as follows (see Li & Montúfar, 2018a).
Theorem 1 (Constrained Wasserstein-2 metric) The constrained Wasserstein-2 metric function dW : Θ×Θ→ R+ has the following formulation:
dW (θ0, θ1) 2 = inf {∫ 1 0 ∫ Rn ‖∇Φ(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] × Rn → R and continuous parameter paths θ : [0, 1] → Rd. Here ∇· and ∇ are the divergence and gradient operators over Rn.
We note that the constrained metric on parameter space can be different from the Wasserstein-2 distance on the full density set. The metric dW 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 ∈ Rd×d is given by
G(θ)ij = ∫ Rn ∇Φi(x)∇Φj(x)ρ(θ, x)dx.
Here for each i ∈ {1, · · · , d}, Φi : Rn → 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
d dt θ(t) = −G(θ(t))−1∇θF (θ(t)). (2)
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
θk+1 = arg min θ∈Θ
F (θ) + 1
2h dW (θ, θ
k)2. (3)
Here, at each step, the distance of the parameter update acts as a regularization to the original loss function.
Computing dW is also often challenging. However, we can approximate the dW 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
Φ ∫ Rn Φ(x)(ρ(θ, x)− ρ(θk, x))− 1 2 (∇Φ(x))2ρ(θk, x)dx,
where the supremum is taken over Φ: Rn → 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 dW .
We implement the Semi-Backward Euler method in implicit generative model as follows. For each parameter θ ∈ Rd, let the generator be given by gθ : Rm → Rn; z 7→ x = g(θ, z). This takes an input noise prior Z ∼ p(z) ∈ P2(Rm) to an output sample with density given by X = g(θ, Z) ∼ ρ(θ, x). Here Rd is the parameter space, Rm is the latent space, and Rn is the sample space.
In this case, the update in Proposition 3 forms
θk+1 = arg min θ∈Θ sup Φ
F (θ) + 1
h EZ∼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.

heading: 2.2 REGULARIZATION ON GENERATORS
text: 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)
dW (θ0, θ1) 2 = inf {∫ 1 0 EZ∼p(z)‖ d dt g(θ(t), Z)‖2dt :
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]×Rn → R and continuous parameter paths θ : [0, 1]→ Rd.
Here the constrained Wasserstein metric requires that the derivative of the generator g w.r.t. θ ∈ Rd be a gradient vector field of Φ w.r.t x ∈ Rn. 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 EZ∼p(z)‖ d dt g(θ(t), Z)‖2dt : θ(0) = θ0, θ(1) = θ1 } .
We approximate the relaxed Wasserstein proximal operator based on the new metric d to obtain
θk+1 = arg min θ∈Θ
F (θ) + 1
2h EZ∼p(z)‖g(θ, Z)− g(θk, Z)‖2, (4)
where the infimum is among all feasible continuous parameter path θ : [0, 1]→ Rd. 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 dis-
criminator). gθ the generator. Require: h proximal step-size, B batch size. Require: OptimizerFω and Optimizergθ . Require: max iterations, and generator iterations
for k = 0 to max iterations do Sample real data {xi}Bi=1 and latent data {zi}Bi=1 ωk ← OptimizerFω ( 1 B ∑B i=1 Fω(gθ(zi)) ) for ` = 0 to generator iterations do
Sample latent data {zi}Bi=1 θk ← Optimizergθ ( 1 B ∑B i=1 Fω(gθ(zi)) + 1 h‖gθ(zi)− gθk−1(zi)‖2 ) end for
end for

heading: 2.3 ILLUSTRATION OF WASSERSTEIN PROXIMAL
text: 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 = (ak, bk). We check the following commonly used statistical distance (divergence) functions d between parameters θ and θk.
1. Wasserstein-2 distance:
dW (θ, θ k)2 = α(a− ak)2 + (1− α)(b− bk)2;
2. Euclidean distance: dE(θ, θ
k)2 = (a− ak)2 + (b− bk)2; 3. Kullback–Leibler divergence:
dKL(ρθ‖ρθk) = ∫ Rn ρ(θ, x) log ρ(θ, x) ρ(θk, x) dx =∞;
4. L2-distance:
dL2(ρθ, ρθk) 2 = ∫ Rn |ρ(θ, x)− ρ(θk, x)|2dx =∞.
Here the KL divergence and L2-distance cannot measure the difference of probability models. The Wasserstein-2 and Euclidean distances still work in these cases. In addition, the Euclidean distance dE does not depend on the structure of model ρ(θ, x), while the constrained Wasserstein-2 metric dW does.
Proposition 5 Given θ∗ = (a∗, b∗) ∈ Θ, consider the Wasserstein-1 metric as the loss function, i.e., FW1(θ) = W1(ρθ, ρθ∗) = α|a− a∗|+ (1− α)|b− b∗|. Denote θk+1W = arg minθ FW1(θ) + 1 2hdW (θ, θ k)2, and θk+1E = arg minθ FW1(θ) + 1 2hdE(θ, θ k)2. For each stepsize h > 0, then FW1(θ k+1 E ) ≥ FW1(θk+1W ).
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.

heading: 3 EXPERIMENTS ON GANS
text: 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.

heading: 3.1 RESULTS OF RELAXED WASSERSTEIN PROXIMAL
text: 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 Fré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.

heading: 3.2 RESULTS OF SEMI-BACKWARD EULER METHOD
text: 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.

heading: 4 RELATED WORKS
text: 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 Malagò et al. (2018).
Compared to previous works, our approach applies the Wasserstein gradient to work on general implicit generative models.

heading: 5 DISCUSSION
text: 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.

heading: A REVIEW OF WASSERSTEIN STATISTICAL MANIFOLD
text: In the full probability set, we consider a metric function W2 : P2(Rn)× P2(Rn)→ R+,
W2(ρ0, ρ1) 2 = inf
Φt {∫ 1 0 ∫ Rn ‖∇Φ(t, x)‖2ρ(t, x)dxdt :
∂tρ(t, x) +∇ · (ρ(t, x)∇Φ(t, x)) = 0, ρ(0, x) = ρ0(x), ρ(1, x) = ρ1(x) } ,
(5)
where the infimum is taken among all feasible Borel potential functions Φ: [0, 1] × Rn → R and continuous density path ρ : [0, 1]× Rn → R+ satisfying the continuity equation. The variational formulation in equation 5 introduces a Riemannian structure in density space. Consider the set of smooth and strictly positive probability densities
P+ = { ρ ∈ C∞(Rn) : ρ(x) > 0, ∫ Rn ρ(x)dx = 1 } ⊂ P2(Rn).
Denote F := C∞(Rn) the set of smooth real valued functions. The tangent space of P+ is given by TρP+ = { σ ∈ F : ∫ Rn σ(x)dx = 0 } . Given Φ ∈ F and ρ ∈ P+, define VΦ(x) := −∇ · (ρ(x)∇Φ(x)).
Thus VΦ ∈ TρP+. The elliptic operator∇·(ρ∇) identifies the function Φ modulo additive constants with the tangent vector VΦ of the space of densities.
Given ρ ∈ P+, σi ∈ TρP+, i = 1, 2, define
gWρ (σ1, σ2) = ∫ Rn (∇Φ1(x),∇Φ2(x))ρ(x)dx,
where Φi(x) ∈ F/R, such that −∇ · (ρ∇Φi) = σi. The inner product gW endows P+ with a Riemannian metric tensor. In other words, the variational problem equation 5 is a geometric action energy in (P+, gW ). Given a loss function F : P+ → R, the Wasserstein gradient operator in (P+, gW ) is given as follows.
gradWF (ρ) = −∇ · (ρ∇ δ
δρ(x) F (ρ)).
Thus the gradient flow satisfies ∂ρ
∂t = −gradWF (ρ) = ∇ · (ρ∇
δ
δρ(x) F (ρ)).
More analytical results on the Wasserstein-2 gradient flow are provided in Ambrosio et al. (2005).
We next consider Wasserstein-2 metric and gradient operator constrained on statistical models. A statistical model is defined by a triplet (Θ,Rn, ρ). For simple presentation of paper, we assume Θ ⊂ Rd and ρ : Θ → P(Rn) is a parameterization function. In this case, ρ(Θ) ⊂ P(Rn). We assume that the parameterization map ρ is locally injective and under suitable regularities. We define a Riemannian metric g on ρ(Θ) by pulling back the Wasserstein-2 metric tensor gW .
Definition 6 (Wasserstein statistical manifold) Given θ ∈ Θ and σi ∈ TθΘ, i = 1, 2, we define
gθ(σ1, σ2) = ∫ Rn ∇Φ1(x)∇Φ2(x)ρ(θ, x)dx,
where −∇ · (ρ(θ, x)∇Φi(x)) = (∇θρ(θ, x), σi).
Here∇θρ = ( ∂∂θi ρ(θ, x)) d i=1 ∈ Rd and (·, ·) is an Euclidean inner product in Rd.
In particular, we denote gθ(σ, σ) = σ TG(θ)σ, where G(θ) = (G(θ)ij)1≤i,j≤d ∈ Rd×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.

heading: B PROOFS OF WASSERSTEIN NATURAL GRADIENT
text: Proof of Theorem 1 The distance dW can be written into the action function in Wasserstein statistical manifold. In other words, consider
dW (θ0, θ1) 2 = inf {∫ 1 0 θ̇(t)TG(θ(t))θ̇(t) : θ(0) = θ0, θ(1) = θ1 } where the infimum is taken over θ(t) ∈ C1([0, 1],Θ). Following the definition of metric tensor in definition 6, we have
θ̇(t)TG(θ(t))θ̇(t) = ∫ Rn (∇Φ(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, θ̇TG(θ)∇Wθ F (θ) = ∇θF (θ)Tσ. Since θ ⊂ Rd 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.
Claim: Denote ‖θ − θk‖ = h, then (θk − θ)TG(θk)(θk − θ) = dW (θ, θk)2 +O(h2), (6)
and 1 2 (θk − θ)TG(θk)(θk − θ) +O(h2) = sup
Φ ∫ Rn Φ(x)(ρ(θ, x)− ρ(θk, x))− 1 2 ‖∇Φ(x)‖2ρ(θk, x)dx.
(7)
Proof of Claim. We next prove the claim. Denote the geodesic path θ∗(t), t ∈ [0, 1], with θ∗(0) = θ, θ∗(1) = θk, s.t.
dW (θ, θ k)2 = ∫ 1 0 ( d dt θ∗(t))TG(θ∗(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 ddτ θ(τ) = θ−θk h +O(τ),
dW (θ, θ k)2 =h ∫ h 0 d dτ θ(τ)TG(θ(τ)) d dτ θ(τ)dτ
=h ∫ h 0 ( θ − θk h +O(h))TG(θk +O(h))( θ − θk h +O(h))dτ =(θ − θk)TG(θk)(θ − θk) +O(h2),
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)TG(θk)(θ − θk) = 1 2 ∫ Rn (∇Φ(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 Φ∗ satisfies
ρ(θ, x)− ρ(θk, x) +∇ · (ρ(θk, x)∇Φ∗(x)) = 0. (8) Applying equation 8 into the R.H.S. of equation 7, we have∫
Rn Φ∗(x)(ρ(θ, x)− ρ(θk, x))− 1 2 ‖∇Φ∗(x)‖2ρ(θk, x)dx
= ∫ Rn Φ∗(x)[−∇ · (ρ(θk, x)∇Φ∗(x)]− 1 2 ∇Φ∗(x)ρ(θk, x)dx
= ∫ Rn ‖∇Φ∗(x)‖2ρ(θk, x)− 1 2 ‖∇Φ∗(x)‖2ρ(θk, x)dx
= 1
2 ∫ Rn ‖∇Φ∗(x)‖2ρ(θk, 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
dW (θ, θ k)2
2
= arg min θ∈Θ
F (θ) + 1
2h (θk − θ)TG(θk)(θk − θ) +O(h)
= arg min θ∈Θ
F (θ) + 1
h sup
Φ ∫ Rn Φ(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 gθ#p(z) = ρ(θ, x), i.e.,∫
Rm f(g(θ, z))p(z)dz = ∫ Rn f(x)ρ(θ, x)dx, for any f ∈ C∞c (Rn). (9)
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
∂ ∂t ρ(θ(t), x) +∇ · (ρ(θ(t), x)∇Φ(t, x)) = 0, (10)
and EZ∼p(z)‖ d
dt g(θ(t), Z)‖2 = ∫ Rn ‖∇Φ(t, x))‖2ρ(θ(t), x)dx. (11)
On the one hand, consider f ∈ C∞c (Rn), then d
dt EZ∼p(z)f(g(θ(t), Z)) =
d
dt ∫ Rm f(g(θ(t), z))p(z)dz
= d
dt ∫ Rn f(x)ρ(θ(t), x)dx
= ∫ Rn f(x) ∂ ∂t ρ(θ(t), x)dx,
(12)
where the second equality holds from the push forward relation in equation 9.
On the other hand, consider d
dt EZ∼p(z)f(g(θ(t), Z)) = lim ∆t→0 EZ∼p(z) f(g(θ(t+ ∆t), Z)− f(g(θ(t), Z)) ∆t
= lim ∆t→0 ∫ Rm f(g(θ(t+ ∆t), z))− f(g(θ(t), z)) ∆t p(z)dz
= ∫ Rm ∇f(g(θ(t), z)) d dt g(θ(t), z)p(z)dz
= ∫ Rm ∇f(g(θ(t), z))∇Φ(t, g(θ(t), z))p(z)dz
= ∫ Rn ∇f(x)∇Φ(t, x)ρ(θ(t), x)dx
=− ∫ Rn f(x)∇ · (∇Φ(t, x)ρ(θ(t), x))dx,
(13)
where ∇, ∇· are gradient and divergence operators w.r.t. x ∈ Rn. 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 (Rn), we prove equation 10. In addition, by the definition of the push forward operator equation 9, we have
EZ∼p(z)‖ d
dt g(θ(t), Z)‖2 = ∫ Rn ‖∇Φ(t, g(θ(t), z))‖2p(z)dz
= ∫ Rn ‖∇Φ(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+1W = (a k+1 W , b k+1 W ) = arg min
θ FW1(θ) +
1
2h dW (θ, θ
k)2
= arg min (a,b) α|a− a∗|+ (1− α)|b− b∗|+ 1 2h (α|a− ak|2 + (1− α)|b− bk|).
I.e.,
aWk+1 = arg min a |a− a∗|+ 1 2h |a− ak|2, bWk+1 = arg min b |b− b∗|+ 1 2h |b− bk|2.
Here
ak+1W = shrinka∗(a k, h) =  ak − h if ak > a∗ + h; ak + h if ak < a∗ − h; a∗ otherwise.
Similarly, bk+1W = shrinkb∗(b k, h).
On the other hand, we calculate the Euclidean proximal operator explicitly:
θk+1E = (a k+1 E , b k+1 E ) = arg min
θ FW1(θ) +
1
2h dE(θ, θ
k)2
= arg min (a,b) α|a− a∗|+ (1− α)|b− b∗|+ 1 2h (|a− ak|2 + |b− bk|2).
I.e.,
ak+1E = arg mina α|a− a∗|+ 1 2h |a− ak|2, bk+1E = arg min b (1− α)|b− b∗|+ 1 2h |b− bk|2.
Here
ak+1E = shrinka∗(a k, αh) =  ak − αh if ak > a∗ + αh; ak + αh if ak < a∗ − αh; a∗ otherwise.
Similarly, bk+1E = shrinkb∗(b k, (1− α)h).
Here we only need to check that for all possible cases, FW1(θ k+1 E ) > FW1(θ k+1 W ). If a k > a∗ + h and bk > b∗ + h, then
FW1(θ k+1 W ) =α[(a k − a∗ − h) + h 2 ] + (1− α)[(bk − b∗ − h) + h 2 ]
=α(ak − a∗) + (1− α)(bk − b∗)− h 2 ,
and
FW1(θ k+1 E ) =α[(a
k − a∗ − αh)] + (αh) 2 2h + (1− α)[(bk − b∗ − αh)] + (1− α) 2h2 2h
=α(ak − a∗) + (1− α)(bk − b∗)− h 2 [α2 + (1− α)2].
Since α ∈ [0, 1], then α2 + (1 − α)2 ≤ [α + (1 − α)]2 = 1, then FW1(θk+1W ) ≤ FW1(θk+1E ). In other cases, the proof follows similarly. We finish the proof.

heading: C HYPERPARAMETERS FOR RELAXED WASSERSTEIN PROXIMAL EXPERIMENTS
text: 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.

heading: D A PRACTICAL DESCRIPTION OF THE RELAXED WASSERSTEIN PROXIMAL
text: 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θ) = Ex∼real[log(Dω(x))] − Ez∼N (0,1)[log(1 − Dω(gθ(z))],
– Choice of optimizers, Adamω and Adamθ,
– Proximal step-sizes h, and generator iterations `, and
– Batch size B.
Then the algorithm follows:
1. Sample real data {xi}Bi=1, and latent data {zi}Bi=1.
2. Update the discriminator:
ωk ← Adamω ( − 1 B B∑ i=1 log(Dω(xi))− 1 B B∑ i=1 log(1−Dω(gθ(zi))) )
3. Sample latent data {zi}Bi=1
4. Perform Adam gradient descent ` number of times:
θk ← Adamθ ( − 1 B B∑ i=1 log(Dω(gθ(zi)))− 1 B B∑ i=1 1 2h ‖gθ(zi)− gθk−1(zi)‖22 ) ,
for ` number of times.
5. Repeat the above until a chosen stopping condition (e.g. maximum number of iterations).
As one can analyze above, the only difference between the standard way of training GANs and using the Relaxed Wasserstein Proximal, are the ‖gθ(zi)− gθk−1(zi)‖22 terms and the number of generator iterations `. Note that in this paper, we call a single loop of updating a discriminator a number of times and then updating the generator a number of a time, an outer-iteration.

heading: E GENERATED SAMPLES FROM THE MODEL
text: In Figure 6, we have samples generated from a Standard GAN with RWP regularization, trained on the CelebA dataset. The FID of these images was 17.105.
In Figure 7, we have samples generated from WGAN-GP with RWP , trained on the CIFAR-10 dataset. The FID for these images is 38.3.

heading: F LATENT SPACE WALK
text: 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.

heading: G ALGORITHM AND PARTICULAR HYPERPARAMETERS FOR THE SEMI-BACKWARD EULER METHOD
text: 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 fully-
connected 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 WGANGP), 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.
Algorithm 2 Semi-backward Euler method, where Fω is a parameterized function to minimize. Require: Fω , a parameterized function to minimize (e.g. Wasserstein-1 with a parameterized dis-
criminator). gθ the generator. Φp the potential. Require: h the proximal step-size, m the batch size. Require: OptimizerFω , Optimizergθ , and OptimizerΦp Require: The number of generator iterations and p iterations to do per update.
1: for k = 0 to max iterations do 2: Sample real data {xi}Bi=1 and latent data {zi}Bi=1. 3: ωk ← OptimizerFω ( 1 B ∑B i=1 Fω(gθ(zi)) ) 4: for s = 0 to phi iterations do 5: Sample latent data {zi}Bi=1 6: pk ← OptimizerΦp ( 1 h 1 B ∑B i=1 Φp(gθ(zi))− Φp(gθk−1(zi))− 12∇Φp(gθk−1(zi))
) 7: end for 8: for ` = 0 to generator iterations do 9: Sample latent data {zi}Bi=1
10: θk ← Optimizergθ ( 1 B ∑B i=1 Fω(gθ(zi)) + 1 h ( Φp(gθ(zi))− Φp(gθk−1(zi))− 12∇Φp(gθk−1(zi)) )) 11: end for 12: end for

Review summaries:
This paper considers natural gradient learning in GAN learning, where the Riemannian structure induced by the Wasserstein-2 distance is employed. More concretely, the constrained Wasserstein-2 metric $d_W$, the geodesic distance on the parameter space induced by the Wasserstein-2 distance in the ambient space, is introduced (Theorem 1). The natural gradient on the parameter space with respect to the constrained Wasserstein-2 metric is then derived (Theorem 2). Since direct evaluation of $G^{-1}$ poses difficulty, the authors go on to considering a backward scheme using the proximal operator (3), yielding:

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.