Gradient Regularization as Approximate Variational Inference

Abstract

We developed Variational Laplace for Bayesian neural networks (BNNs), which exploits a local approximation of the curvature of the likelihood to estimate the ELBO without the need for stochastic sampling of the neural-network weights. The Variational Laplace objective is simple to evaluate, as it is the log-likelihood plus weight-decay, plus a squared-gradient regularizer. Variational Laplace gave better test performance and expected calibration errors than maximum a posteriori inference and standard sampling-based variational inference, despite using the same variational approximate posterior. Finally, we emphasize the care needed in benchmarking standard VI, as there is a risk of stopping before the variance parameters have converged. We show that early-stopping can be avoided by increasing the learning rate for the variance parameters.

1. Introduction

Neural networks are increasingly being used in safety-critical settings, such as self-driving cars [1] and medical diagnosis [2]. In these settings, it is critical to be able to reason about uncertainty in the parameters of the network—for instance, so that the system is able to call for additional human input when necessary [3]. Several approaches to Bayesian inference in neural networks are available, including stochastic gradient Langevin dynamics [4] Laplace’s method [5,6,7] and variational inference [8,9].

Here, we focus on combining the advantages of Laplace’s method [5,6,7] and variational inference (VI; [10]). In particular, Laplace’s method is very fast, as it begins by finding a mode using a standard gradient descent procedure, and then computes a local Gaussian approximate of the mode by performing a second-order Taylor expansion. However, as the mode is discovered by standard gradient descent, it may be a narrow mode that generalizes poorly [11]. In contrast, variational inference (VI; [8]) is slower, as it requires stochastic sampling of the weights, but that stochastic sampling forces it to find a broad, flat mode that presumably generalizes better.

We developed a new Variational Laplace (VL) method that combines the best of both worlds, giving a method that finds broad, flat modes even in the absence of the stochastic sampling. The resulting objective is composed of the log-likelihood, standard weight-decay regularization and a squared-gradient regularizer, which is weighted by the variance of the approximate posterior. VL displayed improved performance over VI and MAP in standard benchmark tasks.

Our squared gradient regularizer relates strongly to the effectiveness of stochastic gradient descent. In particular, recent work has shown that gradient descent implicitly uses a squared gradient regularizer [12], and that full-batch gradient descent with a squared gradient regularizer can recover much of the benefits of implicit regularization from minibatched stochastic gradient descent [13]. Our work implies that these regularizers can be interpreted as a form of approximate inference over the neural-network weights.

2. Background

2.1. Variational Inference (VI) for Bayesian Neural Networks

To perform variational inference for neural networks, we follow the usual approach [8,14], by using independent Gaussian priors, P and approximate posteriors Q for all parameters, $\mathbf{w}$:

$$\begin{array}{cc}\hfill P_{}\left(w_{\lambda }\right)& =\mathcal{N}\left(w_{\lambda };0,s_{\lambda }^2\right)\hfill \end{array}$$ (1)

$$\begin{array}{cccccc}\hfill Q_{}\left(w_{\lambda }\right)& =\mathcal{N}\left(w_{\lambda };{\mathbf{\mu }}_{\lambda },{\sigma }_{\lambda }^2\right)\hfill & \hfill & \mathrm{equivalently}\hfill & \hfill Q_{}\left(\mathbf{w}\right)& =\mathcal{N}\left(\mathbf{w};\mathbf{\mu },\mathbf{\Sigma }\right),\hfill \end{array}$$ (2)

where ${\mathbf{\mu }}_{\lambda }$ and ${\sigma }_{\lambda }^2$ are learned parameters of the approximate posterior, and where $\mathrm{\Sigma }$ is a diagonal matrix, with ${\mathrm{\Sigma }}_{\lambda \lambda }={\sigma }_{\lambda }^2$. We fit the approximate posterior by optimizing the evidence lower bound objective (ELBO) with respect to parameters of the variational posterior, ${\mathbf{\mu }}_{\lambda }$ and ${\sigma }_{\lambda }^2$:

$${\mathcal{L}}_{\mathrm{VI}}=\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)+\beta {\sum }_{\lambda }log\frac{log{P}_{}\left(w_{\lambda }\right)}{log{Q}_{}\left(w_{\lambda }\right)}\right].$$ (3)

Here, $\mathbf{x}$ is all training inputs, $\mathbf{y}$ is all training outputs, and $\beta$ is the tempering parameter which is 1 for a close approximation to Bayesian inference, but is often set to smaller values to “temper” the posterior, which often improves empirical performance [15,16] and has theoretical justification as accounting for the data-curation process [17].

We need to optimize the expectation in Equation (3) with respect to the parameters of $Q_{}\left(\mathbf{w}\right)$, the distribution over which the expectation is taken. To perform this optimization efficiently, the usual approach is to use the reparameterization trick [8,18,19]—we write $\mathbf{w}$ in terms of $ϵ$:

$$\begin{array}{cc}\hfill w_{\lambda }\left({ϵ}_{\lambda }\right)& ={\mathbf{\mu }}_{\lambda }+{\sigma }_{\lambda }{ϵ}_{\lambda }\hfill \end{array}$$ (4)

