Computationally Efficient Laplace Approximations for Neural Networks

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 04:45:38

Abstract: Laplace approximation is arguably the simplest approach for uncertainty quantification using intractable posteriors associated with deep neural networks. While Laplace approximation based methods are widely studied, they are not computationally feasible due to the involved cost of inverting a (large) Hessian matrix. This has led to an emerging line of work which develops lower dimensional or sparse approximations for the Hessian. In this paper, we build upon this work by proposing two novel sparse approximations of the Hessian: (1) greedy subset selection, and (2) gradient based thresholding. We show via simulations that these methods perform well when compared to current benchmarks over a broad range of experimental settings.

Main paper: 
 introduction: 
 Uncertainty quantification for neural networks is critical in understanding the quality of output/estimates provided by the network, and is a crucial component for many standard frameworks (such as contextual bandits) to study exploration-exploitation trade-offs. Such frameworks are useful in many domains, such as recommender systems [Su et al., 2024 , Joachims et al., 2018 , Li et al., 2010] , content moderation [Avadhanula et al., 2022] , healthcare [Durand et al., 2018 , Mintz et al., 2020 , Esteva et al., 2017] , dynamic pricing [Misra et al., 2019] , dialogue systems [Liu et al., 2018] , self-driving cars [Bojarski et al., 2016] , detecting hallucinations in LLMs [Chen and Mueller, 2023 , Felicioni et al., 2024 , Dwaracherla et al., 2024] . It is well known that accurate uncertainty quantification can significantly improve the performance of downstream tasks [Ovadia et al., 2019 , Riquelme et al., 2018] . Despite these benefits, uncertainty quantification is not yet widely used in deep learning due to excessively high computational cost in cutting edge applications. This is in particular true for the Laplace approximation based approach, which approximates the posterior distribution of the parameters in a neural network by a multivariate Gaussian centered at the posterior mode θMAP (with θ denoting the parameters in the network), and with the (negative) observed Fisher information matrix as its precision (inverse-covariance) matrix, see MacKay [1992] , Foong et al. [2019] . This approach is technically well-grounded and is also attractive due to its conceptual simplicity. However, any application of this approach for uncertainty quantification requires inversion of the observed Fisher information matrix whose computational complexity O(p 3 ) can be prohibitive in many settings (here p is the number of parameters in the neural network).To overcome this computational challenge, a nascent strand of the literature has focused on replacing the Laplace precision matrix Ω by a lower dimensional approximation which retains the most relevant components or entries. In Nilsen et al. [2022] , the authors propose replacing the Laplace precision matrix by an approximation which uses only the top k eigenpairs of Ω (identified using the Lanczos algorithm [Cullum and Willoughby, 2002] ). However, the iterative Lanczos algorithm can in general suffer from numerical instability issues when using finite precision arithmetic, especially for large matrices [Chen and Trogdon, 2024] . The PrecisionDiag approach in Riquelme et al. [2018] , Zhang et al. [2020] approximates Ω by setting all off-diagonal entries to zero. This leads to extremely efficient computation, but can suffer from poor statistical performance given that it ignores all cross-parameter correlations.Our particular focus in this paper is a class of methods, termed sub-network LA in Daxberger et al. [2021a] , that focuses on identifying a small subset S of parameters capturing the most variability. The matrix [Ω S,S ] 0 , constructed by setting all entries of Ω outside of the principal submatrix Ω S,S := (Ω rs ) r∈S,s∈S to be zero, is then used as an approximation for Ω. Since the approximation is singular, its Moore-Penrose inverse is utilized for subsequent variance computations. This amounts to ignoring the variability in all other parameters except those in S. To the best of our knowledge, two methods for identifying S have been proposed in the current literature. The Neural-Linear [Snoek et al., 2015 , Riquelme et al., 2018] method chooses S as the subset of parameters in the last layer of the neural network. However, the last layer of the network may not always contain the most relevant parameters in terms of capturing variability (see Section 4). The marginal variance based approach in [Daxberger et al., 2021a, Section 5] relies on an independence assumption (ignoring cross-correlations) for the posterior distribution of the parameters. For a given subset size k, this approach reduces to choosing the parameters corresponding to the k smallest diagonal entries of Ω.In this paper, we propose two methods for selecting the parameter subset S for subnetwork LA which build on these promising initial efforts while aiming to address their deficiencies through a more refined treatment. The first method, Greedy-Laplace pursues a greedy step-wise approach for selecting S. The method starts with the Laplace precision matrix Ω, and at each step the parameter corresponding to the smallest diagonal entry in the current precision matrix is chosen and added to S, and the precision matrix is shrunk and updated by 'removing the effect' of this parameter.The second method, Gradient-Laplace, recognizes that the posterior predictive variance for a future observation also depends on the corresponding gradient vector (see ( 1)), and uses parameters with the highest average absolute gradient evaluated over a reference dataset to construct S. Detailed empirical evaluation based on accuracy of subnetwork choice (Section 4) and coverage of predictive credible intervals (Appendix D) is provided. 
