[Paper Review] On Batch Normalisation for Approximate Bayesian Inference

Jishnu Mukhoti, Puneet K. Dokania+2 more|arXiv (Cornell University)|Dec 24, 2020

Gaussian Processes and Bayesian Inference18 references4 citations

TL;DR

This paper investigates batch normalization (BN) in the context of approximate Bayesian inference in deep neural networks, showing that BN does not affect the evidence lower bound (ELBO) optimum in variational inference. It identifies a critical failure in Monte Carlo Batch Normalization (MCBN): as mini-batch size increases, epistemic uncertainty collapses to zero, rendering models unable to detect out-of-distribution samples. The authors propose treating mini-batch size as a variational parameter to restore uncertainty calibration, especially under large dataset asymptotics.

ABSTRACT

We study batch normalisation in the context of variational inference methods in Bayesian neural networks, such as mean-field or MC Dropout. We show that batch-normalisation does not affect 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 fix 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 infinity, the batch-size must increase towards infinity as well for MCBN to be a valid approximate inference technique.

Motivation & Objective

• To analyze the role of batch normalization (BN) in variational inference (VI) for Bayesian neural networks.
• To investigate the validity of Monte Carlo Batch Normalization (MCBN) as a method for approximate Bayesian inference.
• To identify the failure mode of MCBN when applied with large mini-batch sizes.
• To propose a principled fix by treating mini-batch size as a variational parameter in MCBN.
• To analyze the asymptotic behavior of the ELBO with respect to mini-batch size under large dataset limits.

Proposed method

• The authors formalize the ELBO for Bayesian neural networks with BN layers, showing that BN does not alter the optimal solution of the ELBO objective.
• They analyze MCBN as a stochastic approximation method, where BN's mini-batch statistics provide the source of model uncertainty.
• They derive that the stochasticity in MCBN stems solely from the mini-batch size, making it a key source of epistemic uncertainty.
• They propose treating the mini-batch size as a variational parameter, analogous to dropout rate in MC Dropout, to optimize uncertainty calibration.
• They perform empirical evaluation on CIFAR-10, SVHN, and Fashion-MNIST/MNIST pairs using mutual information as an uncertainty metric.
• They analyze the asymptotic behavior of the ELBO, showing that for MCBN to remain valid, batch size must grow with dataset size.

Experimental results

Research questions

• RQ1 Does batch normalization affect the optimal solution of the evidence lower bound (ELBO) in variational inference?
• RQ2 Can Monte Carlo Batch Normalization (MCBN) reliably capture epistemic uncertainty across different mini-batch sizes?
• RQ3 Why does epistemic uncertainty in MCBN degrade to zero as batch size increases?
• RQ4 What is the asymptotic behavior of the ELBO when the dataset size grows, under MCBN?
• RQ5 Can treating mini-batch size as a variational parameter restore uncertainty calibration in MCBN?

Key findings

• Batch normalization does not affect the optimum of the evidence lower bound (ELBO), meaning BN layers do not alter the fundamental variational inference objective.
• MCBN fails to preserve epistemic uncertainty as mini-batch size increases, with mutual information between predictive and posterior distributions dropping to near zero at batch sizes of 1024 and 2048.
• On OOD detection tasks, models trained with large batch sizes using MCBN lose the ability to distinguish in-distribution from out-of-distribution samples, despite maintaining high test accuracy.
• The failure of MCBN is attributed to the deterministic nature of BN statistics at large batch sizes, eliminating the stochasticity that drives uncertainty estimation.
• The authors show that for MCBN to remain a valid approximate inference method, the mini-batch size must increase in proportion to the dataset size as the dataset grows toward infinity.
• Treating mini-batch size as a variational parameter—similar to how dropout rate is treated—can restore uncertainty calibration in MCBN, suggesting a path toward more robust Bayesian deep learning.