where ${ϵ}_{\lambda }\sim \mathcal{N}\left(0,1\right)$. Thus, the ELBO can be written as an expectation over $ϵ$:

$${\mathcal{L}}_{\mathrm{VI}}=\underset{ϵ}{\mathbb{E}}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}(ϵ)\right)+\beta {\sum }_{\lambda }log\frac{log{P}_{}\left(w_{\lambda }\left({ϵ}_{\lambda }\right)\right)}{log{Q}_{}\left(w_{\lambda }\left({ϵ}_{\lambda }\right)\right)}\right].$$ (5)

where the distribution over $ϵ$ is now fixed. Critically, now the expected gradient of the term inside the expectation is equal to the gradient of ${\mathcal{L}}_{\mathrm{VI}}$, so we can use samples of $ϵ$ to estimate the expectation.

2.2. Laplace’s Method

Laplace’s method [5,6,7] first finds a mode by doing gradient ascent on the log-joint:

$$\begin{array}{cc}\hfill {\mathbf{w}}^{*}& =\underset{\mathbf{w}}{arg\; max}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)+log{P}_{}\left(\mathbf{w}\right)\right]\hfill \end{array}$$ (6)

and uses a Gaussian approximate posterior around that mode,

$$\begin{array}{cc}\hfill Q_{}\left(\mathbf{w}\right)& =\mathcal{N}\left(\mathbf{w};{\mathbf{w}}^{*},{-\mathbf{H}}^{-1}\left({\mathbf{w}}^{*}\right)\right)\hfill \end{array}$$ (7)

where H$\left({\mathbf{w}}^{*}\right)$ is Hessian of the log-joint at ${\mathbf{w}}^{*}$.

3. Related Work

There is past work on Variational Laplace [20,21,22], which learns the mean parameters, $\mathbf{\mu }$, of a Gaussian approximate posterior,

$$\begin{array}{cc}\hfill Q_{\mathbf{\mu }}\left(\mathbf{w}\right)& =\mathcal{N}\left(\mathbf{w};\mathbf{\mu },{-\mathbf{H}}^{-1}\left(\mathbf{\mu }\right)\right)\hfill \end{array}$$ (8)

and obtains the covariance matrix as a function of the mean parameters using the Hessian, as in Laplace’s method. However, instead of taking the approximation to be centered around a MAP solution, ${\mathbf{w}}^{*}$, they take the approximate posterior to be centered on learned mean parameters, $\mathbf{\mu }$. Importantly, they simplify the ELBO by substituting this approximate posterior into Equation (3), and approximating the log-joint using its Taylor series expansion. Ultimately they obtain

$${\mathcal{L}}_{\mathrm{VI}}\approx log{P}_{}\left(\mathbf{y}|\mathbf{w}=\mathbf{\mu },\mathbf{x}\right)+log{P}_{}\left(\mathbf{w}=\mathbf{\mu }\right)-\frac{1}{2}log\left|\mathbf{H}\left(\mathbf{\mu }\right)\right|+\mathrm{const}.$$ (9)

However, there are two problems with this approach when applied to neural networks. First, the algebraic manipulations required to derive Equation (9) require the full $N\times N$ Hessian, $\mathbf{H}\left(\mathbf{\mu }\right)$, for all N parameters, and neural networks have too many parameters for this to be feasible. Second, the $log\left|\mathbf{H}\left(\mathbf{\mu }\right)\right|$ term in Equation (9) cannot be minibatched, as we need the full sum over minibatches inside the log to compute the Hessian:

$$\begin{array}{cc}\hfill log\left|\mathbf{H}\left(\mathbf{\mu }\right)\right|& =log|{\sum }_j{\mathbf{H}}_j\left(\mathbf{\mu }\right)|,\hfill \end{array}$$ (10)

where ${\mathbf{H}}_j\left(\mathbf{\mu }\right)$ is the contribution to the Hessian from an individual minibatch. Due to these issues, past Variational Laplace methods did not scale to large neural networks.

An alternative deterministic approach to variational inference in Bayesian neural networks, approximates the distribution over activities induced by stochasticity in the weights [23]. Unfortunately, it is important to capture the covariance over features induced by stochasticity in the weights. In fully connected networks, this is feasible, as we usually have a small number of features at each layer. However, in convolutional networks, we have a large number of features, $\mathtt{channels}\times \mathtt{height}\times \mathtt{width}$. In the lower layers of a ResNet, we may have 64 channels and a $32\times 32$ feature map, resulting in $64\times {32}^2$ = 65,536 features and a 65,536 × 65,536 covariance matrix. These scalability issues prevented them from applying their approach to convolutional networks. In contrast, our approach is highly scalable and readily applicable to the convolutional setting. Subsequent work such as Haußmann et al. [24] introduced other deterministic approximations, based on decomposing the relu into a linear function and a Heaviside step. However, their approach had errors of ∼30% on CIFAR-10.

Ritter et al. [7] and MacKay [25] used Laplace’s method in Bayesian neural networks, by first finding the mode by doing gradient ascent on the log-joint probability, and expanding around that mode. As usual for Laplace’s method, they risk finding a narrow mode that generalizes poorly. In contrast, we find a mode using an approximation to the ELBO that takes the curvature into account and hence is biased towards broad, flat modes that presumably generalise better.