setup: 
 In this article, our central focus will be on the linearized Laplace Approximation (LLA) [Foong et al., 2019] . We consider a neural network with output f θ (x) for given input x and parameter vector θ ∈ R p . The network is trained using a given datasetD train = {(x n , y n ) : 1 ≤ n ≤ N train },where we assume, y n ∼ N (f θ (x n ), σ 2 0 ). We place independent Gaussian priors on the weights and biases of the network. As in Foong et al. [2019] , under this setup, the LLA provides the following approximation of the posterior predictive distribution for a new observation (x * , y * ),EQUATIONwhere, g(x * ) = ∇ θ f θ (x * )| θ= θMAP ,and Ω is the negative Hessian matrix, which in practice, is replaced by the Gauss-Newton matrix, as the Gauss-Newton matrix is guaranteed to be positive semi-definite. In particular, letting v denote the vector containing all prior variances, we setEQUATIONThe posterior approximation for a classification problem will be similar but with the Gauss-Newton matrix replaced by the Generalized Gauss-Newton matrix, see Appendix C for more details. 
proposed algorithms: 
 Our goal is to develop novel and principled methods for choosing the subset S for subnetwork LA, so that the corresponding approximation [Ω S,S ] 0 serves as an effective surrogate for Ω for the variance computation in (1). We now describe two strategies for this purpose.The Greedy-Laplace Approach. Note that the inverse of the r th diagonal entry of a precision matrix is equal to the conditional variance of the r th variable given all the other variables. Given a user-specified subset-size k, we leverage this fact to sequentially identify the 'best' subset of k parameters (i.e., the subset capturing the most variability) in a greedy fashion. We start with Ω (1) = Ω (as in ( 2) and (C.1)). In the i th step, the parameter with the minimum diagonal entry in Ω (i) is chosen, and the marginal precision matrix of the remaining variables is calculated and setasΩ (i+1) (needs O((p -i) 2 ) computations, seeStep 5 of Algorithm 1). In particular, Ω (i+1) is a (p -i + 1) × (p -i + 1) matrix. We stop after the k th step. We refer to this procedure as the Greedy-Laplace algorithm (Algorithm 1), which has an overall computational complexity of O(kp 2 ). The output of Algorithm 1 is an index set S (of size k), and the surrogate [Ω S,S ] 0 is obtained from Ω by setting all elements outside the principal sub-matrix corresponding to S to zero. Recall that since[Ω S,S ] 0 is singular, its Moore-Penrose inverse is used in (1) instead of Ω -1 .Algorithm 1 Greedy-Laplace Algorithm Require: Laplace Precision matrix Ω as in (2); user specified dimension k 1: Initialize Ω (1) = Ω, collection of indices of selected parameters S = ϕ 2: for t = 1, 2, . . . , k do3: Compute j = argmin i∈{1,2,••• ,p}\S (Ω (t) ) i,i 4:Update S = S ∪ {j} 5:Construct Ω (t+1) by removing row & column for j th variable in Ω (t) , and updating remaining entries: Ω (t+1) = (Ω (t) ) -j,-j -(Ω (t) )-j,j (Ω (t) )j,-j (Ω (t) )j,j 6: end for 7: Return SThe Gradient-Laplace Approach. Note that the variance term in (1) depends on both Ω and g(x * ).While Algorithm 1 exclusively uses the matrix Ω to identify the subset S of the most relevant 4 Empirical Evaluation -Accuracy of Subnetwork Choice Experimental Setup. Daxberger et al. [2021b] propose using the Wasserstein distance between the linearized Laplace posterior distribution of the parameters θ (multivariate normal with covariance matrix Ω -1 ) and its subnetwork LA based surrogate (singular multivariate normal with covariance matrix given by the Moore-Penrose inverse of [Ω S,S ] 0 ) as a measure of accuracy/quality of the subnetwork based approximation using the subset S. We use a derived Wasserstein metric to compare the performance of the proposed methods (Greedy-Laplace and Gradient-Laplace) for choosing S to existing subnetwork LA methods -the Subnet diagonal approach [Daxberger et al., 2021b, Section 5]), the NeuralLinear approach [Snoek et al., 2015 , Riquelme et al., 2018] and the Last k approach (a natural extension of NeuralLinear; see Appendix B). For thoroughness, we also include the (non subnetwork LA based) PresicionDiag approach [Riquelme et al., 2018 , Zhang et al., 2020] for obtaining a surrogate of Ω. For further details on these benchmark methods see Appendix B.For this experiment, we generate synthetic regression data, with 2000 samples, using 5 input variables and a single output variable. The input variable x is generated from N 5 (0, Σ), where Σ i,j = 0.6 |i-j| , i, j = 1, 2, . . . , 5, and a 3-hidden layer fully connected feed-forward Neural Network with random weights is used to generate the output. The error variance σ 2 0 is fixed at 1. such inclusions in the test dataset. We compare our methods with the following Laplace-based algorithms (details in Section B): NeuralLinear [Snoek et al., 2015 , Riquelme et al., 2018] , Presi-cionDiag [Riquelme et al., 2018 , Zhang et al., 2020] , Subnet diagonal [Daxberger et al., 2021b] . For thoroughness, we include two standard (non-Laplace) Bayesian methods in our comparison. The first method, referred to as the Split method [Kuchibhotla et al., 2023] , splits the training data into d parts, uses each of them to fit the model separately and then proposes the highest and the lowest estimated value of f θ (x) among the d fitted models as the upper and lower limits of the confidence interval for f θ (x) respectively. The second method, referred to as the ensemble method [Lakshminarayanan et al., 2017 , Osband et al., 2016] , generates B bootstrap samples using the training data, uses them to train the network separately, and then proposes the 2.5% and 97.5% quantiles to be the lower and upper limit respectively of a 95% confidence interval for f θ (x). and fit 20 models using those samples.As mentioned earlier, for both the regression and classification experiments, we have two working model settings. For both these settings, we consider IID N (0, 1) priors on the weights and biases.For the mis-specified setting, the working model is a Neural network with a single hidden layer with 40 hidden nodes and tanh activation function. For the well-specified setting, as previously mentioned, the working model has 15 hidden nodes in all 3 hidden layers and tanh, sigmoid and tanh activations respectively. To train the models, we use a decaying learning rate to train the models. For regression experiments, the initial learning rate is 0.5 for the mis-specified working model setting, and 0.008 for the well-specified one, whereas for classification experiments, the learning rate is 0.01 for the mis-specified setting, and 0.001 for the well-specified setting. In all 4 cases, the learning rate is multiplied by 0.1 after every 1000 epochs. In the regression setup, we use 3000 epochs to train both the models, and in classification setup, we use 2000 epochs for each model.Once the network is trained, we compute the empirical coverage of the credible interval of f θ (x) for all the benchmark approximation methods, as well as the proposed methods. The coverages are further averaged over 200 replications. The results are depicted in Figure 2 .The linearized Laplace approximation (LLA) provides overall competitive coverage to the two benchmark methods, Split and Ensemble. Among the methods that use a low-dimensional surrogate of the Hessian to further approximate LLA, NeuralLinear clearly provides poor coverage. Its immediate extension, Last k, also provides much lower coverage compared Subnet diagonal, and both the proposed methods, Greedy-Laplace and Gradient-Laplace, especially in the well-specified working model settings. The proposed Gradient-Laplace approach clearly outperforms the Subnet diagonal approach in all but the well-specified working model setting under classification, where the Subnet diagonal provides marginally better coverage. The proposed Greedy-Laplace approach provides similar coverage to the Subnet diagonal approach in mis-specified working model settings, and significantly outperforms it in the well-specified working model setting for the regression setup, while the Subnet diagonal approach provides marginally better coverage in the well-specified setting for the classification setup. Although PrecisionDiag provides better coverage than NeuralLinear, the coverage is still quite low, and is attained by the proposed methods for very small values of k. To summarize, the proposed methods for subnetwork LA overall provide significantly better coverage than existing methods for constructing surrogates for the Laplace precision matrix. 
 
 Appendix:  



