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