Our approach gives a squared-gradient regularizer that is similar to those discovered in past work [12,26]. They showed that squared-gradient regularizers connect to gradient descent, in that approximation errors due to finite-step sizes in gradient-descent imply an effective squared gradient regularization. The similarity of our objectives raises profound questions about the extent to which gradient descent can be said to perform Bayesian inference. That said there are three key differences. First, their approach connects full-batch gradient descent to squared-gradient regularizers. Of course, most neural network training is stochastic gradient descent based on minibatches. Given that the stationary distribution of SGD is isotropic Gaussian (in a quadratic loss-function; Appendix A), we are able to connect stochastic gradient descent to squared gradient regularizers, and hence to approximate variational inference. This is especially important in light of recent work showing full-batch gradient descent gives poor performance, but that performance can be improved by including explicit squared gradient regularization. Our work indicates that explicit squared gradient regularization is mimicking the implicit regularization from stochastic gradient descent. First, our method uses the Fisher, (i.e., the gradients for data sampled from the model) whereas their approach uses the empirical Fisher, (i.e., gradients for the observed data) to form the squared gradient regularizer [27]. Second, our approach gives a principled method to learn a separate weighting for the squared-gradient for each parameter, whereas the connection to SGD forces Barrett and Dherin [12] to use a uniform weighting across all parameters.

Our work differs from, e.g., Khan et al. [28] by explicitly providing a simply implemented loss-function in terms of a squared-gradient regularizer, instead of working with NTK-inspired approximations to the Hessian.

Other approaches include “Broad Bayesian Learning” [29], which optimizes the architecture of a Bayesian neural network, exploiting information from previously trained but different networks. Of course, quantification of uncertainty for Bayesian neural networks is always fraught [30]. As such, we followed standard practice in the literature of reporting OOD detection performance and a measure of calibration accuracy [31].

4. Methods

To combine the best of VI and Laplace’s method, we begin by noting that the ELBO can be rewritten in terms of the KL divergence between the prior and approximate posterior:

$${\mathcal{L}}_{\mathrm{VI}}=\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)\right]-\beta {\sum }_{\lambda }{D}_{\mathrm{KL}}\left(Q_{}\left(w_{\lambda }\right)\left|\right|P_{}\left(w_{\lambda }\right)\right),$$ (11)

where the KL-divergence can be evaluated analytically:

$$\begin{array}{cc}\hfill D_{\mathrm{KL}}\left(Q_{}\left(w_{\lambda }\right)\left|\right|P_{}\left(w_{\lambda }\right)\right)& =\frac{1}{2}\left(\frac{{\sigma }_{\lambda }^2+{\mathbf{\mu }}_{\lambda }^2}{s_{\lambda }^2}-1+log\frac{s_{\lambda }^2}{{\sigma }_{\lambda }^2}\right).\hfill \end{array}$$ (12)

As such, the only term we need to approximate is the expected log-likelihood.

To approximate the expectation, we begin by taking a second-order Taylor series expansion of the log-likelihood around the current settings of the mean parameters, $\mathbf{\mu }$:

$$\begin{array}{c}\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)\right]\approx log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}=\mathbf{\mu }\right)+\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[{\sum }_{j=1}^B{\mathbf{g}}_j^T\left(\mathbf{w}-\mathbf{\mu }\right)\right]\hfill \\ \hfill +\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[\frac{1}{2}{\left(\mathbf{w}-\mathbf{\mu }\right)}^T\mathbf{H}\left(\mathbf{w}-\mathbf{\mu }\right)\right]\end{array}$$ (13)

where B is the number of minibatches, ${\mathbf{g}}_j$ is the gradient for minibatch j and H is the Hessian for the full dataset:

$$\begin{array}{cc}\hfill g_{j;\lambda }& =\frac{\partial }{\partial w_{\lambda }}\left[log{P}_{}\left({\mathbf{y}}_j|{\mathbf{x}}_j,\mathbf{w}\right)\right]\hfill \end{array}$$ (14)

$$\begin{array}{cc}\hfill H_{\lambda ,\nu }& =\frac{log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)}{\partial w_{\lambda }\partial w_{\nu }}.\hfill \end{array}$$ (15)

Here, $\mathbf{x}$ and $\mathbf{y}$ are the the inputs and outputs for the full dataset, whereas ${\mathbf{x}}_j$ and ${\mathbf{y}}_j$ are the inputs and outputs for minibatch j. Now we consider the expectation of each of these terms under the approximate posterior, $Q_{}\left(\mathbf{w}\right)$. The first term is constant and independent of $\mathbf{w}$. The second (linear) term is zero, because the expectation of $\left(\mathbf{w}-\mathbf{\mu }\right)$ under the approximate posterior is zero

$$\begin{array}{cc}\hfill \underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[{\mathbf{g}}_j^T\left(\mathbf{w}-\mathbf{\mu }\right)\right]& ={\mathbf{g}}_j^T\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[\left(\mathbf{w}-\mathbf{\mu }\right)\right]=0.\hfill \end{array}$$ (16)

