Function-Space Regularization for Deep Bayesian Classification

TL;DR

This paper introduces a novel function-space regularization method for deep Bayesian classification that employs a Dirichlet prior in predictive space, enhancing uncertainty estimation and robustness without altering model architecture.

Contribution

It proposes a new approach to Bayesian deep learning by applying a Dirichlet prior in predictive space, enabling flexible, scalable, and interpretable uncertainty quantification across models.

Findings

Improved uncertainty quantification in image classification.

Enhanced adversarial robustness demonstrated on large-scale datasets.

Flexible integration with various neural network architectures.

Abstract

Bayesian deep learning approaches assume model parameters to be latent random variables and infer posterior distributions to quantify uncertainty, increase safety and trust, and prevent overconfident and unpredictable behavior. However, weight-space priors are model-specific, can be difficult to interpret and are hard to specify. Instead, we apply a Dirichlet prior in predictive space and perform approximate function-space variational inference. To this end, we interpret conventional categorical predictions from stochastic neural network classifiers as samples from an implicit Dirichlet distribution. By adapting the inference, the same function-space prior can be combined with different models without affecting model architecture or size. We illustrate the flexibility and efficacy of such a prior with toy experiments and demonstrate scalability, improved uncertainty quantification and adversarial robustness with large-scale image classification experiments.