Uncertainty Quantification with the Empirical Neural Tangent Kernel

I was missing perhaps gradient based MCMC methods which also have been applied to neural networks, see e.g. [1-3]...I would also have liked to see an SGMCMC method in the baselines.

We thank the reviewer for bringing attention to the SGMCMC methods. We originally did not compare to any of these methods due to the often large cost associated with computing them in high-dimensions, and due to some poor performance in previous UQ benchmarks (albiet computed using flawed metrics). However, we appreciate the importance of SGMCMC to the Bayesian framework, in that they are able to compute approximate samples from the posterior. We are actively working toward addressing your comment and will endeavour to include an SGMCMC method in our comparisons before the end of the rebuttal period.

I understand the reasoning to exclude BDE, but do not understand how the same reasoning would not also apply to DE. I think including BDE in the image classification experiments could benefit the clarity and transparency of the paper, and it would also be interesting to see how BDE compares to DE on the newly proposed metric.

This is certainly an interesting point, and we appreciate the opportunity to clarify. We find that the prevailing opinion in current Bayesian UQ works is that DE is still the SOTA, even with the addition of BDE. Further, we find it rare that works will compare with BDE, rather than simply comparing with DE. Hence, we currently rank DE as the highest performing method. A primary goal of our work was to achieve performance similar to DE, whilst reducing the computational cost. We then sought to compare with other methods that seek to achieve a similar goal. Hence, this is why we excluded BDE from the image classification results. However, we agree that a comparison of BDE to both DE and NUQLS on this new metric is enlightening, as BDE may in fact show better performance than DE. Further, the similarity / difference between BDE and NUQLS in this setting will be interesting to observe. Thus, we will aim to evaluate BDE on an image classification result and hope to provide the results before the end of the rebuttal author-reviewer period.

Could the authors elaborate on the new UQ metric and its novelty? Is it the idea of taking the variance of the softmax outputs, or the act of visually inspecting the resulting violin plots?

This is a great question, and we are happy to provide further details on this point that we may have skipped over in the main text. We firstly find that no previous works have employed the variance of the softmax outputs as the quantity used to compare UQ performance. For example, most works simply use the softmax outputs themselves to form evaluation metrics, either by taking the maximum class prediction as the confidence, or the entropy over the classes. The SLU work [4] used the variance of the logits to compute the AUCROC between ID and OOD points (see Appendix G.1 in our work for a small comparison to SLU on this metric). However, we are unaware of other works that have employed the variance of the softmax outputs. Secondly, previous works, such as SWAG (Appendix E.2), have plotted the entropic histograms of the softmax predictions for both ID and OOD points. To our knowledge though, we are the first work which employs violin plots, that compares the median and skew values, and that compares to a poorly performing baseline. Hence, we see that the metric is novel in both the quantity it uses for evaluation, and the manner in which it compares performance between methods. Note that we also provide further details on the VMSP metric in response to a question from reviewer j2t5.

Can the authors elaborate on the computations described in appendix D, perhaps in mathematical forms?

We are happy to provide further mathematical details on the computations in Appendix D. We compare the computational complexity for an epoch of training for the neural network $f_\theta(x)$, for batch $x \in \mathbb{R}^{d \times n}$, and an epoch of training for a single linear network $\hat{f}\_\theta(x)$. We take approximations on the computational complexity of both forward-mode AD and backward-mode AD from the following text [5]. Specifically, we take $[fp]$ as the computational complexity for evaluating $f_\theta(x)$.

