The Deep Weight Prior

1 INTRODUCTION . Bayesian inference is a tool that , after observing training data , allows to transforms a prior distribution over parameters of a machine learning model to a posterior distribution . Recently , stochastic variational inference ( Hoffman et al. , 2013 ) – a method for approximate Bayesian inference – has been successfully adopted to obtain a variational approximation of a posterior distribution over weights of a deep neural network ( Kingma et al. , 2015 ) . Currently , there are two major directions for the development of Bayesian deep learning . The first direction can be summarized as the improvement of approximate inference with richer variational approximations and tighter variational bounds ( Dikmen et al. , 2015 ) . The second direction is the design of probabilistic models , in particular , prior distributions , that widen the scope of applicability of the Bayesian approach . Prior distributions play an important role for sparsification ( Molchanov et al. , 2017 ; Neklyudov et al. , 2017 ) , quantization ( Ullrich et al. , 2017 ) and compression ( Louizos et al. , 2017 ; Federici et al. , 2017 ) of deep learning models . Although these prior distributions proved to be helpful , they are limited to fully-factorized structure . Thus , the often observed spatial structure of convolutional filters can not be enforced with such priors . Convolutional neural networks are an example of the model family , where a correlation of the weights plays an important role , thus it may benefit from more flexible prior distributions . Convolutional neural networks are known to learn similar convolutional kernels on different datasets from similar domains ( Sharif Razavian et al. , 2014 ; Yosinski et al. , 2014 ) . Based on this fact , within ∗Equal contribution a specific data domain , we consider a distribution of convolution kernels of trained convolutional networks . In the rest of the paper , we refer to this distribution as the source kernel distribution . Our main assumption is that within a specific domain the source kernel distribution can be efficiently approximated with convolutional kernels of models that were trained on a small subset of problems from this domain . For example , given a specific architecture , we expect that kernels of a model trained on notMNIST dataset – a dataset of grayscale images – come from the same distribution as kernels of the model trained on MNIST dataset . In this work , we propose a method that estimates the source kernel distribution in an implicit form and allows us to perform variational inference with the specific type of implicit priors . Our contributions can be summarized as follows : 1 . We propose deep weight prior , a framework that approximates the source kernel distribution and incorporates prior knowledge about the structure of convolutional filters into the prior distribution . We also propose to use an implicit form of this prior ( Section 3.1 ) . 2 . We develop a method for variational inference with the proposed type of implicit priors ( Section 3.2 ) . 3 . In experiments ( Section 4 ) , we show that variational inference with deep weight prior significantly improves classification performance upon a number of popular prior distributions in the case of limited training data . We also find that initialization of conventional convolution networks with samples from a deep weight prior leads to faster convergence and better feature extraction without training i.e. , using random weights . 2 DEEP BAYES . In Bayesian setting , after observing a dataset D = { x1 , . . . , xN } of N points , the goal is to transform our prior knowledge p ( ω ) of the unobserved distribution parameters ω to the posterior distribution p ( ω | D ) . However , computing the posterior distribution through Bayes rule p ( ω | D ) = p ( D |ω ) p ( ω ) /p ( D ) may involve computationally intractable integrals . This problem , nonetheless , can be solved approximately . Variational Inference ( Jordan et al. , 1999 ) is one of such approximation methods . It reduces the inference to an optimization problem , where we optimize parameters θ of a variational approximation qθ ( ω ) , so that KL-divergence between qθ ( ω ) and p ( ω | D ) is minimized . This divergence in practice is minimized by maximizing the variational lower bound L ( θ ) of the marginal log-likelihood of the data w.r.t parameters θ of the variational approximation qθ ( W ) . L ( θ ) = LD −DKL ( qθ ( ω ) ‖p ( ω ) ) → max θ ( 1 ) where LD = Eqθ ( ω ) log p ( D |ω ) ( 2 ) The variational lower bound L ( θ ) consists of two terms : 1 ) the ( conditional ) expected log likelihood LD , and 2 ) the regularizerDKL ( qθ ( ω ) ‖p ( ω ) ) . Since log p ( D ) = L ( θ ) +DKL ( qθ ( ω ) ‖p ( ω | D ) ) and p ( D ) does not depend on qθ ( w ) maximizing of L ( θ ) minimizes DKL ( qθ ( ω ) ‖p ( ω | D ) ) . However , in case of intractable expectations in equation 1 neither the variational lower bound L ( θ ) nor its gradients can be computed in a closed form . Recently , Kingma & Welling ( 2013 ) and Rezende et al . ( 2014 ) proposed an efficient mini-batch based approach to stochastic variational inference , so-called stochastic gradient variational Bayes or doubly stochastic variational inference . The idea behind this framework is reparamtetrization , that represents samples from a parametric distribution qθ ( ω ) as a deterministic differentiable function ω = f ( θ , ) of parameters θ and an ( auxiliary ) noise variable ∼ p ( ) . Using this trick we can efficiently compute an unbiased stochastic gradient ∇θL of the variational lower bound w.r.t the parameters of the variational approximation . Bayesian Neural Networks . The stochastic gradient variational Bayes framework has been applied to approximate posterior distributions over parameters of deep neural networks ( Kingma et al. , 2015 ) . We consider a discriminative problem , where dataset D consists of N object-label pairs D = { ( xi , yi ) } Ni=1 . For this problem we maximize the variational lower bound L ( θ ) with respect to parameters θ of a variational approximation qθ ( W ) : L ( θ ) = N∑ i=1 Eqθ ( W ) log p ( yi |xi , W ) −DKL ( qθ ( W ) ‖p ( W ) ) → max θ ( 3 ) where W denotes weights of a neural network , qθ ( W ) is a variational distribution , that allows reparametrization ( Kingma & Welling , 2013 ; Figurnov et al. , 2018 ) and p ( W ) is a prior distribution . In the simplest case qθ ( W ) can be a fully-factorized normal distribution . However , more expressive variational approximations may lead to better quality of variational inference ( Louizos & Welling , 2017 ; Yin & Zhou , 2018 ) . Typically , Bayesian neural networks use fully-factorized normal or log-uniform priors ( Kingma et al. , 2015 ; Molchanov et al. , 2017 ; Louizos & Welling , 2017 ) . Variational Auto-encoder . Stochastic gradient variational Bayes has also been applied for building generative models . The variational auto-encoder proposed by Kingma & Welling ( 2013 ) maximizes a variational lower bound L ( θ , φ ) on the marginal log-likelihood by amortized variational inference : L ( θ , φ ) = N∑ i=1 Eqθ ( zi | xi ) log pφ ( xi | zi ) −DKL ( qθ ( zi |xi ) ‖p ( zi ) ) → max θ , φ , ( 4 ) where an inference model qθ ( zi |xi ) approximates the posterior distribution over local latent variables zi , reconstruction model pφ ( xi | zi ) transforms the distribution over latent variables to a conditional distribution in object space and a prior distribution over latent variables p ( zi ) . The vanilla VAE defines qθ ( z |x ) , pφ ( x | z ) , p ( z ) as fully-factorized distributions , however , a number of richer variational approximations and prior distributions have been proposed ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ; Tomczak & Welling , 2017 ) . The approximation of the data distribution can then be defined as an intractable integral p ( x ) ≈ ∫ pφ ( x | z ) p ( z ) dz which we will refer to as an implicit distribution . 3 DEEP WEIGHT PRIOR . In this section , we introduce the deep weight prior – an expressive prior distribution that is based on generative models . This prior distribution allows us to encode and favor the structure of learned convolutional filters . We consider a neural network with L convolutional layers and denote parameters of l-th convolutional layer as wl ∈ RIl×Ol×Hl×Wl , where Il is the number of input channels , Ol is the number of output channels , Hl and Wl are spatial dimensions of kernels . Parameters of the neural network are denoted as W = ( w1 , . . . wL ) . A variational approximation qθ ( W ) and a prior distribution p ( W ) have the following factorization over layers , filters and channels : qθ ( W ) = L∏ l=1 Il∏ i=1 Ol∏ j=1 q ( wlij | θlij ) p ( W ) = L∏ l=1 Il∏ i=1 Ol∏ j=1 pl ( w l ij ) , ( 5 ) where wlij ∈ RHl×Wl is a kernel of j-th channel in i-th filter of l-th convolutional layer . We also assume that qθ ( W ) allows reparametrization . The prior distribution p ( W ) , in contrast to popular prior distributions , is not factorized over spatial dimensions of the filters Hl , Wl . For a specific data domain and architecture , we define the source kernel distribution – the distribution of trained convolutional kernels of the l-th convolutional layer . The source kernel distribution favors learned kernels , and thus it is a very natural candidate to be the prior distribution pl ( wlij ) for convolutional kernels of the l-th layer . Unfortunately , we do not have access to its probability density function ( p.d.f . ) , that is needed for most approximate inference methods e.g. , variational inference . Therefore , we assume that the p.d.f . of the source kernel distribution can be approximated using kernels of models trained on external datasets from the same domain . For example , given a specific architecture , we expect that kernels of a model trained on CIFAR-100 dataset come from the same distribution as kernels of the model trained on CIFAR-10 dataset . In other words , the p.d.f . of the source kernel distribution can be approximated using a small subset of problems from a specific data domain . In the next subsection , we propose to approximate this intractable probability density function of the source kernel distribution using the framework of generative models . Algorithm 1 Stochastic Variational Inference With Implicit Prior Distribution Require : dataset D = { ( xi , yi ) } Ni=1 Require : variational approximations q ( w | θlij ) and reverse models r ( z |w ; ψl ) Require : reconstruction models p ( w | z ; φl ) , priors for auxiliary variables pl ( z ) while not converged do M̂ ← mini-batch of objects form dataset D ŵlij ← sample weights from q ( w|θlij ) with reparametrization ẑlij ← sample auxiliary variables from r ( z | ŵlij ; ψl ) with reparametrization L̂aux ← LM̂+ ∑ l , i , j − log q ( ŵlij | θlij ) −log r ( ẑlij | ŵlij ; ψl ) +log pl ( ẑlij ) +log p ( ŵlij | ẑlij ; φl ) Obtain unbiased estimate ĝ with E [ ĝ ] = ∇Laux by differentiating L̂aux Update parameters θ and ψ using gradient ĝ and a stochastic optimization algorithm end while return Parameters θ , ψ

Summary: This paper proposes the ‘deep weight prior’: the idea is to elicit a prior on an auxiliary dataset and then use that prior over the CNN filters to jump start inference for a data set of interest.  Both explicit and implicit priors are considered, with the latter having the benefit of increased flexibility but having the drawback of a lack of a parametric form to plug in to the ELBO.  The authors address this last point by extending the ELBO appropriately.  Experiments are performed testing the prior’s ability to capture trained filters (Figure 1), provide a good initialization (Figure 2), improve sample efficiency (Figure 3), improve training speed (Figure 4).