Implicit Normalizing Flows

1 INTRODUCTION . Normalizing flows ( NFs ) ( Rezende & Mohamed , 2015 ; Dinh et al. , 2014 ) are promising methods for density modeling . NFs define a model distribution px ( x ) by specifying an invertible transformation f ( x ) from x to another random variable z . By change-of-variable formula , the model density is ln px ( x ) = ln pz ( f ( x ) ) + ln |det ( Jf ( x ) ) | , ( 1 ) where pz ( z ) follows a simple distribution , such as Gaussian . NFs are particularly attractive due to their tractability , i.e. , the model density px ( x ) can be directly evaluated as Eqn . ( 1 ) . To achieve such tractability , NF models should satisfy two requirements : ( i ) the mapping between x and z is invertible ; ( ii ) the log-determinant of the Jacobian Jf ( x ) is tractable . Searching for rich model families that satisfy these tractability constraints is crucial for the advance of normalizing flow research . For the second requirement , earlier works such as inverse autoregressive flow ( Kingma et al. , 2016 ) and RealNVP ( Dinh et al. , 2017 ) restrict the model family to those with triangular Jacobian matrices . More recently , there emerge some free-form Jacobian approaches , such as Residual Flows ( ResFlows ) ( Behrmann et al. , 2019 ; Chen et al. , 2019 ) . They relax the triangular Jacobian constraint by utilizing a stochastic estimator of the log-determinant , enriching the model family . However , the Lipschitz constant of each transformation block is constrained for invertibility . In general , this is not preferable because mapping a simple prior distribution to a potentially complex data distribution may require a transformation with a very large Lipschitz constant ( See Fig . 3 for a 2D example ) . Moreover , all the aforementioned methods assume that there exists an explicit forward mapping z = f ( x ) . Bijections with explicit forward mapping only covers a fraction of the broad class of invertible functions suggested by the first requirement , which may limit the model capacity . In this paper , we propose implicit flows ( ImpFlows ) to generalize NFs , allowing the transformation to be implicitly defined by an equation F ( z , x ) = 0 . Given x ( or z ) , the other variable can be computed by an implicit root-finding procedure z = RootFind ( F ( · , x ) ) . An explicit mapping z = f ( x ) used in prior NFs can viewed as a special case of ImpFlows in the form of F ( z , x ) = f ( x ) − z = ∗Corresponding Author . 0 . To balance between expressiveness and tractability , we present a specific from of ImpFlows , where each block is the composition of a ResFlow block and the inverse of another ResFlow block . We theoretically study the model capacity of ResFlows and ImpFlows in the function space . We show that the function family of single-block ImpFlows is strictly richer than that of two-block ResFlows by relaxing the Lipschitz constraints . Furthermore , for any ResFlow with a fixed number of blocks , there exists some invertible function that ResFlow has non-negligible approximation error , but ImpFlow can exactly model . On the practical side , we develop a scalable algorithm to estimate the probability density and its gradients , and draw samples from ImpFlows . The algorithm leverages the implicit differentiation formula . Despite being more powerful , the gradient computation of ImpFlow is mostly similar with that of ResFlows , except some additional overhead on root finding . We test the effectiveness of ImpFlow on several classification and generative modeling tasks . ImpFlow outperforms ResFlow on all the benchmarks , with comparable model sizes and computational cost . 2 RELATED WORK . Expressive Normalizing Flows There are many works focusing on improving the capacity of NFs . For example , Dinh et al . ( 2014 ; 2017 ) ; Kingma & Dhariwal ( 2018 ) ; Ho et al . ( 2019 ) ; Song et al . ( 2019 ) ; Hoogeboom et al . ( 2019 ) ; De Cao et al . ( 2020 ) ; Durkan et al . ( 2019 ) design dedicated model architectures with tractable Jacobian . More recently , Grathwohl et al . ( 2019 ) ; Behrmann et al . ( 2019 ) ; Chen et al . ( 2019 ) propose NFs with free-form Jacobian , which approximate the determinant with stochastic estimators . In parallel with architecture design , Chen et al . ( 2020 ) ; Huang et al . ( 2020 ) ; Cornish et al . ( 2020 ) ; Nielsen et al . ( 2020 ) improve the capacity of NFs by operating in a higher-dimensional space . As mentioned in the introduction , all these existing works adopt explicit forward mappings , which is only a subset of the broad class of invertible functions . In contrast , the implicit function family we consider is richer . While we primarily discuss the implicit generalization of ResFlows ( Chen et al. , 2019 ) in this paper , the general idea of utilizing implicit invertible functions could be potentially applied to other models as well . Finally , Zhang et al . ( 2020 ) formally prove that the model capacity of ResFlows is restricted by the dimension of the residual blocks . In comparison , we study another limitation of ResFlows in terms of the bounded Lipschitz constant , and compare the function family of ResFlows and ImpFlows with a comparable depth . Continuous Time Flows ( CTFs ) ( Chen et al. , 2018b ; Grathwohl et al. , 2019 ; Chen et al. , 2018a ) are flexible alternative to discrete time flows for generative modeling . They typically treat the invertible transformation as a dynamical system , which is approximately simulated by ordinary differential equation ( ODE ) solvers . In contrast , the implicit function family considered in this paper does not contain differential equations , and only requires fixed point solvers . Moreover , the theoretical guarantee is different . While CTFs typically study the universal approximation capacity under the continuous time case ( i.e. , “ infinite depth ” limit ) , we consider the model capacity of ImpFlows and ResFlows under a finite number of transformation steps . Finally , while CTFs are flexible , their learning is challenging due to instability ( Liu et al. , 2020 ; Massaroli et al. , 2020 ) and exceedingly many ODE solver steps ( Finlay et al. , 2020 ) , making their large-scale application still an open problem . Implicit Deep Learning Utilizing implicit functions enhances the flexibility of neural networks , enabling the design of network layers in a problem-specific way . For instance , Bai et al . ( 2019 ) propose a deep equilibrium model as a compact replacement of recurrent networks ; Amos & Kolter ( 2017 ) generalize each layer to solve an optimization problem ; Wang et al . ( 2019 ) integrate logical reasoning into neural networks ; Reshniak & Webster ( 2019 ) utilize the implicit Euler method to improve the stability of both forward and backward processes for residual blocks ; and Sitzmann et al . ( 2020 ) incorporate periodic functions for representation learning . Different from these works , which consider implicit functions as a replacement to feed-forward networks , we develop invertible implicit functions for normalizing flows , discuss the conditions of the existence of such functions , and theoretically study the model capacity of our proposed ImpFlow in the function space . 3 IMPLICIT NORMALIZING FLOWS . We now present implicit normalizing flows , by starting with a brief overview of existing work . 3.1 NORMALIZING FLOWS . As shown in Eqn . ( 1 ) , a normalizing flow f : x 7→ z is an invertible function that defines a probability distribution with the change-of-variable formula . The modeling capacity of normalizing flows depends on the expressiveness of the invertible function f . Residual flows ( ResFlows ) ( Chen et al. , 2019 ; Behrmann et al. , 2019 ) are a particular powerful class of NFs due to their free-form Jacobian . ResFlows use f = fL ◦ · · · ◦ f1 to construct the invertible mapping , where each layer fl is an invertible residual network with Lipschitz constraints bounded by a fixed constant κ : fl ( x ) = x+ gl ( x ) , Lip ( gl ) ≤ κ < 1 , ( 2 ) where Lip ( g ) is the Lipschitz constant of a function g ( see Sec . 4.1 for details ) . Despite their free-form Jacobian , the model capacity of ResFlows is still limited by the Lipschitz constant of the invertible function . The Lipschitz constant of each ResFlow block fl can not exceed 2 ( Behrmann et al. , 2019 ) , so the Lipschitz constant of an L-block ResFlow can not exceed 2L . However , to transfer a simple prior distribution to a potentially complex data distribution , the Lipschitz constant of the transformation can be required to be sufficiently large in general . Therefore , ResFlows can be undesirably deep simply to meet the Lipschitz constraints ( see Fig . 3 for a 2D example ) . Below , we present implicit flows ( ImpFlows ) to relax the Lipschitz constraints . 3.2 MODEL SPECIFICATION . In general , an implicit flow ( ImpFlow ) is defined as an invertible mapping between random variables x and z of dimension d by finding the roots of F ( z , x ) = 0 , where F is a function from R2d to Rd . In particular , the explicit mappings z = f ( x ) used in prior flow instances ( Chen et al. , 2019 ; Kingma & Dhariwal , 2018 ) can be expressed as an implicit function in the form F ( z , x ) = f ( x ) − z = 0 . While ImpFlows are a powerful family to explore , generally they are not guaranteed to satisfy the invertibility and the tractability of the log-determinant as required by NFs . In this paper , we focus on the following specific form , which achieves a good balance between expressiveness and tractability , and leave other possibilities for future studies . Definition 1 . Let gz : Rd → Rd and gx : Rd → Rd be two functions such that Lip ( gx ) < 1 and Lip ( gz ) < 1 , where Lip ( g ) is the Lipschitz constant of a function g. A specific form of ImpFlows is defined by F ( z , x ) = 0 , where F ( z , x ) = gx ( x ) − gz ( z ) + x− z . ( 3 ) The root pairs of Eqn . ( 3 ) form a subset in Rd × Rd , which actually defines the assignment rule of a unique invertible function f . To see this , for any x0 , according to Definition 1 , we can construct a contraction hx0 ( z ) = F ( z , x0 ) + z with a unique fixed point , which corresponds to a unique root ( w.r.t . z ) of F ( z , x0 ) = 0 , denoted by f ( x0 ) . Similarly , in the reverse process , given a z0 , the root ( w.r.t . x ) of F ( z0 , x ) = 0 also exists and is unique , denoted by f−1 ( z0 ) . These two properties are sufficient to ensure the existence and the invertibility of f , as summarized in the following theorem . Theorem 1 . Eqn . ( 3 ) defines a unique mapping f : Rd → Rd , z = f ( x ) , and f is invertible . See proof in Appendix A.1 . Theorem 1 characterizes the validness of the ImpFlows introduced in Definition 1 . In fact , a single ImpFlow is a stack of a single ResFlow and the inverse of another single ResFlow , which will be formally stated in Sec 4 . We will investigate the expressiveness of the function family of the ImpFlows in Sec 4 , and present a scalable algorithm to learn a deep generative model built upon ImpFlows in Sec . 5 .

Summary: This work is about a new architecture for normalizing flows inspired by implicit neural networks. In particular, the authors show that a specific implicit neural network build from residual blocks defines a bijective map. From this insight, they show how to efficiently train such architecture by estimating the log Jacobian and solving the inverse problem related to the implicit architecture. On top of this achievement, the authors thoroughly study the attractiveness of their architecture compared to i-res-flow. They demonstrate analytically and empirically that implicit normalizing flows are strictly more expressive than i-res-flow.