We note from [5] that both a JVP and a VJP cost roughly $2$ $\times$ a forward pass in computational complexity and memory. Further, both a JVP and a JVP return a function evaluation. Now, a standard epoch of training for the neural network involves a forward-pass to compute the error, and then a backward-pass (VJP), hence the complexity is approx. $3[fp]$. The linearised network involves a JVP (which includes a function evaluation) to form the linear network, and a VJP to compute the gradient. Hence, the complexity for an epoch of training for the linearized network is approx. $4[fp]$. So we observe that each epoch for a linearized network is only $4/3~\times$ as expensive as for a neural network. In regards memory, we see that the memory requirement for both the linearized network and the neural network are similar. However, we generally train all linearized networks in parallel. For computational complexity, this will incur some additional cost, thought it will not be linear in number of networks, that is dependent upon the specific software and hardware. However, memory cost will scale linearly by number of networks, i.e. $2S \times M([fp])$, where $M([fp])$ is the memory cost for a forward-pass, and $S$ is the number of linear networks. As an example, we employ a batch size of $56$ for ImageNet on ResNet50 with $10$ ensemble members when using an 80GB H100 GPU, due to the large parameter count and large number of classes.

These improvements have been made to the working copy of our paper (pending additional experiments). Please let us know if you have any other questions or points that need clarification.

---

References:

[1] Wenzel, Florian, et al. "How good is the bayes posterior in deep neural networks really?." arXiv preprint arXiv:2002.02405 (2020).

[2] Izmailov, Pavel, et al. "What are Bayesian neural network posteriors really like?." International conference on machine learning. PMLR, 2021

[3] Garriga-Alonso, Adrià, and Vincent Fortuin. "Exact langevin dynamics with stochastic gradients." arXiv preprint arXiv:2102.01691 (2021).

[4] Sketched Lanczos uncertainty score: a low-memory summary of the Fisher information, Miani et. al. 2024

[5] The Elements of Differentiable Programming, Blondel et. al. 2024, Chapter 8

We thank the reviewer for their careful reading of our work, and constructive criticism.

The comparison to prior work is limited. Several efficient ensemble methods, such as BatchEnsemble [1], Snapshot Ensembles [2] were proposed. SNGP [3] are highly relevant but not included in the evaluation. Additionally, LLA is not a particularly strong Bayesian baseline.

We thank the reviewer for bringing attention to the importance of these methods. Unfortunately due to the limited time of the rebuttal period, we were only able to impliment one of the referenced methods, of which we chose SNGP, as there exists a Lightning Pytorch implementation of the method [5]. We tested SNGP using VMSP on both ResNet9 / FashionMNIST and ResNet50 / SVHN. The median results are presented in the following table. We write $(>)$ to denote values which we want to be large, and $(<)$ for values we want to be small (synonymous with the up and down arrows in the paper, respectively).

and the skew results:

We see that on ResNet9/FashionMNIST, NUQLS far out-performs SNGP. For ResNet50 / SVHN, $3$ learning rates and $2$ optimizers were tested; unfortunately SNGP did not properly train for the tested combinations, and hence the best results are very poor.

Note that we focused our comparison on LLA variants as they have become popular recently, with the release of the Laplace package [5], and through scalable variants such as Sampling-LLA and VaLLA.

We believe that the efficient ensemble methods are also important methods to consider, and hence we are working toward including results comparing these methods to NUQLS before the end of the review period.

Although the paper evaluates both regression and classification tasks, it would strengthen the contribution to demonstrate that improved uncertainty estimates lead to better downstream performance (e.g., in active learning or decision-making under uncertainty).

We appreciate the suggestion of the reviewer to display the performance of NUQLS on a downstream task, as we believe the results would be interesting and would strengthen the value of the paper. Unfortunately we are unable to run a well-thought out, large scale experiment in the rebuttal period due to time constraints. However, we will include such an evaluation in future work.

Prior work has studied the connection between uncertainty quantification and the neural tangent kernel (NTK). It would be helpful for the authors to clarify how their method differs from or extends these works, such as [4].

The referenced work PNC is very interesting, and the connection of NUQLS to PNC will be included in the paper. We refer the reviewer to our response to Reviewer Z6xt, where we discuss the inherent differences between NUQLS and PNC, as well as empirically compare them on an experiment from [4]. To summarise, PNC is developed under the asymptotic NTK theory and performs well for very wide networks near initialization (i.e., poorly trained networks), where NUQLS struggles. In contrast, NUQLS does not rely on the asymptotic NTK regime and instead leverages the trained, data-dependent NTK to construct meaningful linear ensembles. For finite-width networks that are well trained—which we believe to be the practically relevant setting—NUQLS is more appropriate and clearly outperforms PNC.

