Reducing the Computational Cost of Deep Generative Models with Binary Neural Networks

1 INTRODUCTION . As machine learning models continue to grow in number of parameters , there is a corresponding effort to try and reduce the ever-increasing memory and computational requirements that these models incur . One method to make models more efficient is to use neural networks with weights and possibly activations restricted to be binary-valued ( Courbariaux et al. , 2015 ; Courbariaux et al. , 2016 ; Rastegari et al. , 2016 ; McDonnell , 2018 ; Gu et al. , 2018 ) . Binary weights and activations require significantly less memory , and also admit faster low-level implementations of key operations such as linear transformations than when using the usual floating-point precision . Although the application of binary neural networks for classification is relatively well-studied , there has been no research that we are aware of that has examined whether binary neural networks can be used effectively in unsupervised learning problems . Indeed , many of the deep generative models that are popular for unsupervised learning do have high parameter counts and are computationally expensive ( Vaswani et al. , 2017 ; Maaløe et al. , 2019 ; Ho et al. , 2019a ) . These models would stand to benefit significantly from converting the weights and activations to binary values , which we call binarization for brevity . In this work we focus on non-autoregressive models with explicit densities . One such class of density model is the variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) , a latent variable model which has been used to model many high-dimensional data domains accurately . The state-of-the-art VAE models tend to have deep hierarchies of latent layers , and have demonstrated good performance relative to comparable modelling approaches ( Ranganath et al. , 2016 ; Kingma et al. , 2016 ; Maaløe et al. , 2019 ) . Whilst this deep hierarchy makes the model powerful , the model size and compute requirements increases with the number of latent layers , making very deep models resource intensive . Another class of density model which has shown promising results are flow-based generative models ( Dinh et al. , 2014 ; Rezende & Mohamed , 2015 ; Dinh et al. , 2017 ) . These models perform a series of invertible transformations to a simple density , with the transformed density approximating the data-generating distribution . Flow models which achieve state-of-the-art performance compose many transformations to give flexibility to the learned density ( Kingma & Dhariwal , 2018 ; Ho et al. , 2019a ) . Again the model computational cost increases as the number of transformations increases . To examine how to binarize hierarchical VAEs and flow models successfully , we take two models which have demonstrated excellent modelling performance - the ResNet VAE ( Kingma et al. , 2016 ) and the Flow++ model ( Ho et al. , 2019a ) - and implement the majority of each model with binary neural networks . Using binary weights and activations reduces the computational cost , but also decreases the representational capability of the model . Therefore our aim is to strike a balance between reducing the computational cost and maintaining good modelling performance . We show that it is possible to decrease the model size drastically , and allow for significant speed ups in run time , with only a minor impact on the achieved loss value . We make the following key contributions : • We propose an efficient binary adaptation of weight normalization , a reparameterization technique often used in deep generative models to accelerate convergence . Binary weight normalization is the generative-modelling alternative to the usual batch normalization used in binary neural networks . • We show that we can binarize the majority of weights and activations in deep hierarchical VAE and flow models , without significantly hurting performance . We demonstrate the corresponding binary architecture designs for both the ResNet VAE and the Flow++ model . • We perform experiments on different levels of binarization , clearly demonstrating the trade-off between binarization and performance . 2 BACKGROUND . In this section we give background on the implementation and training of binary neural networks . We also describe the generative models that we implement with binary neural networks in detail . 2.1 BINARY NEURAL NETWORKS . In order to reduce the memory and computational requirements of neural networks , there has been recent research into how to effectively utilise networks which use binary-valued weights wB and possibly also activations αB rather than the usual real-valued1 weights and activations ( Courbariaux et al. , 2015 ; Courbariaux et al. , 2016 ; Rastegari et al. , 2016 ; McDonnell , 2018 ; Gu et al. , 2018 ) . In this work , we use the convention of binary values being in B : = { −1 , 1 } . Motivation . The primary motivation for using binary neural networks is to decrease the memory and computational requirements of the model . Clearly binary weights require less memory to be stored : 32× less than the usual 32-bit floating-point weights . Binary neural networks also admit significant speed-ups . A reported 2× speed-up can be achieved by a layer with binary weights and real-valued inputs ( Rastegari et al. , 2016 ) . This can be made an additional 29× faster if the inputs to the layer are also constrained to be binary ( Rastegari et al. , 2016 ) . With both binary weights and inputs , linear operators such as convolutions can be implemented using the inexpensive XNOR and bit-count binary operations . A simple way to ensure binary inputs to a layer is to have a binary activation function before the layer ( Courbariaux et al. , 2016 ; Rastegari et al. , 2016 ) . Optimization . Taking a trained model with real-valued weights and binarizing the weights has been shown to be lead to significant worsening of performance ( Alizadeh et al. , 2019 ) . So instead the binary weights are optimized . It is common to not optimize the binary weights directly , but instead optimize a set of underlying real-valued weights wR which can then be binarized in some fashion for inference . In this paper we will adopt the convention of binarizing the underlying weights using the sign function ( see Equation 2 ) . We also use the sign function as the activation function when we use binary activations ( see Equation 5 , where αR are the real-valued pre-activations ) . We define the sign 1We use real-valued throughout the paper to be synonymous with `` implemented with floating-point precision '' . function as : sign ( x ) : = { −1 , if x < 0 1 , if x ≥ 0 ( 1 ) Since the derivative of the sign function is zero almost everywhere2 , the gradients of the underlying weights wR and through binary activations are zero almost everywhere . This makes gradient-based optimization challenging . To overcome this issue , the straight-through estimator ( STE ) ( Bengio et al. , 2013 ) can be used . When computing the gradient of the loss L , the STE replaces the gradient of the sign function ( or other discrete output functions ) with an approximate surrogate . A straightforward and widely used surrogate gradient is the identity function , which we use to calculate the gradients of the real-valued weights wR ( see Equation 3 ) . It has been shown useful to cancel the gradients when their magnitude becomes too large ( Courbariaux et al. , 2015 ; Alizadeh et al. , 2019 ) . Therefore we use a clipped identity function for the gradients of the pre-activations ( see Equation 6 ) . This avoids saturating a binary activation . Lastly , the loss value only depends on the sign of the real-valued weights . Therefore , the values of the weights are generally clipped to be in [ −1 , 1 ] after each gradient update ( see Equation 4 ) . This restricts the magnitude of the weights and thus makes it easier to flip the sign . Forward pass : Backward pass : After update : Weights wB = sign ( wR ) ( 2 ) ∂L ∂wR : = ∂L ∂wB ( 3 ) wR ← max ( −1 , min ( 1 , wR ) ) ( 4 ) Activations αB = sign ( αR ) ( 5 ) ∂L ∂αR : = ∂L ∂αB ∗ 1|αR|≤1 ( 6 ) − 2.2 DEEP GENERATIVE MODELS . Hierarchical VAEs . The variational autoencoder ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) is a latent variable model for observed data x conditioned on unobserved latent variables z . It consists of a generative model pθ ( x , z ) and an inference model qφ ( z|x ) . The generative model can be decomposed into the prior on the latent variables pθ ( z ) and the likelihood of our data given the latent variables pθ ( x|z ) . The inference model is a variational approximation to the true posterior , since the true posterior is usually intractable in models of interest . Training is generally performed by maximization of the evidence lower bound ( ELBO ) , a lower bound on the log-likelihood of the data : log pθ ( x ) ≥ Eqφ ( z|x ) [ log pθ ( x , z ) − log qφ ( z|x ) ] ( 7 ) To give a more expressive model , the latent space can be structured into a hierarchy of latent variables z1 : L. In the generative model each latent layer is conditioned on deeper latents pθ ( zi|zi+1 : L ) . A common problem with hierarchical VAEs is that the deeper latents can struggle to learn , often `` collapsing '' such that the layer posterior matches the prior : qφ ( zi|zi+1 : L , x ) ≈ pθ ( zi|zi+1 : L ) 3 . One method to help prevent posterior collapse is to use skip connections between latent layers ( Kingma et al. , 2016 ; Maaløe et al. , 2019 ) , turning the layers into residual layers ( He et al. , 2016 ) . We focus on the ResNet VAE ( RVAE ) model ( Kingma et al. , 2016 ) . In this model , both the generative and inference model structure their layers as residual layers . The ResNet VAE uses a bi-directional inference structure with both a bottom-up and top-down residual channel . This is a similar structure to the BIVA model ( Maaløe et al. , 2019 ) , which has demonstrated state-of-the-art results for a latent variable model . We give a more detailed description of the model in Appendix C. Flow models . Flow models consist of a parameterized invertible transformation , z = fθ ( x ) , and a known density pz ( z ) usually taken to be a unit normal distribution . Given observed data x we obtain the objective for θ by applying a change-of-variables to the log-likelihood : log pθ ( x ) = log pz ( fθ ( x ) ) + log ∣∣∣∣det dfθdx ∣∣∣∣ ( 8 ) For training to be possible , it is required that computation of the Jacobian determinant det ( dfθ/dx ) is tractable . We therefore aim to specify of flow model fθ which is sufficiently flexible to model the 2Apart from at 0 , where it is non-differentiable . 3We have assumed here that the inference model is factored `` top-down '' . data distribution well , whilst also being invertible and having a tractable Jacobian determinant . One common approach is to construct fθ as a composition of many simpler functions : fθ = f1 ◦ f2 ◦ ... ◦ fL , with each fi invertible and with tractable Jacobian . So the objective becomes : log pθ ( x ) = log pz ( fθ ( x ) ) + L∑ i=1 log ∣∣∣∣det dfidfi−1 ∣∣∣∣ ( 9 ) There are many approaches to construct the fi layers ( Dinh et al. , 2014 ; Rezende & Mohamed , 2015 ; Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ; Ho et al. , 2019a ) . In this work we will focus on the Flow++ model ( Ho et al. , 2019a ) , which has state-of-the-art results for flow models . In the Flow++ model , the fi are coupling layers which partition the input into x1 and x2 , then transform only x2 : fi ( x1 ) = x1 , fi ( x2 ) = σ −1 ( MixLogCDF ( x2 ; t ( x1 ) ) ) · exp ( a ( x1 ) ) + b ( x1 ) ( 10 ) Where MixLogCDF is the CDF for a mixture of logistic distributions . This is an iteration on the affine coupling layer ( Dinh et al. , 2014 ; 2017 ) . Note that keeping part of the input fixed ensures that the layer is invertible . To ensure that all dimensions are transformed in the composition , adjacent coupling layers will keep different parts of the input fixed , often using an alternating checkerboard or stripe pattern to choose the fixed dimensions ( Dinh et al. , 2017 ) . The majority of parameters in this flow model come from the functions t , a and b in the coupling layer , and in Flow++ these are parameterized as stacks of convolutional residual layers . In this work we will focus on how to binarize these functions whilst maintaining good modelling performance . We give a more detailed description of the full flow model in Appendix D .

Summary: The paper introduces the first way (to the best of authors knowledge) of building generative models with binary weights. Also the case with binary activations is considered. The authors consider two SOTA generative models (flow++ and RVAE) and develop technique to binarize all weights (and possibly activcations) in residual layers. They show that residual layers can be binarized with relatively small drop in performance and further binarization of remaining blocks in computational graph leads to significant degradation. Binary modification of weight normalization is suggested although no ablation is study is performed so it is unclear how crucial is BWN for robust learning. The training process itself is pretty standard way of training binary DNNs - they use STE + truncation of real-valued weights counter-parts.