Optimality of Sub-network Laplace Approximations: New Results and Methods

Authors: Swarnali Raha, Kshitij Khare, Rohit K Patra

arXiv · PDF

Summary

Laplace approximation provides a way to quantify uncertainty in neural networks by approximating the posterior with a Gaussian centered at the MAP estimate, but inverting the full Hessian is computationally prohibitive. Sub-network Laplace methods restrict attention to a subset of parameters to make it tractable. This paper proves that all such methods systematically underestimate predictive variance and that the bias decreases monotonically as you include more parameters. They propose two principled subset selection methods: Gradient-Laplace (selects parameters with largest average squared output gradients) and Greedy-Laplace (iteratively adds parameters accounting for correlations). They prove Gradient-Laplace provably beats existing heuristics and demonstrate strong empirical performance.

Main takeaways:

• Sub-network Laplace restricts the Hessian to a parameter subset to make uncertainty quantification tractable
• All such methods underestimate predictive variance; bias decreases monotonically as the subset grows
• Gradient-Laplace selects parameters with largest average squared gradients of model output (provably optimal in a certain sense)
• Greedy-Laplace iteratively refines selection by accounting for off-diagonal Hessian terms
• Extensive experiments show these methods outperform existing heuristics (diagonal, layer-wise, etc.)

Relevance

Potentially relevant if I ever need uncertainty quantification for fine-tuned models, but not directly connected to my persona or behavioral installation work—this is about Bayesian posterior approximation for neural networks rather than steering or prompting.

Abstract

arXiv:2605.09075v1 Announce Type: new Abstract: Although the Laplace approximation offers a simple route to uncertainty quantification in deep neural networks, its reliance on inverting large Hessian matrices has motivated a range of computationally feasible low-dimensional or sparse approximations. A prominent class of such methods - sub-network Laplace approximations, constructs surrogates by restricting attention to a small subset of parameters. Existing approaches in this family typically rely on diagonal, layer-wise, or other architectural heuristics for subset selection, which ignore cross-parameter interactions and lack formal optimality guarantees. In this paper, we provide a rigorous theoretical analysis of the sub-network Laplace paradigm. We prove that all sub-network Laplace methods systematically underestimate the predictive variance of the full Laplace posterior, and that this bias decreases monotonically as the retained sub-matrix expands. Leveraging this insight, we propose two principled, analytically grounded sub-network Hessian approximations: \textit{Gradient-Laplace} selects parameters with the largest average squared gradients of the model output with respect to the parameters over a reference dataset; while \textit{Greedy-Laplace} iteratively refines this selection by accounting for off-diagonal interactions in the precision matrix. We establish theoretical guarantees characterizing their optimality properties and show that Gradient-Laplace provably outperforms existing heuristic approaches. Extensive numerical studies across diverse settings indicate that these methods perform strongly relative to existing benchmarks.