Since all linearized models are initialized around the same local optimum, is there a risk of mode collapse—i.e., all models converging to similar behavior? A discussion of this potential limitation and how it is mitigated would be helpful.

This is an excellent question, and we're glad to offer a concrete answer: mode collapse cannot occur—almost surely.

To see this, note that the solution to equation (3) in our main paper for each linearized model consists of a unique row-space component (under our assumptions), plus a projection of the initialization $z_i$ onto the null space of the Jacobian. For mode collapse to occur, the projections for two i.i.d. initializations $z_1, z_2 \sim \mathcal{N}(0, \gamma^2 I)$ would need to coincide—that is, we would require $(I - J^{\dagger}J)(z_1 - z_2) = 0$, even though $z_1 - z_2 \neq 0$ almost surely. This implies $z_1 - z_2 \in \text{Range}(J^T)$. However, $\text{Range}(J^T)$ is a low-dimensional subspace of $\mathbb{R}^p$, while $z_1$ and $z_2$ are drawn iid from full-rank (i.e., non-degenerate) $p$-dimensional Gaussian distribution. The probability that $z_1 - z_2$ lies in a fixed lower-dimensional subspace is zero. Therefore, mode collapse occurs with probability zero, and we can safely rule it out as a concern.

We hope that this has alleviated your concerns with our work. Please let us know if there are any further questions of points that require clarification, as we would like to engage further if it will improve your assessment of our paper.

---

References:

[1] BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning, ICLR 2020

[2] Training Independent Subnetworks for Robust Prediction, ICLR 2021

[3] Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness, NeurIPS 2022

[4] Efficient Uncertainty Quantification and Reduction for Over-Parameterized Neural Networks, NeurIPS 2023

[5] Lightning UQ Box: A Comprehensive Framework for Uncertainty Quantification in Deep Learning, Lehmann et. al. 2024

[6] Laplace Redux – Effortless Bayesian Deep Learning, Daxberger 2021

Unlike the NeurIPS 2023 PNC paper, the submission lacks rigorous statistical theory (e.g., asymptotic coverage guarantees, bias-variance decomposition) to back the empirical results.

We appreciate the opportunity to highlight the differences between the PNC paper and ours. Firstly, PNC is formed from a frequentist framework, and hence, similar to conformal predictions, they are able to get statistical guarantees on confidence interval coverage. Further, PNC relies on the NTK theory, where for infinitely wide networks, the network behaves according to kernel regression. From this, the authors are able to estimate possible function values. Instead, NUQLS relies on the network being trained to a local minima, and then employs the feature-learning of finite-width networks to create meaningful ensembles.

To test the difference between the two methods, we tested NUQLS on the confidence intervals problem from [3]. In this experiment, an MLP with a single hidden-layer is trained on $n$ datapoints from $U([0,0.2]^d)$, where the target is $y = \sum_{i=1}^d \sin(x_i) + \epsilon$, and $\epsilon$ is a small noise term. The method is then asked to form a confidence interval around the prediction for $x = (0.1, \dots, 0.1)$. This setup is repeated several times, and the coverage and width of the confidence intervals are recorded, as well as the mean prediction. For more details, we refer the reviewer to [3]. When the width of the network is $32 \times n$, and the network is only partially trained with $80$ epochs and a learning rate of $0.01$, NUQLS performs poorly compared to PNC, as can be seen in the following table. Note that bolded numbers indicate intervals that have reached or exceeded the expected coverage. Further, smaller width intervals are preferred.

