Practical Integration via Separable Bijective Networks

1 INTRODUCTION . Most supervised learning problems operate by training a model on a finite collection , T , of N ( typically paired ) examples , ( x , y ) . The model is updated by comparing its predicted output to the expected output and performing some flavor of stochastic gradient descent based on the comparison and various regularizers . The process is repeated by reexamining the elements of T in a random order , possibly with augmentation , until the model parameters converge or an iteration budget is exceeded . This relatively simple procedure has proven to be remarkably effective in a variety of domains and these models have begun to permeate every aspect of modern science and everyday life ( Brown et al. , 2020 ; He et al. , 2016 ; Silver et al. , 2017 ) . The deep learning revolution has also resulted in highly effective generative models such as VAEs ( Kingma & Welling , 2014 ) , GANs ( Goodfellow et al. , 2014 ) , and tractable likelihood models ( Dinh et al. , 2017 ; Grathwohl et al. , 2019 ) . These models are largely used to create novel samples of impressive quality . In addition to sampling , tractable likelihood models also provide an estimate of the probability density function of the data which can be used for additional , downstream processes . We augment the training process by constructing neural networks that allow for tractable integration over the input domain . This differs from implicit layers which utilize integration over a parameterized variable ( Chen et al. , 2018 ; Grathwohl et al. , 2019 ) . Access to fast and differentiable integrals allows us to regularize a model ’ s average behavior using metrics that may not be available otherwise . Integration over the input space also allows us to supervise how the model behaves in continuous regions that are not directly observed in T and may even be out-of-distribution ( OOD ) . Alternative methods attempt to supervise examples outside of T by performing random perturbations ( Gutmann & Hyvärinen , 2010 ) , along the line between known examples ( Zhang et al. , 2018 ) , or via a generative process ( Zenati et al. , 2018 ; Akcay et al. , 2018 ) . However , these methods are only capable of observing a small quantity of the total space . By integrating over the entire regions , it is possible to observe a large portion of the space based on statistical relevance . Main Contributions We propose a general architecture that enables tractable integration over the input space , enabling supervision and custom regularization over continuous regions . We demonstrate how to construct this network and how it allows for a reduction in the computation cost required for dense numeric integration from exponential in the number of dimensions to linear . We derive several useful integrated formulations over continuous regions . Finally , we explore the impact of this architecture and regularizers on the standard accuracy and robustness to OOD examples on several standard classification datasets . Notation Throughout this work , we consider theM dimensional input features , x = [ x1 , ... , xM ] ∈ RM ; the latent features , z ∈ Z ⊆ RM ; and the K-wise classification probability , y . The input features x the training set , Tx , are drawn from the in-distribution data , D ⊆ RM . Subscripts represent a particular dimension , e.g. , Dm corresponds to the mth dimension of the space . Paranthetical superscripts represent the subspace corresponding to a particular class , e.g. , D ( c ) is the subset of D where the data belongs to class c. The bijective network is given as h such that h : D → Z . Probability distributions overD and Z are given by p with the corresponding subscript . Classification networks are given as f and g. Variables with a “ hat , ” ŷ , are predictions of the true quantity , y . 2 MOTIVATION . Neural networks are highly effective function approximators between two ( typically ) continuous spaces : f : X → Y . However , networks are typically trained and evaluated using a finite collection of points without any explicit assessment of the complete hypervolume spanned by the data . This omission is understandable from an implementation perspective as the number of samples required to obtain a reasonable estimate over a volume scales exponentially with data dimensionality . However , human beings often have an understanding of how a process should behave on average . Ideally , we would like to embed this intuition into the model but currently can not assess the average performance of a trained model outside of the held-out test set . Specifically , we would like to regularize the model by estimating the expected behavior of some metric , Ω , produced by the model over the training data Ex∼p ( x ) [ Ω ( ŷ ( x ) ) ] = ∫ X Ω ( ŷ ( x ) ) p ( x ) dx . ( 1 ) There are many useful choices of Ω over a variety of applications . If it is known what the model output should be on average ( ȳ ) , we can construct Ω to encourage that behavior , e.g. , Ω ( ŷ ) = ( ȳ− ŷ ) 2 . Minimizing consistency metrics ( Xie et al. , 2020 ) are a common method to improve learning in label-starved problems . These encourage the model to produce similar outputs over neighborhoods around ( labeled ) examples from Tx where neighborhoods are created by random or domain-specific augmentations . This process can be viewed as an approximation to an integral over the neighborhood , E ∼p ( ) [ L ( y , ŷ ( x + ) ) ] = ∫ L ( y , ŷ ( x + ) ) p ( ) d ( 2 ) where L is a distance-like metric , and is the neighborhood . Equation 2 can be generalized to other neighborhoods . We can recast the standard classification problem as a discrete approximation to an integral . Typically , we minimize the cross-entropy between ŷ ( x ; θ ) and y over the model parameters , θ , for all ( x , y ) ∈ T which becomes an integral over class-conditioned distributions , p ( x|c ) , min θ − ∑ x , y∈T ∑ k yk log ( ŷk ( x ; θ ) ) ⇒ min θ − ∑ k ∫ D ( k ) yk log ( ŷk ( x ; θ ) ) p ( x|k ) dx . ( 3 ) Unfortunately , integration in high dimension is difficult . Naive gridded solutions require an exponential number of points with error decreasing as O ( G−M ) , for G , M -dimensional points . Monte Carlo ( MC ) methods theoretically have better performance with error that decreases asO ( G−1/2 ) . However , the rate of convergence for MC methods depends on the variance of samples ( Veach , 1998 ) , which may make for poor approximations in practice . Importance sampling ( Bishop , 2006 ) can improve the performance of Monte Carlo methods to adapt to the regions with the greatest contribution . We choose to model the data using a separable function . Separable functions have the key benefit of decomposing M -dimensional integrals into a combination of M one-dimensional integrals . Each of these one-dimensional integrals can then be solved using any number of highly-accurate solvers ( e.g. , Runge-Kutta ( Shampine , 2005 ) , Dormand-Prince ( Dormand & Prince , 1980 ) , Tsit ( Tsitouras , 2011 ) , etc . ) that have error rates better than O ( G−4 ) but are unavailable in high dimensions . The use of separable functions is a component of the VEGAS algorithm ( Peter Lepage , 1978 ) and is utilized in conjunction with adaptive Monte Carlo sampling to approximate high-dimensional integrals . The use of a separable function over the input space may make estimation of integrals over the model more accessible ; however , they impose a strong , inappropriate , inductive-bias . The obvious solution of using a standard neural network feature extractor and integrating over the learned features means that we would no longer have a valid integrator over the input space . Fortunately , we can solve this problem by including a bijective transform prior to the separable network as a feature extractor . The bijective transform allows us to decouple the data into a latent space where the data can be modeled using a separable network and guarantees equality between integrals in the latent space and integrals in the input space . The bijective model has the added benefit of allowing us to create a tractable likelihood model ( Dinh et al. , 2017 ) over the data which gives us access to p ( x ) ( or possibly p ( x|y ) ) and makes it easier to draw new samples . We discuss bijective and separable functions in Sec . 3 , how they come together in Sec . 4 , and the specific choices we use in Sec . 5 . 3 BACKGROUND . We perform practical integration over the input space by breaking neural network models down into two key components : ( 1 ) a bijective feature extractor , ( 2 ) a separable task network , see Fig . 1 . For the sake of simplicity , we only consider classification tasks in this work . This makes our total network analogous with the common network breakdown where a classifier , often a linear layer or an MLP , is constructed on a feature extractor , such as a CNN . However , unlike the typical process , we must place constraints on both networks so that we can integrate over the input domain . This network breakdown is similar to hybrid networks ( Chen et al. , 2019 ; Nalisnick et al. , 2019b ) except for the separability requirement on the classifier . 3.1 BIJECTIVE NETWORKS . Bijective networks are the key component in flow-based likelihood models . A bijective network , h : D → Z , has a known forward and inverse operation so that data can exactly be reconstructed after the transformation . This allows for exact likelihood estimation via the change of variables formula : z = h ( x ; θ ) , x = h−1 ( z ; θ ) , pX ( x ) = pZ ( h ( x ; θ ) ) ∣∣∣∣∂h∂x ∣∣∣∣ ( 4 ) where pZ is a predefined distribution over the latent space , often a standard Gaussian . These models are trained via Eq . 4 to maximize the likelihood of the examples in T . Once trained , flow-based likelihood models are commonly used as a generative process where samples are drawn from pZ and are then inverted through h to arrive at an example in D. Our goal is not generative , however it is possible for us to use this process once training is complete . Instead , we will take advantage of the fact that we can choose pZ and then utilize it in downstream tasks . Normalizing flow bijectors provide rich feature extractors that can represent a distribution of complicated inputs with simply-distributed , independent features , z . Given the independence and expressibility of these learned independent features , we build estimators using separable functions over z , which enables us to integrate over the data ’ s domain while retaining expressibility . The requirement for bijectivity places a strong constraint on network design , eliminating many common choices due to the need to maintain dimension or invert element-wise activations . Even naive convolutional operations become unavailable since they are not generally invertible . Modern advances have demonstrated methods to work around these limitations through the use of clever partitioning and coupling tricks ( Dinh et al. , 2017 ) or the use of constraints ( Chen et al. , 2019 ) . However , the field of learnable bijective functions is less advanced then its unconstrained counterparts which often results in reduced performance on auxiliary tasks . Flow models are well-studied and We take advantage of several recent architectures to test our methods . In particular , we consider three different bijective models . The simplest models are constructed via stacked affine-coupling layers ( Dinh et al. , 2017 ) and invertible , fully-connected layers . We primarily utilize Glow ’ s ( Kingma & Dhariwal , 2018 ) invertible 1× 1 convolutions to process image data .

Summary: For input-output pair $(x,y)$, the paper undertakes the task of estimating (*) $E_{x \sim p(x)} \Omega(\hat y(x))$. When $x$ is high dimensional, integration is difficult. Standard ML/DL applies MC to estimate this integral based on an iid sample $(x_i,y_i), i=1,\ldots,n$. This paper suggests there is a better way.   The main idea is to recognize that if $f$ is a separable function and $h$ is a bijective function, then (**) $E_{x \sim p(x)} f(h(x))$ is easy to evaluate.