The third (quadratic) term might at first appear difficult to evaluate because it involves, the $N\times N$ matrix of second derivatives, where N is the number of parameters in the model. However, using properties of the trace, and noting that the expectation of $\left(\mathbf{w}-\mathbf{\mu }\right){\left(\mathbf{w}-\mathbf{\mu }\right)}^T$ is the covariance of the approximate posterior, we obtain

$$\begin{array}{cc}\hfill \underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}[\frac{1}{2}(\mathbf{w}& {-\mathbf{\mu })}^T(\mathbf{w}-\mathbf{\mu })]=\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[\frac{1}{2}Tr\left(\left(\mathbf{w}-\mathbf{\mu }\right){\left(\mathbf{w}-\mathbf{\mu }\right)}^T\right)\right]=\frac{1}{2}Tr\left(\mathrm{\Sigma }\right)\hfill \end{array}$$ (17)

writing the trace in index notation, and substituting for the (diagonal) posterior covariance, $\mathrm{\Sigma }$:

$$\frac{1}{2}Tr\left(\mathrm{\Sigma }\right)=\frac{1}{2}\sum _{\lambda \nu }{H}_{\lambda \nu }{\mathrm{\Sigma }}_{\lambda \nu }=\frac{1}{2}\sum _{\lambda }{H}_{\lambda \lambda }{\sigma }_{\lambda }^2.$$ (18)

Thus, our first approximation of the expected log-likelihood is

$$\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)\right]\approx log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}=\mathbf{\mu }\right)+\frac{1}{2}{\sum }_{\lambda }{\sigma }_{\lambda }^2{H}_{\lambda \lambda },$$ (19)

and substituting this into Equation (11) gives

$$\begin{array}{c}\hfill {\mathcal{L}}_{\mathrm{VI}}\approx {\mathcal{L}}_{\mathrm{VL}\left(\mathrm{H}\right)}=log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}=\mathbf{\mu }\right)+\frac{1}{2}{\sum }_{\lambda }{\sigma }_{\lambda }^2{H}_{\lambda \lambda }-\beta {\sum }_{\lambda }{D}_{\mathrm{KL}}\left(Q_{}\left(w_{\lambda }\right)\left|\right|P_{}\left(w_{\lambda }\right)\right).\end{array}$$ (20)

This resolves most of the issues with the original Variational Laplace method: it requires only the diagonal of the Hessian, it can be minibatched and it does not blow up if $H_{\lambda \lambda }$ is zero.

4.1. Pathological Optima When Using the Hessian

However, a new issue arises: $H_{\lambda \lambda }$ is usually negative, in which case the approximation in Equation (20) can be expected to work well. However there is nothing to stop $H_{\lambda \lambda }$ from becoming positive. Usually if we, e.g., took the log-determinant of the negative Hessian, this would immediately break the optimization process (as we would be taking the logarithm of a negative number). However, in our context, there is no immediate issue as Equation (20) takes on a well-defined value even when one or more $H_{\lambda \lambda }$’s are positive. That said, we rapidly encounter similar issues as we get pathological optimal values of ${\sigma }_{\lambda }^2$. In particular, picking out the terms in the objective that depend on ${\sigma }_{\lambda }^2$, absorbing the other terms into the constant, and taking $\beta =1$ for simplicity, we have

$$\begin{array}{c}\hfill {\mathcal{L}}_{\mathrm{VL}\left(\mathrm{H}\right)}=\frac{1}{2}{\sum }_{\lambda }\left(-\left(\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }\right){\sigma }_{\lambda }^2+log{\sigma }_{\lambda }^2\right)+\mathrm{const}.\end{array}$$ (21)

Thus, the gradient wrt a single variance parameter is

$$\begin{array}{c}\hfill \frac{\partial }{\partial {\sigma }_{\lambda }^2}{\mathcal{L}}_{\mathrm{VL}\left(\mathrm{H}\right)}=\frac{1}{2}\left(-\left(\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }\right)+\frac{1}{{\sigma }_{\lambda }^2}\right).\end{array}$$ (22)

In the typical case, $H_{\lambda \lambda }$ is negative so $\left(\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }\right)$ is positive, and we can find the optimum by solving for the value of ${\sigma }_{\lambda }^2$ where the gradient is zero:

$$\begin{array}{cc}\hfill {\sigma }_{\lambda }^2& =\frac{1}{\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }}.\hfill \end{array}$$ (23)

However, if $H_{\lambda \lambda }$ is positive and sufficiently large, $H_{\lambda \lambda }>\frac{1}{s_{\lambda }^2}$, then $\left(\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }\right)$ becomes negative, and not only is the mode in Equation (23) undefined, but the gradient is always positive:

$$\begin{array}{c}\hfill 0<\frac{\partial }{\partial {\sigma }_{\lambda }^2}{\mathcal{L}}_{\mathrm{VL}\left(\mathrm{H}\right)}=\frac{1}{2}\left(-\left(\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }\right)+\frac{1}{{\sigma }_{\lambda }^2}\right).\end{array}$$ (24)

as both terms in the sum: $-\left(\frac{1}{s_{\lambda }^2}-H_{\lambda \lambda }\right)$ and $\frac{1}{{\sigma }_{\lambda }^2}$ are positive. As such, when $H_{\lambda \lambda }>\frac{1}{s_{\lambda }^2}$, the variance, ${\sigma }_{\lambda }^2$ grows without bound.

4.2. Avoiding Pathologies with the Fisher

To avoid pathologies arising from the fact that the Hessian is not necessarily negative definite, a common approach is to approximate the Hessian using the Fisher information matrix:

$$\begin{array}{cc}\hfill -\mathbf{H}\approx \mathbf{F}& =\sum _{j=1}^B\underset{P_{}\left({\tilde{\mathbf{y}}}_j|{\mathbf{x}}_j,\mathbf{w}=\mathbf{\mu }\right)}{\mathbb{E}}\left[{\tilde{\mathbf{g}}}_j\left({\tilde{\mathbf{y}}}_j\right){\tilde{\mathbf{g}}}_j^T\left({\tilde{\mathbf{y}}}_j\right)\right].\hfill \end{array}$$ (25)

Importantly, $\tilde{\mathbf{g}}$ is the gradient of the log-likelihood for data sampled from the model, ${\tilde{\mathbf{y}}}_j$, not for the true data:

$$\begin{array}{cc}\hfill {\tilde{g}}_{j;\lambda }\left({\tilde{\mathbf{y}}}_j\right)& =\frac{\partial }{\partial w_{\lambda }}\left[log{P}_{}\left({\tilde{\mathbf{y}}}_j|{\mathbf{x}}_j,\mathbf{w}\right)\right].\hfill \end{array}$$ (26)

This gives us the Fisher, which is a commonly used and well-understood approximation to the Hessian [27]. Importantly, this contrasts with the empirical Fisher [27], which uses the gradient conditioned on the actual data (and not data sampled from the model):

$$\begin{array}{cc}\hfill {\mathbf{F}}_{\mathrm{emp}}& =\sum _{j=1}^B{\mathbf{g}}_j{\mathbf{g}}_j^T,\hfill \end{array}$$ (27)

which is problematic, because there is a large rank-1 component in the direction of the mean gradient, which disrupts the estimated matrix specifically in the direction of interest for problems such as optimization [27].

Using the Fisher information (Equation (25)) in Equation (19), we obtain an approximate expected log-likelihood:

$$\underset{Q_{}\left(\mathbf{w}\right)}{\mathbb{E}}\left[log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}\right)\right]\approx log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}=\mathbf{\mu }\right)-\frac{1}{2}{\sum }_{\lambda }{\sigma }_{\lambda }^2{\sum }_{j=1}^B{\tilde{g}}_{j;\lambda }^2.$$ (28)

Substituting this into Equation (11) gives us the final VL objective, ${\mathcal{L}}_{\mathrm{VL}}$, which is an approximation of the ELBO:

$${\mathcal{L}}_{\mathrm{VI}}\approx {\mathcal{L}}_{\mathrm{VL}}=log{P}_{}\left(\mathbf{y}|\mathbf{x},\mathbf{w}=\mathbf{\mu }\right)-\frac{1}{2}{\sum }_{\lambda }{\sigma }_{\lambda }^2{\sum }_{j=1}^B{\tilde{g}}_{j;\lambda }^2-\beta {\sum }_{\lambda }{D}_{\mathrm{KL}}\left(Q_{}\left(w_{\lambda }\right)\left|\right|P_{}\left(w_{\lambda }\right)\right).$$ (29)

In practice, we typically take the objective for a minibatch, divided by the number of datapoints in a minibatch, S:

$$\begin{array}{c}\hfill \frac{1}{S}{\mathcal{L}}_{\mathrm{VL};j}=\frac{1}{S}log{P}_{}\left({\mathbf{y}}_j|{\mathbf{x}}_j,\mathbf{w}=\mathbf{\mu }\right)-\frac{S}{2}{\sum }_{\lambda }{\sigma }_{\lambda }^2{\left(\frac{1}{S}{\tilde{g}}_{j;\lambda }\right)}^2-\frac{\beta }{2SB}{\sum }_{\lambda }\left(\frac{{\sigma }_{\lambda }^2+{\mathbf{\mu }}_{\lambda }^2}{s_{\lambda }^2}-1+log\frac{s_{\lambda }^2}{{\sigma }_{\lambda }^2}\right),\end{array}$$ (30)

where $\left(\frac{1}{S}{\tilde{g}}_{j;\lambda }\right)$ are the gradients of the log-likelihood for the minibatch averaged across datapoints, i.e., the gradient of $\frac{1}{S}log{P}_{}\left({\tilde{\mathbf{y}}}_j|{\mathbf{x}}_j,\mathbf{w}=\mathbf{\mu }\right)$. Remember B is the number of minibatches so $SB$ is the total number of training datapoints.

4.3. Constraints on the Network Architecture

Importantly, here the regularizer is the squared gradient of the loss with respect to the parameters. As such, computing the loss implicitly involves a second-derivative of the log-likelihood, and we therefore cannot use piecewise linear activation functions such as ReLU, which have pathological second derivatives. In particular, the second derivative has a delta-function “spike” at zero:

$$\begin{array}{cc}\hfill \frac{d^2}{dx}\varphi \left(x\right)& =\frac{d}{dx}\left[\frac{d}{dx}\varphi \left(x\right)\right]=\frac{d}{dx}\mathrm{\Theta }\left(x\right)=\delta \left(x\right)\hfill \end{array}$$ (31)

where $\varphi$ is the relu nonlinearity, $\mathrm{\Theta }\left(x\right)$ is the Heaviside step function which is zero for $x<0$ and one for $0<x$, and $\delta \left(x\right)$ is the Dirac delta function. As the function is almost never evaluated at exactly zero, it is not possible to sensibly take into account the contribution of the infinitely high spike in the second derivative at zero. Interestingly, this issue is very similar to the one that turns up when differentiating step (i.e., $\mathrm{\Theta }\left(x\right)$) activations—the derivative is well-defined and zero almost everywhere. The issue is that there are delta-function spikes in the gradient at zero that gradient descent cannot reasonably work with. Instead, we used a softplus activation function, but any activation with well-behaved second derivatives is admissible.

5. Results

We compared MAP, VI and our method (VL) on four different datasets (CIFAR-10, CIFAR-100 [32], SVHN [33] and fashion-MNIST [34] MIT Licensed) using a PreactResNets-18 [35] with an initial learning rate of 1E-4, which decreased by a factor of 10 after 100 and 150 epochs and a batch size of 128 with all the other optimizer hyperparameters set to their default values. We tried two variants of variational inference: evaluating test-performance using the mean network, VI (mean), and evaluating test performance by drawing 10 samples from the approximate posterior, VI (sampled). We swept across different degrees of posterior tempering, $\beta$. Using $\beta <1$ is normatively justified in the Bayesian framework as accounting for the effect of data curation [17]. For many values of $\beta$ VL gave better test accuracies, test log-likelihoods and expected calibration errors [31,36] than VI or MAP inference (Figure 1). Importantly though, for the optimal value of $\beta$, VL almost always gave better performance on these metrics (Table 1). These experiments took ∼480 GPU hours, and were run on a mixture of nVidia 1080 and 2080 GPUs in an internal cluster.

Figure 1.

Training a PreactResNet-18 on various datasets, displaying the test accuracy, test log-likelihood, expected calibration error (ECE) [31,36] and OOD detection metric (AUROC) for CIFAR-10, CIFAR-100, SVHN and fashion MNIST. Downsampled Imagenet [37] was used as OOD data. See Appendix B.

Table 1.

Best values test NLL, test accuracy and ECE for a variety of datasets, as we used different values of the tempering parameter, $\beta$.

Dataset	Method	Test NLL	Test Acc.	ECE
CIFAR-10	VL	0.23	92.4%	0.017
	VI (Mean)	0.37	91.1%	0.053
	VI (10 Samples)	0.35	90.2%	0.044
	MAP	0.43	90.8%	0.058
CIFAR-100	VL	1.00	71.4%	0.024
	VI (Mean)	1.29	68.8%	0.100
	VI (10 Samples)	1.49	67.3%	0.026
	MAP	1.61	67.5%	0.159
SVHN	VL	0.14	97.1%	0.009
	VI (Mean)	0.16	96.3%	0.012
	VI (10 Samples)	0.22	95.5%	0.022
	MAP	0.24	95.7%	0.028
Fashion MNIST	VL	0.16	94.6%	0.010
	VI (Mean)	0.23	94.0%	0.034
	VI (10 Samples)	0.29	93.6%	0.016
	MAP	0.29	93.6%	0.096

The runtimes of the methods are listed in Table 2. VL or gradient regularization was around a factor of three slower than either VI or MAP due to the need to compute second-derivatives. It is still eminently feasible, especially in comparison to past methods for deterministic variational inference that have fundamental difficulties in scaling to convolutional networks [23]. Furthermore, we did not find that increasing the number of epochs improved performance either for VI or MAP, as we were already training for convergence.

Table 2.

Time per epoch for different methods on CIFAR-10.

Method	Time per Epoch (s)
VL	114.9
VI	43.2
MAP	41.8

Early-Stopping and Poor Performance in VI

Before performing comparisons where we learn the approximate posterior variance, it is important to understand the pitfalls when optimizing variational Bayesian neural networks using adaptive optimizers such as Adam. In particular, there is a strong danger of stopping the optimization before the variances have converged. To illustrate this risk, note that Adam [38] updates take the form

$$\begin{array}{cc}\hfill \mathsf{\Delta }\theta & =\eta \frac{m}{\sqrt{v}+ϵ}\hfill \end{array}$$ (32)

where $\eta$ is the learning rate, m is an unbiased estimator of the mean gradient, $\langle g\rangle$, v is an unbiased estimator of the squared gradient, $\langle g^2\rangle$ and $ϵ$ is a small positive constant to avoid divide-by-zero. The magnitude of the updates, $\left|\mathsf{\Delta }\theta \right|$, is maximized by having exactly the same gradient on each step, in which case, neglecting $ϵ$, we have $\left|\mathsf{\Delta }\theta \right|=\eta$. As such, with a learning rate of $\eta ={10}^{-4}$, a training set of 50,000 and a batch size of 128 parameters can move at most 50,000/128 $\times {10}^{-4}\approx 0.04$ per epoch. Doing 100 epochs at this learning rate, a parameter can change by at most 4 over the 100 epochs before the first learning rate step.

This is fine for the weights, which typically have very small values. However, the underlying parameters used for the variances typically take on larger values. In our case, we will use $log{\sigma }_{\lambda }$ as the parameter, and initialize it to three less than the prior standard deviation, $log{s}_{\lambda }-3$. To ensure reasonable convergence, $log{\sigma }_{\lambda }$ should be able to revert back to the prior, implying that it must be able to change by at least three during the course of training. Unfortunately, 3 is very close to the maximum possible change of 4, raising the possibility that the variance parameters will not actually converge. To check whether early-stopping was indeed an issue, we plotted the (tempered) ELBO for VI (Figure 2A) and VL (Figure 2B). For VI (Figure 2A) with the standard setup (lightest line with a learning rate multiplier of 1), the ELBO clearly has not converged at 100 epochs, indicating early-stopping. Notably, this was still an issue with VL (Figure 2B), especially if we were to train for fewer epochs. However, the effect was smaller for VL, which may have been because the gradients were more consistent, as it did not sample the weights. These issues can be rectified by increasing the learning rate specifically for the $log{\sigma }_{\lambda }$ parameters (darker lines).

Figure 2.

Analysis of early stopping in VI and VL. The first row is untempered ($\beta =1$), and the second row is tempered ($\beta =0.1$). (A) ELBO over epochs 0–100 (with the highest initial learning rate) for VI. Different lines correspond to networks with learning rate multipliers for $log{\sigma }_{\lambda }$ of 1, 3, 10 and 30. (B) As (A), but for VL. CDE Final test-log-likelihood (C), test accuracy (D) and ELBO (E) after 200 epochs for different learning rate multipliers.

We then plotted the test log-likelihood (Figure 2C), test accuracy (Figure 2D) and ELBO (Figure 2E) against the learning rate multiplier. Again, the performance for VL (orange) was reasonably robust to changes in the learning rate multiplier. However, the performance of VI (blue) was very sensitive to the multiplier: as the multiplier increased, test performance fell but the ELBO rose. As we ultimately care about test performance, these results would suggest that we should use the lowest multiplier (1), and accept the possibility of early-stopping. That may be a perfectly good choice in many cases. However, VI is supposed to be an approximate Bayesian method, and using an alternative form for the ELBO,

$$\begin{array}{cc}\hfill {\mathcal{L}}_{\mathrm{VI}}& =log{P}_{}\left(\mathbf{y}|\mathbf{x}\right)-D_{\mathrm{KL}}\left(Q_{}\left(\mathbf{w}\right)\left|\right|P_{}\left(\mathbf{w}|\mathbf{y},\mathbf{x}\right)\right),\hfill \end{array}$$ (33)

we can see that the ELBO measures KL-divergence between the true and approximate posterior, and hence the quality of our approximate Bayesian inference. As such, very poor ELBOs imply that the KL-divergence between the true and approximate posterior is very large, and hence the “approximate posterior” is no longer actually approximating the true posterior. As such, if we are to retain a Bayesian interpretation of VI, we need to use larger learning rate multipliers which give better values for the ELBO (Figure 2E). However, in doing that, we get worse test performance (Figure 2C,D). This conflict between approximate posterior quality and test performance is very problematic: the Bayesian framework would suggest that as Bayesian inference becomes more accurate, performance should improve, whereas for VI, performance gets worse. Concretely, by initializing $log{\sigma }_{\lambda }$ to a small value and then early-stopping, we leave $log{\sigma }_{\lambda }$ at a small value through training, in which case VI becomes equivalent to MAP inference with a negligibly small amount of noise added to the weights. We would therefore expect early-stopped VI to behave (and be) very similar to MAP inference.

In subsequent experiments, we chose to use a learning rate multiplier of 10, as this largely eliminated early-stopping (though see VI with $\beta =0.1$; Figure 2E). Overall, this indicates that we have to be very careful to avoid early stopping when running standard, sampling-based variational inference.

6. Conclusions

We gave a novel Variational Laplace approach to inference in Bayesian neural networks which combines the best of previous approaches based on variational inference and Laplace’s Method. This method gave excellent empirical performance compared to VI.

Appendix A. The Stationary Distribution of SGD

We sought to relate these gradient regularizers back to work on SGD. In particular, we looked to work on the stationary distribution of SGD, which noted that under quadratic losses functions, SGD samples an isotropic Gaussian (i.e., with covariance proportional to the identity matrix). In particular, consider a loss function which is locally closely approximated by a quadratic. Without loss of generality, we consider a mode at $\mathbf{w}=\mathbf{0}$,

$$log{P}_{}\left({\left\{y_i\right\}}_{i=1}^P|{\left\{x_i\right\}}_{i=1}^P,\mathbf{w}\right)=-\frac{P}{2}{\mathbf{w}}^T\mathbf{H}\mathbf{w}+const.$$ (A1)

where P is the total number of datapoints. Typically, the objective used in SGD is the loss for a minibatch of size $S_{\mathrm{SGD}}$. Following Mandt et al. [39], we use the Fisher information to identify the noise in the minibatch gradient,

