# Last Layer Marginal Likelihood for Invariance Learning

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 45 references that are in the map.
- First posted on arXiv 14 Jun 2021.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Rotation Equivariant CNNs for Digital Pathology](https://mlsanity.com/paper/1806.03962/) (2018)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Scalable Variational Gaussian Process Classification](https://mlsanity.com/paper/1411.2005/) (2014)
- [Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules](https://mlsanity.com/paper/1905.05393/) (2019)
- [GPflow: A Gaussian process library using TensorFlow](https://mlsanity.com/paper/1610.08733/) (2016)
- [Efficiently Sampling Functions from Gaussian Process Posteriors](https://mlsanity.com/paper/2002.09309/) (2020)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning](https://mlsanity.com/paper/2104.04975/) (2021)
- [Stochastic Variational Deep Kernel Learning](https://mlsanity.com/paper/1611.00336/) (2016)
- [On Feature Collapse and Deep Kernel Learning for Single Forward Pass Uncertainty](https://mlsanity.com/paper/2102.11409/) (2021)
- [The Promises and Pitfalls of Deep Kernel Learning](https://mlsanity.com/paper/2102.12108/) (2021)
- [Manifold Gaussian Processes for Regression](https://mlsanity.com/paper/1402.5876/) (2014)
- [Meta-Learning Symmetries by Reparameterization](https://mlsanity.com/paper/2007.02933/) (2020)
- [On Sparse variational methods and the Kullback-Leibler divergence between stochastic processes](https://mlsanity.com/paper/1504.07027/) (2015)
- [Convergence of Sparse Variational Inference in Gaussian Processes Regression](https://mlsanity.com/paper/2008.00323/) (2020)
- [Learning Invariances in Neural Networks](https://mlsanity.com/paper/2010.11882/) (2020)
- [Learning Invariances using the Marginal Likelihood](https://mlsanity.com/paper/1808.05563/) (2018)
- [Adversarial Examples, Uncertainty, and Transfer Testing Robustness in Gaussian Process Hybrid Deep Networks](https://mlsanity.com/paper/1707.02476/) (2017)
- [A Framework for Interdomain and Multioutput Gaussian Processes](https://mlsanity.com/paper/2003.01115/) (2020)
- [Global inducing point variational posteriors for Bayesian neural networks and deep Gaussian processes](https://mlsanity.com/paper/2005.08140/) (2020)
- [Data augmentation in Bayesian neural networks and the cold posterior effect](https://mlsanity.com/paper/2106.05586/) (2021)
- [Deep Neural Networks as Point Estimates for Deep Gaussian Processes](https://mlsanity.com/paper/2105.04504/) (2021)
- [Bayesian Image Classification with Deep Convolutional Gaussian Processes](https://mlsanity.com/paper/1902.05888/) (2019)

## Built on by (most cited first, 21 of 21)

- [Variational Bayesian Last Layers](https://mlsanity.com/paper/2404.11599/) (2024)
- [Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI](https://mlsanity.com/paper/2402.00809/) (2024)
- [Bayesian Model Selection, the Marginal Likelihood, and Generalization](https://mlsanity.com/paper/2202.11678/) (2022)
- [Invariance Learning in Deep Neural Networks with Differentiable Laplace Approximations](https://mlsanity.com/paper/2202.10638/) (2022)
- [Relaxing Equivariance Constraints with Non-stationary Continuous Filters](https://mlsanity.com/paper/2204.07178/) (2022)
- [Learning Layer-wise Equivariances Automatically using Gradients](https://mlsanity.com/paper/2310.06131/) (2023)
- [Gaussian Stochastic Weight Averaging for Bayesian Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2405.03425/) (2024)
- [Learning Invariant Weights in Neural Networks](https://mlsanity.com/paper/2202.12439/) (2022)
- [Improving Neural Additive Models with Bayesian Principles](https://mlsanity.com/paper/2305.16905/) (2023)
- [Learning Augmentation Distributions using Transformed Risk Minimization](https://mlsanity.com/paper/2111.08190/) (2021)
- [Modulated Neural ODEs](https://mlsanity.com/paper/2302.13262/) (2023)
- [Promises and Pitfalls of the Linearized Laplace in Bayesian Optimization](https://mlsanity.com/paper/2304.08309/) (2023)
- [Stochastic Marginal Likelihood Gradients using Neural Tangent Kernels](https://mlsanity.com/paper/2306.03968/) (2023)
- [A Generative Model of Symmetry Transformations](https://mlsanity.com/paper/2403.01946/) (2024)
- [Hyperparameter Optimization through Neural Network Partitioning](https://mlsanity.com/paper/2304.14766/) (2023)
- [Noether's razor: Learning Conserved Quantities](https://mlsanity.com/paper/2410.08087/) (2024)
- [Evaluation of machine learning architectures on the quantification of epistemic and aleatoric uncertainties in complex dynamical systems](https://mlsanity.com/paper/2306.15159/) (2023)
- [A tradeoff between universality of equivariant models and learnability of symmetries](https://mlsanity.com/paper/2210.09444/) (2022)
- [On Equivariant Model Selection through the Lens of Uncertainty](https://mlsanity.com/paper/2506.18629/) (2025)
- [Uncertainty in Graph Contrastive Learning with Bayesian Neural Networks](https://mlsanity.com/paper/2312.00232/) (2023)
- [Twin: Tuning Learning Rate and Weight Decay of Deep Homogeneous Classifiers without Validation](https://mlsanity.com/paper/2403.05532/) (2024)

## Abstract (arXiv)

Data augmentation is often used to incorporate inductive biases into models. Traditionally, these are hand-crafted and tuned with cross validation. The Bayesian paradigm for model selection provides a path towards end-to-end learning of invariances using only the training data, by optimising the marginal likelihood. Computing the marginal likelihood is hard for neural networks, but success with tractable approaches that compute the marginal likelihood for the last layer only raises the question of whether this convenient approach might be employed for learning invariances. We show partial success on standard benchmarks, in the low-data regime and on a medical imaging dataset by designing a custom optimisation routine. Introducing a new lower bound to the marginal likelihood allows us to perform inference for a larger class of likelihood functions than before. On the other hand, we demonstrate failure modes on the CIFAR10 dataset, where the last layer approximation is not sufficient due to the increased complexity of our neural network. Our results indicate that once more sophisticated approximations become available the marginal likelihood is a promising approach for invariance learning in neural networks.