Reviews: review: title: title: This paper introduces two innovative methods for improving computational efficiency in uncertainty quantification for neural networks, with strong performance and relevance in various applications, but it could benefit from clearer explanations and more real-world validation
main_review: main_review: Pros:
This paper Introduces Greedy-Laplace and Gradient-Laplace, improving computational efficiency for uncertainty quantification.
It reduces computational complexity, making it feasible for large neural networks.
It has Rigorous benchmarking against state-of-the-art methods using multiple datasets.
It addresses important applications like healthcare and autonomous systems.
The outperforms existing methods in terms of accuracy and credible interval coverage.

Cons:
Its complex mathematical details may be difficult for non-experts.
It lacks extensive real-world application demonstrations.
The minimal exploration of where the methods might fail.
And the gradient-based method may introduce computational overhead.

In summary, the paper presents an important contribution to computationally efficient uncertainty quantification in neural networks. Its originality, practical relevance, and rigorous evaluation make it a valuable piece of research in the field of Bayesian neural networks and Laplace approximations. However, further clarity in presentation and real-world validation would enhance its impact.
title: title: Official review for Submission93
main_review: main_review: This paper proposes two novel methods for approximating the Hessian matrix in Laplace approximations for neural networks: Greedy-Laplace and Gradient-Laplace. These methods aim to address the computational bottlenecks associated with the inversion of large Hessian matrices by selecting subsets of parameters that contribute the most to the posterior variance. The paper presents a clear motivation for sparse Hessian approximations and provides empirical evidence on synthetic datasets showing that the proposed methods outperform existing subnetwork-based Laplace approximations. The topic is timely and relevant to the workshop, the use of Laplace approximations is well-grounded in Bayesian statistics, and efforts to make these methods computationally tractable are both timely and valuable.
My comments are as follows:
1.	How to consider the trade-offs between computational complexity and the quality of the approximation (in terms of posterior variance)?
2.	Are there formal error bounds for the approximation using the Moore-Penrose inverse versus the true inverse of the precision matrix?
3.	Additionally, while the Greedy-Laplace algorithm intuitively selects parameters based on the smallest diagonal entries, it would be valuable to discuss the theoretical properties of this subset selection approach in more detail, such as whether this strategy can lead to suboptimal selections in some cases.
4.	The paper demonstrates improvements in computational efficiency, however, the current experiments are primarily limited to synthetic data with a relatively small number of parameters (5 inputs and 3 hidden layers). To claim practical relevance, it would be important to test the methods on real-world datasets and/or larger neural network architectures (e.g., convolutional neural networks or transformers). 
5.	I’m wondering: Can the proposed methods scale to models with millions of parameters?
6.	How does the runtime compare to other efficient uncertainty quantification methods? Please discuss.
7.	How do Greedy-Laplace and Gradient-Laplace compare to popular Bayesian neural network approaches (e.g., variational inference or MCMC-based methods)?
8.	When Gradient-Laplace is computationally preferable (versus the Greedy-Laplace)? I think adding more details on this point would help practitioners understand when to use each method.\