FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models

Reversible generative models map points from a simple distribution to a complex distribution through an easily invertible neural network . Likelihood-based training of these models requires restricting their architectures to allow cheap computation of Jacobian determinants . Alternatively , the Jacobian trace can be used if the transformation is specified by an ordinary differential equation . In this paper , we use Hutchinson ’ s trace estimator to give a scalable unbiased estimate of the logdensity . The result is a continuous-time invertible generative model with unbiased density estimation and one-pass sampling , while allowing unrestricted neural network architectures . We demonstrate our approach on high-dimensional density estimation , image generation , and variational inference , improving the state-ofthe-art among exact likelihood methods with efficient sampling . 1 INTRODUCTION Reversible generative models use cheaply invertible neural networks to transform samples from a fixed base distribution . Examples include NICE ( Dinh et al. , 2014 ) , Real NVP ( Dinh et al. , 2017 ) , and Glow ( Kingma & Dhariwal , 2018 ) . These models are easy to sample from , and can be trained by maximum likelihood using the change of variables formula . However , this requires placing awkward restrictions on their architectures , such as partitioning dimensions or using rank one weight matrices , in order to avoid an O ( D3 ) cost determinant computation . Recently , Chen et al . ( 2018 ) introduced continuous normalizing flows ( CNF ) , defining the mapping from latent variables to data using ordinary differential equations ( ODE ) . In their model , the likelihood can be computed using trace operations costing only O ( D2 ) . This allows a more flexible , but still restricted , family of network architectures to be used . Extending this work , we introduce an unbiased stochastic estimator of the likelihood that has O ( D ) time cost , allowing completely unrestricted architectures . Furthermore , we have implemented GPU-based adaptive ODE solvers to train and evaluate these models on modern hardware . We call our approach Free-form Jacobian of Reversible Dynamics ( FFJORD ) . Figure 1 shows FFJORD smoothly transforming a Gaussian distribution into a multi-modal distribution . ∗Equal contribution . Order determined by coin toss . { wgrathwohl , rtqichen } @ cs.toronto.edu †University of Toronto and Vector Institute . ‡OpenAI . 2 BACKGROUND : GENERATIVE MODELS AND CHANGE OF VARIABLES . In contrast to directly parameterizing a normalized distribution ( e.g . Oord et al . ( 2016 ) ; Germain et al . ( 2015 ) ) , the change of variables formula allows one to specify a complex normalized distribution px ( x ) implicitly by warping a normalized base distribution pz ( z ) through an invertible function f : RD → RD . Given a random variable z ∼ pz ( z ) the log density of x = f ( z ) follows log px ( x ) = log pz ( z ) − log det ∣∣∣∣∂f ( z ) ∂z ∣∣∣∣ ( 1 ) where ∂f ( z ) /∂z is the Jacobian of f . In general , computing the log determinant has a time cost of O ( D3 ) . Much work has gone into developing restricted neural network architectures which make computing the Jacobian ’ s determinant more tractable . These approaches broadly fall into three categories : Normalizing flows . By restricting the functional form of f , various determinant identities can be exploited ( Rezende & Mohamed , 2015 ; Berg et al. , 2018 ) . These models can not be trained as generative models from data because they do not have a tractable inverse f−1 . However , they are useful for specifying approximate posteriors for variational inference ( Kingma & Welling , 2014 ) . Autoregressive transformations . By using an autoregressive model and specifying an ordering of the dimensions , the Jacobian of f is enforced to be lower triangular ( Kingma et al. , 2016 ; Oliva et al. , 2018 ) . These models excel at density estimation for tabular datasets ( Papamakarios et al. , 2017 ) , but require D sequential evaluations of f to invert , which is prohibitive when D is large . Partitioned transformations . Partitioning the dimensions and using affine transformations makes the determinant of the Jacobian cheap to compute , and the inverse f−1 computable with the same cost as f ( Dinh et al. , 2014 ; 2017 ) . This method allows the use of convolutional architectures , excelling at density estimation for image data ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ) . Throughout this work , we refer to reversible generative models as those which use the change of variables to transform a base distribution to the model distribution while maintaining both efficient density estimation and efficient sampling capabilities using a single pass of the model . 2.1 OTHER GENERATIVE MODELS . There exist several approaches to generative modeling approaches which do not use the change of variables equation for training . Generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) use large , unrestricted neural networks to transform samples from a fixed base distribution . Lacking a closed-form likelihood , an auxiliary discriminator model must be trained to estimate divergences or density ratios in order to provide a training signal . Autoregressive models ( Germain et al. , 2015 ; Oord et al. , 2016 ) directly specify the joint distribution p ( x ) as a sequence of explicit conditional distributions using the product rule . These models require at least O ( D ) evaluations to sample from . Variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ) use an unrestricted architecture to explicitly specify the conditional likelihood p ( x|z ) , but can only efficiently provide a stochastic lower bound on the marginal likelihood p ( x ) . 2.2 CONTINUOUS NORMALIZING FLOWS . Chen et al . ( 2018 ) define a generative model for data x ∈ RD similar to those based on ( 1 ) , but replace the warping function with an integral of continuous-time dynamics . The generative process first samples from a base distribution z0 ∼ pz0 ( z0 ) . Then , given an ODE whose dynamics are defined by the parametric function ∂z ( t ) /∂t = f ( z ( t ) , t ; θ ) , we solve the initial value problem with z ( t0 ) = z0 to obtain a data sample x = z ( t1 ) . These models are called Continous Normalizing Flows ( CNF ) . The change in log-density under this model follows a second differential equation , called the instantaneous change of variables formula ( Chen et al. , 2018 ) : ∂ log p ( z ( t ) ) ∂t = −Tr ( ∂f ∂z ( t ) ) . ( 2 ) We can compute total change in log-density by integrating across time : log p ( z ( t1 ) ) = log p ( z ( t0 ) ) − ∫ t1 t0 Tr ( ∂f ∂z ( t ) ) dt . ( 3 ) Given a datapoint x , we can compute both the point z0 which generates x , as well as log p ( x ) under the model by solving the combined initial value problem : [ z0 log p ( x ) − log pz0 ( z0 ) ] ︸ ︷︷ ︸ solutions = [ x 0 ] + ∫ t0 t1 [ f ( z ( t ) , t ; θ ) Tr ( ∂f ∂z ( t ) ) ] dt︸ ︷︷ ︸ dynamics , [ z ( t1 ) log p ( x ) − log p ( z ( t1 ) ) ] = [ x 0 ] ︸ ︷︷ ︸ initial values ( 4 ) which integrates the combined dynamics of z ( t ) and the log-density of the sample backwards in time from t1 to t0 . We can then compute log p ( x ) using the solution of ( 4 ) and adding log pz0 ( z0 ) . The existence and uniqueness of ( 4 ) require that f and its first derivatives be Lipschitz continuous ( Khalil , 2002 ) , which can be satisfied in practice using neural networks with smooth Lipschitz activations , such as softplus or tanh . 2.2.1 BACKPROPAGATING THROUGH ODE SOLUTIONS WITH THE ADJOINT METHOD . CNFs are trained to maximize ( 3 ) . This objective involves the solution to an initial value problem with dynamics parameterized by θ . For any scalar loss function which operates on the solution to an initial value problem L ( z ( t1 ) ) = L ( ∫ t1 t0 f ( z ( t ) , t ; θ ) dt ) ( 5 ) then Pontryagin ( 1962 ) shows that its derivative takes the form of another initial value problem dL dθ = − ∫ t0 t1 ( ∂L ∂z ( t ) ) T ∂f ( z ( t ) , t ; θ ) ∂θ dt . ( 6 ) The quantity −∂L/∂z ( t ) is known as the adjoint state of the ODE . Chen et al . ( 2018 ) use a black-box ODE solver to compute z ( t1 ) , and then a separate call to a solver to compute ( 6 ) with the initial value ∂L/∂z ( t1 ) . This approach is a continuous-time analog to the backpropgation algorithm ( Rumelhart et al. , 1986 ; Andersson , 2013 ) and can be combined with gradient-based optimization to fit the parameters θ by maximum likelihood . 3 SCALABLE DENSITY EVALUATION WITH UNRESTRICTED ARCHITECTURES . Switching from discrete-time dynamics to continuous-time dynamics reduces the primary computational bottleneck of normalizing flows fromO ( D3 ) toO ( D2 ) , at the cost of introducing a numerical ODE solver . This allows the use of more expressive architectures . For example , each layer of the original normalizing flows model of Rezende & Mohamed ( 2015 ) is a one-layer neural network with only a single hidden unit . In contrast , the instantaneous transformation used in planar continuous normalizing flows ( Chen et al. , 2018 ) is a one-layer neural network with many hidden units . In this section , we construct an unbiased estimate of the log-density with O ( D ) cost , allowing completely unrestricted neural network architectures to be used . 3.1 UNBIASED LINEAR-TIME LOG-DENSITY ESTIMATION . In general , computing Tr ( ∂f/∂z ( t ) ) exactly costs O ( D2 ) , or approximately the same cost as D evaluations of f , since each entry of the diagonal of the Jacobian requires computing a separate derivative of f ( Griewank & Walther , 2008 ) . However , there are two tricks that can help . First , vector-Jacobian products vT ∂f∂z can be computed for approximately the same cost as evaluating f using reverse-mode automatic differentiation . Second , we can get an unbiased estimate of the trace of a matrix by taking a double product of that matrix with a noise vector : Tr ( A ) = Ep ( ) [ TA ] . ( 7 ) The above equation holds for any D-by-D matrix A and distribution p ( ) over D-dimensional vectors such that E [ ] = 0 and Cov ( ) = I . The Monte Carlo estimator derived from ( 7 ) is known as Hutchinson ’ s trace estimator ( Hutchinson , 1989 ; Adams et al. , 2018 ) . To keep the dynamics deterministic within each call to the ODE solver , we can use a fixed noise vector for the duration of each solve without introducing bias : log p ( z ( t1 ) ) = log p ( z ( t0 ) ) − ∫ t1 t0 Tr ( ∂f ∂z ( t ) ) dt = log p ( z ( t0 ) ) − ∫ t1 t0 Ep ( ) [ T ∂f ∂z ( t ) ] dt = log p ( z ( t0 ) ) − Ep ( ) [ ∫ t1 t0 T ∂f ∂z ( t ) dt ] ( 8 ) Typical choices of p ( ) are a standard Gaussian or Rademacher distribution ( Hutchinson , 1989 ) .

Summary: This paper discusses an advance in the framework of normalizing flows for generative modeling, named FFJORD. The authors consider normalizing flows in the form of ordinary differential equations, as also discussed in [1]. Their contributions are two-fold: (1) they use an unbiased estimator of the likelihood of the model by approximating the trace of the jacobian with Hutchinson’s trace estimator, (2) they have implemented the required ODE solvers on GPUs.