We note that scaling the width of a network by $32$ times the number of training points is rare in practice, and that such a wide-network is difficult to train. If we instead form an MLP with width equal to the number of training points, and increase epochs to $1000$ and learning rate to $0.5$, so that the network is properly trained, NUQLS far outperforms PNC, as can be seen in the following table. We note that for these experiments we tuned the $\gamma$ hyper-parameter for NUQLS on a small validation set.

We observe that NUQLS and PNC are in fact complementary methods. For infinite-width networks near initialization, PNC performs well while NUQLS struggles. Conversely, for finite-width networks trained to a minimum, NUQLS excels while PNC performs poorly. We believe the latter regime is more representative of practical scenarios, where NUQLS offers significantly better performance.

Finally, we note that rigorous statistical guarantees that exist for PNC do not currently exist for Bayesian methods, a point that is raised in the PNC paper [3].

Compared to the SWAG paper, insufficient empirical evidence is provided to demonstrate how Equation (9) (the main result) relates to uncertainty quantification

In regards to the experiments in the submitted paper, we find that under a variety of regression and classification tasks, our method does in fact accurately quantify the uncertainty in the predictions of a neural network. Further, we have now also provided an extra regression experiment, i.e. the confidence intervals experiment from [3], and a large-scale image classification experiment, i.e. the ResNet50/ImageNet result that has been provided in the response to Reviewer j2t5. From this, we hope that the reviewer finds adequate empirical evidence that the variance term derived in (9) does in fact accurately quantify the uncertainty in a neural networks predictions.

The method relies on approximate Gaussianity of the hidden-layer representations' output distribution and the strong convexity of the loss function. These assumptions may break down for complex real-world classification tasks (especially with multimodal outputs), yet this limitation is not critically analyzed...How do you handle classification tasks where predictive distributions are inherently non-Gaussian (e.g., multimodal posteriors)? Does your method underperform in such settings?

We would like to respectfully clarify a potential misunderstanding: we do not assume Gaussianity for the original neural network at any point. Rather, if the network is approximated by its linearization—a reasonable assumption near a local minimum—we show that the predictive distribution of the linearized model follows a Gaussian Process. Consequently, the predictive distribution of the original network is approximately Gaussian in this regime.

Even in real-world classification tasks such as ImageNet with ResNet-50, where the posterior is likely multimodal, we observe excellent uncertainty quantification. While we have no guarantee that the predictive distribution is exactly Gaussian in these settings, we find that the linearized networks form meaningful ensembles that effectively explore the local function space and capture predictive uncertainty.

There is little discussion on when or why the method might fail (e.g., poorly trained base networks, overfitting, high-dimensional noise in embeddings)...How sensitive is your method to the quality of the learned representations? Would randomly initialized or undertrained networks still yield calibrated uncertainty?

This is an insightful question, and has prompted an important addition to our paper. As has been previously noted, we can see from the attached confidence interval results that untrained or undertrained networks do not necessarily yield well-calibrated uncertainty estimates. Instead, if the network is very wide and poorly trained, a method like PNC would be more appropriate. However, we again note that this would be uncommon in practice, as most networks will be well-trained. In the case of well-trained networks, we find that NUQLS excels.

The method’s performance is evaluated primarily on small to medium-scale datasets, limiting the generalizability of its claims to larger or more diverse real-world scenarios.

We kindly refer the reviewer to our response to Reviewer j2t5, where we show results for ResNet50 trained on ImageNet. This setting is generally considered large-scale in the Bayesian UQ commmunity.

Please let us know if you have any questions or there is something else that needs clarification, as we are eager to increase your assessment of our work.

---

References:

[1] Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel, NeurIPS2020

[2] Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks, Arora et. al. 2019.

[3] Efficient Uncertainty Quantification and Reduction for Over-Parameterized Neural Networks, Huang et. al. 2023

We thank the reviewer for their feedback and constructive criticism. We hope that the following aids in alleviating your concerns.

The core mathematical construction in equation (4) is presented as general but only holds under specific loss function assumptions (strong/strict convexity)...The meaningful deep learning experiments (classification on MNIST/CIFAR) operate entirely outside the theoretical framework. The remaining empirical evaluation consists only of nine UCI regression tasks and toy experiments, severely limiting the practical relevance of the work...

