# Variational Inference on the Final-Layer Output of Neural Networks

arXiv: https://arxiv.org/abs/2302.02420 · HTML: https://mlsanity.com/paper/2302.02420/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 50 references that are in the map.
- First posted on arXiv 5 Feb 2023.

## Builds on (most cited first, 29 of 29)

- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Predictive Uncertainty Estimation via Prior Networks](https://mlsanity.com/paper/1802.10501/) (2018)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://mlsanity.com/paper/2006.10108/) (2020)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [What Are Bayesian Neural Network Posteriors Really Like?](https://mlsanity.com/paper/2104.14421/) (2021)
- [Size-Independent Sample Complexity of Neural Networks](https://mlsanity.com/paper/1712.06541/) (2017)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Posterior Network: Uncertainty Estimation without OOD Samples via Density-Based Pseudo-Counts](https://mlsanity.com/paper/2006.09239/) (2020)
- [Functional Variational Bayesian Neural Networks](https://mlsanity.com/paper/1903.05779/) (2019)
- [Repulsive Deep Ensembles are Bayesian](https://mlsanity.com/paper/2106.11642/) (2021)
- [PAC-Bayesian Theory Meets Bayesian Inference](https://mlsanity.com/paper/1605.08636/) (2016)
- [Pitfalls of Epistemic Uncertainty Quantification through Loss Minimisation](https://mlsanity.com/paper/2203.06102/) (2022)
- [Do Bayesian Neural Networks Need To Be Fully Stochastic?](https://mlsanity.com/paper/2211.06291/) (2022)
- [All You Need is a Good Functional Prior for Bayesian Deep Learning](https://mlsanity.com/paper/2011.12829/) (2020)
- [On Last-Layer Algorithms for Classification: Decoupling Representation from Uncertainty Estimation](https://mlsanity.com/paper/2001.08049/) (2020)
- [Stochastic Expectation Propagation](https://mlsanity.com/paper/1506.04132/) (2015)
- [Distributed Bayesian Learning with Stochastic Natural-gradient Expectation Propagation and the Posterior Server](https://mlsanity.com/paper/1512.09327/) (2015)
- [Direct loss minimization algorithms for sparse Gaussian processes](https://mlsanity.com/paper/2004.03083/) (2020)
- [On the Performance of Direct Loss Minimization for Bayesian Neural Networks](https://mlsanity.com/paper/2211.08393/) (2022)

## Abstract (arXiv)

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.