Marginal Likelihood Gradient for Bayesian Neural Networks
Abstract
Bayesian learning of neural networks is attractive as it can protecting against over-ﬁtting and provide automatic methods for inferring important hyperparameters by maximizing the marginal probability of the data.
However, existing approaches in this vein, such as those based on variational inference, do not perform well.
In this paper, we take a diﬀerent approach and directly derive a practical estimator of the gradient of the marginal loglikelihood for BNNs by combining local reparametrization of the network w.r.t.
the prior distribution with the self-normalized importance sampling estimator.
We show promising preliminary results on a toy example and on vectorized MNIST classiﬁcation where the new method results in signiﬁcantly improved performance of variational inference compared to relying on predeﬁned hyperparameters.
1. Introduction
We consider a supervised learning task given a data set consisting of input-output pairs D =
( xxx i ,yyy i ) } Ni =1 .
We denote the observation model for data as p ( yyy | xxx,θθθ,βββ ) deﬁning the likelihood p ( D| θθθ,βββ ) =
(cid:81)
Ni =1 p
( yyy i |
xxx i ,θθθ,βββ ) and prior p ( θθθ | ααα ), where ααα and βββ denote hyperparameters.
Marginal likelihood of data is defnied as p ( D| ααα,βββ )
=
p ( { yyy i } Ni =1 |{ xxx i } Ni =1 ,ααα,βββ ).
Given a new data point ( xxx ∗ ,yyy ∗ ), a fully Bayesian approach to modeling makes a prediction by marginalizing out posterior over parameters θθθ and hyperparameters ααα,βββ (MacKay, 1992a):
[Equation]
The integrals in Equation (1) inevitably carry high computational burden for non-trivial models, requiring us to recourse to approximate inference.
We subsequently discuss approximations to posterior distributions p ( θθθ |D ,ααα,βββ ) and p ( ααα,βββ |D ).
Variational inference (VI).
Variational inference (Barber and Bishop, 1998; Hinton and van Camp, 1993) reformulates ﬁnding an approximation to a posterior p ( θθθ |D ,ααα,βββ ) into an optimization problem ﬁnding a a simpler distribution q ( θθθ | λλλ ) approximating p ( θθθ |D ).
This is done by minimizing the divergence D KL (cid:0) q ( θθθ | [Equation]
) || p ( θθθ |D ,ααα,βββ ) (cid:1) , which is equivalent to maximizing the Evidence Lower BOund (ELBO), a lower bound to log p ( D| ααα,βββ ), deﬁned as:
[Equation]
Evidence framework.
The evidence framework (MacKay, 1992a,b) provides a systematic framework for adaptation of hyperparameters ααα,βββ .
Equation (1) indicates that p ( ααα,βββ |D ), (and the marginal likelihood p ( D| ααα,βββ ) if the prior over ααα,βββ is uniform) embodies the inﬂuence of predictions arising from hyperparameters ααα,βββ on marginalized test predictions
p
( yyy ∗ | xxx ∗ , D ).
If the integral in Equation (1) has no closed form, it can be approximated as p ( yyy ∗ | xxx ∗ , D )
≈ (cid:82) p ( yyy ∗ |
xxx ∗ ,θθθ ) p ( θθθ |D ,ααα ∗ ,βββ ∗ ) dθθθ where ααα ∗ ,βββ ∗ = argmax ααα,βββ p ( ααα,βββ |D ).
Relying only on the mode ααα ∗ ,βββ ∗ can lead to poor approximation of p ( yyy ∗ | xxx ∗ , D ), depending on the concentration of density p ( ααα,βββ |D ).
Estimating marginal log-likelihood log p ( D| ααα,βββ ) for non trivial models is a well-known problem (Barber and Bishop, 1998).
Poorly chosen prior can result in ELBO not correlating with the predictive performance of the learned posterior (Trippe and Turner, 2018; Turner and Sahani, 2011).
Bayesian Neural Networks.
We consider a neural network F WWW represented as composition of L fully-connected layers interleaved with activation functions F WWW =
F WWW L ◦ φ ◦ F WWW L − 1 . . .
φ ◦ F WWW 1 where φ denotes an activation function.
We denote a forward pass with input xxx l through l -th layer as F WWW l ( xxx l ) =
WWW l xxx l where biases are handled by expanding the dimension of xxx l .
We restrict the attention to Gaussian mean-ﬁeld priors p ( WWW l | µµµ ααα ,σσσ 2 ααα ) and mean-ﬁeld approximate posteriors q ( WWW l | µµµ l ,σσσ 2 l ) independent over layers (see Foong et al. for discussion of limitations of these) and use the notation WWW l ∼
N ( µµµ l ,σσσ 2 l ) meaning w l,ij ∼ N ( µ l,ij , σ 2 l,ij ).
We adopt the notation WWW =
( WWW 1 ,WWW 2 , . . .
,WWW l ) for weights, µµµ =
( µµµ 1 ,µµµ 2 , . .
. ,µµµ l ), σσσ = ( σσσ 21 ,σσσ 22 , . .
. ,σσσ 2 L ) for variational parameters and µµµ ααα =
( µµµ ααα 1 ,µµµ ααα 2 , . . .
,µµµ ααα L ), σσσ 2 ααα
=
( σσσ 2 ααα 1 ,σσσ 2 ααα 1 , . .
,σσσ 2 ααα L ) for prior hyperarameters.
Hyperparameter βββ belongs to an observation model, for instance, the scale of Gaussian observation noise
N ( yyy | 000 , diag[ βββ 2 ]).
Local Reparametrization Trick (LRT).
The forward pass F WWW l ( xxx l ) through a fullyconnected layer with mean-ﬁeld Gaussian distribution over weights WWW l ∼ N ( µµµ l ,σσσ 2 l ), can be reparametrized as (Kingma et al., 2015; Tomczak et al., 2020):
[Equation]
where √ is applied element-wise and the Equation (3) denotes the equality between two random variables (LHS and RHS follow the same distribution).
We denote the reparametrization noise as (cid:15)(cid:15)(cid:15) =
( (cid:15)(cid:15)(cid:15) 1 ,(cid:15)(cid:15)(cid:15) 2 , . .
,(cid:15)(cid:15)(cid:15) L ) and reparametrization of the network f WWW resulting from applying Equation (3) to every layer l by g µµµ,σσσ 2 ,(cid:15)(cid:15)(cid:15) . Then we have that E q ( WWW | µµµ,σσσ
) h ( f WWW ( xxx )) =
E (cid:15)(cid:15)(cid:15) ∼N (000 ,III ) h ( g µµµ,σσσ 2 ,(cid:15)(cid:15)(cid:15) ( xxx )) for a continuous function h .
Approaches to adjust priors for BNNs. Hyperparameters ααα ∗ ,βββ ∗ can be also learned by optimizing L ( λλλ ; ααα,βββ ) as proxy for marginal log-likelihood since log p ( D ; ααα,βββ ) ≥ L ( λλλ ; ααα,βββ ) (MacKay, 1999; Wu et al., 2019). The local optima of log p ( D ; ααα,βββ ) and L ( λλλ ; ααα,βββ ) can signiﬁcantly vary. This is manifested when a ﬂexible prior distribution p ( θθθ | ααα ) learned by optimizing L ( λλλ ; ααα,βββ ) collapses to the approximate posterior q ( θθθ | λλλ ), resulting in no regularization (Blundell et al., 2015).
Laplace approximation assumes log p ( D ,θθθ | ααα,βββ ) is a quadratic w.r.t.
θθθ enabling approximating log p ( D| ααα,βββ ) in closed form (MacKay, 1992c), which can be used to compare diﬀerent hyperparameters ααα,βββ .
Hyperparameters ααα,βββ can be also chosen using cross-validation (Farquhar et al., 2020; Tomczak et al., 2018), but this requires increased computational eﬀort, limiting the practicability of this approach.
First, we discuss the problems with the naive application of self-normalized importance sampling (SNIS) to estimate the gradient of marginal log-likelihood of data ∇ ααα log p ( D| ααα,βββ ) for BNNs. We then demonstrate how taking advantage of local reparametrization of the network allows us to mitigate these problems and derive a practical estimator of the gradient of ∇ ααα log p ( D| ααα,βββ ) given a BNN with mean-ﬁeld Gaussian prior.
Recall that the gradient ∇ ααα log p ( D| ααα,βββ ) can be derived as (Bornschein and Bengio, 2014; Tucker et al., 2019):
[Equation]
See Appendix A for similar expression for ∇ βββ log p ( D| ααα,βββ ) and associated derivations.
A naive attempt to apply SNIS to Equation (4) results in estimator:
[Equation]
[Equation]
where w ( WWW ( s ) )
= p ( D| WWW ( s ) ,βββ ) p ( WWW ( s ) | ααα ) /q
( WWW ( s ) ) and weights are proposed from WWW ( s ) ∼ q ( WWW ).
There are several challenges associated with applying the estimator given by Equation (5) making it impractical to use with BNNs.
First, sampling multiple weights WWW ( s ) is excessively costly even for networks of moderate size as this corresponds to sampling separate copies of the network f WWW .
Second, RHS of Equation (5) has high variance being a score gradient estimator (Williams, 1992).
Third, a separate objective needs to be optimized to derive proposal q ( WWW | ααα ), e.g. by minimizing divergence between q ( WWW ) and p ( WWW |D ,ααα,βββ ).
Lastly, calculating weights w ( WWW ( s ) ) requires iterating over whole data set D .
We now demonstrate that reparametrizing the network using LRT leads to the estimator free of the two ﬁrst described issues.
More precisely, rather than operating in the space of random weights WWW l , we can reparametrize the network using LRT and operate in the space of preactivation reparametrization noise (cid:15)(cid:15)(cid:15) l ∼ N (000 ,III ).
This reduces the cost of sampling, as noise (cid:15)(cid:15)(cid:15) l has signiﬁcantly lower dimensionality than weight matrices WWW l .
In addition, the derived estimator is an improvement over one given by Equation (5) in a similar way as reparametrizing (Kingma and Welling, 2014; Rezende et al., 2014; Kingma et al., 2015) improves upon the score function estimator applied to the ELBO (Paisley et al., 2012;
Ranganath et al., 2014; Titsias and L´azaro-Gredilla, 2014), since path derivative usually exhibits lower variance than score function gradient (Roeder et al., 2017).
We subsequently discuss how to derive better proposal distributions, reduce the cost of computing IS weights, and derive an approximation using one optimization objective in Appendix A. We now state
Lemma deriving the reparametrized gradient of marginal log-likelihood of data log p ( D| ααα,βββ ).
Lemma 1 (Marginal log-likelihood gradient for BNN with MF Gaussian prior)
Given a neural network f WWW with layer-wise mean-ﬁeld Gaussian prior N ( WWW l | µµµ ααα l ,σσσ 2 ααα l ) where ααα,βββ are hyperparameters, the marginal log-likelihood gradient can be computed as:
[Equation]
where p ( (cid:15)(cid:15)(cid:15) |D ,ααα,βββ ) is a posterior distribution over layer’s preactivations resulting from applying reparametrization in Equation (3) to layers of a network f WWW .
Substituting ∇ ααα with ∇ βββ in Equation (6) yields the gradient for hyperparameters of observation model ∇ βββ log p ( D| ααα,βββ ) .
Note that Equation (6) still requires sampling from posterior distribution p ( (cid:15)(cid:15)(cid:15) |D ,ααα,βββ ), where (cid:15)(cid:15)(cid:15)
=
( (cid:15)(cid:15)(cid:15) 1 ,(cid:15)(cid:15)(cid:15) 2 , . .
,(cid:15)(cid:15)(cid:15) l ) denotes the local reparametrization noise for consecutive layers F WWW l .
While analyzing the posterior distribution over reparametrization noise might be uncommon, it is well-deﬁned: for ﬁxed ααα,βββ and for every (cid:15)(cid:15)(cid:15) we deﬁne the likelihood of data (cid:81)
Ni =1 p
( yyy i | g µµµ ααα ,σσσ 2 ααα ,(cid:15)(cid:15)(cid:15) ( xxx ) ,βββ ) and hence we can calculate posterior distribution over reparametrization noise p ( (cid:15)(cid:15)(cid:15) |D ,ααα,βββ ) given prior p ( (cid:15)(cid:15)(cid:15) ).
The gradient in the RHS of Equation (6) resembles gradients derived for generative models (Bornschein and Bengio, 2014;
Burda et al., 2015; Tucker et al., 2019). However, it is important to note not all derivations for generative models transfer to BNNs if local reparametrization is desired, as LRT requires the expression under the expectation to be only a function of layer’s outputs {F ( xxx l ) } l .
For instance, the expectation E q ( WWW ) p ( yyy | xxx, f WWW ) p ( WWW ) /q
( WWW ) cannot be locally reparametrized due to the dependency on p ( WWW ) /q ( WWW ). Applying SNIS to Equation (6) by proposing from distribution q ( (cid:15)(cid:15)(cid:15) ) gives us the following estimator of ∇ ααα log p ( D| µµµ ααα ,σσσ 2 ααα ):
[Equation]
where w ( (cid:15)(cid:15)(cid:15) ( i,s ) )
= (cid:81) Ni =1 p
( yyy i | g µµµ ααα ,
[Equation]
[Equation]
,(cid:15)(cid:15)(cid:15) ( i
[Equation]
s ) ( xxx ) ,βββ ) p ( (cid:15)(cid:15)(cid:15) ( i,s ) ) /q
( (cid:15)(cid:15)(cid:15) ( i,s ) ) are normalized IS weights.
We will employ the estimator of ∇ ααα log p ( D| µµµ ααα ,σσσ 2 ααα ) in Equation (7) to optimize marginal loglikelihood of data log p ( D| µµµ ααα ,σσσ 2 ααα ) w.r.t.
prior means µµµ ααα and variances σσσ 2 ααα .
Better proposal distributions.
We now focus on deriving better proposal distributions q ( (cid:15)(cid:15)(cid:15) ) for estimator in Equation (7).
This requires ﬁnding a distribution over noise (cid:15)(cid:15)(cid:15) closer to posterior distribution p ( (cid:15)(cid:15)(cid:15) |D ,ααα,βββ ).
We deﬁne means and variances for Gaussian distributions over layer’s preactivations as µµµ (cid:15)(cid:15)(cid:15) =
( µµµ 1 ,(cid:15)(cid:15)(cid:15) ,µµµ 2 ,(cid:15)(cid:15)(cid:15) , . .
,µµµ L,(cid:15)(cid:15)(cid:15) )
and σσσ 2 (cid:15)(cid:15)(cid:15) =
( σσσ 21 ,(cid:15)(cid:15)(cid:15) ,σσσ 22 ,(cid:15)(cid:15)(cid:15) , . .
. ,σσσ 2 L,(cid:15)(cid:15)(cid:15) ).
We assume factors q ( (cid:15)(cid:15)(cid:15) l | µµµ l,(cid:15)(cid:15)(cid:15) ,σσσ 2 l,(cid:15)(cid:15)(cid:15) ) to be mean-ﬁeld Gaussian distributions and construct q ( (cid:15)(cid:15)(cid:15) | µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) by assuming independence between layers q ( (cid:15)(cid:15)(cid:15) | µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) =
(cid:81)
Ll q ( (cid:15)(cid:15)(cid:15) l
| µµµ (cid:15)(cid:15)(cid:15) l ,σσσ 2 (cid:15)(cid:15)(cid:15) l ).
We learn the approximate posterior distribution q ( (cid:15)(cid:15)(cid:15) | µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) by minimizing D KL (cid:0) q ( (cid:15)(cid:15)(cid:15) | µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) )
|| p ( (cid:15)(cid:15)(cid:15) |D ,µµµ ααα ,σσσ 2 ααα ) (cid:1) leading to ELBO-like expression:
[Equation]
The optimal posterior distribution q ∗ ( (cid:15)(cid:15)(cid:15) | µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) derived by optimizing the L ( µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) depends on the prior hyperparameters µµµ ααα ,σσσ 2 ααα , so in practice we will optimize L ( µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) by performing one (or possibly a few) updates while we simultaneously learning µµµ ααα ,σσσ 2 ααα .
This setup of learning jointly learning model parameters µµµ ααα ,σσσ 2 ααα and auxiliary parameters µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) resembles algorithms developed for generative models (Kingma and Welling, 2014; Bornschein and Bengio, 2014; Burda et al., 2015; Rezende et al., 2014).
Computing weights w ( (cid:15)(cid:15)(cid:15) ) .
The cost of calculating weights w ( (cid:15)(cid:15)(cid:15) ) grows linearly with the size of data set D .
To ease the computational cost of iterating over whole data set D to calculate w ( (cid:15)(cid:15)(cid:15) ), we propose to employ the “average” likelihood calculated based on subsample set of data points, similarly as Li and Turner (2016).
While the theoretical properties of this approximation are not very well studied, it appears to be a sensible heuristic.
Preventing over-ﬁtting.
Adjusting hyperparameters to data by optimizing log p ( D| ααα,βββ ) (type II MLE) can be prone to over-ﬁtting.
The standard approach to mitigate over-ﬁtting hyperparameters is to introduce hyperpriors p ( ααα ) and p ( βββ ).
In addition, the variances of units σσσ 2 l can be shared across the layer l , so that σσσ 2 l = σ 2 l 1 .
The gradients of log hyperprior ∇ ααα log p ( ααα ) and ∇ βββ log p ( βββ ) are then added to the Equation (7).
3. Experimental demonstration
We will now empirically investigate the diﬀerences between learning hyperparameters ααα,βββ via optimizing log p ( D| ααα,βββ ) and L ( µµµ,σσσ 2 ; µµµ ααα ,σσσ 2 ααα ).
We demonstrate that relying on ﬁxed hyperparameters and optimizing ELBO w.r.t.
hyperparameters ααα,βββ leads to underconﬁdent and overconﬁdent predictions, respectively.
Employing the derived estimator of ∇ ααα log p ( D| µµµ ααα ,σσσ 2 ααα ) allows us to mitigate this behavior and leads to ααα yielding better generalization.
In our experiments we assume layer-wise Gaussian hyperprior over shared (within every layer) prior log variances p (log σ 2 α l ) =
N (log σ 2 α l | log[50 / ( N in + 1)] , γ 2 l ), where N in is the size of the layer’s input
and we set γ 2 l = 500.
[Figure]
Fixed
hyperparameters       0.
( b ) optimizing ELBO       0.
p r
i o r      3 0.
p o s t e r i o r ( c ) Log marginal gradient One dimensional regression.
We ﬁrst consider one dimensional regression where we try to learn function f ( x ) =
( x + 1) sin(2 x + 2)
+ 0 .
3 (cid:15), (cid:15) ∼ N (0 , 1) after observing onlydata points.
We learn a variational posterior for two hidden layer network with 128 hidden units, tanh activations and assume Gaussian observation model log p ( y | f WWW ( x ) , β ) ∝ − ( y − f WWW ( x )) 2 / 2 β 2 − log β .
We use ADAM optimizer (Kingma and Ba, 2014) to optimize L ( µµµ,σσσ 2 ; µµµ ααα ,σσσ 2 ααα ).
We consider three approaches: (a) use ﬁxed value of observation noise β = 1, zero prior unit means µµµ and ﬁxed variances σ 2 α = 1 /N in (Neal, 1996), (b) optimize L ( µµµ,σσσ 2 ; µµµ ααα ,σσσ 2 ααα , β ) w.r.t.
both variational parameters µµµ,σσσ 2 , and prior hyperparameters µµµ ααα ,σσσ 2 ααα , β , and (c) learn µµµ ααα ,σσσ 2 ααα , β
using evidence gradient given by Equation (7) and then learn variational posterior q ( µµµ,σσσ 2 ) by optimizing L ( µµµ,σσσ 2 ; µµµ ααα ,σσσ 2 ααα , β ).
We report prior and posterior predictive distribution in Figure 1.
First, in Figure 1( a ) we observe that MF-VI with predeﬁned variances of units can severely underﬁt data if prior variances of the units are not tuned.
This happens when the optimization of L ( µµµ,σσσ 2 ; µµµ ααα ,σσσ 2 ααα , β ) is driven by reducing complexity penalty D KL (cid:0) q ( θθθ | λλλ ) || p ( θθθ | ααα ) (cid:1) and explains data D by noise.
Next, in Figure 1( b ) we see that learning prior hyperparameters by optimizing L ( µµµ,σσσ 2 ; µµµ ααα ,σσσ 2 ααα , β ) leads to an improved performance, but results in overconﬁdent predictive distribution.
This happens as optimizing prior p ( WWW | µµµ ααα ,σσσ 2 ααα ) using Equation (2) can set it very close to posterior q ( WWW | µµµ,σσσ 2 ), yielding almost no regularization.
Lastly, we see from Figure 1( c ) that using the derived marginal log-likelihood gradient results in a prior explaining the data D with subsequent application of VI yielding reasonable performance.
Vectorized MNIST classiﬁcation.
Next we investigate if adjusting hyperparameters leads to improved performance of Gaussian MF-VI on a vectorized MNIST classiﬁcation task, a standard benchmark for Bayesian deep learning (Blundell et al., 2015).
We learn a two hidden layer network with 400 hidden units and ReLU activations (Nair and Hinton, 2010) and report learning curves as optimization proceeds in Figure 2.
We use the same hyperpriors for all approaches and compare using predeﬁned hyperparameters with learning them via optimizing ELBO and using the derived marginal log-likelihood gradient.
Overall, we observe similar patterns as in 1D regression experiment.
Using predeﬁned set of hyperparameters yields underﬁtting with test NLL 0 .
095 and test error rate 2 .
49%, as optimizing ELBO constrains q ( WWW | µµµ,σσσ 2 ) to prior distribution incapable of modeling data.
Tuning prior by optimizing ELBO leads to over-ﬁtting (test NLL 0 .
177 and test error rate 1 . 57%), as prior distribution p ( WWW | µµµ ααα ,σσσ 2 ααα ) collapse to posterior q ( WWW | µµµ,σσσ 2 ) leading to insufﬁcient regularization resulting in over-ﬁtting.
Using marginal log-likelihood gradient given by Equation (6) optimizes prior p ( WWW | µµµ ααα ,σσσ 2 ααα ) directly to ﬁt the data and hence introduces suﬃcient regularization when optimizing ELBO, leading to the best test log likelihood 0 .and
test error rate 1 .
[Figure]
4. Conclusions
We have derived a pragmatic algorithm adjusting hyperparameters of BNNs with meanﬁeld Gaussian priors by optimizing marginal log-likelihood of data.
To this end, we applied local reparametrization of the network combined with self-normalized importance sampling.
The optimized priors yielded predictive distributions explaining the data and subsequent application of VI resulted in ﬁtting posterior distributions with good predictive performance.
We hope our work will catalyze research on beneﬁts of reparametrization and algorithms inferring hyperparameters for BNNs.
Appendix A.
We provide derivations of ∇ ααα log p ( D| ααα,βββ ) and ∇ βββ log p ( D| ααα,βββ ) in Equation (6)
[Equation]
We now present the proof of Lemma 1.
The derivation uses the reparametrization of the network f WWW into g µµµ,σσσ 2 ,(cid:15)(cid:15)(cid:15) together with applying log trick ∇ f ( x ) =
f ( x ) ∇ log f ( x ).
Proof Derivation of Lemma 1.
[Equation]
[Equation]
Following the same derivations for ∇ βββ yields
[Equation]
Approximate scheme weights tribution be used to formulate the objective be
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Equation]
Alternatively, an approximation
[Equation]
(cid:88)
[Equation]
34) where (cid:15)(cid:15)(cid:15) ( s ) ∼ q ( (cid:15)(cid:15)(cid:15) | µµµ (cid:15)(cid:15)(cid:15) ,σσσ 2 (cid:15)(cid:15)(cid:15) ) can be used
[Equation]