While our theory does not currently extend to cross-entropy loss, as is common in NTK-related work and many aspects of theory in machine learning more broadly, we would like to highlight the following two points. Firstly, it is practical extensions of the theory to settings where the assumptions fail that are often the most valuable to empirically test. As is seen in our results, we find that our framework also correctly quantifies the uncertainty in a range of image classification tasks. Secondly, we note that our theory does hold for classification using the Brier loss, and models trained using this loss have been shown to perform similarly to cross-entropy based models [1]. Testing NUQLS with the Brier loss is therefore expected to yield comparable performance and presents an interesting direction for future work. We have clarified these points in the updated version of our paper.

The authors...could better address the theory-practice disconnect and provide more analysis of when the linearization approximation might fail... How does the method perform when the linearization approximation quality degrades?

We kindly refer the review to the confidence intervals experimental results we have provided in our response to Reviewer Z6xt. We observe from the first given table that, in line with its theoretical construction, our method does not perform very well for poorly trained networks. This is because far away from a local minima, the linearization approximation may be poor. Further, we rely on the 'data-dependent' NTK that is learned during training [2]. While this is a space where PNC becomes a good method to choose, we do not view it as a limitation of our work. This is because (useful) networks in practice are well-trained; we observe in the second table of results that the performance of NUQLS is excellent once the network has been trained.

The paper introduces VMSP (Variance of Maximum Softmax Probability) as a UQ metric without explicit definition or proper theoretical justification (even taking into account appendix C).

We are happy to provide a more explicit definition of VMSP: for a Bayesian method, we generally have a mean predictor $\mu : \mathbb{R}^d \to \mathbb{R}^c$ and a covariance function $\Sigma : \mathbb{R}^d \to \mathbb{R}^{c \times c}$, for example the mean and covariance of the linearized ensemble in the case of NUQLS, that output in the probit space. To compute VMSP for a given test point $\mathbf{x}^\star$, we first find $\hat{c} = \text{argmax}\_k ~\mu(\mathbf{x}^{*})\_{k}$, where $\mu(\mathbf{x}^\star)\_k$ denotes the $k$-th output of $\mu(\mathbf{x}^\star)$. That is, we find the class that the Bayesian method predicts, given $\mathbf{x}^\star$. We then define $\text{VMSP} := \Sigma(\mathbf{x}^\star)_{\hat{c},\hat{c}} = \sigma^2(\mathbf{x}^\star)\_{\hat{c}}$, that is, the variance of this prediction. We will include this description in the paper, to aid comprehension of the metric.

As for theoretical justification, we can also be more explicit: we have detailed in Appendix C how the uncertainty of the predictions is captured by their variance. For a given dataset, we want a UQ method to provide low uncertainty for correctly predicted test points, and high uncertainty for incorrectly predicted or OOD test points. We then compare these distributions pictorially using a violin plot, and quantitatively using the median and skew values for the respective distributions. With the addition of a poorly-performing baseline model, we are able to easily compare the ability of UQ models to quantify uncertainty.

The core ideas of network linearization, NTK connections, and sample-then-optimize frameworks have been extensively explored. Particularly, Immer et al. (2021) already studied linearization, refinement through additional training, and NTK-GP connections. The contribution appears incremental rather than fundamental...No comparison with important recent UQ methods, e.g. Immer et al. (2021).

To be clear, we have in fact compared to the work of Immer et al. (2021) (i.e. LLA) in most of the experiments, and kindly refer the reviewer to Appendix B.1, where we describe in detail the connection to, and differences with, the LLA framework methods. To summarize, the work of Immer et. al. showed how forming a Laplace Approximation over the parameter posterior of a linearized model resulted in a noisy NTK-GP predictive posterior. This approach incurs several challenges: forming the posterior requires explicitly forming and inverting the NTK gram matrix, independent GP's are required for each class in classification, and approximations such as last-layer and KFAC are required for scaling. In comparison, our method skips the parameter posterior, and uses the randomness in the null-space of the Jacobian to form a noise-free GP. Due to our derivation, we are able to easily sample from this GP posterior using linearized networks, without requiring covariance structure approximations. This enables us to out-perform LLA in larger image classification tasks (see Figure 3 in main text).

