Variational Inference on the Final-Layer Output of Neural Networks

TL;DR

This paper introduces VIFO, a method that applies Variational Inference directly to the neural network's output layer, balancing training simplicity with improved uncertainty estimation, especially for out-of-distribution data.

Contribution

It proposes a novel approach combining variational inference with neural network output layers, enhancing uncertainty quantification while maintaining efficient training.

Findings

VIFO improves uncertainty estimation for out-of-distribution data.

The method offers a good tradeoff between computational efficiency and uncertainty quality.

VIFO allows risk bounds via Rademacher complexity similar to non-Bayesian models.

Abstract

Traditional neural networks are simple to train but they typically produce overconfident predictions. In contrast, Bayesian neural networks provide good uncertainty quantification but optimizing them is time consuming due to the large parameter space. This paper proposes to combine the advantages of both approaches by performing Variational Inference in the Final layer Output space (VIFO), because the output space is much smaller than the parameter space. We use neural networks to learn the mean and the variance of the probabilistic output. Using the Bayesian formulation we incorporate collapsed variational inference with VIFO which significantly improves the performance in practice. On the other hand, like standard, non-Bayesian models, VIFO enjoys simple training and one can use Rademacher complexity to provide risk bounds for the model. Experiments show that VIFO provides a good tradeoff in terms of run time and uncertainty quantification, especially for out of distribution data.