Wasserstein-2 Generative Networks

1 INTRODUCTION . Generative learning framework has become widespread over the last couple of years tentatively starting with the introduction of generative adversarial networks ( GANs ) by Goodfellow et al . ( 2014 ) . The framework aims to define a stochastic procedure to sample from a given complex probability distribution Q on a space Y ⊂ RD , e.g . a space of images . The usual generative pipeline includes sampling from tractable distribution P on space X and applying a generative mapping g : X → Y that transforms P into the desired Q . In many cases for probability distributions P , Q , there may exist several different generative mappings . For example , the mapping in Figure 1b seems to be better than the one in Figure 1a and should be preferred : the mapping in Figure 1b is straightforward , wellstructured and invertible . Existing generative learning approaches mainly do not focus on the structural properties of the generative mapping . For example , GAN-based approaches , such as f -GAN by Nowozin et al . ( 2016 ) ; Yadav et al . ( 2017 ) , WGAN by Arjovsky et al . ( 2017 ) and others Li et al . ( 2017 ) ; Mroueh & Sercu ( 2017 ) , approximate generative mapping by a neural network with a problem-specific architecture . The reasonable question is how to find a generative mapping g ◦ P = Q that is well-structured . Typically , the better the structure of the mapping is , the easier it is to find such a mapping . There are many ways to define what the well-structured mapping is . But usually , such a mapping is expected to be continuous and , if possible , invertible . One may note that when P and Q are both one-dimensional ( X , Y ⊂ R1 ) , the only class of mappings g : X → Y satisfying these properties are monotone mappings1 , i.e . ∀x , x′ ∈ X ( x 6= x′ ) satisfying ( g ( x ) − g ( x′ ) ) · ( x − x′ ) > 0 . The intuition of 1-dimensional spaces can be easily extended to X , Y ⊂ RD . We can require the similar condition to hold true : ∀x , x′ ∈ X ( x 6= x′ ) 〈g ( x ) − g ( x′ ) , x− x′〉 > 0 . ( 1 ) The condition ( 1 ) is called monotonicity , and every surjective function satisfying this condition is invertible . In one-dimensional case , for any pair of continuous P , Q with non-zero density there exists a unique monotone generative map given by g ( x ) = F−1Q ( FP ( x ) ) McCann et al . ( 1995 ) , where F ( · ) is the cumulative distribution function of P or Q . However , for D > 1 there might exist more than one generative monotone mapping . For example , when P = Q are standard 2-dimensional Gaussian distributions , all rotations by angles −π2 < α < π 2 are monotone and preserve the distribution . One may impose uniqueness by considering only maximal Peyré ( 2018 ) monotone mappings g : X → Y satisfying ∀N = 2 , 3 . . . and N distinct points x1 , . . . , xN ∈ X ( N + 1 ≡ 1 ) : N∑ n=1 〈g ( xn ) , xn − xn+1〉 > 0 . ( 2 ) The condition ( 2 ) is called cycle monotonicity and also implies `` usual '' monotonicity ( 1 ) . Importantly , for almost every two continuous probability distributions P , Q on X = Y = RD there exists a unique cycle monotone mapping g : X → Y satisfying g ◦ P = Q , see McCann et al . ( 1995 ) . Thus , instead of searching for arbitrary generative mapping , one may significantly reduce the considered approximating class of mappings by using only cycle monotone ones . According to Rockafellar ( 1966 ) , every cycle monotone mapping g is contained in a sub-gradient of some convex function ψ : X → R. Thus , every convex class of functions may produce cycle monotone mappings ( by considering sub-gradients of these functions ) . In practice , deep input convex neural networks ( ICNNs , see Amos et al . ( 2017 ) ) can be used as a class of convex functions . Formally , to fit a cycle monotone generative mapping , one may apply any existing approach , such as GANs Goodfellow et al . ( 2014 ) , with the set of generators restricted to gradients of ICNN . However , GANs typically require solving a minimax optimization problem . It turns out that the cycle monotone generators are strongly related to Wasserstein-2 distance ( W2 ) . The approaches by Taghvaei & Jalali ( 2019 ) ; Makkuva et al . ( 2019 ) use dual form of W2 to find the optimal generative mapping which is cycle monotone . The predecessor of both approaches is the gradient-descent algorithm for computing W2 distance by Chartrand et al . ( 2009 ) . The drawback of all these methods is similar to the one of GANs – their optimization objectives are minimax . Cyclically monotone generators require that both spaces X and Y have the same dimension , which poses no practical limitation . Indeed , it is possible to combine a generative mapping with a decoder of a pre-trained autoencoder , i.e . train a generative mapping into a latent space . It should be also noted that the cases with equal dimensions ofX and Y are common in computer vision . The typical example is image-to-image style transfer when both the input and the output images have the same size and number of channels . Other examples include image-to-image color transfer , domain adaptation , etc . In this paper , we develop the concept of cyclically monotone generative learning . The main contributions of the paper are as follows : 1 . Developing an end-to-end non-minimax algorithm for training cyclically monotone generative maps , i.e . optimal maps for quadratic transport cost ( Wasserstein-2 distance ) . 2 . Proving theoretical bound on the approximation properties of the transport mapping fitted by the developed approach . 3 . Developing a class of Input Convex Neural Networks whose gradients are used to approximate cyclically monotone mappings . 1We consider only monotone increasing mappings . Decreasing mappings have analogous properties . 4 . Demonstrating the performance of the method in practical problems of image-to-image color transfer , mass transport in latent spaces , image-to-image style translation and domain adaptation . Our algorithm extends the approach of Makkuva et al . ( 2019 ) , eliminates minimax optimization imposing cyclic regularization and solves non-minimax optimization problem . At the result , the algorithm scales well to high dimensions and converges up to 10x times faster than its predecessors . The paper is structured as follows . Section 2 is devoted to Related Work . In Section 3 , we give the necessary mathematical tools on Wasserstein-2 optimal transport . In Section 4 , we derive our algorithm and state our main theoretical results . In Section 5 , we provide the results of computational experiments . In Appendix A , we prove our theoretical results . In Appendix B , we describe the particular architectures of ICNN that we use for experiments . In Appendix C , additional experiments and training details are provided . 2 RELATED WORK . Modern generative learning is mainly associated with Generative Adversarial Networks ( GANs ) Goodfellow et al . ( 2014 ) ; Arjovsky et al . ( 2017 ) . Basic GAN model consists of two competing networks : generator g and discriminator d. Generator g takes as input samples x from given distribution P and tries to produce realistic samples from real data distribution Q. Discriminator d attempts to distinguish between generated and real distributions g ◦P and Q respectively . Formally , it approximates a dissimilarity measure between g ◦ P and Q ( e.g . f -divergence Nowozin et al . ( 2016 ) or Wasserstein-1 distance Arjovsky et al . ( 2017 ) ) . Although superior performance is reported for many applications of GANs Karras et al . ( 2017 ) ; Mirza & Osindero ( 2014 ) , training such models is always hard due to the minimax nature of the optimization objective . Another important branch of generative learning is related to the theory of Optimal Transport ( OT ) Villani ( 2008 ) ; Peyré et al . ( 2019 ) . OT methods seek generative mapping2 g : X → Y , optimal in the sense of the given transport cost c : X × Y → R : Cost ( P , Q ) = min g◦P=Q ∫ X c ( x , g ( x ) ) dP ( x ) . ( 3 ) Equation ( 3 ) is also known as Monge ’ s formulation of optimal transportation Villani ( 2008 ) . The principal OT generative method Seguy et al . ( 2017 ) is based on optimizing the regularized dual form of the transport cost ( 3 ) . It fits two potentials ψ , ψ ( primal and conjugate ) and then uses the barycentric projection to establish the desired ( third ) generative network g. Although the method uses non-minimax optimization objective , it is not end-to-end ( consists of two sequential steps ) . In the case of quadratic transport cost c ( x , y ) = ‖x−y‖ 2 2 , the value ( 3 ) is known as the square of Wasserstein-2 distance : W22 ( P , Q ) = min g◦P=Q ∫ X ‖x− g ( x ) ‖2 2 dP ( x ) . ( 4 ) It has been well studied in literature Brenier ( 1991 ) ; McCann et al . ( 1995 ) ; Villani ( 2003 ; 2008 ) and has many useful properties which we discuss in Section 3 in more detail . The optimal mapping for the quadratic cost is cyclically monotone . Several algorithms exist Lei et al . ( 2019 ) ; Taghvaei & Jalali ( 2019 ) ; Makkuva et al . ( 2019 ) for finding this mapping . The recent approach by Taghvaei & Jalali ( 2019 ) uses the gradient-descent-based algorithm by Chartrand et al . ( 2009 ) for computing W2 . The key idea is to approximate the optimal potential ψ∗ by an ICNN Amos et al . ( 2017 ) , and extract the optimal generator g∗ from its gradient ∇ψ∗ . The method is impractical due to high computational complexity : during the main optimization cycle , it solves an additional optimization sub-problem . The inner problem is convex but computationally costly . This was noted in the original paper and de-facto confirmed by the lack of experiments with complex distributions . A refinement of this approach is proposed by Makkuva et al . ( 2019 ) . The inner optimization sub-problem is removed , and a network is used to approximate its solution . This speeds up the computation , but the problem is still minimax . 2Commonly , in OT it is assumed that dimX = dimY . 3 PRELIMINARIES . In the section , we recall the properties of W2 distance ( 4 ) and its relation to cycle monotone mappings . Throughout the paper , we assume that P and Q are continuous distributions on X = Y = RD with finite second moments.3 This condition guarantees that ( 3 ) is well-defined in the sense that the optimal mapping g∗ always exists . It follows from ( Villani , 2003 , Brenier ’ s Theorem 2.12 ) that its restriction to the support of P is unique ( up to the values on the small sets ) and invertible . The symmetric characteristics apply to its inverse ( g∗ ) −1 , which induces symmetry to definition ( 4 ) for quadratic cost . According to Villani ( 2003 ) , the dual form of ( 4 ) is given by W22 ( P , Q ) = ∫ X ‖x‖2 2 dP ( x ) + ∫ Y ‖y‖2 2 dQ ( y ) ︸ ︷︷ ︸ Const ( P , Q ) − min ψ∈Convex [ ∫ X ψ ( x ) dP ( x ) + ∫ Y ψ ( y ) dQ ( y ) ] ︸ ︷︷ ︸ Corr ( P , Q ) , ( 5 ) where the minimum is taken over all the convex functions ( potentials ) ψ : X → R ∪ { ∞ } , and ψ ( y ) = maxx∈X ( 〈x , y〉 − ψ ( x ) ) is the convex conjugate Fenchel ( 1949 ) to ψ , which is also a convex function , ψ : Y → R ∪ { ∞ } . We call the value of the minimum in ( 5 ) cyclically monotone correlations and denote it by Corr ( P , Q ) . By equating ( 5 ) with ( 4 ) , one may derive the formula Corr ( P , Q ) = max g◦P=Q ∫ X 〈x , g ( x ) 〉dP ( x ) . ( 6 ) Note that ( − Corr ( P , Q ) ) can be viewed as an optimal transport cost for bilinear cost function c ( x , y ) = −〈x , y〉 , see McCann et al . ( 1995 ) . Thus , searching for optimal transport map g∗ for W2 is equivalent to finding the mapping which maximizes correlations ( 6 ) . It is known for W2 distance that the gradient g∗ = ∇ψ∗ of optimal potential ψ∗ readily gives the minimizer of ( 4 ) , see Villani ( 2003 ) . Being a gradient of a convex function , it is necessarily cycle monotone . In particular , the inverse mapping can be obtained by taking the gradient w.r.t . input of the conjugate of optimal potential ψ∗ ( y ) McCann et al . ( 1995 ) . Thus , we have ( g∗ ) −1 ( y ) = ( ∇ψ∗ ) −1 ( y ) = ∇ψ∗ ( y ) . ( 7 ) In fact , one may approximate the primal potential ψ by a parametric class Θ of input convex functions ψθ and optimize correlations min θ∈Θ Corr ( P , Q | ψθ ) = min θ∈Θ [ ∫ X ψθ ( x ) dP ( x ) + ∫ Y ψθ ( y ) dQ ( y ) ] ( 8 ) in order to extract the approximate optimal generator gθ† : X → Y from the approximate potential ψθ† . Note that in general it is not true that gθ† ◦ P will be equal to Q . However , we prove that if Corr ( P , Q | ψθ† ) is close to Corr ( P , Q ) , then gθ† ◦ P ≈ Q , see our Theorem A.3 in Appendix A.2 . The optimization of ( 8 ) can be performed via stochastic gradient descent . It is possible to get rid of conjugate ψθ and extract an analytic formula for the gradient of ( 8 ) w.r.t . parameters θ by using ψθ only , see the derivations in Taghvaei & Jalali ( 2019 ) ; Chartrand et al . ( 2009 ) : ∂Corr ( P , Q | ψθ ) ∂θ = ∫ X ∂ψθ ( x ) ∂θ dP ( x ) − ∫ Y ∂ψθ ( x̂ ) ∂θ dQ ( y ) , where ∂ψθ∂θ in the second integral is computed at x̂ = ( ∇ψθ ) −1 ( y ) , i.e . inverse value of y for ∇ψθ . In practice , both integrals are replaced by their Monte Carlo estimates over random mini-batches from P and Q . Yet to compute the second integral , one needs to recover the inverse values of the current mapping∇ψθ for all y ∼ Q in the mini batch . To do this , the following optimization sub-problem has to be solved x̂ = ( ∇ψθ ) −1 ( y ) ⇔ x̂ = arg max x∈X ( 〈x , y〉 − ψθ ( x ) ) ( 9 ) 3In practice , the continuity condition can be assumed to hold true . Indeed , widely used heuristics , such as adding small Gaussian noise to data Sønderby et al . ( 2016 ) , make considered distributions to be continuous . for each y ∼ Q in the mini batch . The optimization problem ( 9 ) is convex but complex because it requires computing the gradient of ψθ multiple times . It is computationally costly since ψθ is in general a large neural network . Besides , during iterations over θ , each time a new independent batch of samples arrives . This makes it hard to use the information on the solution of ( 9 ) from the previous gradient descent step over θ in ( 8 ) .

Summary: From my perspective, this is a much needed and love-to-see work for the line of neural generative modeling. Previous approaches were dominated by GAN based approaches which require solving a minimax optimization which has technical hurdles in practice. This paper reviews the literature of OT theory, and provides a comprehensive explanation about the relation between W_2 and cyclical monotonic map. In particular, it motivates the opportunity to approximate Eq. (8) with Eq.(12), as which a non-minimax formulation exist.