On Batch Normalisation for Approximate Bayesian Inference
Abstract
We study batch normalisation in the context of variational inference methods in Bayesian neural networks, such as mean-ﬁeld or MC Dropout.
We show that batch-normalisation does not aﬀect the optimum of the evidence lower bound (ELBO).
Furthermore, we study the Monte Carlo Batch Normalisation (MCBN) algorithm, proposed as an approximate inference technique parallel to MC Dropout, and show that for larger batch sizes, MCBN fails to capture epistemic uncertainty.
Finally, we provide insights into what is required to ﬁx this failure, namely having to view the mini-batch size as a variational parameter in MCBN.
We comment on the asymptotics of the ELBO with respect to this variational parameter, showing that as dataset size increases towards inﬁnity, the batch-size must increase towards inﬁnity as well for MCBN to be a valid approximate inference technique.
1. Introduction
Deep neural networks have achieved tremendous success in several ﬁelds and are now being applied to various important real-life applications like medical diagnosis (Esteva et al., 2017) and autonomous driving (Levinson et al., 2011), where it is particularly important to produce reliable uncertainty estimates in addition to predictions.
The Bayesian formalism (Neal, 2012; MacKay, 1992) provides a principled approach for modeling uncertainty in deep neural networks.
However, performing exact Bayesian inference in deep neural networks is computationally intractable.
In this regard, one of the well-known and accepted methods of performing approximate Bayesian inference is variational inference (VI) (Graves, 2011;
Hinton and Van Camp, 1993) in which an approximating distribution is used instead of the true Bayesian posterior over model parameters.
There have been several works (Graves, 2011; Hern´andez-Lobato and Adams, 2015;
Blundell et al., 2015; Gal and Ghahramani, 2016) which have tried to apply VI in deep neural networks.
Among them, (Gal and Ghahramani, 2016) provided a practical solution by showing that any neural network trained using dropout layers is performing VI and can be treated as a Bayesian neural network.
Their method, MC Dropout, uses dropout at test time to sample from the approximate posterior in a deep neural network.
Following on (Gal and Ghahramani, 2016), the paper (Teye et al., 2018) claimed that batch normalisation (BN) layers can also be used to cast deep neural networks as Bayesian models.
Their method, Monte Carlo Batch Normalisation (MCBN) leverages the randomness of mini-batch statistics computed by BN layers due to the stochasticity of mini-batches drawn from the training set.
The running mini-batch statistics (mean and variance) of BN layers are stochastic estimators of the true population statistics .
This observation has direct implications on how batch normalisation can or cannot be used to perform VI in deep neural networks.
In this paper, we use this observation to ﬁrst show that adding BN layers in deep neural networks does not aﬀect the probabilistic inference of existing VI methods like MC Dropout (Gal and Ghahramani, 2016) or Bayes by Backprop (Blundell et al., 2015).
Secondly, we also expose a failure case of MCBN.
In particular, we see that by increasing the mini-batch size, the epistemic uncertainty (Kendall and Gal, 2017) estimates produced by a model reduce to zero as BN layers become deterministic.
We empirically see how this renders a model completely incapable of distinguishing between in-distribution and out-of-distribution samples.
Finally, we provide insight into what is required to ﬁx this failure, namely having to view the mini-batch size as a variational parameter in MCBN.
We comment on the asymptotics of the ELBO with respect to this variational parameter, showing that as dataset size increases towards inﬁnity, the batch-size must increase towards inﬁnity as well for MCBN to be a valid approximate inference technique.
2. Batch Normalisation with Variational Inference
Background on Variational Inference (VI):
In a Bayesian framework, we are interested in ﬁnding the posterior distribution p
( θ | X , Y ) over model parameters θ given data
D =
{ X , Y } =
{ ( x i , y i )
Ni =1 } .
However, ﬁnding the true posterior for modern neural networks is computationally intractable.
Hence, in variational inference, we ﬁnd an approximating parameterised distribution q φ ( θ ) (also called the variational distribution with parameters φ ) over model weights θ to approximate the true posterior.
We want the variational distribution to be as close to the true posterior as possible.
This is done by minimising the
KL divergence KL( q φ ( θ ) || p ( θ | X , Y ) between q φ ( θ ) and the true posterior p ( θ | X , Y ).
Minimising this KL divergence is equivalent to maximising the evidence lower bound (ELBO) (or minimising its negative):
[Equation]
In variational inference, we minimise L ELBO ( φ ) to ﬁnd the optimal variational distribution q ∗ φ ( θ ) which explains our data.
Note that the ELBO objective ( L ELBO ( φ )) is over variational parameters φ and not model weights θ .
Does Batch Normalisation (BN) impact VI?
The batch normalisation (Ioﬀe and
Szegedy, 2015; LeCun et al., 2012) operation is used to standardise each layer’s input to have zero mean and unit variance.
Given input h l for a fully connected layer l , the BN operation uses two data-dependent statistics: the mini-batch mean µ lB and the mini-batch variance ( σ 2 ) lB to compute:
[Equation]
where γ l and β l are the scale and shift parameters which are learnt through backpropagation.
If we compute L ELBO ( φ ) using the entire dataset D
= { ( x i , y i )
Ni =1 } as shown in equation 1, the mini-batch statistics { µ lB , ( σ 2 ) lB } are exactly equal to the population statistics { µ l , ( σ 2 ) l } and hence, are ﬁxed.
However, in practical situations, the entire dataset is not used while optimising L ELBO ( φ ) and mini-batch optimisation methods (Hoﬀman et al., 2013) are employed.
In such cases, the ˆ L ELBO ( φ ) computed from a minibatch B is:
[Equation]
During training, BN layers keep a running average of means and variances (batch statistics) which get updated with each incoming mini-batch from the training set (Ioﬀe and Szegedy, 2015).
Following the law of large numbers, in expectation, the running batch statistics converge to the population statistics { µ l , ( σ 2 ) l } and hence, ˆ L ELBO ( φ ) computed from a mini-batch forms an unbiased estimator of L ELBO ( φ ) computed using the entire dataset, i.e.,
E B [ ˆ L ELBO ( φ )]
=
L ELBO ( φ ).
There is no change to the optimum of the ELBO L ELBO ( φ ) due to the presence of BN layers.
Thus, batch normalisation does not impact the probabilistic inference of other VI methods like Mean Field VI (Graves, 2011; Blundell et al., 2015) or
MC Dropout (Gal and Ghahramani, 2016).
3. Stochastic Batch Normalisation (MCBN) as VI
What is MCBN?
In the work (Teye et al., 2018), the authors propose using batch normalisation as a way of performing VI in deep neural networks.
In a nutshell, they categorise the parameters of a neural network into:
Learnable Parameters: θ = { ( W l , β l , γ l ) Ll =1 } , weight matrices and ( β, γ ) BN parameters learnable through backpropagation for a neural network with L layers.
Stochastic Parameters: w = { ( µ lB , ( σ 2 ) lB ) Ll =1 } , the running average batch statistics (mean and variance) for all BN layers.
In their method, they ﬁrst train a network with BN layers using a conventional optimiser (like SGD or Adam etc.) and loss function (like cross entropy).
During test time, they keep the BN layers in training mode and for each test sample x ∗ , they perform multiple stochastic forward passes using diﬀerent, randomly chosen mini-batches from the training set.
Each such forward pass corresponds to a particular instantiation of w = { ( µ lB , ( σ 2 ) lB ) Ll =1 } and provides a single output ˆy ∗ t (for the t th stochastic forward pass, say).
The mean E [ y ] ≈ 1 T (cid:80) Tt =1 ˆy ∗ t is then used as the prediction and some measure of variance of these outputs can be interpreted as an uncertainty estimate.
It is important to note here that the mini-batch size used during stochastic forward passes at test time is the same as the mini-batch size used for training.
This method was named Monte Carlo Batch Normalisation (MCBN).
Failure case of MCBN:
In MCBN, the randomness of model outputs at test time is solely dependent on the stochastic parameters w = { ( µ l B , ( σ 2 ) l B ) L l =1 } , which in turn are dependent on the mini-batch size.
As mentioned before, in case of larger mini-batch sizes, the batch statistics will converge to the population statistics, which is in fact desirable [Figure] for BN layers following their original motivation.
When the batch size equals the dataset size [Figure] there will be no stochasticity at all, and ‘full BN’ is recovered (the optimal case for regularisation).
However, in this case, any measure of epistemic uncertainty which depends on variation of model outputs over stochastic forward passes will reduce to zero and hence, the resulting model won’t be able to capture epistemic uncertainty at all
[Figure]
We empirically validate this claim by showing the eﬀect of increasing the mini-batch size in MCBN on the epistemic uncertainty of the model.
As epistemic uncertainty should be high on unseen data (Kendall and Gal, 2017), we choose a notably diﬃcult pair of in-distribution and OOD datasets (Nalisnick et al., 2019), (CIFAR-10 (Krizhevsky et al., 2009) vs. SVHN (Netzer et al., 2011)) and (Fashion-MNIST (Xiao et al., 2017) vs MNIST (LeCun)).
We train a ResNet-50 on CIFAR-10 and Fashion-MNIST and we train a VGG16 on Fashion-MNIST using diﬀerent training batch sizes from 32 to 2048 (see Table 1 for speciﬁc batch sizes).
We use SGD with a momentum of 0.9 as the optimiser and crossentropy as the objective function.
We train the ResNet-50 on CIFAR-10 for 350 epochs with an initial learning rate of 0.1 and drop the learning rate by a factor of 10 at epochs 150 and 250.
For the models trained on Fashion-MNIST we use 100 training epochs and drop the learning rate at epochs 40 and 60.
During test time, we use 20 stochastic forward passes to make predictions for each model.
In order to compute the epistemic uncertainty, we use the mutual information (Gal, 2016) between the predictive and the posterior distributions as the uncertainty estimate of choice.
In Table 1, we show the test set errors of the models for every batch size.
Note that we did not train ResNet-50 on CIFAR-10 with batch sizes 1024 and 2048 due to compute constraints.
Finally, we show the variation of epistemic uncertainty with batch size for both in-distribution and OOD datasets in Figure 1.
It is very clear from Figure 1 that with increase in batch sizes, the model’s epistemic uncertainty lowers to the point where it can no longer distinguish between in and out of distribution samples .
In fact, for the models trained on Fashion MNIST, we ﬁnd the epistemic uncertainty to almost become 0 for batch sizes 1024 and 2048.
Thus, although the test set accuracy remains more or less the same, MCBN fails to capture epistemic uncertainty with growing mini-batch size. [Table]
Insight into what is required to ﬁx this failure:
The model’s stochasticity in MCBN is clearly dependent on the mini-batch size used for training.
Hence, the mini-batch size itself should be considered as a variational parameter much like the dropout rate is a variational parameter in MC Dropout (Gal and Ghahramani, 2016).
One can ﬁnd the best dropout rate either by learning it directly (Gal et al., 2017) or performing grid-search over possible dropout rates (Mukhoti et al., 2018).
Similarly, one could think of performing a grid-search over possible mini-batch sizes and ﬁnd the one which maximises the ELBO.
Comment on ELBO asymptotics:
We brieﬂy comment on the asymptotics of the
ELBO with respect to this variational parameter.
As the dataset size increases towards inﬁnity, we expect epistemic uncertainty of a parametric model to diminish to zero.
However, as we commented above, for the epistemic uncertainty of MCBN to diminish to zero, the batch-size must increase towards inﬁnity as well.
This may be impractical in certain applications.
4. Conclusion
In this paper we have studied how batch normalisation layers ﬁt into the context of variational inference in deep neural networks.
Mini-batch statistics in BN layers are stochastic estimators of the population statistics.
We use this observation to show that batch normalisation does not aﬀect the probabilistic inference of VI approaches.
Furthermore, we use the same observation to expose a failure case of the MCBN (Teye et al., 2018) algorithm and see that for bigger batch sizes, MCBN is unable to produce sensible epistemic uncertainty and cannot distinguish in and out of distribution samples.
Finally, we provide some insights into how to ﬁx the problem and leave the development of more involved approaches on MCBN as future work.