Tackling covariate shift with node-based Bayesian neural networks
Bayesian neural networks (BNNs) promise improved generalization under
covariate shift by providing principled probabilistic representations of
epistemic uncertainty. However, weight-based BNNs often struggle with high
computational complexity of large-scale architectures and datasets. Node-based
BNNs have recently been introduced as scalable alternatives, which induce
epistemic uncertainty by multiplying each hidden node with latent random
variables, while learning a point-estimate of the weights. In this paper, we
interpret these latent noise variables as implicit representations of simple
and domain-agnostic data perturbations during training, producing BNNs that
perform well under covariate shift due to input corruptions. We observe that
the diversity of the implicit corruptions depends on the entropy of the latent
variables, and propose a straightforward approach to increase the entropy of
these variables during training. We evaluate the method on out-of-distribution
image classification benchmarks, and show improved uncertainty estimation of
node-based BNNs under covariate shift due to input perturbations. As a side
effect, the method also provides robustness against noisy training labels.