Experiments limited to MNIST/CIFAR-scale datasets and outdated architectures rather than modern benchmarks and model architectures...While computationally faster than deep ensembles... the linearization approximation quality at scale is unclear...Missing evaluation on uncertainty-specific benchmarks (e.g., diabetic retinopathy detection tasks by Band (2022) or others)...Can you demonstrate scalability to larger, more modern architectures? And additional datasets with additional epistemic uncertainty metrics?

We thank the reviewer for raising this point, as the evaluation of our method on a larger problem will add to the value of our paper. We have now trained and evaluated NUQLS on ResNet50 trained on ImageNet, and using the ImageNet-o dataset for out-of-distribution testing [3]. We employed the pre-trained weights found on torch.hub for ResNet50 / ImageNet. While we cannot provide an image of the violin plot of the VMSP, we can still provide the median values (where (>) denotes values which we want to be large, and (<) those we want to be small):

and the skew values:

For the in-distribution test set, we see excellent separation between the correct and incorrect predictions. As was stated in [4], approximately $21$% of the images in ImageNet-o are in-distribution. Hence, we would expect that the OOD VMSP values to skew slightly more towards zero than the in-distribution incorrect predictions, as we observe. We note that ImageNet/ResNet50 is generally considered a large-scale setting for Bayesian methods.

Related work discussion relegated entirely to appendix

We have relegated this discussion due to the large amount of related works, and the detail the discussion required. However, we can include a condensed version in the main body, given the additional space allowed in the camera-ready version.

Pseudo-inverse specification (left vs. right) not clarified

As we assume $J_X$ to be full-row rank, the pseudo-inverse is computed as the right inverse.

Insufficient explanation of why objective perturbation (standard in sample-then-optimize?) is unnecessary

We appreciate the opportunity to clarify this point, as this is a contribution of the paper. Standard sample-then-optimize works rely on randomness in the loss function. Instead, NUQLS relies on the fact that we are in an underdetermined linear system, and hence we can form an ensemble of solutions simply through the variability of random projections onto the null-space of the Jacobian. Thus, we do not require any variability in the loss function to sample from the posterior, in contrast to other sample-then-optimize methods.

The approach forgoes having a principled prior and predictive variance formulation, potentially making it more prone to overfitting on validation data compared to proper Bayesian methods... Hyperparameter $\gamma$ tuning may lead to validation overfitting without principled prior specification

We respectfully disagree with the reviewer that we do not have a principled predictive variance formulation, as the predictive distribution follows a Gaussian Process. In regards the prior, we do not find evidence of overfitting on validation data, and instead show very competitive results. Further, we find that it is common for Bayesian methods to choose a Guassian isotropic prior.

Please let us know if there are any other points that we can clarify, as we are eager to engage further to improve your assessment of our work.

---

References: [1] Evaluation of Neural Architectures Trained with Square Loss vs Cross-Entropy in Classification Tasks, Hui et. al. 2020

[2] Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel, Fort et. al. 2020

[3] Natural Adversarial Examples, Hendrycks et. al. 2017

[4] In or Out? Fixing ImageNet Out-of-Distribution Detection Evaluation, Bitterwolf et. al. 2023

Thank you for the response. I'm increasing my score accordingly.

PS: Have the authors seen Miani, Marco, Hrittik Roy, and Søren Hauberg. "Bayes without Underfitting: Fully Correlated Deep Learning Posteriors via Alternating Projections." arXiv preprint arXiv:2410.16901 (2024)? It also projects onto the kernel of the GGN to avoid underfitting.