FLUID FLOW MASS TRANSPORT FOR GENERATIVE NETWORKS

1 INTRODUCTION . Generative Networks have been intensively studied in recent years yielding some impressive results in different fields ( see for example Salimans et al . ( 2016 ) ; Karras et al . ( 2017 ) ; Brock et al . ( 2018 ) ; Karras et al . ( 2018 ) ; Zhu et al . ( 2017 ) and reference within ) . The most common technique used to train a generative network is by formulating the problem as a Generative Adversarial Networks ( GANs ) . Nonetheless , GANs are notoriously difficult to train and in many cases do not converge , converge to undesirable points , or suffer from problems such as mode collapse . Our goal here is to better understand the generative network problem and to investigate a new formulation and numerical optimization techniques that do not suffer from similar shortcomings . To be more specific , we consider a data set made up of two sets of vectors , template vectors ( organized as a matrix ) T = [ T1 , . . . , Tn ] ∈ T and reference vectors R = [ R1 , . . . , Rn ] ∈ R. The goal of our network is to find a transformation f ( T , θ ) that generates reference vectors , that is vectors from the spaceR , from template vectors in the space T , where θ are parameters that control the function f . For simplicity , we write our generator as T ( θ ) = f ( T , θ ) ( 1.1 ) Equation equation 1.1 defines a generator that depends on the template data and some unknown parameters to be learned in the training process . We will use a deep residual network to approximate the function with θ being the weights of the network . In order to find the parameters for the generator we need to minimize some loss that captures how well the generator works . Assume first that we have correspondence between the data points in T and R , that is , we know that the vector Tj ( θ ) = Rj , ∀j . In other words , our training set consists of paired input and output vectors . In this case we can find θ by simply minimizing the sum of squares difference ( and may add some regularization term such as weight decay ) . E ( θ ) = 1 2n ∑ j ‖Tj ( θ ) −Rj‖2 . However , such correspondence is not always available and therefore , a different approach is needed if we are to estimate the generator parameters . The most commonly used approach to solve the lack of correspondence is Generative Adversarial Networks ( GANs ) and more recently , the Wasserstein GANs ( WGANs ) Arjovsky et al . ( 2017 ) . In these approaches one generates a discriminator network , or a critic in the case of WGANs , c ( r , η ) that gives a score for a vector , r , to be in the space R. In the original formulation c ( · , · ) yields the probability and in more recent work on WGANs it yields a real number . The discriminator/critic depends on the parameters η to be estimated from the data in the training process . One then defines the function J ( η , θ ) = g ( c ( R , η ) ) + h ( c ( T ( θ ) , η ) ) ( 1.2 ) In the case of a simple GAN we have that g ( x ) = log ( sm ( x ) ) and h ( x ) = log ( 1− sm ( x ) ) where the sm ( · ) function , is a soft max function that converts the score to a probability . In the case of WGANs a simpler expression is derived where we use the score directly setting g to the identity and h = −g , and require some extra regularity on the score function . Minimizing J with respect to η detects the `` fake '' vectors generated by the generator while maximizing J with respect to θ ” fools ” the discriminator , thus generating vectors T ( θ ) that are similar to vectors that are drawn fromR . Training GANs is a minimax problem where J ( η , θ ) is minimized with respect to η and maximize with respect to θ. Minimax problems are very difficult to solve and are typically unstable . Furthermore , the solution is based on gradient d/ascent which is known to be slow , especially when considering a saddle point problem Nocedal & Wright ( 1999 ) , and this can be demonstrated by solving the following simple quadratic problem . Example 1.1 Let J ( η , θ ) = 1 2 θ2 − 1 20 η2 + 10θ η The d/ascent algorithm used for the solution of the problem reads θk+1 = θk − µ ( θk + 10ηk ) ηk+1 = ηk − µ ( 1 10 ηk − 10θk+1 ) The convergence path of the method is plotted in Figure 1 . It is evident that the algorithm takes an inefficient path to reach its destination . This inefficient path is a known property of the method and avoiding it requires much more complex algorithms that , for example , eliminate one of the unknowns at least locally ( see a discussion in Nocedal & Wright ( 1999 ) ) . While such approaches have been derived for problems in fluid dynamics and constrained optimization , they are much more difficult to derive for deep learning due to the non-linearity of the learning problem . Besides the slowly converging algorithm , the simple GAN approach has a number of known fundamental problems . It has been shown in Zhang et al . ( 2017 ) that a deep network can classify vectors with random labels . This implies that given sufficient capacity in the classifier , it is always possible to obtain 0 loss even if R 6= T ( θ ) ∈ R , implying that there is no real metrics to stop the process ( see further discussion in Arjovsky & Bottou ( 2017 ) ) . This problem also leads to mode collapse , as it is possible to obtain the saddle point of the function when T ( θ ) = R , however this yields a local solution that is redundant . Therefore , it is fairly well known that the discriminator needs to be heavily regularized in order to effectively train the generator Gulrajani et al . ( 2017 ) , such as by using a very small learning rate or by weight clipping in the case of WGANs . A number of improvements have been proposed for GANs and WGANs , however these techniques still involve a minimax problem that can be difficult to solve . Another common property of all GANs is that they minimize some distance between the two vector spaces that are spanned by some probabilities . While a simple GAN is minimizing the JS-Divergence , the Wasserstein GAN minimizes the Wasserstein Distance . Minimizing the distance for probabilities makes sense when the probabilities are known exactly . However , as we discuss in the next section , such a process is not reasonable when the probabilities are estimated , that is , sampled and are therefore noisy . Since both spaces , R and T are only sampled , we only have noisy estimations of the probabilities and this has to be taken into consideration when solving the problem . In this paper we therefore propose a different point of view that allows us to solve the problem without the need of minimax . Our point of view stems from the following observation . Assume for a moment that the vectors R and T ( θ ) are in R2 or R3 . Then , the problem described above is nothing but a registration of point clouds under a non-common transformation rule ( i.e . not affine transformation ) . Such problem has been addressed in computer graphics and image registration for decades ( see for example Eckart et al . ( 2018 ) ; Myronenko & Song ( 2010 ) and reference within ) yielding successful software packages , numerical analysis and computational treatment . This observation is demonstrated in the following example , that we use throughout the paper . Example 1.2 Assume that the space T is defined by vectors Ti that are in R2 and that each vector is drawn from a simple Gaussian distribution with 0 mean and standard deviation of 1 . To generate the spaceR we use a known generator that is a simple resnet with fully connected layers xj+1 = xj + hσ ( Kjxj ) j = 1 , . . . , n x1 = T Here σ ( · ) is the tanh activation function and we choose Kj and save them for later use . The original points as well as points from the same distribution that are transformed using this simple resenet are plotted in Figure 2 . Finding the transformation parameters ( the matrices Kj ) is a kin to registering the red and blue points , when no correspondence map is given . We therefore propose to extend the ideas behind cloud point registration to higher dimensions , adapting them to Generative Models . Furthermore , recent techniques for such registration are strongly connected to the fluid flow formulation of optimal mass transport . As we explore next , there is a strong connection between our formulation and the Wasserstein GAN . Similar connections have been proposed in Lei et al . ( 2019 ) . Our approach can be viewed as a discretization of the fluid flow formulation of the optimal mass transport problem proposed in Benamou & Brenier ( 2003 ) and further discussed in Haber & Horesh ( 2014 ) ; Ryu et al . ( 2017 ) . It has some commonalities with normalized flow generators Kingma & Dhariwal ( 2018 ) with two main differences . First , our flow become regular due to the addition of regularization that explicitly keep the flow normalized and second , and more importantly , it solves the correspondence problem between the two data sets . The fluid formulation is known to be easier to solve and less nonlinear than the standard approach and ( at least in 3D ) leads to better and faster algorithms . We name our algorithm Mass Transport Generative Networks ( MTGN ) since our algorithm is based on mass transport that tries to match distributions but not adversarial networks . Similar to the fluid flow formulation for the OMT problem , our algorithm leads to a simple minimization problem and does not require the solution of a minimax . The rest of this paper is organized as follows . In Section 2 we lay out the foundation of the idea used to solve the problem , including discretization and numerical optimization . In Section 3 we demonstrate the idea on very simple example that help us gain insight into the method . In Section 4 we perform numerical experiments in higher dimensions and discuss how to effectively use the method and finally in Section 5 we summarize the paper and suggest future work . 2 GENERATIVE MODELS , MASS TRANSPORT AND CLOUD POINT REGISTRATION . In this section we discuss our approach for the solution of the problem and the connection to the registration of cloud of points and the fluid flow formulation of optimal mass transport . 2.1 GENERATIVE MODELS AND FLUID FLOW MASS TRANSPORT . We start by associating the spaces T andR with two probability density functions pT ( x ) and pR ( x ) . The goal is to find a transformation such that the probability , pT is transformed to the probability pR , minimizing some distance ( see the review paper Evans ( 1989 ) for details ) . An L2 distance leads to the Monge Kantorovich problem but it is possible to use different distances to obtain different transformations Burger et al . ( 2013 ) . The computation of such a transformation has been addressed by vast amount of literature . Solution techniques range from linear programming , to the solution of the notoriously nonlinear Monge Ampre equation Evans ( 1989 ) . However , in a seminal paper Benamou & Brenier ( 2003 ) , it was shown that the problem can be formed as minimizing the energy of a simple flow min E ( v ( x , t ) ) = ∫ 1 0 ∫ Ω ρ ( x , t ) |v ( x , t ) |2 dx dt ( 2.3a ) s.t ρt +∇ · ( vρ ) = 0 ( 2.3b ) ρ ( 0 , x ) = pT ( x ) ρ ( 1 , x ) = pR ( x ) ( 2.3c ) Here E is the total energy that depends on the velocity and density of the flow . The idea was extended in Chen et al . ( 2017 ) to solve the problem with different distances on vector spaces . The problem is also commonly solved in fields such as computational flow of fluids in porous media , where different energy and more complex transport equations are considered Sarma et al . ( 2007 ) . A simple modification which we use here , is to relax the constraint ρ ( 1 , x ) = pR ( x ) and to demand it holds only approximately . This formulation is better where we have noisy realizations of the distributions and a perfect fit may lead to overfitting . The formulation leads to the optimization problem min E ( v ( x , t ) ) = 1 2 ∫ Ω ( ρ ( 1 , x ) − pR ( x ) ) 2 dx+ α ∫ 1 0 ∫ Ω ρ ( x , t ) |v ( x , t ) |2 dx dt ( 2.4a ) s.t ρt +∇ · ( vρ ) = 0 ρ ( 0 , x ) = pT ( x ) ( 2.4b ) Here , the first term M = 1 2 ∫ Ω ( ρ ( 1 , x ) − pR ( x ) ) 2 dx ( 2.5 ) can be viewed as a data misfit , while the second term S = ∫ 1 0 ∫ Ω ρ ( x , t ) |v ( x , t ) |2 dxδt ( 2.6 ) can be thought of as regularization . Typical to all learning problems , the regularization parameter α needs to be chosen , usually by using some cross validation set . Such a formulation is commonly solved in applied inverse transport problems Dean & Chen ( 2011 ) . Here we see that our formulation equation 2.4a differs from both WGAN and normalized flow . It allows us to have a tradeoff between the regularity of the transformation as expressed in equation 2.6 to fitting the data as expressed in equation 2.5 . This is different from both WGAN and normalized flow when such a choice is not given .

Summary: The paper addresses the task of constructing a generative model for data using a novel optimal transport-based method. The paper proposes an alternative view of obtaining generative models by viewing the generation process as a transport problem (specifically, fluid flow mass transport) between two point clouds living in high-dimensional space. To solve the transport problem, a discretization scheme is proposed, which gives rise to a variant of point cloud registration problem, which is solved using numerical optimization. The results are provided on synthetic data and a real MNIST data.