Wasserstein proximal of GANs

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 . 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 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 . 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 ∈ 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 . 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 ) 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

Summary: 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: