HWA: Hyperparameters Weight Averaging in Bayesian Neural Networks
Abstract
Bayesian neural networks attempt to combine the strong predictive performance of neural networks with formal quantiﬁcation of uncertainty of the predicted output in the Bayesian framework.
In deterministic deep neural network, the conﬁdence of the model and the predictions at inference time are left alone.
Applying randomness and Bayes Rule to the weights of a deep neural network is a step towards achieving this goal.
Current state of the art optimization methods for training Bayesian Neural Networks are relatively slow and ineﬃcient, compared to their deterministic counterparts.
In this paper, we propose HWA (Hyperparameters Weight Averaging) algorithm that exploits an averaging procedure in order to optimize faster and achieve better accuracy.
We develop our main algorithm using the simple averaging heuristic and demonstrate its eﬀectiveness on the space of the hyperparameters of the networks random weights.
Numerical applications are presented to conﬁrm the empirical beneﬁts of our method.
1. Introduction
While Deep Learning methods have shown increasing eﬃciency in various domains such as natural language processing, computer vision or robotics, sensible areas including autonomous driving or medical imaging not only require accurate predictions but also uncertainty quantiﬁcation.
In (Neal, 2012), authors develop a bayesian variant of plain feedforward multilayer neural networks in which weights and biases are considered as random variables.
For supervised learning tasks, deterministic models are prone to overﬁtting and are not capable of estimating uncertainty in the training data which results in making overly conﬁdent decisions about the correct class, also known as miscalibration (Guo et al., 2017;
Kendall and Gal, 2017).
Nevertheless, representing that aforementioned uncertainty is crucial for decision making.
Bayesian methods display a hierarchical probabilistic model that assume a (prior) random distribution over the parameters of the parameters and are useful for assessing the uncertainty of the model via posterior predictive distribution quantiﬁcation (Blundell et al., 2015; Kingma et al., 2015).
Current training methods for Bayesian
Neural Networks (BNN) (Neal, 2012) include Variational Inference (Graves, 2011; Hoﬀman et al., 2013) or BayesByBackprop (Blundell et al., 2015) based on the Evidence Lower
Bound (ELBO) maximization task.
Naturally, Bayesian methods, and in particular BNNs, are highly sensitive to choice of the prior distribution parameters.
Besides, current stateof-the-art models are not as eﬃcient and robust as traditional deep learning models.
In this paper, we introduce a new optimization algorithm to alleviate those challenges.
Our main contributions read as follows:
• We introduce Hyperparameter Weight Averaging (HWA), a training algorithm that leverages stochastic averaging techniques (Polyak and Juditsky, 1992) and posterior sampling methods to eﬃciently train bayesian neural networks.
• Given the high nonconvexity of the loss landscape, our method ﬁnds heuristic explanation from theoretical works on averaging and generalization such as (Keskar et al., 2016;
He et al., 2019) and more practical work on Deep Neural Networks (DNN) optimization such as (Izmailov et al., 2018).
• We provide numerical examples showcasing the eﬀectiveness of our method on simple and complex supervised classiﬁcation tasks.
The remaining of the paper is organized as follows.
Section 2 presents the related works in the ﬁelds of optimization, Variational Inference and posterior sampling.
Section 3 introduces our main contribution, namely the HWA algorithm.
Section 4 highlights the beneﬁts of our procedure on various classiﬁcation tasks.
Section 5 concludes our work.
2. Related Work
Posterior Prediction.
Due to the nonconvexity of the loss landscapes involved in modern deep learning tasks, sampling directly from the posterior distribution of the weights is not an option.
Depending on the nature and the dimensionality of the problem, Markov
Chain Monte Carlo (MCMC) methods have been employed to overcome this intractability issue.
The samples drawn at convergence of the Markov chain are guaranteed to be drawn from the target distribution.
Hamiltonian Monte Carlo (HMC) (Neal et al., 2011) or Metropolis Hastings (MH) (Hastings, 1970) are two standard solutions used in practice.
Their stochastic gradients counterpart are extensively studied in (Ma et al., 2015).
Variational Inference (VI).
When tackling an optimization problem, exact posterior sampling may be computationally involved and not even required.
Variational inference was proposed in (Graves, 2011), in the particular case of BNNs, in order to ﬁt a Gaussian variational posterior approximation over the weights of neural networks.
Via a simple reparameterization trick (Blundell et al., 2015), several methods have emerged to train
BNNs leveraging the ease of use and implementation of VI (Kingma et al., 2015; Blundell et al., 2015; Molchanov et al., 2017).
Though, those methods appear to be ineﬃcient for large datasets and newer ones were proposed to alleviate this issue such as normalizing ﬂows (Louizos and Welling, 2017), deterministic VI (Wu et al., 2018) or dropout VI (Gal and Ghahramani, 2016).
Stochastic Averaging.
Averaging methods include the seminal papers of (Polyak, 1990) and (Ruppert, 1988), both based on the combination of past iterates along a stochastic approximation trajectory.
For nonconvex objectives, this averaging procedure has been adapted to Stochastic Gradient Descent (SGD) trajectory in (Zhou and Cong, 2017).
In particular, for recent deep learning examples, Izmailov et al.
(2018) develops a method that averages snapshots of a DNN through the iterations leading to a better empirical generalization.
Those experimental discoveries are then backed by theoretical understanding of the multilayer neural network loss landscape in (Keskar et al., 2016;
He et al., 2019).
3. Hyperparameters Averaging in Bayesian Neural Networks
In this section, we introduce the basic concepts of bayesian neural networks and their associated loss function which plays a key role in this paper.
From an optimization perspective, we ﬁrst review the Stochastic Weight Averaging (SWA) (Izmailov et al., 2018) procedure, which can be seen as an approximation of the mean trajectory of the SGD iterates and then introduce our method – namely HWA .
While SWA averages snapshots of the weights of the neural networks from successive past iterations, our method HWA only captures snapshots of the hyperparameters, and not of the weights that are sampled at each training iteration.
We then discuss the uncertainty estimation prediction of such method and how our proposed extra step, combining posterior sampling and optimization , can lead to better generalization of the trained model on test sets.
Based on the idea of ensemble learning, as in (Garipov et al., 2018), averaging procedure leads to a similar ensemble eﬀect.
Indeed, popular Bayesian methods for training BNNs tend to focus on a single mode, leading in general to mode collapse, whereas, as stated and exhaustively tested in (Fort et al., 2019), ensembles tend to explore diverse modes in function space, so are SWA and HWA.
3.1. Bayesian Neural Networks and ELBO Maximization
Let (( x i , y i ) , i ∈ [1 , n ]) be i.i.d.
input-output pairs and w ∈ W ⊆
R d be a latent variable.
When conditioned on the input data x = ( x i , i ∈ [ n ]), the joint distribution of y =
( y i , i ∈ [ n ]) and w is given by:
[Equation]
In the particular case of BNN, this likelihood function is parameterized by a multilayer neural network, which can be convolutional or not.
The latent variables w are thus the weights and the biases of the model and are considered as latent (and random) variables.
Training such hierarchical models involves sampling from the posterior distribution of the weights w conditioned on the data ( x, y ), noted p ( w | y, x ).
In most cases, this posterior distribution p ( w | y, x ) is intractable and is approximated using a family of parametric distributions, { q ( w, θ ) , θ ∈ Θ } .
The variational inference (VI) problem (Blei et al., 2017) boils down to minimizing the Kullback-Leibler (KL) divergence between q ( w, θ ) and the posterior distribution p ( w | y, x ).
The objective is the ELBO (Evidence Lower BOund) and reads:
[Equation]
Directly optimizing the objective function in (2) can be diﬃcult.
First, with n (cid:29) 1, evaluating the objective function L ( θ ) requires a full pass over the entire dataset.
Second, for some complex models, the expectations in (2) can be intractable even if we assume a simple parametric model for q ( w ; θ ).
Thus, the computation of the gradient requires an approximation step usually invoking a Monte Carlo (MC) approximation step.
Training solutions simply include using SGD (Bottou and Bousquet, 2008) where the gradient of the individual ELBO (2) is computed using Automatic Diﬀerentiation (Kucukelbir et al., 2017).
The ﬁnal update goes in the opposite direction of that gradient up to a learning rate factor.
In the sequel, we develop an improvement over baseline SGD, invoking averaging virtue of several successive snapshots of the gradients.
The method, called Hyperparameters Weight Averaging (HWA), aims at improving the generalization property of the trained model on unseen data.
3.2. Averaging model snapshots through hyperparameters loss landscapes
Consider a deterministic deep neural network, the idea behind the Stochastic Weight Averaging (SWA) procedure, developed in (Izmailov et al., 2018), is to run several iterates of the classical SGD procedure, starting from a pre-trained model.
At each timestep noted T avg , the model estimate is set to be the the average of the last T avg iterates.
After establishing the connectivity between several modes (point estimates of minimal loss) of the same deep neural network (after diﬀerent training procedures) in (Garipov et al., 2018), the ability to average over all those iterates probably traversing several models, or at least model estimates that belong to low loss region, would make the resulting trained model more robust and thus generalize better to unseen data.
Several theoretical papers such as (He et al., 2019) or (Keskar et al., 2016) provide justiﬁcations of this empirical phenomena.
– Hyperparameters Weight Averaging:
Based
on the probabilistic model devel-
oped
Section
.1, the loss function
2) is deﬁned
on
the space of the hyperparameters
i.e
the mean and
the variance
of the
variational candidate distribution.
Regarding the parameterization choice
of
the variational candidate q ( w θ ) we choose for simplicity a scalar mean
µ (cid:96) depending on the layer (cid:96) ∈
1 , L
and constant between each neuron of the same layer.
Classically the covariance
of
this variational distribution is diagonal, see (Kirkpatrick et al., 2017; Blundell et al., 2015),
yet
this assumption can be too restrictive.
We follow the direction
taken
in (Maddox et al., 2019), where the covariance of q ( w, θ ) is a diagonal matrix.
As a result, the averaging procedure practically occurs on the set of hyperparameters and requires updating the mean and the variance of the variational candidate distribution, at iteration k + 1, if k , the iteration index, is a multiple of the cycle length c , as below: µ HWA(cid:96) =
[Equation]
n m +and
[Equation]
where for all (cid:96) ∈ [1 , L ], µ k +1 (cid:96) and σ k +1 are obtained via Stochastic VI (Hoﬀman et al., 2013).
1: Input: Iteration index k .
Trained hyperparameters ˆ µ (cid:96) and ˆ σ .
LR γ k .
Cycle length c .
Gradient vector ∇L i k
( θ k ) 2: γ ← γ ( k ) (Cyclical LR for the iteration) 3: SVI updates: 4: µ k +1 (cid:96) ← µ k (cid:96) − γ k ∇ µ (cid:96) L i k ( µ k (cid:96) ) 5: σ k +1 ← σ k − γ k ∇ σ L i k ( σ k ) 6:
if mod( k, c ) = 0 then 7: n m ← k/c (Number of models to average over) µ HWA(cid:96) ← n m µ HWA(cid:96) + µ k +1 (cid:96) n m +and σ HWA ← n m σ HWA + ( µ k +1 (cid:96) )n m +− ( µ HWA(cid:96) )8: end if 9: Return: if mod ( k, c ) = 0, ( { µ HWA(cid:96) } Ll =1 , σ HWA ) else , ( { µ k(cid:96) } Ll =1 , σ k )
Our main contribution lies in Algorithm 1.
Stochastic Variational updates are executed
Line 4.
The stochastic averaging procedure happens every c iterations, and consists in computing the weighted sum between the latest model estimate and the running average noted by the superscript HWA.
Once the parameter estimates are updated via (3), the network weights are then sampled according to the updated variational candidate distributions in order to compute the next iteration approximate stochastic gradient, see Alg. 2 in supplementary material for more details on the end-to-end VI procedure embedding HWA.
Note that in the above procedure, the variational candidate q ( w, θ ) has a diagonal covariance matrix where the scalar standard deviations are obtained through Algorithm 1.
Yet, it is also possible build a non diagonal proposal covariance to bypass the restriction of such structure.
Besides, given the nonconvexity and high dimensionality of the true posterior distribution, adding a low rank non diagonal structure to the covariance of our proposal would yield a gain in eﬃciency in the VI procedure.
Of course the ideal option would be to construct a curvature-informed covariance for our proposal but at a higher cost.
The low-rank plus diagonal posterior approximation matrix, noted Σ of q ( w, θ ) introduced in (Maddox et al., 2019) reads:
[Equation]
where µ HWA = ( µ HWA(cid:96) , (cid:96) ∈ [1 , L ]), R is the maximum number of columns in the low rank deviation matrix (cid:98) D and Σ diag is the diagonal covariance deﬁned above.
The r -th component of the low rank deviation matrix (cid:98) D is deﬁned as the gap between the current estimate and the running average: (cid:98) D r = θ r − θ HWAr .
It quantiﬁes how far the current estimate parameter deviate from the current average.
Several hyperparameters are worth highlighting here.
The standard learning rate γ k plays a key role and is either a constant or cyclical, see (Zhang et al., 2019).
The cycle length c , monitoring the number of times snapshots of the model estimates are being averaged, is also of utmost importance and needs careful tuning.
4. Numerical Experiments
We provide experiments on classiﬁcation tasks with various neural network architectures and datasets to demonstrate the eﬀectiveness of our method, namely HWA.
Methods.
We consider three baselines: vanilla BayesByBackprop (BBB) developed in (Blundell et al., 2015), the Stochastic Gradient Langevin Dynamics (SGLD) method introduced in (Welling and Teh, 2011) and its cyclical variant in (Zhang et al., 2019).
The algorithms are initialized at the same point and the results are averaged over 5 repetitions.
Datasets.
We compare the diﬀerent algorithms on MNIST (LeCun, 1998) and CIFAR10 (Krizhevsky and Hinton, 2009) datasets.
Network architectures.
(MNIST)
We train a Bayesian variant of LeNet-5 convolutional neural network (LeCun et al., 1998) on the MNIST dataset.
Under the prior distribution π , see (1), the weights are assumed independent and identically distributed according to N (0 , 1).
We also assume a Gaussian variational candidate distribution such that q ( · ; θ ) ≡ N ( µ, σ 2 I ), where I is the identity matrix.
The variational posterior parameters are thus θ =
( µ, σ ) where µ = ( µ (cid:96) , (cid:96) ∈ [ d ]) with d the number of weights in the neural network.
(CIFAR-10)
We train the Bayesian variant of the VGG neural network (Simonyan and Zisserman, 2014) on the CIFAR-10 dataset.
As in the previous example, the weights are assumed i.i.d.
according to N (0 , I ).
Standard hyperparameters values found in the literature, such as the annealing constant or the number of MC samples, were used for the benchmark methods.
For better eﬃciency and lower variance, the Flipout estimator (Wen et al., 2018) is used to compute the MC approximation of the gradient of the loss function.
Results.
Results for both datasets and network architectures are reported Figure 1.
While for the MNIST dataset, the runs for HWA and SGLD are comparable both in terms of train and testing loss and accuracy, they both highlights better convergence properties compared to BayesByBackprop (BBB).
It is worth mentioning that our method HWA displays a similar behavior as a gradient based method, such as SGLD , by only leveraging the average of past snapshots of the variational candidate hyperparameters.
Regarding the
CIFAR10 experiment, our method shows the lowest training loss and generalize better to unseen data (cf.
last ﬁgure on bottom line in Figure 1).
In conclusion, HWA achieves stateof-the-art results for either small or large bayesian variants of standard network architectures while using a simple and eﬃcient averaging update at each cycle .
[Figure]
5. Conclusion
We present in this paper an averaging procedure on the hyperparameters of the weights of a bayesian neural network architecture.
Based on both empirical and theoretical results regarding stochastic averaging, we propose the HWA algorithm in order to increase the generalization ability of a BNN.
The procedure is easily implementable on top of any vanilla optimizer with standard design choices for prior and candidate distributions, crucial quantities in variational inference.
Numerical experiments show the advantage of our method matching and sometimes surpassing baselines such as SGLD or CSGMCMC, which require additional expensive gradient computation.
Appendix for HWA: Hyperparameters Weight Averaging in Bayesian Neural Networks A Comparison with other classical averaging procedures in nonconvex optimization
From Alg. 1, we note that the averaging procedure happens once at each cycle c , a tuning hyperparameter, on the parameter estimates resulting from a simple stochastic gradient descent update.
Yet, another natural averaging step would be to keep K snapshots of the past stochastic gradients and compute an aggregated sum used as the drift term in the general update rule, see (Zhou and Cong, 2017).
Nevertheless, in our setting, the objective function while being nonconvex is (possibly) parameterized by a high dimensional neural network making it computationally involved to store those K gradients.
Incremental Aggregated Gradients methods: Popular optimization methods, such as SAG (Schmidt et al., 2017) or SAGA (Defazio et al., 2014), make use of the past individual gradient and compute a moving average of those vectors as the ﬁnal drift term.
Those methods are proven to be faster than plain SGD in both convex and nonconvex cases, leveraging among other reasons variance reduction eﬀect, but suﬀer from a high storage cost.
Indeed the drift term is composed of the sum of the n past individual gradient where n is equal to the size of the training set.
MISO (Mairal, 2015): Another important method invoking variance reduction through incremental update of the drift term in a gradient descent step is the Minimization by
Incremental Surrogate Optimization method, namely MISO, developed in Mairal (2015).
Contrary to the method mentionned above, the accumulation does not happen on the gradient but on the sum of individual surrogate objective functions.
While this framework is more general than SAG or SAGA, and also does not require storing n past gradients, it is still computationally heavy to store those n past objective functions, rather their model parameter estimates, when tackling deep neural networks training.
For all those reasons, HWA surely combines the virtue of the accumulation/aggregation eﬀect and the low computing cost of vanilla SGD.
B Embedding HWA in Variational Inference
Note that in the above procedure, the variational candidate q ( w, θ ) has a diagonal covariance matrix where the scalar standard deviations are obtained through Alg. 1.
One parameter estimates are updated via (3), the neural network weights are then sampled according to the updated variational candidate distributions in order to compute the next iteration approximate stochastic gradient ∇L i k ( θ ).
Yet it is also possible build a non diagonal proposal covariance to bypass the restriction of such structure.
Besides, given the nonconvexity and high dimensionality of the true posterior distribution, adding even a low rank non diagonal structure to the covariance of our proposal would yield a gain in eﬃciency in the variational inference procedure.
Of course the ideal option would be to build a posterior curvature-informed covariance but at a higher cost.
The trade-oﬀ between computational costs and proposal eﬃciency is detailed in the following.
The low-rank plus diagonal posterior approximation matrix, noted Σ of q ( w, θ ) introduced in (Maddox et al.,
2019) reads:
[Equation]
where µ HWA = ( µ HWA(cid:96) , (cid:96) ∈ [1 , L ]), R is the maximum number of columns in the low rank deviation matrix (cid:98) D and Σ diag is the diagonal covariance deﬁned above.
The r -th component, where r ∈ [ R ], of the low rank deviation matrix (cid:98) D is deﬁned as the gap between the current estimate and the running average: (cid:98) D r = θ r − θ HWAr .
It quantiﬁes how far the current estimate parameter deviate from the current averaged parameter.
Then the covariance of the proposal q ( · ) in Alg.
2 is either set to (3) or (4).
Discussion on the choice of the variational candidate distribution: The general aim of the update rules presented above is to construct an eﬃcient variational candidate distribution that would provide an approximate shape of the true posterior.
Our method acts on the mean and covariance of a simple Gaussian distribution where the covariance matrix is either diagonal or low rank.
Nevertheless, other choice of proposal can be employed such as the spike and slab variational distribution, in (Gal and Ghahramani, 2016), leveraging dropout mechanism in VI.
The other similar idea, namely concrete dropout in (Gal et al., 2017) not only optimizes the hyperparameters of the weights but also the dropout probabilities.
We do not consider those variants as our work focuses on Gaussian approximations of the posterior distribution and how their parameters are updated, see Section 4 for a description of baseline methods used in our numerical experiments.
We now give in Alg. 2, the overall training algorithm of the bayesian neural network using the proposed HWA algorithm to update the parameters.
Cycle length c
K iterations do
[Equation]
6: Update the vector of parameter estimates calling
Alg
[Equation]
HWA
[Equation]
7: end for 8: