Original: Gradient Regularisation as Approximate Variational Inference
Revised: Gradient Regularisation as Approximate Variational Inference

Original: Anonymous Authors

Original: Anonymous Institution

Original: Abstract
Revised: Abstract

Original: Variational inference in Bayesian neural networks is usually performed using stochastic sampling which gives very high-variance gradients, and hence slow learning.
Revised: Variational inference in Bayesian neural networks is usually performed using stochastic sampling which gives very high-variance gradients, and hence slow learning.

Original: Here, we show that it is possible to obtain a deterministic approximation of the ELBO for a Bayesian neural network by doing a Taylor-series expansion around the mean of the current variational distribution.
Revised: Here, we show that it is possible to obtain a deterministic approximation of the ELBO for a Bayesian neural network by doing a Taylor-series expansion around the mean of the current variational distribution.

Original: The resulting approximate ELBO is the training-log-likelihood plus a squared gradient regulariser.
Revised: The resulting approximate ELBO is the training-log-likelihood plus a squared gradient regulariser.

Original: In addition to learning the approximate posterior variance, we also consider a uniform-variance approximate posterior, inspired by the stationary distribution of SGD.
Revised: In addition to learning the approximate posterior variance, we also consider a uniform-variance approximate posterior, inspired by the stationary distribution of SGD.

Original: The corresponding approximate ELBO has a simple form, as the log-likelihood plus a simple squared-gradient regulariser.
Revised: The corresponding approximate ELBO has a simple form, as the log-likelihood plus a simple squared-gradient regulariser.

Original: We argue that this squared-gradient regularisation may at the root of the excellent empirical performance of SGD.
Revised: We argue that this squared-gradient regularisation may at the root of the excellent empirical performance of SGD.

Original: 1. Introduction
Revised: 1. Introduction

Original: Neural networks are increasingly being used in safety-critical settings such as self-driving cars (Bojarski et al., 2016) and medical diagnosis (Amato et al., 2013).
Revised: Neural networks are increasingly being used in safety-critical settings such as self-driving cars (Bojarski et al., 2016) and medical diagnosis (Amato et al., 2013).

Original: 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 (McAllister et al., 2017).
Revised: 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 (McAllister et al., 2017).

Original: Several approaches to Bayesian inference in neural networks are available, including from stochastic gradient Langevin dynamics (Welling and Teh, 2011) but here we focus on variational inference (Blundell et al., 2015).
Revised: Several approaches to Bayesian inference in neural networks are available, including from stochastic gradient Langevin dynamics (Welling and Teh, 2011) but here we focus on variational inference (Blundell et al., 2015).

Original: Despite the theoretical advantages of Bayesian neural networks (e.g. their relationship with PAC-Bayesian methods that provide bounds on generalisation error Germain et al., 2016; Rivasplata et al., 2019), they are rare amongst networks claiming to give state-of-theart performance on tasks such as image classiﬁcation.
Revised: Despite the theoretical advantages of Bayesian neural networks (e g. their relationship with PAC-Bayesian methods that provide bounds on generalisation error Germain et al., 2016; Rivasplata et al., 2019), they are rare amongst networks claiming to give state-of-theart performance on tasks such as image classiﬁcation.

Original: Instead, much simpler techniques such as stochastic gradient descent (SGD) (Krizhevsky et al., 2012) are found to give excellent performance in practice.
Revised: Instead, much simpler techniques such as stochastic gradient descent (SGD) (Krizhevsky et al., 2012) are found to give excellent performance in practice.

Original: Indeed, SGD is also found to give better performance than adaptive optimizers such as Adam (Kingma and Ba, 2014; Keskar and Socher, 2017; Loshchilov and
Revised: Indeed, SGD is also found to give better performance than adaptive optimizers such as Adam (Kingma and Ba, 2014; Keskar and Socher, 2017; Loshchilov and

Original: Hutter, 2017; Wilson et al., 2017).
Revised: Hutter, 2017; Wilson et al., 2017).

Original: This is usually understood as regularisation that is implicitly embodied in SGD (Keskar et al., 2016;
Revised: This is usually understood as regularisation that is implicitly embodied in SGD (Keskar et al., 2016;

Original: Wu et al., 2017; Lei et al., 2018; Roberts, 2018).
Revised: Wu et al., 2017; Lei et al., 2018; Roberts, 2018).

Original: While SGD’s implicit regularisation gives excellent practical performance, it becomes problematic when trying to improve the optimizer used in deep learning.
Revised: While SGD’s implicit regularisation gives excellent practical performance, it becomes problematic when trying to improve the optimizer used in deep learning.

Original: In the ideal case, we would be able to separate the objective (including regularisation terms) from the optimizer, such that e.g. improvements in the optimizer’s convergence rate can always be expected to improve performance.
Revised: In the ideal case, we would be able to separate the objective (including regularisation terms) from the optimizer, such that e.g. improvements in the optimizer’s convergence rate can always be expected to improve performance.

Original: In contrast, at the moment, methods that converge faster (such as Adam; Kingma and Ba, 2014) are generally believed to have worse performance at convergence, due to the lack of the implicit regularisation embodied in SGD (Wilson et al., ; Keskar and Socher, 2017).
Revised: In contrast, at the moment, methods that converge faster (such as Adam; Kingma and Ba, 2014) are generally believed to have worse performance at convergence, due to the lack of the implicit regularisation embodied in SGD (Wilson et al., 2017; Keskar and Socher, 2017).

Original: Therefore, it is important to understand the implicit regularisation embodied in SGD, such that we can use that regularisation in combination with other optimizers.
Revised: Therefore, it is important to understand the implicit regularisation embodied in SGD, such that we can use that regularisation in combination with other optimizers.

Original: Here, we start by noting that variational inference for Bayesian neural networks typically involves stochastic sampling, which can give rise to high-variance gradients.
Revised: Here, we start by noting that variational inference for Bayesian neural networks typically involves stochastic sampling, which can give rise to high-variance gradients.

Original: We note that we can form a deterministic approximation of the ELBO by doing a second-order Taylor expansion around the mode of the approximate posterior.
Revised: We note that we can form a deterministic approximation of the ELBO by doing a second-order Taylor expansion around the mode of the approximate posterior.

Original: This gives an approximate
Revised: This gives an approximate

Original: ELBO that is composed of the log-likelihood, standard weight-decay regularisation and a squared-gradient regulariser, which is weighted by the variance of the approximate posterior.
Revised: ELBO that is composed of the log-likelihood, standard weight-decay regularisation and a squared-gradient regulariser, which is weighted by the variance of the approximate posterior.

Original: Next, we noted that the implicit regularisation eﬀects of SGD can be understood using the isotropic Gaussian stationary distribution under locally quadratic loss functions (Mandt et al., 2017).
Revised: Next, we noted that the implicit regularisation eﬀects of SGD can be understood using the isotropic Gaussian stationary distribution under locally quadratic loss functions (Mandt et al., 2017).

Original: Inspired by this stationary distribution, we considered uniform-variance approximate posteriors, which correspond to a simpler squared-gradient regulariser.
Revised: Inspired by this stationary distribution, we considered uniform-variance approximate posteriors, which correspond to a simpler squared-gradient regulariser.

Original: Incorporating this regulariser improved performance on standard benchmark tasks, and provided a potential explanation for the excellent performance of stochastic gradient descent.
Revised: Incorporating this regulariser improved performance on standard benchmark tasks, and provided a potential explanation for the excellent performance of stochastic gradient descent.

Original: 2. Background 2.1. Variational inference for Bayesian neural networks
Revised: 2. Background 2.1. Variational inference for Bayesian neural networks

Original: Following the usual convention (Blundell et al., 2015), we use independent Gaussian priors and approximate posteriors for all parameters,
Revised: Following the usual convention (Blundell et al., 2015), we use independent Gaussian priors and approximate posteriors for all parameters,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: equivalently
Revised: equivalently

Original: [Equation]
Revised: [Equation]

Original: where µ λ and σ 2 λ are learned parameters of the approximate posterior, and where Σ is a diagonal matrix, with Σ λλ = σ 2 λ .
Revised: where µ λ and σ 2 λ are learned parameters of the approximate posterior, and where Σ is a diagonal matrix, with Σ λλ = σ 2 λ .

Original: Please see Appendix A for more information.
Revised: Please see Appendix A for more information.

Original: 2.2. The stationary distribution of SGD
Revised: 2.2. The stationary distribution of SGD

Original: When we solve for steady-state in which Σ
Revised: When we solve for steady-state in which Σ

Original: = C [ w ( t + 1)]
Revised: = C [ w ( t + 1)]

Original: = C [ w ( t )], we derive that Σ ≈ η 2 S
Revised: = C [ w ( t )], we derive that Σ ≈ η 2 S

Original: I .
Revised: I .

Original: For the derivation, please see Appendix B.
Revised: For the derivation, please see Appendix B.

Original: 3. Methods 3.1. Deterministic approximations to variational inference
Revised: 3. Methods 3.1. Deterministic approximations to variational inference

Original: We begin by noting that the ELBO can be rewritten in terms of the KL divergence between the prior and approximate posterior,
Revised: We begin by noting that the ELBO can be rewritten in terms of the KL divergence between the prior and approximate posterior,

Original: [Equation]
Revised: [Equation]

Original: And the KL-divergence can be evaluated analytically,
Revised: And the KL-divergence can be evaluated analytically,

Original: [Equation]
Revised: [Equation]

Original: As such, the only term we need to approximate is the expected log-likelihood.
Revised: As such [Equation] the only term we need to approximate is the expected log-likelihood.

Original: To evaluate this expectation
Revised: To evaluate this expectation

Original: [Equation]
Revised: [Equation]

Original: we begin by taking a second-order Taylor series expansion of the log-likelihood around the current setting of the mean parameters, µ ,
Revised: we begin by taking a second-order Taylor series expansion of the log-likelihood around the current setting of the mean parameters, µ ,

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Now we consider the expectation of each of these terms under the approximate posterior,
Revised: Now we consider the expectation of each of these terms under the approximate posterior,

Original: Q ( w ).
Revised: Q ( w ).

Original: The ﬁrst term is constant and independent of w .
Revised: The ﬁrst term is constant and independent of w .

Original: The second (linear) term is zero, because the expectation of ( w − m ) under the approximate posterior is zero
Revised: The second (linear) term is zero, because the expectation of ( w − m ) under the approximate posterior is zero

Original: [Equation]
Revised: [Equation]

Original: the third (quadratic) term is diﬃcult to evaluate because it involves H , the N × N matrix of second derivatives, where N is the number of parameters in the model.
Revised: the third (quadratic) term is diﬃcult to evaluate because it involves H , the N × N matrix of second derivatives, where N is the number of parameters in the model.

Original: Nonetheless, we begin by using properties of the trace [Equation] and noting that the expectation is the covariance of the approximate posterior,
Revised: Nonetheless, we begin by using properties of the trace, and noting that the expectation is the covariance of the approximate posterior,

Revised: E

Original: [Equation]
Revised: [Equation]

Original: writing the trace in index notation, and substituting for the (diagonal) posterior covariance, Σ ,
Revised: writing the trace in index notation, and substituting for the (diagonal) posterior covariance, Σ ,

Original: [Equation]
Revised: [Equation]

Original: Evaluating this term requires the diagonal of the Hessian, and as such our methods can be used with any drop-in estimate of this Hessian.
Revised: Evaluating this term requires the diagonal of the Hessian, and as such our methods can be used with any drop-in estimate of this Hessian.

Original: Note that this regulariser, explicitly in terms of the Hessian, can also be justiﬁed by empirical work on generalisation in neural networks (Wu et al., 2017), which found that networks with smaller Hessian norm generalised better.
Revised: Note that this regulariser, explicitly in terms of the Hessian, can also be justiﬁed by empirical work on generalisation in neural networks (Wu et al., 2017), which found that networks with smaller Hessian norm generalised better.

Original: In our case, we use the Fisher Information (Kunstner et al., 2019), for three reasons.
Revised: In our case, we use the Fisher Information (Kunstner et al., 2019), for three reasons.

Original: First, it is extremely stable (the Fisher Information matrix is always positive deﬁnite, and hence H λλ is always positive).
Revised: First, it is extremely stable (the Fisher Information matrix is always positive deﬁnite, and hence H λλ is always positive).

Original: Second the FI will allow us to relate to recently published work on implicit gradient regularisation.
Revised: Second the FI will allow us to relate to recently published work on implicit gradient regularisation.

Original: Third, use of the FI is standard practice in a variety of work (Khan and Lin, 2017; Khan et al., 2017, 2018; Aitchison, 2018).
Revised: Third, use of the FI is standard practice in a variety of work (Khan and Lin, 2017; Khan et al., 2017, 2018; Aitchison, 2018).

Original: Nonetheless, it is interesting future work to establish whether any other estimates of the diagonal of the Hessian (e.g. from Dangel et al., 2019) improve performance.
Revised: Nonetheless, it is interesting future work to establish whether any other estimates of the diagonal of the Hessian (e.g. from Dangel et al., 2019) improve performance.

Original: The Fisher Information identity allows us to approximate the Hessian,
Revised: The Fisher Information identity allows us to approximate the Hessian,

Original: [Equation]
Revised: [Equation]

Original: Note that the S 2 arises because we deﬁned g j as the average gradient for the minibatch
Revised: Note that the S 2 arises because we deﬁned g j as the average gradient for the minibatch

Original: Eq. 6, whereas the Fisher Information requires the raw log probability, which is formed by the sum.
Revised: Eq. 6, whereas the Fisher Information requires the raw log probability, which is formed by the sum.

Original: Substituting this approximation of the diagonal of the Hessian we can write the quadratic term in the Taylor expansion of the expected log-likelihood as,
Revised: Substituting this approximation of the diagonal of the Hessian we can write the quadratic term in the Taylor expansion of the expected log-likelihood as,

Original: [Equation]
Revised: [Equation]

Original: Thus, our deterministic approximation to the expected log-likelihood is,E
Revised: Thus, our deterministic approximation to the expected log-likelihood is,

Original: [Equation]
Revised: [Equation]

Original: Critically, this quantity is a sum over datapoints, and we can therefore use minibatches of data to give unbiased estimates of the objective in stochastic gradient descent.
Revised: Critically, this quantity is a sum over datapoints, and we can therefore use minibatches of data to give unbiased estimates of the objective in stochastic gradient descent.

Original: Our full approximation to the ELBO is thus,
Revised: Our full approximation to the ELBO is thus,

Original: [Equation]
Revised: [Equation]

Original: We now have to be extremely careful when using minibatch estimates of the gradient.
Revised: We now have to be extremely careful when using minibatch estimates of the gradient.

Original: The easiest approach is to consider x i and y i as minibatches, where there are M minibatches, each with S examples in the full dataset.
Revised: The easiest approach is to consider x i and y i as minibatches, where there are M minibatches, each with S examples in the full dataset.

Original: In that case, the objective for one minibatch can be written
Revised: In that case, the objective for one minibatch can be written

Original: [Equation]
Revised: [Equation]

Original: where, remember, S is the minibatch size and P is the total number of datapoints in the full dataset.
Revised: where, remember, S is the minibatch size and P is the total number of datapoints in the full dataset.

Original: 3.2. Connections to SGD
Revised: 3.2. Connections to SGD

Original: If we set the approximate posterior to that suggested by the stationary distribution of SGD,
Revised: If we set the approximate posterior to that suggested by the stationary distribution of SGD,

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: and we use ﬁxed prior variance, s 2 λ , then the ELBO simpliﬁes further, simply including a log-likelihood, a squared-gradient regulariser and weight-decay,
Revised: and we use ﬁxed prior variance, s 2 λ , then the ELBO simpliﬁes further, simply including a log-likelihood, a squared-gradient regulariser and weight-decay,

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: 3.3. Network architecture
Revised: 3.3. Network architecture

Original: Importantly, here the regulariser is the squared gradient of the loss with respect to the parameters.
Revised: Importantly, here the regulariser is the squared gradient of the loss with respect to the parameters.

Original: As such, computing the loss implicitly involves a second-derivative of the loss, and we therefore should not use piecewise linear activation functions such as ReLU, which have pathological second derivatives.
Revised: As such, computing the loss implicitly involves a second-derivative of the loss, and we therefore should not use piecewise linear activation functions such as ReLU, which have pathological second derivatives.

Original: Instead, we used a softplus activation function, but any activation with well-behaved second derivatives is admissible.
Revised: Instead, we used a softplus activation function, but any activation with well-behaved second derivatives is admissible.

Original: 4. Results
Revised: 4. Results

Original: We trained a PreactResNets-18 (He et al., 2016b) on CIFAR-10 with ﬁxed approximate posteriors inspired by SGD.
Revised: We trained a PreactResNets-18 (He et al., 2016b) on CIFAR-10 with ﬁxed approximate posteriors inspired by SGD.

Original: We used the Adam optimizer (Kingma and Ba, 2014), with an initial learning rate of 1E-4, which decreased by a factor of 10 after 100 epochs and a batch size of 128 with all the other optimizer parameters set to their default values.
Revised: We used the Adam optimizer (Kingma and Ba, 2014), with an initial learning rate of 1E-4, which decreased by a factor of 10 after 100 epochs and a batch size of 128 with all the other optimizer parameters set to their default values.

Original: We have also run experiments whereby Beta-annealing is used with β = 0 .
Revised: We have also run experiments whereby Beta-annealing is used with β = 0 .

Original: 1 (Huang et al., 2018).
Revised: 1 (Huang et al., 2018).

Original: We compared the performance of adding noise to the parameters using “MCVI” (i.e. using samples to evaluate the expectation in Eq. 18) against using our approximate ELBO, and against MAP inference.
Revised: We compared the performance of adding noise to the parameters using “MCVI” (i.e. using samples to evaluate the expectation in Eq. 18) against using our approximate ELBO, and against MAP inference.

Original: We found that our approximate ELBO gave superior performance both to MCVI and MAP.
Revised: We found that our approximate ELBO gave superior performance both to MCVI and MAP.

Original: We also tried MCVI for four times as many epochs (“rescaled”), to compensate for the additional compute required to backpropagate gradients.
Revised: We also tried MCVI for four times as many epochs (“rescaled”), to compensate for the additional compute required to backpropagate gradients.

Original: Remarkably, this actually performed worse, than the standard settings of MCVI and we hypothesise that the additional training time allows it to overﬁt more strongly.
Revised: Remarkably, this actually performed worse, than the standard settings of MCVI and we hypothesise that the additional training time allows it to overﬁt more strongly.

Original: Our approximate ELBO displayed optimal performance with a variance of around σ 2 = 10 − 4 or 10 − 5 (Fig. 1).
Revised: Our approximate ELBO displayed optimal performance with a variance of around σ 2 = 10 − 4 or 10 − 5 (Fig. 1).

Original: To understand whether this is sensible, in comparison with SGD learning rates, we solved for the implied learning rate,
Revised: To understand whether this is sensible, in comparison with SGD learning rates, we solved for the implied learning rate,

Original: [Equation]
Revised: [Equation]

Original: if we were to use, a batch size of S = 50.
Revised: if we were to use, a batch size of S = 50.

Original: These settings correspond to an SGD learning rate of 10 − 2 or 10 − 3 , which is a common ﬁnal learning rate (e.g. if we start at a learning rate of 10 − 1 and decay the learning rate by a factor of 10 once or twice;
Revised: These settings correspond to an SGD learning rate of 10 − 2 or 10 − 3 , which is a common ﬁnal learning rate (e.g. if we start at a learning rate of 10 − 1 and decay the learning rate by a factor of 10 once or twice;

Original: He et al., 2016a,b).
Revised: He et al., 2016a,b).

Original: In addition, our approximate ELBO can be used to learn the approximate posterior variance.
Revised: In addition, our approximate ELBO can be used to learn the approximate posterior variance.

Original: We used a single scalar posterior variance for each convolutional weight matrix, which would correspond to a using a diﬀerent SGD learning rate for each convolutional weight matrix.
Revised: We used a single scalar posterior variance for each convolutional weight matrix, which would correspond to a using a diﬀerent SGD learning rate for each convolutional weight matrix.

Original: Sampling-based MCVI performed poorly, even with small initial variances (initialized to e − 6 times their prior value).
Revised: Sampling-based MCVI performed poorly, even with small initial variances (initialized to e − 6 times their prior value).

Original: In past work on MCVI it was found that achieving even 80% performance on CIFAR-10 required choosing an unusually small network (Ober and Aitchison, 2020) with 32 channels in all layers.
Revised: In past work on MCVI it was found that achieving even 80% performance on CIFAR-10 required choosing an unusually small network (Ober and Aitchison, 2020) with 32 channels in all layers.

Original: Here by contrast, we used a standard-size PreActResNet-18 with 64 channels in the ﬁrst layers, increasing tochannels in the ﬁnal layers.
Revised: Here by contrast, we used a standard-size PreActResNet-18 with 64 channels in the ﬁrst layers, increasing tochannels in the ﬁnal layers.

Original: Remarkably, our approximate ELBO in these larger standard networks was gave superior performance even to MCVI in networks speciﬁcally tuned for variational inference (Table 1).
Revised: Remarkably, our approximate ELBO in these larger standard networks was gave superior performance even to MCVI in networks speciﬁcally tuned for variational inference (Table 1).

Original: 5. Related work
Revised: 5. Related work

Original: Recent work also showed that gradient-descent implies an implicit squared-gradient regularisation (Barrett and Dherin, 2020) with a very similar dependence of the strength of regularisation on the SGD step-size.
Revised: Recent work also showed that gradient-descent implies an implicit squared-gradient regularisation (Barrett and Dherin, 2020) with a very similar dependence of the strength of regularisation on the SGD step-size.

Original: Remarkably, this arose despite their taking a radically diﬀerent approach.
Revised: Remarkably, this arose despite their taking a radically diﬀerent approach.

Original: In particular, they showed that squared-gradient regularisation can arise through errors when using ﬁnite step size Euler integration, as compared to following the true underlying gradient ﬂow.
Revised: In particular, they showed that squared-gradient regularisation can arise through errors when using ﬁnite step size Euler integration, as compared to following the true underlying gradient ﬂow.

Original: However, this approach was derived for gradient descent, and does not take into account the fundamental stochasticity of stochastic gradient descent.
Revised: However, this approach was derived for gradient descent, and does not take into account the fundamental stochasticity of stochastic gradient descent.

Original: The emergence of such similar results from fundamentally very diﬀerent analytical approaches would suggest that gradient regularisation is a fundamental component of the implicit regularisation inherent in SGD.
Revised: The emergence of such similar results from fundamentally very diﬀerent analytical approaches would suggest that gradient regularisation is a fundamental component of the implicit regularisation inherent in SGD.

Original: [Table]
Revised: [Table]

Original: 6. Conclusions
Revised: 6. Conclusions

Original: We showed that it is possible to use a second-order Taylor expansion to compute a deterministic approximation of the ELBO in Bayesian neural networks.
Revised: We showed that it is possible to use a second-order Taylor expansion to compute a deterministic approximation of the ELBO in Bayesian neural networks.

Original: While any drop-in estimate of the diagonal of the Hessian could be used, here we consider the squared gradient (via the Fisher Information)
Revised: While any drop-in estimate of the diagonal of the Hessian could be used, here we consider the squared gradient (via the Fisher Information)

Original: [Table]
Revised: [Table]

Original: as this allows eﬃcient optimization and allows us to connect to gradient regularisation (Barrett and Dherin, 2020).
Revised: as this allows eﬃcient optimization and allows us to connect to gradient regularisation (Barrett and Dherin, 2020).

Original: Further, we connect to SGD by noting that the stationary distribution for SGD is isotropic Gaussian (Mandt et al., 2017).
Revised: Further, we connect to SGD by noting that the stationary distribution for SGD is isotropic Gaussian (Mandt et al., 2017).

Original: This approximate ELBO gives superior performance to unbiased sampling-based estimation of the ELBO, and to MAP inference.
Revised: This approximate ELBO gives superior performance to unbiased sampling-based estimation of the ELBO, and to MAP inference.

Original: Finally, our framework makes a prediction: that the optimal posterior variance for the ELBO corresponds to standard settings of the SGD learning rate.
Revised: Finally, our framework makes a prediction: that the optimal posterior variance for the ELBO corresponds to standard settings of the SGD learning rate.

Original: Remarkably, we found that this prediction held, with the optimal posterior variance in ResNets on CIFAR-10 being 10 − 4 or 10 − 5 corresponds to standard values for the ﬁnal learning rate in SGD used in these models.
Revised: Remarkably, we found that this prediction held, with the optimal posterior variance in ResNets on CIFAR-10 being 10 − 4 or 10 − 5 corresponds to standard values for the ﬁnal learning rate in SGD used in these models.

Original: Appendix A. Variational Inference for Bayesian Neural Networks
Revised: Appendix A. Variational Inference for Bayesian Neural Networks

Original: We optimize the evidence lower bound objective with respect to parameters of the variational posterior,
Revised: We optimize the evidence lower bound objective with respect to parameters of the variational posterior,

Original: [Equation]
Revised: [Equation]

Original: where x i is the training inputs and y i is the training outputs.
Revised: where x i is the training inputs and y i is the training outputs.

Original: Unfortunately, we need to optimize this expectation with respect to parameters of Q ( w ), the distribution over which the expectation is taken.
Revised: Unfortunately, we need to optimize this expectation with respect to parameters of Q ( w ), the distribution over which the expectation is taken.

Original: We therefore use the reparameterisation trick (Kingma and
Revised: We therefore use the reparameterisation trick (Kingma and

Original: Welling, 2013; Rezende et al., 2014) — we write w in terms of (cid:15) ,
Revised: Welling, 2013; Rezende et al., 2014) — we write w in terms of (cid:15) ,

Original: [Equation]
Revised: [Equation]

Original: 19) where (cid:15) λ ∼ N (0 , 1).
Revised: 19) where (cid:15) λ ∼ N (0 , 1).

Original: Thus, the ELBO can be written as an expectation over (cid:15) ,
Revised: Thus, the ELBO can be written as an expectation over (cid:15) ,

Original: [Equation]
Revised: [Equation]

Original: However, exactly evaluating this objective remains intractable, as it still involves a highdimensional expectation.
Revised: However, exactly evaluating this objective remains intractable, as it still involves a highdimensional expectation.

Original: Instead, we typically form an unbiased estimator, by drawing a few (often one) sample of w or (cid:15) .
Revised: Instead, we typically form an unbiased estimator, by drawing a few (often one) sample of w or (cid:15) .

Original: While this approach can be eﬀective, it gives high-variance estimates of the gradient, which can cause slow-learning (Wu et al., 2018; Rainforth et al., 2018).
Revised: While this approach can be eﬀective, it gives high-variance estimates of the gradient, which can cause slow-learning (Wu et al., 2018; Rainforth et al., 2018).

Original: Appendix B. The Stationary Distribution of SGD
Revised: Appendix B. The Stationary Distribution of SGD

Original: We then sought to relate these gradient regularisers back to the excellent empirical performance of SGD.
Revised: We then sought to relate these gradient regularisers back to the excellent empirical performance of SGD.

Original: 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).
Revised: 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).

Original: In particular, consider a loss function which is locally closely approximated by a quadratic.
Revised: In particular, consider a loss function which is locally closely approximated by a quadratic.

Original: Without loss of generality, we consider a mode at w = 0 ,
Revised: Without loss of generality, we consider a mode at w = 0 ,

Original: [Equation]
Revised: [Equation]

Original: where P is the total number of datapoints.
Revised: where P is the total number of datapoints.

Original: Typically, the objective used in SGD is the loss for a minibatch of size S with indices S .
Revised: Typically, the objective used in SGD is the loss for a minibatch of size S with indices S .

Original: Following Mandt et al.
Revised: Following Mandt et al.

Original: (2017), we use the Fisher
Revised: (2017), we use the Fisher

Original: Information to identify the noise in the minibatch gradient,
Revised: Information to identify the noise in the minibatch gradient,

Original: [Equation]
Revised: [Equation]

Original: (22) where ξ ( t ) is sampled from a standard IID Gaussian.
Revised: (22) where ξ ( t ) is sampled from a standard IID Gaussian.

Original: For SGD, this gradient is multiplied by a learning rate, η ,
Revised: For SGD, this gradient is multiplied by a learning rate, η ,

Original: [Equation]
Revised: [Equation]

Original: This is an multivariate Gaussian autoregressive process, so we can solve for the stationary distribution of the weights.
Revised: This is an multivariate Gaussian autoregressive process, so we can solve for the stationary distribution of the weights.

Original: In particular, we note that the covariance at time t + 1 is
Revised: In particular, we note that the covariance at time t + 1 is

Original: [Equation]
Revised: [Equation]

Original: Following Mandt et al.
Revised: Following Mandt et al.

Original: (2017), when the learning rate is small, the quadratic term can be neglected.
Revised: (2017), when the learning rate is small, the quadratic term can be neglected.

Original: [Equation]
Revised: [Equation]

Original: We then solve for steady-state in which Σ
Revised: We then solve for steady-state in which Σ

Original: = C [ w ( t + 1)]
Revised: = C [ w ( t + 1)]

Original: =
Revised: =

Original: C [ w ( t )],
Revised: C [ w ( t )],

Original: [Equation]
Revised: [Equation]

Original: so,
Revised: so,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: We therefore considered a family of approximate posteriors matching these stationary distributions, by setting all variances, σ 2 λ , to a single, ﬁxed value, η 2 S .
Revised: We therefore considered a family of approximate posteriors matching these stationary distributions, by setting all variances, σ 2 λ , to a single, ﬁxed value, η 2 S .