Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The paper proposes ELLA, a Nyström approximation, as an alternative posterior (predictive) approximation for linearized Laplace. Their method uses the NTK-based formulation of linearized Laplace and improves over other common Laplace approximations in terms of quality of the predictive while maintaining a relatively low computational complexity.
strengths_and_weaknesses: ## Strengths
- Can reduce the cubic complexity in data points $N$ or parameters $P$ of linearized Laplace in neural networks down to much lower complexity that can be controlled by the number of samples used for Nyström. This way, a performance-accuracy trade-off can be made.
- Performance of the proposed posterior predictive improves over cheap last-layer Laplace and is better than the full-network variants.
- Interesting and open discussion of the overfitting issue of the proposed method ELLA. 
- Exhaustive experiments with ablations on the number of factors $K$ and several metrics of interest. Overall convinces that the proposed method is a valuable addition to the family of Laplace approximations.

## Weaknesses
- Only posterior predictive is discussed but not marginal likelihood estimation. Since the proposed method is rather simple ("take linearized Laplace in its kernel variant and apply Nyström, which is well-known from GP literature"), it would greatly strengthen the paper to at least show not only predictives but also the marginal likelihood, which should in principle be easy to add.
- Misunderstanding of relationship between linearized Laplace and GGN. ll 3-6 and 30-32: linearized Laplace and Laplace with GGN are equivalent as shown in reference [22]. The equivalent variant with (neural tangent) kernels has been used before as well in [22,28] but without Nyström approximations. ll 92-99 discusses the scalability issues only of the parametric LLA which requires approximations to the $P \times P$ GGN. The problem that is tackled in the present paper, however, is the approximation of the equivalent $N \times N$ kernel, which is for example defined. Clearly, $N \times N$ is still intractable so the proposed Nyström approximation is well justified.
- ELLA does not seem to approximate full LA, which it should in theory. Figure 1: discrepancy of ELLA and LLA, it seems like ELLA does not get very close to LLA in Figure 1 and there is no explanation as to why that would be. In line 50, it is said though that the approximation becomes accurate.
- Relation to prior work could be pointed out more clearly. For example, it is unclear to what extent the Theorems are novel/original or if they follow in a straightforward way from previous theoretical contributions. 
questions: - ll 31 says that LLA has a more benign covariance. Can you elaborate on what that means or in what way it is benign?
- Figure 1: how comes that ELLA gives better "in-between" uncertainty? Are the hyperparameters equivalent to the LLA method? Generally, ELLA approximates the LLA variant so it is unclear why/how it is better. This would be worth investigating more. 
- Could you elaborate on the simplification in Equation (6), for example, by putting a detailed derivation in the appendix? It is not clear to me where the equivalence comes from.
- Regarding computational overhead (ll. 162 ff): Can you relate this to the prevelant alternative of KFAC-LLA? It would be interesting to know the complexity of computing the posterior predictive for an unseen data point after inference.
- The paper at hand only discusses the posterior predictive. Have the authors considered using the Nyström approximation for marginal likelihoods as well. In [*scalable marginal likelihood estimation for model selection in deep learning*](https://arxiv.org/abs/2104.04975), the NTK variant for linearized Laplace was used but without further approximations. It would be interesting to experiment if Nyström can provide an alternative scalable method. Potentially, this could alleviate overfitting in Figure 2c) by selecting the number of bases $K$?
- Figure 2 is hard to parse since the colorbars are unclear. What is the differene between Nyström and ELLA there in particular?
- Table 2 has no error bars but bold figures. Which policy is used to make numbers bold? Did the authors run for several seeds? Same Q for Table 3.
- Regarding the overfitting issue of ELLA: to what extent can changing the prior fix such overfitting issue? For example, the prior precision is often changed after training to improve the performance of LA [6]. Is this an option here, too?
limitations: There are no immediate societal impacts of the present work. Computational limitations are discussed in the paper, for example, the authors discuss the issue of overfitting present in their method.
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This work proposes a new, scalable variant of the (linearized) Laplace approximation (LLA) for inferring posterior distributions over the weights of deep learning models. Leveraging connections between LLA and neural tangent kernels (NTKs) -- namely that the LLA can essentially be viewed in function space as a Gaussian process with the NTK -- the paper describes how to speed up the LLA via a Nyström approximation to the NTK, effectively yielding an implicit sparse approximation to the LLA-GP predictive distribution. The work then elaborates on how to efficiently implement the proposed approach using automatic differentiation frameworks, and provides a theoretical analysis of the induced approximation error for the kernel. The resulting ELLA method is then empirically shown to outperform several baselines (both other LLA variants and a variational method) across a range of benchmarks -- even large-scale results with a Vision Transformer architecture on ImageNet are reported. It is also empirically found that ELLA (and LLA methods more generally) tends to "overfit" in some settings, which is proposed to be remedied by only using a subset of the data for estimating the covariance matrix, similar in spirit to early stopping methods often used in standard neural network training.
strengths_and_weaknesses: **Summary of Review**

This work proposes a novel, theoretically-justified and clever variant of the Laplace approximation that is empirically shown to be performant and scalable to fairly modern deep learning settings, and that could therefore have significant impact by enabling practical uncertainty quantification for deep neural networks. Overall, there are many things that I liked about this paper; however, I also see some significant issues with the empirical evaluation (see details below). All in all, I am inclined to recommend acceptance of the manuscript, but would be more enthusiastic in my judgement if the authors can convincingly address the concerns raised.

**Strengths**
- The paper is overall well-written, easy-to-follow, and clearly-structured.
- The proposed ELLA method for accelerating LLA is novel, theoretically sound, and intuitively sensible; it is therefore a timely and welcome addition to the fast-growing Laplace toolbox for Bayesian deep learning.
- I appreciate the implementation-specific details in Section 3.4 including the code snippets in Algorithm 1-2, which make it fairly clear how one would practically implement ELLA (without needing to look at the actual Python code).
- The theoretical results on the approximation error induced by ELLA are nice-to-have (even though they appear to be somewhat straightforwardly inherited from previous analyses of Nyström-like kernel approximations).
- A main advantage of the proposed ELLA method is that a practitioner can actively control the performance vs. cost trade-off by choosing the hyperparameters $M$ and $K$, which is in contrast to most other methods (e.g. LLA-Diag and LLA-KFAC) that have a fixed cost and performance (however, as mentioned in the weaknesses below, this trade-off should be assessed more thoroughly in the empirical evaluation).
- The empirical evaluation demonstrates that ELLA can outperform relevant baselines on image-classification-based uncertainty calibration and out-of-distribution detection tasks; I particularly appreciate the inclusion of experiments on ImageNet, which is (unfortunately) still not so common in the Bayesian deep learning literature; I was especially impressed by the results on the recent Vision Transformer architecture, which promisingly demonstrates that the method can be applied in fairly modern deep learning settings (which, again, is not typically the case -- Bayesian methods normally lack behind significantly when it comes to adoption of advances in deep learning).
- I really liked the experiment showing that ELLA (and LLA methods more generally) have a tendency to "overfit" to larger datasets (in the sense that the test loss starts increasing again past a certain number of data points used for fitting the covariance matrix), demonstrating that LLA methods can be significantly sped-up by subsampling the data in some settings.

**Weaknesses**
- As all experiments seem to have been run with just a single random seed, the results do not come with any error bars, which makes it difficult to reason about the statistical significance of the reported conclusions; while I understand that repeating experiments for multiple random seeds linearly increases the required computational effort, reporting some sort of error statistic would significantly aid the credibility of the drawn conclusions; for the particularly expensive experiments (i.e. on ImageNet) I could perhaps accept lack of compute as an excuse (although even there, a minimum number of 3 seeds would be desirable), but on the cheaper experiments (i.e. on CIFAR-10) I would really expect some repetition of the experiments.
- The empirical comparison does not seem entirely fair, as different methods have different memory and compute requirements; in particular, for the hyperparameters chosen, ELLA seems to be more expensive than some of the other baselines considered (in terms of compute and/or memory cost, as e.g. also shown in Fig. 4c); therefore, it might be more appropriate and insightful to plot performances in 2D with memory and/or compute effort on the x-axis and performance on the y-axis, so that it becomes clear how the different methods can trade-off cost vs. performance; for ELLA, it would be great to then report performance for different values of $M$ and/or $K$ (i.e. different costs), yielding a Pareto curve; a practitioner can then choose the method/hyperparameters on the Pareto front that best fulfills their needs, i.e. which is either 1) as cheap as possible for a given desired performance X, or 2) as performant as possible for a given desired cost Y; without such a plot it is difficult to draw conclusions on how superior ELLA really is -- although ELLA's ability to cater to a wide range of performance/cost trade-offs by sweeping over $M$ and $K$ should already make it look better than most of the baselines; you already show several plots that are related in spirit (Fig 2. a-b, Fig. 4 b-c), but I believe the specific kind of plot I described would be much more insightful.
- For the ImageNet experiments, I would expect that at least the last-layer KFAC variant should be feasible to run and not yield an OOM error (with either a ResNet or ViT) -- e.g. the Eschenhagen et al. (2021) paper you cite seems to have managed to run this Laplace variant on ImageNet with a ResNet; I think it would significantly strengthen the paper to at least have this one Laplace baseline to compare with; if this is really an issue with the Laplace library, I would encourage you to reach out to the authors / raise an issue on the GitHub repo -- I also had issues with this library in the past, and in my experience, the authors are typically happy to help and fairly quick to respond.
- The proposed approach can be viewed as limited in novelty in the sense that it effectively combines the LLA with a Nyström approximation to the NTK, which both have been studied (fairly extensively) before. That being said, I think the idea is pretty neat (even though somewhat obvious in hindsight) and appreciate that the authors made the effort of actually getting this to work, so I don't view this to be a major issue.

