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: The paper proposes a new method for learning an optimal pushforward (for the quadratic cost) from a distribution to another distribution based on samples of both distributions. The optimal transport problem is first written equivalently as a minimax problem over set of convex functions, as in Makkuva et al. 19. Then, the convex functions are parametrized as Input Convex Neural Networks (ICNN). To my understanding, the novelty of the approach is to replace the minimax problem over the parameters of the ICNN by an easier to solve a regularized minimization problem. The authors prove the consistency of their approach: loosely speaking, they show that true minimizers of their approximate problem are epsilon minimizers of the original optimal transport problem. They also provide promising numerical experiments.