Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization

1 INTRODUCTION . Normalizing flows ( Dinh et al. , 2014 ; Rezende & Mohamed , 2015 ) have recently gathered much interest within the machine learning community , ever since its recent breakthrough in modelling high dimensional image data ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ) . They are characterized by an invertible mapping that can reshape the distribution of its input data into a simpler or more complex one . To enable efficient training , numerous tricks have been proposed to impose structural constraints on its parameterization , such that the density of the model can be tractably computed . We ask the following question : “ what is the natural way to parameterize a normalizing flow ? ” To gain a bit more intuition , we start from the one-dimension case . If a function f : R ! R is continuous , it is invertible ( injective onto its image ) if and only if it is strictly monotonic . This means that if we are only allowed to move the probability mass continuously without flipping the order of the particles , then we can only rearrange them by changing the distance in between . In this work , we seek to generalize the above intuition of monotone rearrangement in 1D . We do so by motivating the parameterization of normalizing flows from an optimal transport perspective , which allows us to define some notion of rearrangement cost ( Villani , 2008 ) . It turns out , if we want the output of a flow to follow some desired distribution , under mild regularity conditions , we can characterize the unique optimal mapping by a convex potential ( Brenier , 1991 ) . In light of this , we propose to parameterize normalizing flows by the gradient map of a ( strongly ) convex potential . Owing to this theoretical insight , the proposed method is provably universal and optimal ; this means the proposed flow family can approximate arbitrary distributions and requires the least amount of transport cost . Furthermore , the parameterization with convex potentials allows us to formulate model inversion and gradient estimation as convex optimization problems . As such , we make use of existing tools from the convex optimization literature to cheaply and efficiently estimate all quantities of interest . In terms of the benefits of parameterizing a flow as a gradient field , the convex potential is an Rd ! R function , which is different from most existing discrete-time flows which are Rd ! Rd . This makes CP-Flow relatively compact . It is also arguably easier to design a convex architecture , as we do not need to satisfy constraints such as orthogonality or Lipschitzness ; the latter two usually require a direct or an iterative reparameterization of the parameters . Finally , it is possible to incorporate additional structure such as equivariance ( Cohen & Welling , 2016 ; Zaheer et al. , 2017 ) into the flow ’ s parameterization , making CP-Flow a more flexible general purpose density model . 2 BACKGROUND : NORMALIZING FLOWS AND OPTIMAL TRANSPORT . Normalizing flows are characterized by a differentiable , invertible neural network f such that the probability density of the network ’ s output can be computed conveniently using the change-ofvariable formula pY ( f ( x ) ) = pX ( x ) @ f ( x ) @ x 1 ( ) pY ( y ) = pX ( f 1 ( y ) ) @ f 1 ( y ) @ y ( 1 ) where the Jacobian determinant term captures the local expansion or contraction of the density near x ( resp . y ) induced by the mapping f ( resp . f 1 ) , and pX is the density of a random variable X . The invertibility requirement has led to the design of many special neural network parameterizations such as triangular maps , ordinary differential equations , orthogonality or Lipschitz constraints . Universal Flows For a general learning framework to be meaningful , a model needs to be flexible enough to capture variations in the data distribution . In the context of density modeling , this corresponds to the model ’ s capability to represent arbitrary probability distributions of interest . Even though there exists a long history of literature on universal approximation capability of deep neural networks ( Cybenko , 1989 ; Lu et al. , 2017 ; Lin & Jegelka , 2018 ) , invertible neural networks generally have limited expressivity and can not approximate arbitrary functions . However , for the purpose of approximating a probability distribution , it suffices to show that the distribution induced by a normalizing flow is universal . Among many ways to establish distributional universality of flow based methods ( e.g . Huang et al . 2018 ; 2020b ; Teshima et al . 2020 ; Kong & Chaudhuri 2020 ) , one particular approach is to approximate a deterministic coupling between probability measures . Given a pair of probability densities pX and pY , a deterministic coupling is a mapping g such that g ( X ) ⇠ pY if X ⇠ pX . We seek to find a coupling that is invertible , or at least can be approximated by invertible mappings . Optimal Transport Let c ( x , y ) be a cost function . The Monge problem ( Villani , 2008 ) pertains to finding the optimal transport map g that realizes the minimal expected cost Jc ( pX , pY ) = inf eg : eg ( X ) ⇠pY EX⇠pX [ c ( X , eg ( X ) ) ] ( 2 ) When the second moments of X and Y are both finite , and X is regular enough ( e.g . having a density ) , then the special case of c ( x , y ) = ||x y||2 has an interesting solution , a celebrated theorem due to Brenier ( 1987 ; 1991 ) : Theorem 1 ( Brenier ’ s Theorem , Theorem 1.22 of Santambrogio ( 2015 ) ) . Let µ , ⌫ be probability measures with a finite second moment , and assume µ has a Lebesgue density pX . Then there exists a convex potential G such that the gradient map g : = rG ( defined up to a null set ) uniquely solves the Monge problem in eq . ( 2 ) with the quadratic cost function c ( x , y ) = ||x y||2 . Some recent works are also inspired by Brenier ’ s theorem and utilize a convex potential to parameterize a critic model , starting from Taghvaei & Jalali ( 2019 ) , and further built upon by Makkuva et al . ( 2019 ) who parameterize a generator with a convex potential and concurrently by Korotin et al . ( 2019 ) . Our work sets itself apart from these prior works in that it is entirely likelihood-based , minimizing the ( empirical ) KL divergence as opposed to an approximate optimal transport cost . 3 CONVEX POTENTIAL FLOWS . Given a strictly convex potential F , we can define an injective map ( invertible from its image ) via its gradient f = rF , since the Jacobian of f is the Hessian matrix of F , and is thus positive definite . In this section , we discuss the parameterization of the convex potential F ( 3.1 ) , and then address gradient estimation for CP-Flows ( 3.2 ) . We examine the connection to other parameterization of normalizing flows ( 3.3 ) , and finally rigorously prove universality in the next section . 3.1 MODELING . Input Convex Neural Networks We use L ( x ) to denote a linear layer , and L+ ( x ) to denote a linear layer with positive weights . We use the ( fully ) input-convex neural network ( ICNN , Amos et al . ( 2017 ) ) to parameterize the convex potential , which has the following form F ( x ) = L+ K+1 ( s ( zK ) ) + LK+1 ( x ) zk : = L + k ( s ( zk 1 ) ) + Lk ( x ) z1 : = L1 ( x ) where s is a non-decreasing , convex activation function . In this work , we use softplus-type activation functions , which is a rich family of activation functions that can be shown to uniformly approximate the ReLU activation . See Appendix B for details . Algorithm 1 Inverting CP-Flow . 1 : procedure INVERT ( F , y , CvxSolver ) 2 : Initialize x y 3 : def closure ( ) : 4 : Compute loss : l F ( x ) y > x 5 : return l 6 : x CvxSolver ( closure , x ) 7 : return x Invertibility and Inversion Procedure If the activation s is twice differentiable , then the Hessian HF is positive semi-definite . We can make it strongly convex by adding a quadratic term F↵ ( x ) = ↵ 2 ||x|| 2 2 + F ( x ) , such that HF↵ ⌫ ↵I 0 . This means the gradient map f↵ = rF↵ is injective onto its image . Furthermore , it is surjective since for any y 2 Rd , the potential x 7 ! F↵ ( x ) y > x has a unique minimizer1 satisfying the first order condition rF↵ ( x ) = y , due to the strong convexity and differentiability . We refer to this invertible mapping f↵ as the convex potential flow , or the CP-Flow . The above discussion also implies we can plug in a black-box convex solver to invert the gradient map f↵ , which we summarize in Algorithm 1 . Inverting a batch of independent inputs is as simple as summing the convex potential over all inputs : since all of the entries of the scalar l in the minibatch are independent of each other , computing the gradient all l ’ s wrt all x ’ s amounts to computing the gradient of the summation of l ’ s wrt all x ’ s . Due to the convex nature of the problem , a wide selection of algorithms can be used with convergence guarantees ( Nesterov , 1998 ) . In practice , we use the L-BFGS algorithm ( Byrd et al. , 1995 ) as our CvxSolver . 1The minimizer x⇤ corresponds to the gradient map of the convex conjugate of the potential . See Appendix A for a formal discussion . Estimating Log Probability Following equation ( 1 ) , computing the log density for CP-Flows requires taking the log determinant of a symmetric positive definite Jacobian matrix ( as it is the Hessian of the potential ) . There exists numerous works on estimating spectral densities ( e.g . TalEzer & Kosloff , 1984 ; Silver & Röder , 1994 ; Han et al. , 2018 ; Adams et al. , 2018 ) , of which this quantity is a special case . See Lin et al . ( 2016 ) for an overview of methods that only require access to Hessian-vector products . Hessian-vector products ( hvp ) are cheap to compute with reverse-mode automatic differentiation ( Baydin et al. , 2017 ) , which does not require constructing the full Hessian matrix and has the same asymptotic cost as evaluating F↵ . In particular , the log determinant can be rewritten in the form of a generalized trace tr logH . Chen et al . ( 2019a ) limit the spectral norm ( i.e . eigenvalues ) of H and directly use the Taylor expansion of the matrix logarithm . Since our H has unbounded eigenvalues , we use a more complex algorithm designed for symmetric matrices , the stochastic Lanczos quadrature ( SLQ ; Ubaru et al. , 2017 ) . At the core of SLQ is the Lanczos method , which computes m eigenvalues of H by first constructing a symmetric tridiagonal matrix T 2 Rm⇥m and computing the eigenvalues of T . The Lanczos procedure only requires Hessian-vector products , and it can be combined with a stochastic trace estimator to provide a stochastic estimate of our log probability . We chose SLQ because it has shown theoretically and empirically to have low variance ( Ubaru et al. , 2017 ) . 3.2 O ( 1 ) -MEMORY UNBIASED r log detH ESTIMATOR We would also like to have an estimator for the gradient of the log determinant to enable variants of stochastic gradient descent for optimization . Unfortunately , directly backpropagating through the log determinant estimator is not ideal . Two major drawbacks of directly differentiating through SLQ are that it requires ( i ) differentiating through an eigendecomposition routine and ( ii ) storing all Hessian-vector products in memory ( see fig . 2 ) . Problem ( i ) is more specific to SLQ , because the gradient of an eigendecomposition is not defined when the eigenvalues are not unique ( Seeger et al. , 2017 ) . Consequently , we have empirically observed that differentiating through SLQ can be unstable , frequently resulting in NaNs due to the eigendecomposition . Problem ( ii ) will hold true for other algorithms that also estimate log detH with Hessian-vector products , and generally the only difference is that a different numerical routine would need to be differentiated through . Due to these problems , we do not differentiate through SLQ , but we still use it as an efficient method for monitoring training progress . Instead , it is possible to construct an alternative formulation of the gradient as the solution of a convex optimization problem , foregoing the necessity of differentiating through an estimation routine of the log determinant . We adapt the gradient formula from Chen et al . ( 2019a , Appendix C ) to the context of convex potentials . Using Jacobi ’ s formula⇤ and the adjugate representation of the matrix inverse† , for any invertible matrix H with parameter ✓ , we have the following identity : @ @ ✓ log detH = 1detH @ @ ✓ detH ⇤ = 1detH tr adj ( H ) @ H @ ✓ † = tr H 1 @ H @ ✓ = Ev ⇥ v > H 1 @ H @ ✓ v ⇤ . ( 3 ) Notably , in the last equality , we used the Hutchinson trace estimator ( Hutchinson , 1989 ) with a Rademacher random vector v , leading to a O ( 1 ) -memory , unbiased Monte Carlo gradient estimator . Computing the quantity v > H 1 in eq . ( 3 ) by constructing and inverting the full Hessian requires d calls to an automatic differentiation routine and is too costly for our purposes . However , we can recast this quantity as the solution of a quadratic optimization problem argmin z ⇢ 1 2 z > Hz v > z ( 4 ) which has the unique minimizer z⇤ = H 1v since H is symmetric positive definite . Algorithm 2 Surrogate training objective . 1 : procedure SURROGATEOBJ ( F , x , CG ) 2 : Obtain the gradient f ( x ) , rxF ( x ) 3 : Sample Rademacher random vector r 4 : def hvp ( v ) : 5 : return v > @ @ x f ( x ) 6 : z stop gradient ( CG ( hvp , r ) ) 7 : return hvp ( z ) > r We use the conjugate gradient ( CG ) method , which is specifically designed for solving the unconstrained optimization problems in eq . ( 4 ) with symmetric positive definite H . It uses only Hessian-vector products and is straightforward to parallelize . Conjugate gradient is guaranteed to return the exact solution z⇤ within d iterations , and the error of the approximation is known to converge exponentially fast ||zm z⇤||H 2 m||z0 z⇤||H , where zm is the estimate after m iterations . The rate of convergence < 1 relates to the condition number of H . For more details , see Nocedal & Wright ( 2006 , Ch . 5 ) . In practice , we terminate CG when ||Hzm v||1 < ⌧ is satisfied for some user-controlled tolerance . Empirically , we find that stringent tolerance values are unnecessary for stochastic optimization ( see appendix F ) . Estimating the full quantity in eq . ( 3 ) is then simply a matter of computing and differentiating a scalar quantity ( a surrogate objective ) involving another Hessian-vector product : d d✓ ( zm ) > Hv , where only H is differentiated through ( since zm is only used to approximate v > H 1 as a modifier of the gradient ) . We summarize this procedure in Algorithm 2 . Similar to inversion , the hvp can also be computed in batch by summing over the data index , since all entries are independent .

Summary: This paper proposes the flow based representation of a probability distribution so that the corresponding density remains tractable. In particular, the push-forward map that generates the desired distribution is characterized by the gradient of a strongly convex potential function. The invertability of the mapping as well as the Jacobian of the mapping is hence guaranteed by such a convexity property of the potential function. The proposed CP-flows are proved to be universal density approximators and are optimal in the OT (2-Wasserstein) sense.