**Minor issues**
- l. 226-227: to back up the claim that "ELLA delivers a closer approximation to LLA" it might be worth explicitly quantifying the discrepancies between the predictive distributions using some metric; qualitatively (i.e. by just looking at the plots), it seems like LLA-KFAC also comes pretty close (but I agree that even LLA might underestimate in-between uncertainty in this setting, so the ELLA fit could be more desirable).
- Generally, it would be useful to have short take-away messages in the captions of each table/figure for clarity and convenience.
questions: - l. 234-235: in the footnote, you mention that "storing 2000 vectors of size $P$ is costly", but do you not just need to store $K=20$ instead of $M=2000$ vectors of size $P$ here (cf. l. 163-164)?
- l. 239-240: it would be good to understand why LLA-Diag and LLA-KFAC fail in this setting, given that previous works have been able to make them work fairly well, e.g. Daxberger et al. (2021); what have you tried to make these work and what is your intuition on why they fail?

limitations: Yes, the authors have adequately addressed this.
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: The paper targets the improvement of probabilistic predictions in Bayesian deep learning with the (linearized) Laplace approximation. In particular, it proposes to apply the Nystrom kernel approximation to the neural-tangent-like covariance matrix of the corresponding Gaussian process that is the function-space dual of the linearized Laplace in weight-space. Using the Nystrom approximation allows the authors to scale the linearized Laplace to ImageNet and ViT scale. They compare their method to other Laplace approximations in weight space and find that their method performs better, though in some cases only marginally, at a similar or lower computational cost.

