Stable uncertainty transfer speeds up training of large neural networks
$σ$Transfer: Uncertainty Transfer from Small to Large Networks under $μ\mathrm{P}$
Machine Learning
Summary
Predicting how sure a large neural network is about its predictions usually requires a lot of expensive calculations. The authors show a way to pick an important setting, called prior precision, on a small network and reuse it directly on a much bigger one without extra searching. This speeds up the process by thousands of times in some cases, while keeping predictions nearly as good. Their method works across different tasks like image recognition and language models.
What this means in practice
- •For machine learning engineers: Speed up hyperparameter tuning of uncertainty for very large neural networks by selecting it once on a smaller model.
- •For ai model deployment teams: Enable reliable uncertainty estimates and related decisions like abstention on large models without costly additional tuning.
Authors
Richard Bergna, Fernando Ruiz Mazo, Nicolò Felicioni, José Miguel Hernández-Lobato, Kamil Ciosek
Abstract
Reliable predictive uncertainty in Laplace approximations depends critically on the prior precision, yet selecting it requires a posterior sweep that is prohibitively expensive for neural networks with billions of parameters. Under the Maximal Update Parametrization ($μ\mathrm{P}$), we derive a rescaling of the prior covariance that makes the selected precision stable as model width grows. This leads to $σ\mathrm{Transfer}$: we select the precision on a smaller model and zero-shot transfer it to the much larger model, i.e., without searching for the precision on the larger model at all. We show convergence of the prior kernel, posterior covariance, selected precision, and posterior-derived decisions under explicit conditions, and verify $σ\mathrm{Transfer}$ across regression, image classification, and Transformer readouts. For example, measured precision-sweep speedups reach $\sim 5000\times$ when transferring from width 128 to 4096 on MNIST, at a target-NLL degradation of $0.002$; transferring from a public 1B to 7B model gives a median search speedup of $\sim 2.3\times$ (up to $\sim 330\times$), with a mean measured target-NLL increase below $10^{-4}$ across ten tasks. The same posterior stability also enables transfer of acquisition, OOD-detection, and abstention decisions without constructing a target posterior.