$$\frac{\partial }{\partial \mathbf{w}}\left[\frac{1}{S_{\mathrm{SGD}}}log{P}_{}\left({\mathbf{y}}_j|x_j,\mathbf{w}\right)\right]=-\mathbf{H}\mathbf{w}+{\frac{1}{\sqrt{S_{\mathrm{SGD}}}}\mathbf{H}}^{1/2}\xi \left(t\right),$$ (A2)

where $\xi \left(t\right)$ is sampled from a standard IID Gaussian. For SGD, this gradient is multiplied by a learning rate, ${\eta }_{\mathrm{SGD}}$,

$$\begin{array}{cc}\hfill \mathbf{w}(t+1)& =\mathbf{w}\left(t\right)-{\eta }_{\mathrm{SGD}}\mathbf{H}\mathbf{w}\left(t\right)+{\frac{{\eta }_{\mathrm{SGD}}}{\sqrt{S_{\mathrm{SGD}}}}\mathbf{H}}^{1/2}\xi \left(t\right),\hfill \end{array}$$ (A3)

This is an multivariate Gaussian autoregressive process, so we can solve for the stationary distribution of the weights. In particular, we note that the covariance at time $t+1$ is

$$\begin{array}{cc}\hfill \mathbb{C}\left[\mathbf{w}\left(t+1\right)\right]& =\mathbb{E}\left[{\left(\mathbf{w}\left(t\right)-{\eta }_{\mathrm{SGD}}\mathbf{H}\mathbf{w}\left(t\right)-{\frac{{\eta }_{\mathrm{SGD}}}{\sqrt{S_{\mathrm{SGD}}}}\mathbf{H}}^{1/2}\xi \left(t\right)\right)}^T\left(\mathbf{w}\left(t\right)-{\eta }_{\mathrm{SGD}}\mathbf{H}\mathbf{w}\left(t\right)-{\frac{{\eta }_{\mathrm{SGD}}}{\sqrt{S_{\mathrm{SGD}}}}\mathbf{H}}^{1/2}\xi \left(t\right)\right)\right]\hfill \\ \hfill \mathbb{C}\left[\mathbf{w}\left(t+1\right)\right]& ={\left(\mathbf{I}-{\eta }_{\mathrm{SGD}}\mathbf{H}\right)}^T\mathbb{C}\left[\mathbf{w}\left(t\right)\right]\left(\mathbf{I}-{\eta }_{\mathrm{SGD}}\mathbf{H}\right)+\frac{{\eta }_{\mathrm{SGD}}^2}{S_{\mathrm{SGD}}}\mathbf{H}\hfill \end{array}$$ (A4)

Following Mandt et al. [39], when the learning rate is small, the quadratic term can be neglected.

$$\begin{array}{cc}\hfill \mathbb{C}\left[\mathbf{w}\left(t+1\right)\right]& \approx \mathbb{C}\left[\mathbf{w}\left(t\right)\right]-{\eta }_{\mathrm{SGD}}\mathbf{H}\mathbb{C}\left[\mathbf{w}\left(t\right)\right]-{\eta }_{\mathrm{SGD}}\mathbb{C}\left[\mathbf{w}\left(t\right)\right]\mathbf{H}+\frac{{\eta }_{\mathrm{SGD}}^2}{S_{\mathrm{SGD}}}\mathbf{H}\hfill \end{array}$$ (A5)

We then solve for steady-state in which $\mathrm{\Sigma }=\mathbb{C}\left[\mathbf{w}(t+1)\right]=\mathbb{C}\left[\mathbf{w}\left(t\right)\right]$,

$$\begin{array}{cc}\hfill 0& \approx -{\eta }_{\mathrm{SGD}}\left({\mathbf{H}}^T{\mathrm{\Sigma }}_{\mathrm{SGD}}+{\mathrm{\Sigma }}_{\mathrm{SGD}}\mathbf{H}\right)+\frac{{\eta }_{\mathrm{SGD}}^2}{S_{\mathrm{SGD}}}\mathbf{H}\hfill \end{array}$$ (A6)

so,

$$\begin{array}{cc}\hfill \mathrm{\Sigma }& \approx \frac{{\eta }_{\mathrm{SGD}}}{2S_{\mathrm{SGD}}}\mathbf{I}.\hfill \end{array}$$ (A7)

Appendix B. Varying the Batch Size

Here, we vary the batch size. We used a batch size of 128 in the main text. We found that batch sizes of 64 and 256 have no effect on the relative performance of the methods.

Figure A1.

Replication of Figure 1 with a batch size of 64.

Figure A2.

Replication of Figure 1 with a batch size of 256.

Author Contributions

Conceptualization, L.A.; methodology, A.U.; software, A.U.; validation, L.A.; formal analysis, L.A.; investigation, A.U.; resources, L.A.; data curation, A.U.; writing—original draft preparation, L.A.; writing—review and editing, L.A.; visualization, A.U.; supervision, L.A.; project administration, L.A.; funding acquisition, L.A. All authors have read and agreed to the published version of the manuscript.

Funding

This research received no external funding.

Institutional Review Board Statement

Not applicable.

Informed Consent Statement

Not applicable.

Data Availability Statement

Code is available at https://github.com/LaurenceA/fitr (accessed on 2 December 2021).

Conflicts of Interest

The authors declare no conflict of interest.