strengths_and_weaknesses: Strengths:
* The Laplace approximation for larger scale machine learning has recently seen increasing attention and this paper further improves scalability by introducing the Nystrom approximation in its dual (function space) formulation.
* The paper seems to be technically correct, though I haven't checked all the details.

Weaknesses:
* Generality + Novelty: The novelty and originality presented in this paper is limited, especially when compared to other papers at ICML or NeurIPS in this space (e.g. ref [6] provides a library for Laplace approximation and a thorough review/benchmark, or ref [9] introduces kernel approximations more generally and not only for large-scale Laplace predictions). While this paper extends the Laplace approximation to ImageNet and ViTs, the paper ultimately discusses one particular approximation (Nystrom) to one particular covariance function (NTK) for one particular task (probabilistic predictions on pre-trained models); it largely builds on existing libraries and improvements over MAP predictions are often marginal (especially on some of the larger-scale problems) and provided without errorbars. The authors also claim that no other methods can be applied to ViT or ImageNet, yet last-layer Laplace (LLA*) should be applicable in this case as well (at least when subsampling the dataset; e.g. [6] apply last layer Laplace to transformer models).
* Contributions + relation to prior work: 
  - Several results in Sec 3.2 and 3.3 are, in my opinion, not novel and have been similarly derived in the kernel literature (where the Nystrom approximation is commonplace) as well as in Appendix A.2 and A.3 of NeuralEF (ICML2022, reference [9] in the paper under review) for the NTK kernel. While "NeuralEF" is referenced (as ref [9]) in the paper, it is not mentioned in Sec 3.2 and 3.3 of this paper and it is not clarified, which results in Sec 3.2 and 3.3 are novel and which are already known.
  - The Nystrom approximation is a well known and well studied approximation to kernel functions, yet a discussion/presentation of this is missing from this paper.
  - NeuralEF [9] highlights several limitations of the Nystrom approximation when applied to NTK mostly to do with scalability (see Sec 2.2 in [9] in particular). They highlight the cost of the eigendecomposition as well as prediction costs at test time. I agree with [9] that this could be problems and scalability of the Nystrom should be discussed more prominently.
* also see "Limitations" below
* I found the presentation/clarity to be not clear:
	- The setting (probabilistic prediction with Bayesian neural networks), ultimate goal/objective (improve linearized(?) posterior predictive), as well as the baseline methods (discussion + equations for e.g. Laplace with KFAC, ...) should be introduced more clearly. E.g. the posterior predictive distributions  evaluated are only mentioned in passing/do not have clear equations/equation numbers
	- The paper ultimately presents a technical contribution (Nystrom for this particular kernel) in an arguably very technical way. In my opinion, more emphasis should be put on explanation and some of the technical details (e.g. indexing with $i$ in Sec 3.3) should be left to the appendix.
	- Alg 1 and 2 should be explained in much more detail; comments + references to the text as well as from the text to the algorithms are missing

questions: * I found it hard to understand/follow the discussion of the overfitting issue/behaviour of ELLA in Sec 5.2. As far as I understood, ELLA uses pre-trained models and only deals with the GP predictive distribution. As such, "early stopping" does not really make sense to me as it typically refers to early stopping of training. Could you please elaborate?
limitations: * see Weaknesses above, in particular, the authors only investigate ELLA for predictions.
* several works (e.g. ref [22], [6]) point out that the choice of prior precision is important to achieve good predictive performance (in particular for uncertainties/NLL/ECE); however, it looks like the values here have simply been taken from the pre-trained model regularisation, which may explain the poor performance of LLA and LLA+KFAC. For this reason, [6] (as far as I understand) proposes to use last-layer Laplace by default when these hyperparameters are not tuned/the MAP model itself does not overfit.
* the illustrative regression example (sec 5.3) is not discussed in enough detail. In particular, it should be investigated and further discussed why ELLA can be better than the full GGN model (LLA), to which it should be equivalent (except for the Nystrom approximation, of course). The authors are "surprised" (see caption to Fig 1) by ELLA performing better than the full model (LLA) as well, yet they do not investigate or discuss this. For example, the full GP predictive (without Nystrom approximation) as well as different/larger values for M and K should be considered in this case to build further intuition for the model and to investigate this behaviour.
review_title: null

Meta-review: metareview: This paper introduces an approach to accelerating linearized Laplace approximations to Bayesian neural network posteriors, particularly considering prediction tasks, by performing a Nyström approximation to the neural tangent kernel.

The reviewers all recommended acceptance, eventually — one reviewer was initially quite critical but revised after a rather extensive discussion with the authors revised to a borderline accept. In particular some additional experiments analyzing overfitting in these models, in general, were appreciated.

While this is perhaps borderline on the scores (5,6,7), given the overall quality of the work and the extent to which this was updated and improved during the rebuttal period, I would recommend acceptance.
confidence: Less certain
Decision: Accept