Sketched Lanczos uncertainty score: a low-memory summary of the Fisher information

Authors: marco  miani
lorenzo  beretta
søren  hauberg
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:47:58

Abstract: Current uncertainty quantification is memory and compute expensive, which hinders practical uptake. To counter, we develop Sketched Lanczos Uncertainty (SLU): an architecture-agnostic uncertainty score that can be applied to pre-trained neural networks with minimal overhead. Importantly, the memory use of SLU only grows logarithmically with the number of model parameters. We combine Lanczos' algorithm with dimensionality reduction techniques to compute a sketch of the leading eigenvectors of a matrix. Applying this novel algorithm to the Fisher information matrix yields a cheap and reliable uncertainty score. Empirically, SLU yields well-calibrated uncertainties, reliably detects out-of-distribution examples, and consistently outperforms existing methods in the low-memory regime.

Main paper: 
 introduction: 
 The best-performing uncertainty quantification methods share the same problem: scaling. Practically this prevents their use for deep neural networks with high parameter counts. Perhaps, the simplest way of defining such a score is to independently train several models, perform inference, and check how consistent predictions are across models. The overhead of the resulting 'Deep Ensemble' [Lakshminarayanan et al., 2017] notably introduces a multiplicative overhead equal to the ensemble size. Current approaches aim to reduce the growth in training time costs by quantifying uncertainty through local information of a single pre-trained model. This approach has shown some success for methods like Laplace's approximation [MacKay, 1992 , Immer et al., 2021 , Khan et al., 2019] , SWAG [Maddox et al., 2019] , SCOD [Sharma et al., 2021] or Local Ensembles [Madras et al., 2019] . They avoid the need to re-train but still have impractical memory needs.A popular approach to characterizing local information is the empirical Fisher information matrix, which essentially coincides with the Generalized Gauss-Newton (GGN) matrix [Kunstner et al., 2019] . Unfortunately, for a p-parameter model, the GGN is a p × p matrix, yielding such high memory costs that it cannot be instantiated for anything but the simplest models. The GGN is, thus, mostly explored through approximations, e.g. block-diagonal [Botev et al., 2017] , Kronecker-factorized [Ritter et al., 2018a , Lee et al., 2020 , Martens and Grosse, 2015] or even diagonal [Ritter et al., 2018b , Miani et al., 2022] . An alternative heuristic is to only assign uncertainties to a subset of the model parameters [Daxberger et al., 2021b , Kristiadi et al., 2020] , e.g. the last layer.Instead, we approximate the GGN with a low-rank matrix. A rank-k approximation of the GGN can be computed using Lanczos algorithm [Madras et al., 2019 , Daxberger et al., 2021a] or truncated Our contribution is twofold: (1) we prove that orthogonalization approximately commutes with sketching in Section 3, which makes it possible to sketch Lanczos vectors on the fly and orthogonalize them post-hoc, with significant memory savings; (2) we empirically show that, in the lowmemory-budget regime, the disadvantage of introducing noise through sketching is outweighed by a higherrank approximation, thus performing better than baselines when the same amount of memory is used. 
background: 
 Let f θ : R d → R t denote a neural network with parameter θ ∈ R p , or equivalently f : R p ×R d → R t . Let J θ ′ (x) = ∇ θ f θ (x)| θ=θ ′ ∈ R t×p be its Jacobian with respect to the parameters, evaluated at a datapoint x ∈ R d . Given a training dataset D = {(x i , y i )} i=1,...,n and a loss function L(θ) = (x,y)∈D ℓ(y, f (x, θ), the Generalized Gauss-Newton matrix (GGN) is defined asEQUATIONwhere H(x i ) = ∇ 2 f θ (xi) ℓ(y i |f θ (x i )) ∈ R t×t is the Hessian of the loss with respect to the neural network output. We reduce the notational load by stacking the per-datum Jacobians into J θ = [J θ (x 1 ); . . . ; J θ (x n )] ∈ R nt×p and similarly for the Hessians, and write the GGN matrix as G θ = J ⊤ θ HJ θ . For extended derivations and connections with the Fisher matrix we refer to the excellent review by Kunstner et al. [2019] . In the following, we assume access to a pre-trained model with parameter θ * and omit the dependency of G on θ * .Computationally we emphasize that Jacobian-vector products can be performed efficiently when f is a deep NN. Consequently, the GGN-vector product has twice the cost of a gradient backpropagation, at least for common choices of ℓ like MSE or cross-entropy [Khan and Rue, 2021] .We measure the uncertainty at a datapoint x as the variance of the prediction f θ (x) with respect to a distribution over parameter θ defined at training time and independently of x. This general scheme has received significant attention. For example, Deep Ensemble [Lakshminarayanan et al., 2017] uses a sum of delta distribution supported on independently trained models, while methods that only train a single network θ * generally use a Gaussian N (θ|θ * , M ) with covariance M ∈ R p×p . In the latter case, a first-order approximation of the prediction variance is given by a M -norm of the JacobianEQUATIONwhere f L θ (x) = f θ (x) + J θ * (x) • (θ * -θ) is a linearization of θ → f θ (x) around θ * . The GGN matrix (or empirical Fisher; Kunstner et al. [2019] ) is notably connected to uncertainty measures and, more specifically, to the choice of the matrix M . Different theoretical reasoning leads to different choices and we focus on two of them:EQUATIONwhere Π G is the projection onto the non-zero eigenvectors of G and α > 0 is a constant. Madras et al. [2019] justify M noeig through small perturbation along zero-curvature directions. 2 Immer et al. [2021] justify M eig in the Bayesian setting where α is interpreted as prior precision.We do not question these score derivations and refer to the original works, but we highlight their similarity. Given an orthonormal eigen-decomposition of G = i λ i v i v ⊤ i we see thatEQUATIONThus both covariances are higher in the directions of zero-eigenvalues, and the hyperparameter α controls how many eigenvectors are relevant in M eig .Practical choices. The matrix G is too big to even be stored in memory and approximations are required, commonly in a way that allows access to either the inverse or the projection. A variety of techniques are introduced like diagonal, block diagonal, and block KFAC which also allow for easy access to the inverse. Another line of work, like SWAG [Maddox et al., 2019] , directly tries to find an approximation of the covariance M based on stochastic gradient descent trajectories. Alternatively, low-rank approximations use the eigendecomposition relative to the top k eigenvalues, which allows direct access to both inverse and projection.Is low-rank a good idea? The spectrum of the GGN has been empirically shown to decay exponentially [Sagun et al., 2017 , Papyan, 2018 , Ghorbani et al., 2019] (see Figure 2 ). We investigate this phenomenon further in Appendix C.1 with an ablation over Lanczos hyperparameters. This fast decay implies that the quality of a rank-k approximation of G improves exponentially w.r.t. k if we measure it with an operator or Frobenius norm, thus supporting the choice of low-rank approximation. Moreover, in the overparametrized setting p ≫ nt, the rank of the GGN is by construction at most nt, which is closely linked with functional reparametrizations [Roy et al., 2024] .The low-rank approach has been applied intensively [Madras et al., 2019 , Daxberger et al., 2021a , Sharma et al., 2021] with success, although always limited by memory footprint of pk. Madras et al. [2019] argue in favor of a not-too-high k also from a numerical perspective, as using the "small" eigenvectors appears sensitive to noise and consequently is not robust.The Lanczos algorithm is an iterative method for tridiagonalizing a symmetric matrix G. If stopped at iteration k, Lanczos returns a column-orthonormal matrixV = [V 1 | . . . |V k ] ∈ R p×k and a tridiagonal matrix T ∈ R k×k such that V ⊤ GV = T .The range space of V corresponds to the Krylov subspace K k = span{v, Gv, . . . , G k-1 v}, where v = V 1 is a randomly chosen vector. Provably, K k 2 To be precise, Madras et al. [2019] uses the Hessian rather than the GGN, although their reasoning for discarding the eigenvalues applies to both matrices. Thus, while the score with Mnoeig is technically novel, it is a natural connection between Laplace approximations and local ensembles. Also note that some work uses the Hessian in the eigenvalues setting [MacKay, 2003] despite this requires ad-hoc care for negative values.approximates the eigenspace spanned by the top-k eigenvectors of G, i.e. those corresponding to the eigenvalues of largest values [Meurant, 2006] . Thus, V T V ⊤ approximates the projection of G onto its top-k eigenspace. Notice that projecting G onto its top-k eigenspace yields the best rank-k approximation of G under any unitarily-invariant norm [Mirsky, 1960] . Once the decomposition G ≈ V T V ⊤ is available, we can retrieve an approximation to the top-k eigenpairs of G by diagonalizing T into T = W ΛW ⊤ , which can be done efficiently for tridiagonal matrices [Dhillon, 1997] . It has both practically and theoretically been found that this eigenpairs' approximation is very good. We point the reader to Meurant [2006] and Cullum and Willoughby [2002] for a comprehensive survey on this topic.The benefits of Lanczos. Lanczos has two features that make it particularly appealing. First, Lanczos does not need explicit access to the input matrix G, but only access to an implementation of G-vector product u → Gu. Second, Lanczos uses a small working space: only 3p floating point numbers, where the input matrix is p × p. Indeed, we can think of Lanczos as releasing its output in streaming and only storing a small state consisting of the last three vectors V i-1 , V i and V i+1 .The downsides of Lanczos. Unfortunately, the implementation of Lanczos described above is prone to numerical instability, causing V 1 . . . V k to be far from orthogonal. A careful analysis of the rounding errors causing this pathology was carried out by Paige [1971 Paige [ , 1976 Paige [ , 1980]] . To counteract, a standard technique is to re-orthogonalize V i+1 against all {V j } j≤i , at each iteration. This technique has been employed to compute the low-rank approximation of huge sparse matrices [Simon and Zha, 2000] , as well as by Madras et al. [2019] to compute an approximation to the top-k eigenvectors. Unfortunately, this version of Lanczos loses one of the two benefits described above, in that it must store a larger state consisting of the entire matrix V . Therefore, we dub this version of the algorithm hi-memory Lanczos and the memory-efficient version described above low-memory Lanczos. See Appendix C for a comprehensive discussion on these two versions.Post-hoc orthogonalization Lanczos. Alternatively, instead of re-orthogonalizing at every step as in hi-memory Lanczos, we can run low-memory Lanczos, store all the vectors, and orthogonalize all together at the end. Based on the observations of Paige [1980] , we expect that orthogonalizing the output of low-memory Lanczos post-hoc should yield an orthonormal basis that approximately spans the top-k eigenspace, similar to hi-memory Lanczos. This post-hoc version of Lanczos is however insufficient. It avoids the cost of orthogonalizing at every iteration but still requires storing the vectors, thus losing again the benefit of memory requirement. Or at least it does unless we find an efficient way to store the vectors.Sketching is a key technique in randomized numerical linear algebra [Martinsson and Tropp, 2020] to reduce memory requirements. Specifically, sketching embeds high dimensional vectors from R p into a lower-dimensional space R s , such that the expected norm error of vector dot-products is bounded, and, as a result, also the score in Equation 2. Here we give a concise introduction to this technique. Definition 2.1 (Subspace embedding). Fix ε > 0. A (1 ± ε) ℓ 2 -subspace embedding for the column space of an p × k matrix U is a matrix S for which for all y ∈ R k∥SU y∥ 2 = (1 ± ε)∥U y∥ 2 .(5) Here ω is such that the current best matrixmultiplication algorithm runs in time n ω .The goal is to design an oblivious subspace embedding, that is a random matrix S such that, for any matrix U , S is a subspace embedding for U with sufficiently high probability. In our method, we use a Subsampled Randomized Fourier Transform (SRFT) to achieve this goal [Ailon and Chazelle, 2009] .A SRFT is a s × p matrix defined by the product 1 / √ spP HD, where D is a diagonal matrix where each diagonal entry is an independent Rademacher random variable, H is the discrete Fourier transform, and P is a diagonal matrix where s random diagonal entries are set to one and every other entry is set to zero. Thanks to the Fast Fourier Transform algorithm, SRFT can be evaluated in O(p log p) time, and its memory footprint is only p + s.The following theorem shows that, as long as the sketch size s is big enough, SRFT is an oblivious subspace embedding with high probability.Theorem 2.2 (Essentially, Theorem 7 in Woodruff et al. [2014] ). For any p × k matrix U , SRFT is a (1 ± ε)-subspace embedding for the column space of U with probability 1 -δ as long as s = Ω((k + log p)ε -2 log(k/δ)).We stress that, although several other random projections work as subspace embeddings, our choice is not incidental. Indeed other sketches, including the Sparse JL transform [Kane and Nelson, 2014, Nelson and Nguyên, 2013] or the Dense JL transform (Theorem 4, Woodruff et al. [2014] ), theoretically have a larger memory footprint or a worse trade-off between s and k, as clarified in Table 1 . From such a comparison, it is clear that SRFT is best if our goal is to minimize memory footprint. At the same time, evaluation time is still quasilinear.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The limitations are related to memory requirements and scaling of the error induced by the sketching. Both of them are fully addressed in the Method Section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
method: 
 We now develop the novel SKETCHED LANCZOS algorithm by combining the 'vanilla' Lanczos algorithm (Section 2.2) with sketching (Section 2.3). Pseudo-code is presented in Algorithm 1. Next, we apply this algorithm in the uncertainty quantification setting and compute an approximation of the score in Equation 2 due to Madras et al. [2019] . Our motivation is that given a fixed memory budget, the much lower memory footprint induced by sketching allows for a higher-rank approximation of G.We find the best way to explain our algorithm is to first explain a didactic variant of it, where sketching and orthogonalization happen in reverse order.Running low-memory Lanczos for k iterations on a p × p matrix iteratively constructs the columns of a p × k matrix V . Then, post-hoc, we re-orthogonalize the columns of V in a matrix U ∈ R p×k . Such a matrix is expensive to store due to the value of p, but if we sample a SRFT sketch matrix S ∈ R s×p , we can then store a sketched version SU ∈ R s×k , saving memory as long as s < p. In other words, this is post-hoc orthogonalization Lanczos with a sketching at the end.We observe that sketching the columns of U is sufficient to ε-preserve the norm of matrix-vector products, with high probability. In particular, the following lemma holds (proof in Appendix A). Lemma 3.1 (Sketching low-rank matrices). Fix 0 < ε, δ < 1 /2 and sample a random s × p SRFT matrix S. Then, for any v ∈ R p and any matrix U ∈ R p×k with ||v|| 2 , ||U || 2 = O(1) we haveEQUATION)as long as s = Ω(kε -2 • log p • log(k/δ)).This algorithm may initially seem appealing since we can compute a tight approximation of ∥U ⊤ v∥ by paying only s × k in memory (plus the neglectable cost of storing S). However, as an intermediate step of such an algorithm, we still need to construct the matrix U , paying p × k in memory. Indeed, we defined U as a matrix whose columns are an orthonormal basis of the column space of V and we would like to avoid storing V explicitly and rather sketch each column V i on the fly, without ever paying p × k memory. This requires swapping the order of orthogonalization and sketching.This motivates us to prove that if we orthonormalize the columns of SV and apply the same orthonormalization steps to the columns of V , then we obtain an approximately orthonormal basis. Essentially, this means that sketching and orthogonalization approximately commute. As a consequence, we can use a matrix whose columns are an orthonormal basis of the column space of SV as a proxy for SU while incurring a small error. Formally, the following holds (proof in Appendix A). Lemma 3.2 (Orthogonalizing the sketch). Fix 0 < ε, δ < 1 /2 and sample a random s × p SRFT matrix S. As long as s = Ω(kε -2 • log p • log(k/δ)) the following holds with probability 1 -δ.Given any p × k full-rank matrix V , decompose V = U R and SV = U S R S so that U ∈ R p×k , U S ∈ R s×k and both U and U S have orthonormal columns. For any unit-norm v ∈ R p we haveEQUATIONIn Lemma 3.1, we proved that given a matrix U with orthonormal columns we can store SU instead of U to compute v → ||U ⊤ v|| 2 while incurring a small error. However, in our use case, we do not have explicit access to U . Indeed, we abstractly define U as a matrix whose columns are an orthonormal basis of the column space of V , but we only compute U S as a matrix whose columns are an orthonormal basis of the column space of SV , without ever paying p • k memory. Lemma 3.2 implies that the resulting error is controlled.Algorithm 1 lists the pseudocode of our new algorithm: SKETCHED LANCZOS.Algorithm 1 SKETCHED LANCZOS 1: Input: Rank k, sketch matrix S∈R s×p , matrix-vector product function v → Gv for G∈R p×p . 2: Initialize v 0 ∈ R p as a uniformly random unit-norm vector.3: for i in 1, . . . k do 4:v i ← LANCZOSITERATION(v i-1 , v → Gv)(glossing over the tridiagonal detail)5:Sketch and store v S i ← Sv i 6: end for 7: Construct the matrixV S = [v S 1 , . . . , v S k ] ∈ R s×k 8:Orthogonalize the columns of V S and return U S ∈ R s×k Preconditioned SKETCHED LANCZOS. We empirically noticed that low-memory Lanczos' stability is quite dependent on the conditioning number of the considered matrix. From this observation, we propose a slight modification of SKETCHED LANCZOS that trades some memory consumption for numerical stability.The idea is simple, we first run hi-memory Lanczos for k 0 iterations, obtaining an approximation of the top-k 0 eigenspectrum U 0 ∈ R p×k0 , Λ 0 ∈ R k0×k0 . Then we define a new matrix-vector productEQUATIONand run SKETCHED LANCZOS for k 1 iterations on this new, better-conditioned, matrix Ḡ. This results in a U S ∈ R s×k1 with sketched orthogonal columns. With k = k 1 + k 0 , the simple concatenation[SU 0 |U S ] ∈ R (k0+k1)×s is a sketched orthogonal k-dimensional base of the top-k eigenspace of G, analogous to non-preconditioned SKETCHED LANCZOS. The extra stability comes at a memory cost of k 0 • p, thus preventing k 0 from being too large.The uncertainty score in Equation 2 is computed by first approximating the Generalized Gauss-Newton matrix G. The approach of constructing an orthonormal basis U ∈ R p×k of the top-k eigenvectors of G with relative eigenvalues Λ, leads to the low-rank approximation G ≈ U ΛU ⊤ . This step is done, for example, by Madras et al. [2019] through hi-memory Lanczos and by Sharma et al. [2021] through truncated randomized SVD. In this step, we employ our novel SKETCHED LANCZOS. Similar to Madras et al. [2019] , we focus on the score with M noeig and thus we neglect the eigenvalues.Having access to U , we can compute the score for a test datapointEQUATIONwhich clarifies that computing ||J θ * (x)U || F is the challenging bit to retrieve the score in Equation 2. Note that ∥J θ * (x)∥ 2 F can be computed exactly with t Jacobian-vector products. Computing the uncertainty score through sketching. Employing the novel SKETCHED LANCZOS algorithm we can ε-approximate the score in Equation 9, with only minor modifications. Running the algorithm for k iterations returns a matrix U S ∈ R s×k , which we recall is the orthogonalization of SV where V ∈ R p×k are the column vectors iteratively computed by Lanczos.Having access to U S , we can then compute the score for a test datapoint x ∈ R d asEQUATION)where parentheses indicate the order in which computations should be performed: first a sketch of J θ * (x) is computed, and then it is multiplied by U ⊤ S . Algorithm 2 summarizes the pipeline. Approximation quality. Recall that the Jacobian J θ * (x) is a t × p matrix, where t is the output size of the neural network. A slight extension of Lemma 3.2 (formalized in Lemma A.1) implies that the Algorithm 2 SKETCHED LANCZOS Uncertainty score (SLU) 1: Input: Observed data D, trained MAP parameter θ * ∈ R p , approximation rank k, sketch size s. 2: Initialize SRFT sketch matrix S∈R s×p 3: Initialize GGN-vector product v → Gv = xi∈D J θ * (x i ) ⊤ H(x i )J θ * (x i )v 4: U S ← SKETCHED LANCZOS(k, S, v → Gv) 5: At query time, for a test point x return ∥J θ * (x)∥ 2 F -∥U ⊤ S (SJ θ * (x) ⊤ )∥ 2 Fscore in Equation 9 is guaranteed to be well-approximated by the score in Equation 10 up to a factor 1 ± ε with probability 1 -δ as long as the sketch size is big enough s = Ω(ktε -2 log p log(kt/δ)).Neglecting the log terms to develop some intuition, we can think of the sketch size to be s ≈ ktε -2 , thus resulting in an orthogonal matrix U S of size ≈ ktε -2 × k, which also correspond to the memory requirement. From an alternative view, we can expect the error induced by the sketching to scale as ε ≈ kt/s, naturally implying that a larger k will have a larger error, and larger sketch sizes s will have lower error. Importantly, the sketch size s (and consequently the memory consumption and the error bound) depends only logarithmically on the number of parameters p, while the memory-savingratio s /p clearly improves a lot for bigger architectures.Memory footprint. The memory footprint of our algorithm is at most 4p + s(k + 1) floating point numbers. Indeed, the SRFT sketch matrix uses p + s numbers to store S, whereas low-memory Lanczos stores at most 3 size-p vectors at a time. Finally, U S is a s × k matrix. At query time, we only need to store S and U S , resulting in a memory footprint of p + s(k + 1). Therefore, our method is significantly less memory-intensive than the methods of Madras et al. [2019] , Sharma et al. [2021] , and all other low-rank Laplace baselines, that use Ω(kp) memory.Time footprint. The time requirement is comparable to Vanilla Lanczos. We need k SRFT sketch-vector products and each takes O(p log p) time, while the orthogonalization of SV = U S R through QR decomposition takes O(pk 2 ) time. Both Vanilla and Sketched algorithm performs k GGN-vector products and each takes O(pn) time, where n is the size of the dataset, and that is expected to dominate the overall O(pk(log p + k + n)). Query time is also fast: sketching, Jacobian-vector product and U S -vector product respectively add up to O(tp(log p + 1) + tsk).Preprocessing 4p + s(k + 1) O(pk(log p + k + n)) Query p + s(k + 1) O(tp(log p + 1) + tsk)Table 2 : Recall that through the whole paper p is number of parameters, n is dataset size, t is output dimensionality, k is rank approximation and s is sketch size.Note that the linear scaling with output dimension t can slow down inference for generative models. We refer to Immer et al. [2023] for the effect of considering a subset on dataset size n (or on output dimension t).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: The paper is predominantly theoretical. However, uncertainty quantification is a potential remedy for several issues with currently deployed machine learning models, so there is a potential for positive societal impact. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
related work: 
 Modern deep neural networks tend to be more overconfident than their predecessors [Guo et al., 2017] . This motivated intensive research on uncertainty estimation. Here we survey the most relevant work.Perhaps, the simplest technique to estimate uncertainty over a classification task is to use the softmax probabilities output by the model [Hendrycks and Gimpel, 2016] . A more sophisticated approach, combine softmax with temperature scaling (also Platt scaling) [Liang et al., 2017 , Guo et al., 2017] .The main benefit of these techniques is their simplicity: they do not require any computation besides inference. However, they do not extend to regression and do not make use of higher-order information. Moreover, this type of score relies on the extrapolation capabilities of neural networks since they use predictions made far away from the training data. Thus, their poor performance is not surprising.To alleviate this issue, an immense family of methods has been deployed, all sharing the same common idea of using the predictions of more than one model, either explicitly or implicitly. A complete review of these methods is unrealistic, but the most established includes Variational inference [Graves, 2011 , Hinton and Van Camp, 1993 , Liu and Wang, 2016] , Deep ensembles [Lakshminarayanan et al., 2017] , Monte Carlo dropout [Gal and Ghahramani, 2016, Kingma et al., 2015] and Bayes by Backprop [Blundell et al., 2015] .More closely related to us, the two uncertainty quantification scores introduced in Equation 2with covariances M eig and M noeig (Equation 3) have been already derived from Bayesian (with Laplace's approximation) and frequentist (with local perturbations) notions of underspecification, respectively:Laplace's approximation. By interpreting the loss function as an unnormalized Bayesian logposterior distribution over the model parameters and performing a second-order Taylor expansion, Laplace's approximation [MacKay, 1992] results in a Gaussian approximate posterior whose covariance matrix is the loss Hessian. The linearized Laplace approximation [Immer et al., 2021 , Khan et al., 2019] further linearize f θ at a chosen weight θ * , i.e.f θ (x) ≈ f θ * (x) + J θ * (x)(θ -θ * ).In this setting the posterior is exactly Gaussian and the covariance is exactly M eig .Local perturbations. A different, frequentist, family of methods studies the change in optimal parameter values induced by change in observed data. A consequence is that the parameter directions corresponding to functional invariance on the training set are the best candidates for OoD detection. This was directly formalized by Madras et al. [2019] which derives the score with covariance M noeig , which they call a local ensemble. A similar objective is approximated by Resampling Under Uncertainty (RUE, Schulam and Saria [2019] ) perturbing the training data via influence functions, and by Stochastic Weight Averaging Gaussian (SWAG, Maddox et al. [2019] ) by following stochastic gradient descent trajectories.Sketching in the deep neural networks literature. Sketching techniques are not new to the deep learning community. Indeed several works used randomized SVD [Halko et al., 2011] , which is a popular algorithm [Tropp and Webber, 2023 ] that leverages sketch matrices to reduce dimensionality and compute an approximate truncated SVD faster and with fewer passes over the original matrix. It was used by Antorán et al. [2023] to compute a preconditioner for conjugate gradient and, similarly, by Mishkin et al. [2018] to extend Variational Online Gauss-Newton (VOGN) training [Khan et al., 2018] . More related to us, Sketching Curvature for OoD Detection (SCOD, Sharma et al. [2021] ) uses Randomized SVD to compute exactly the score in Equation 2with M eig , thus serving as a baseline with an alternative to Lanczos. We compare to it more extensively in Appendix B.To the best of our knowledge no other work in uncertainty estimation for deep neural networks uses sketching directly to reduce the size of the data structure used for uncertainty estimation. Nonetheless, a recent line of work in numerical linear algebra studies how to apply sketching techniques to Krylov methods, like Lanczos or Arnoldi [Balabanov and Grigori, 2022 , Timsit et al., 2023 , Simoncini and Wang, 2024 , Güttel and Schweitzer, 2023] . We believe sketching is a promising technique to be applied in uncertainty estimation and, more broadly, in Bayesian deep learning. Indeed, all the techniques based on approximating the GGN could benefit from dimensionality reduction. 
experiments: 
 With a focus on memory budget, we benchmark our method against a series of methods, models, and datasets. The code for both training and testing is implemented in JAX [Bradbury et al., 2018] and it is publicly available3 . Details, hyperparameters and more experiments can be found in Appendix D.To evaluate the uncertainty score we measure the performance of out-of-distribution (OoD) detection and report the Area Under Receiver Operator Curve (AUROC). We choose models with increasing complexity and number of parameters: MLP, LeNet, ResNet, VisualAttentionNet and SwinTransformer architectures, with the number of parameters ranging from 15K to 200M. We train such models on 5 different datasets: MNIST [Lecun et al., 1998] , FASHIONMNIST [Xiao et al., 2017] , CIFAR-10 [ Krizhevsky et al., 2009] , CELEBA [Liu et al., 2015] and IMAGENET [Deng et al., 2009] . We test the score performance on a series of OoD datasets, including rotations and corruptions [Hendrycks, 2019] of ID datasets, as well as SVHN [Netzer et al., 2011] and FOOD101 [Bossard et al., 2014] .For CELEBA and IMAGENET we hold out some classes from training and use them as OoD datasets.We compare to the most relevant methods in literature: Linearized Laplace Approximation with a low-rank structure (LLA) [Immer et al., 2021] and Local Ensemble (LE) are the most similar to us. We also consider Laplace with diagonal structure (LLA-D) and Local Ensemble Hessian variant (LE-H) [Madras et al., 2019] . Another approach using randomized SVD instead of Lanczos is Sketching Curvature for OoD Detection (SCOD) [Sharma et al., 2021] , which also serves as a Lanczos baseline. Lastly, we include SWAG [Maddox et al., 2019] Effect of different sketch sizes. We expect the error induced by the sketching to scale as ϵ ≈ k/s, thus a larger k will have a larger error, and larger sketch sizes s will have a lower error, independently on number of parameters up to log-terms. On the other hand, a larger parameter count will have a better memory-saving-ratio s /p, leading to an advantage for bigger architectures as shown in Figure 3 . Results for all OoD datasets are illustrated more extensively in Figure 4 . Mean values and standard deviations for each table block obtained over 10, 10, 5, 3, 1 independently trained models, respectively.model MLP p = 15K LeNet p = 40K ResNet p = 200KSummary of the experiments. For most of the ID-OoD dataset pairs we tested, our SKETCHED LANCZOS Uncertainty outperforms the baselines, as shown in Table 3 and more extensively in Figure 4 where we fix the memory budget to be 3p. Deep Ensemble performs very well in the small architecture but progressively deteriorates for bigger parameter sizes. SWAG outperforms the other methods on some specific choices of CIFAR-10 corruptions, but we found this method to be extremely dependent on the choice of hyperparameters. SCOD is also a strong baseline in some settings, but we highlight that it requires instantiating the full Jacobian with an actual memory requirement of tp. These memory requirements make SCOD inapplicable to ImageNet. In this setting, given the significant training time, also DEEP ENSEMBLE becomes not applicable.The budget of 3p is an arbitrary choice, but the results are consistent with different values. More experiments, including a 10p memory budget setting, a study on the effect of preconditioning, and a synthetic-data ablation on the trade-off sketch size vs low rank, are presented in Appendix D.3. 
conclusion: 
 We have introduced SKETCHED LANCZOS, a powerful memory-efficient technique to compute approximate matrix eigendecompositions. We take a first step in exploiting this technique showing that sketching the top eigenvectors of the Generalized Gauss-Newton matrix leads to high-quality SLU outperforms the baselines on several choices of ID (4a, 4b, 4c, 4d, 4e) and OoD (x-axis) datasets pairs. Dashed lines are for improved visualization only; see Table 3 for values and standard deviations. Plots 4a, 4b, 4c, 4d, 4e are averaged respectively over 10, 10, 5, 3, 1 independently trained models.scalable uncertainty measure. We empirically show the superiority of the SKETCHED LANCZOS Uncertainty score (SLU) among a variety of baselines in the low-memory-budget setting, where the assumption is that the network has so many parameters that we can only store a few copies.Limitations. The data structure produced by Sketched Lanczos is sufficiently rich to evaluate the predictive variance, and consequently the uncertainty score. However, from a Bayesian point of view, it is worth noting that the method does not allow us to sample according to the posterior. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: we do not deem this to be relevant.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: we only consider well-established benchmark data, which we cite appropriately.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Proof of "Sketching low-rank matrices" LemmaIn this section, we prove Lemma 3.1. We restate it below for convenience. Lemma 3.1 (Sketching low-rank matrices). Fix 0 < ε, δ < 1 /2 and sample a random s × p SRFT matrix S. Then, for any v ∈ R p and any matrix U ∈ R p×k with ||v|| 2 , ||U || 2 = O(1) we haveEQUATION)as long as s = Ω(kε -2 • log p • log(k/δ)).Proof. In order to prove that ||(SU ) ⊤ (Sv)|| 2 = ||U ⊤ v|| 2 ± ε with probability ≥ 1 -δ we prove the following stronger statement. Denote with [u] i the i-th coordinae of u, then for each i = 1 . . . k we have[(SU ) ⊤ (Sv)] i = [U ⊤ v] i ± ε/ √k with probability at least 1 -δ/k. Apparently, the former follows from the latter by union bound.Notice that [(SU ) ⊤ (Sv)] i = (SU i ) ⊤ (Sv), where U i is the i-th column of U . Moreover, as long ass = Ω(kε -2 • log p • log(k/δ)), then S is a (1 ± ε/ √ k)-subspace embedding for the 2-dimensional subspace spanned by v and U i with probability 1 -δ/k (Theorem 2.2). Therefore, conditioning on this event:(SU i ) ⊤ (Sv) = 1 4 ||S(U i + v)|| 2 2 -||S(U i -v)|| 2 2 = 1 4 ||(U i + v)|| 2 2 -||(U i -v)|| 2 2 ± O(ε/ √ k) = (U i ) ⊤ v ± O(ε/ √ k) = [U ⊤ v] i ± O(ε/ √ k).The second equality sign holds because of the subspace embedding property, together with||U i + v|| 2 , ||U i -v|| 2 = O(1).A.2 Proof of "Orthogonalizing the sketch" LemmaIn this section, we prove Lemma 3.2. We restate it below for convenience. Lemma 3.2 (Orthogonalizing the sketch). Fix 0 < ε, δ < 1 /2 and sample a random s × p SRFT matrix S. As long as s = Ω(kε -2 • log p • log(k/δ)) the following holds with probability 1 -δ.Given any p × k full-rank matrix V , decompose V = U R and SV = U S R S so that U ∈ R p×k , U S ∈ R s×k and both U and U S have orthonormal columns. For any unit-norm v ∈ R p we haveEQUATIONEssentially, the proof of Lemma 3.2 was already known, see for example Corollary 2.2 in Balabanov and Grigori [2022] . Nonetheless, we include a short proof for completeness' sake.Proof. Since V is full-rank, it is easy to verify that SV is also full rank with high probability. Thus, R is non-singular and we can defineŪ = V R -1 S . Notice that U S = S Ū . Now, we prove that ∥U U ⊤ -Ū Ū ⊤ ∥ 2 ≤ ε (11) where ∥•∥ 2 is the operator norm. By definition of U S , we have I k = U ⊤ S U S = Ū S ⊤ S Ū ⊤ By subspace embedding property of S, we have that ∥ Ū Ū ⊤ -I k ∥ 2 = ∥ Ū Ū ⊤ -Ū S ⊤ S Ū ⊤ ∥ 2 ≤ ε with probability 1 -δ. Conditioning on this event, the singular values of Ū lie in [1 -ε, 1 + ε].Moreover, Ran(U ) = Ran( Ū ). Therefore, taking the singular value decomposition (SVD) of U and Ū yields Equation 11. Now we are left to prove that ∥ Ū ⊤ v∥ 2 = ∥U ⊤ S Sv∥ 2 ±ε. We have U ⊤ S S = Ū ⊤ S ⊤ S and by Lemma 3.1 we have ∥ Ū ⊤ S ⊤ Sv∥ 2 = ∥ Ū ⊤ v∥ 2 ± ε with probability 1 -δ. Thus, ∥ Ū ⊤ v∥ 2 = ∥U ⊤ S Sv∥ 2 ± ε and combined with Equation 11gives Equation 7, up to constant blow-ups in ε and δ.In this section we extend Lemma 3.2 to the case where instead of having a query vector v ∈ R p we have a query matrix J ∈ R p×t . In our application, J is the Jacobian J θ * (x). Lemma A.1 (Orthogonalizing the sketch, for matrix queries). Fix 0 < ε, δ < 1/2 and sample a random s × p SRFT matrix S. As long as s = Ω(tkε -2 • log p • log(tk/δ)) the following holds with probability 1 -δ.Given any p × k full-rank matrix V , decompose V = U R and SV = U S R S so that U ∈ R p×k , U S ∈ R s×k and both U and U S have orthonormal columns. For any J ∈ R p×t we haveEQUATIONProof. Let J i the i-th column of J. Suppose that the following intermediate statement holds: for each i,∥U ⊤ S (SJ i )∥ 2 = ∥U ⊤ J i ∥ 2 ± ∥J∥ F • ε/ √ t with probability 1 -δ/t.If the statement above holds, then we have|∥U ⊤ S (SJ)∥ 2 F -∥U ⊤ J∥ 2 F | = (∥U ⊤ S (SJ)∥ F + ∥U ⊤ J∥ F ) • |∥U ⊤ S (SJ)∥ F -∥U ⊤ J∥ F | ≤ O(∥J∥ F ) • i (∥U ⊤ S (SJ i )∥ 2 -∥U ⊤ J i ∥ 2 ) 2 ≤ O(∥J∥ 2 F • ε).The penultimate inequality holds because of triangle inequality, whereas the last inequality holds with probability 1 -δ by union bound.To prove the intermediate statement above, it is sufficient to apply Lemma 3.2 with ε ′ = ε/ √ t and δ ′ = δ/t.Local ensembles. In Madras et al. [2019] , they consider local ensembles, a computationally efficient way to detect underdetermination without training an actual ensemble of models. A local ensemble is a distribution over models obtained by perturbing the parameter along zero-curvature directions. Small parameter perturbations along zero-curvature directions cause small perturbation of the model prediction. Given a test point, Madras et al. [2019] define the local ensemble score as the first-order term of the predictive variance for small perturbations along zero-curvature directions.Moreover, in Madras et al. [2019] they prove that the aforementioned score is equivalent to ∥Π k J θ * (x)∥ F where Π k is a projection onto a subspace defined as the orthogonal of the zerocurvature subspace at θ * , J θ * (x) is the Jacobian of the test datapoint and ∥•∥ F is the Frobenius norm. Computationally, they approximate their score by settingΠ k = I -U k U ⊤ k ,where U k is a p × k matrix which columns are the top-k eigenvectors of the Hessian of the loss L(θ).Their method uses Lanczos algorithm to compute the top-k eigenvectors of the Hessian and uses k • p memory. This is costly, indeed Madras et al. [2019] write "The main constraint of our method is space rather than time -while estimating the first k eigenvectors enables easy caching for later use, it may be difficult to work with these eigenvectors in memory as k and model size p increase. [. . .] we note in some cases that increasing k further could have improved performance [. . .] This suggests that further work on techniques for mitigating this tradeof."Comparison with SCOD. SCOD, the uncertainty estimation score introduced in Sharma et al.[2021], has gained traction lately and it has been used as an uncertainty estimation baseline in applications Qiao et al. [2023] , Elhafsi et al. [2023] . SCOD can be interpreted as a Bayesian method that approximates the posterior with a Gaussian having a covariance matrix of the form identity + low-rank. Computing such approximation of the variance boils down to computing a low-rank approximation of the inverse of F + αI, where F is the Fisher information matrix and α depends on the variance of the isotropic Gaussian prior on θ. Ultimately, the SCOD score is defined as∥(I -U k ΛU ⊤ k )J θ * (x)∥ F, where U k is a p × k matrix which columns are the top-k eigenvectors of F and Λ is a k × k diagonal matrix such that Λ ii = λ i /(λ i + α).SCOD and Local Ensembles scores are extremely similar. Indeed, the only difference between their scores is that SCOD weighs projection onto different eigenvectors according to Λ. However, in Sharma et al. [2021] they observe empirically that different values of α yield the same score accuracy, hinting that replacing Λ with the identity (that essentially gives local ensembles score) should have a modest impact on accuracy. From a computational point of view, both methods compute the top-k eigenvectors of the Hessian / Fisher accessing it solely through matrix-vector product. Local ensembles achieves that using Lanczos algorithm, whereas SCOD uses truncated randomized singular values decomposition (SVD). In both cases the space complexity k • p. Similarly to Madras et al. [2019] , also Sharma et al. [2021] observed that the most prominent limiting factor for their method was memory usage4 : "The memory footprint of the offline stage of SCOD is still linear, i.e., O(kp). As a result, GPU memory constraints can restrict k [. . .] substantially for large models." GGN / Fisher vs Hessian. Local Ensembles uses computes the top-k eigenvalue of the Hessian of L(θ * ), whereas Sharma et al. [2021] , Immer et al. [2021] employs the Fisher / GGN. Implementing matrix-vector product for both Hessian and Fisher /GGN is straightforward, thanks to the power of modern software for automatic differentiation. However, we believe that using the GGN / Fisher is more appropriate in this context. Indeed, the latter is guaranteed to be positive semi-definite (PSD), whereas the Hessian is not necessarily PSD for non-convex landscapes.Lanczos algorithm is an iterative method for tridiagonalizing an Hermitian5 matrix G. If stopped at iteration k, Lanczos returns a column-orthogonal matrixV = [V 1 | . . . |V k ] ∈ R p×k and a tridiagonal matrix T ∈ R k×k such that V ⊤ GV = T .The range space of V corresponds to the Krylov subspace K k = span{v, Gv, . . . , G k-1 v}, where v = V 1 is a randomly chosen vector. Provably K k approximates the eigenspace spanned by the top-k eigenvectors (those corresponding to eigenvalues of largest modulus) of G. Thus, V T V ⊤ approximates the projection of G onto its top-k eigenspace. Notice that projecting G onto its top-k eigenspace yields the best rank-k approximation of G under any unitarily-invariant norm Mirsky [1960] . Moreover, as observed in the previous section the spectrum of the GGN decays rapidly, making low-rank decomposition particularly accurate.Once the decomposition G ≈ V T V ⊤ is available, we can retrieve an approximation to the top-k eigenpairs of G by diagonalizing T into T = W ΛW ⊤ , which can be done efficiently for tridiagonal matrices Dhillon [1997] . The quality of eigenpairs' approximation has been studied in theory and practice. We point the reader to Meurant [2006] , Cullum and Willoughby [2002] for a comprehensive sourvey on this topic.Lanczos, in a nutshell. Here we give a minimal description of Lanczos algorithm. Lanczos maintains a set of vectorsV = [V 1 | . . . |V k ],where V 1 is initialized at random. At iteration i + 1, Lanczos performs a matrix-vector product V i+1 ← G • V i , orthogonalizes V i+1 against V i and V i-1 and normalizes it. The tridiagonal entries of T are given by coefficients computed during orthogonalization and normalization. See Chapter 1 of Meurant [2006] for a full description of Lanczos algorithm.The benefits of Lanczos. Lanczos has two features that make it particularly appealing for our uses case. First, Lanczos does not need explicit access to the input matrix G, but only access to an implementation of u → Gu. Second, Lanczos uses a small working space: only 3p floating point numbers, where the input matrix is p × p. Indeed, we can think of Lanczos as releasing its output in streaming and only storing a tiny state consisting of the last three vectors V i-1 , V i and V i+1 .The curse of numerical instability. Unfortunately, the implementation of Lanczos described above is prone to numerical instability, causing V 1 . . . V k to be far from orthogonal. A careful analysis of , and Λ L be the eigenvectors and eigenvalues computed by the two algorithms respectively. We sort both sets of eigenvectors in decreasing order of corresponding eigenvalues. In position (i, j) we plot ⟨H i , L j ⟩. It is apparent that multiple eigenvectors L j correspond to the same eigenvector H i .the rounding errors causing this pathology was carried out by Paige in his PhD thesis as well as a series of papers Paige [1971 Paige [ , 1976 Paige [ , 1980]] .To counteract this, a standard technique is to re-orthogonalize V i+1 against all {V j } j≤i , at each iteration. This technique has been employed to compute the low-rank approximation of huge sparse matrices Simon and Zha [2000] , as well as in Madras et al. [2019] to compute an approximation to the top-k eigenvectors. Unfortunately, this version of Lanczos loses one the two benefits described above, in that it must store a larger state consisting of the whole matrix V . Therefore, we dub this version of thew algorithm hi-memory Lanczosand the cheaper version described above low-memory Lanczos.Paige Paige [1980] proved that the loss of orthogonality in low-memory Lanczosis strictly linked to the convergence of some of the approximate eigenpairs to the corresponding eigenpairs of G. Moreover, the vectors V 1 . . . V k are not orthogonal because they are tilted towards such eigenpairs. Accordingly, he observed that among the eigenpairs computed by low-memory Lanczosthere are multiple eigenpairs approximating the same eigenpair of G. This can easily be observed while running both low-memory Lanczosand hi-memory Lanczosand computing the dot products of the their retrieved eigenvectors; see Figure 5 .Post-hoc orthogonalization. Based on the observations of Paige Paige [1980] , we expect that orthogonalizing the output of low-memory Lanczos post-hoc should yield an orthonormal basis that approximately spans the top-k eigenspace. Since in our method we only need to project vectors onto the top-k eigenspace, this suffices to our purpose.We confirm this expectation empirically. Indeed, using the same setting of Figure 5 we define Π LM as the projection onto the top-10 principal components of LΛ L , and define Π HM likewise. Then, we measure the operator norm6 (i.e., the largest singular value) of Π LM -Π HM and verify that it is fairly low, only 0.03. Taking principal components was necessary, because there is no clear one-to-one correspondence between L i s and H i s, in that as observed in Figure 5 many L i s can correspond to a single H i . Nonetheless, this proves that low-memory Lanczos is capable of approximating the top eigenspace.Spectral property of the GGN and Hessian have been studied both theoretically and empirically. In theory, it was observed that in the limit of infinitely-wide NN, the Hessian is constant through training Jacot et al. [2018] . However, this is no longer true if we consider NN with finite depth Lee et al. [2019] . In Sagun et al. [2017] , Papyan [2018] , Ghorbani et al. [2019] , they show empirically that the spectrum of the GGN / Hessian is composed of two components: a bulk of near-zero eigenvalues, and a small fraction of outliers away from the bulk. Moreover, in Ghorbani et al. [2019] , they observe a large concentration of the gradients along the eigenspace corresponding to the outlier eigenvalues.In summary, the spectrum of the GGN / Hessian of a deep NN is expected to decay rapidly, having a few outlier eigenvalues, along which most gradients are oriented. Therefore, we would expect a low-rank approximation of the GGN to capture most relevant information.We test this phenomenon by using hi-memory Lanczosand performing an ablation over both number of iterations and random initialization of the first vector. Figure 6 shows the results over three different model trained on different datasets. The plots cointains means and std bars over 5 random initializations of the first Lanczos vector. It is interesting to observe how, varying the number of iterations, approximately the first 90% values are the same, while the tail quickly drop down to numerical precision. Based on this observation we choose to only use the top 90% vectors returned by Lanczos.Experiments are reproducible and the code is available at https://github.com/ IlMioFrizzantinoAmabile/uncertainty_quantification. As explained later in greater details, calling the scripts train_model.py and score_model.py with the right parameters is enough to reproduce all the numbers we presented.Reproducibility. The script bash/plot_figure_4.sh collects the entire pipeline of setup, training, scoring and plotting and finally reproduce Figure 4 and save them in the folder figures. The first three plots takes approximately 3 days to complete on a NVIDIA H100 80GB GPU.The code is implemented in JAX [Bradbury et al., 2018] in order to guarantee max control on the randomness, and, given a trained model and a random seed, the scoring code is fully deterministic (and thus fully reproducible). Nonetheless, the dataloader shuffling used at training time depends on PyTorch (bad) randomness management, and consequently it can be machine-dependent.All experiments were conducted on models trained on one of the following datasets: MNIST [Lecun et al., 1998] , FASHIONMNIST [Xiao et al., 2017] , CIFAR-10 [ Krizhevsky et al., 2009] , CELEBA [Liu et al., 2015] and IMAGENET [Deng et al., 2009] . Specifically we train:• on MNIST with a MLP consisting of one hidden layer of size 20 with tanh activation functions, we trained with ADAM for 50 epochs with a batch size 128 and a learning rate 10 -3 ; parameter size is p = 15910; • on FASHIONMNIST with a LeNet consisting of two convolution layers with maxpool followed by 3 fully connected layers with tanh activation functions, we trained with ADAM for 50 epochs with a batch size 128 and a learning rate 10 -3 ; parameter size is p = 44426; • on CIFAR-10 with a ResNet with (3, 3, 3) blocks of channel sizes (16, 32, 64) with relu activation functions, we trained with SGD for 200 epochs with a batch size 128 and a learning rate 0.1, momentum 0.9 and weight decay 10 -4 ; parameter size is p = 272378; • on CELEBA with a VisualAttentionNetwork with blocks of depths (3, 3, 5, 2) and embedded dimensions of (32, 64, 160, 256) with relu activation functions, we trained with Adam for 50 epochs with a batch size 128 and a learning rate decreasing from 10 -3 to 10 -5 ; parameter size is p = 3858309. • on IMAGENET with a SWIN model with an embed dimension of 192, blocks of depths (2, 2, 18, 2) and number of heads (6, 12, 24, 48), we trained with Adam with weight decay for 60 epochs plus a 10 epoch warmup, with a batch size 128 and a learning rate decreasing from 10 -3 to 10 -5 ; parameter size is p = 196517106.For training and testing the default splits were used. All images were normalized to be in the [0, 1] range. For the CELEBA dataset we scale the loss by the class frequencies in order to avoid imbalances. The cross entropy loss was used as the reconstruction loss in all models but the VisualAttentionNetwork one, on which multiclass binary cross entropy was used.We train all models with 10 different random seed, (except for ImageNet, that we only trained with one seed) which can be reproduced by running bash/setup.sh source virtualenv/bin/activate echo "Train all models" for seed in {1..10} do python train_model.py --dataset MNIST --likelihood classification --model MLP --seed $seed --run_name good --default_hyperparams python train_model.py --dataset FMNIST --likelihood classification --model LeNet --seed $seed --run_name good --default_hyperparams python train_model.py --dataset CIFAR-10 --likelihood classification --model ResNet --seed $seed --run_name good --default_hyperparams python train_model.py --dataset CelebA --likelihood binary_multiclassification --model VAN_tiny --seed $seed --run_name good --default_hyperparams python train_model.py --dataset ImageNet --likelihood classification --model SWIN_large --seed $seed --run_name good --default_hyperparams doneThe trained model test set accuracies are: 0.952±0.001, 0.886±0.003, 0.911±0.003, 0.889±0.002, 0.672, respectively for the 5 different settings.The standard deviations of the scores presented are obtained scoring the indepentently trained models. For all methods we use: 10 seeds for MNIST, 10 seeds for FMNIST, 5 seeds for CIFAR-10, 3 seeds for CELEBA and 1 seeds for IMAGENET. The only exception is Deep Ensemble for which we used only 3 independent scores for the budget3 experiment (for a total of 9 models used) and only 1 score for the budget10 experiments (using all the 10 trained model).For the fixed memory budget experiments we fix the rank k to be equal to the budget (3 or 10) for all baseline, while for our method we were able to have a higher rank thanks to the memory saving induced by sketching, specifically:• for experiments with MNIST as in-distribution we used a sketch size s = 1000 and a rank k = 45 for the budget3 and k = 150 for the budget10;• for experiments with FASHIONMNIST as in-distribution we used a sketch size s = 1000 and a rank k = 132 for the budget3 and k = 440 for the budget10;• for experiments with CIFAR-10 as in-distribution we used a sketch size s = 10000 and a rank k = 81 for the budget3 and a sketch size s = 50000 and a rank k = 50 for the budget10;• for experiments with CELEBA as in-distribution we used a sketch size s = 10000 and a rank k = 100 both for the budget3 and for the budget10.• for experiments with IMAGENET as in-distribution we used a sketch size s = 20000000 and a rank k = 30.Baseline-specific details:• SWAG: we set the parameter collecting frequency to be once every epoch; we perform a grid search on momentum and learning rate hyperparameter: momentum in [0.9, 0.99] resulting in the latter being better, and learning rate in [0.1, 1, 10] times the one used in training resulting in 1 being better. Although we spent a decent amount of effort in optimizing these, it is likely that a more fine search may yield slighlty better results;• SCOD: the truncated randomized SVD has a parameter T and we refer to the paper for full explanation. Importantly to us, the default value is T = 6k + 4 and the memory usage at preprocessing is T p, which is greater than kp. Nonetheless at query time the memory requirement is pk and we consider only this value in order to present the score in the most fair way possible;• LOCAL ENSEMBLE: the method presented in the original paper actually makes use of the Hessian, rather than the GGN. And we denoted them as LE-H and LE respectively, which is not fully respectful of the original paper but we think this approach helps clarity. Anyway, we extensively test both variant and the resulting tables show that they perform very similarly;• DEEP ENSEMBLE: each ensemble consists of either 3 or 10 models, depending on the experiment, each initialized with a different random seed;• DIAGONAL LAPLACE: has a fixed memory requirement of 1p.All the scores relative to the budget3 experiment on CIFAR-10 reported in Table 3 can be reproduced by running the following bash code bash/setup.sh source ./virtualenv/bin/activate dataset="CIFAR-10" ood_datasets="SVHN CIFAR-100 CIFAR-10-C" model_name="ResNet" for seed in {1..5} do echo "Sketched Lanczos" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --subsample_trainset 10000 --lanczos_hm_iter 0 --lanczos_lm_iter 81 --lanczos_seed 1 --sketch srft --sketch_size 10000 echo "Linearized Laplace" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --subsample_trainset 10000 --lanczos_hm_iter 3 --lanczos_lm_iter 0 --lanczos_seed 1 --use_eigenvals echo "Local Ensembles" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --subsample_trainset 10000 --lanczos_hm_iter 3 --lanczos_lm_iter 0 --lanczos_seed 1 echo "Local Ensemble Hessian" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --subsample_trainset 10000 --lanczos_hm_iter 3 --lanczos_lm_iter 0 --lanczos_seed 1 --use_hessian echo "Diag Laplace" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --score diagonal_lla --subsample_trainset 10000 echo "SCOD" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --score scod --n_eigenvec_hm 3 --subsample_trainset 10000 echo "SWAG" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --score swag --swag_n_vec 3 --swag_momentum 0.99 --swag_collect_interval 1000 done for seed in (1, 4, 7) echo "Deep Ensembles" python score_model.py --ID_dataset $dataset --OOD_datasets $ood_datasets --model $model_name --model_seed $seed --run_name good --score ensemble --ensemble_size 3 do done Disclaimer. We did not include the KFAC approximation as a baseline, although its main selling point is the memory efficiency. The reason is that it is a layer-by-layer approximation (and so it neglects correlation between layers) and its implementation is layer-dependent. There exist implementations for both linear layers and convolutions, which makes the method applicable to MLPs and LeNet.But, to the best of our knowledge, there is no implementation (or even a formal expression) for skip-connections and attention layers, consequently making the method inapplicable to ResNet, Visual Transformer, SWIN, or more complex architectures. Disclaimer 2. We did not include the max logit values [Hendrycks and Gimpel, 2016 ] as a baseline. The reason is that it is only applicable for classification tasks, so for example would not be applicable for CelebA which is a binary multiclassfication task.Out-of-Distribution datasets. In an attempt to evaluate the score performance as fairly as possible, we include a big variety of OoD datasets. For models trained on MNIST we used the rotated versions of MNIST, and similarly for FASHIONMNIST. We also include KMNIST as an MNIST-Out-of-Distribution. For models trained on CIFAR-10 we used CIFAR-100, SVHN and the corrupted versions of CIFAR-10, of the 19 corruption types available we only select and present the 14 types for which at least one method achieve an AUROC≥ 0.5. For models trained on CELEBA we used the three subsets of CELEBA corresponding to faces with eyeglasses, mustache or beard. These images were of course excluded from the In-Distribution train and test dataset. Similarly for models trained on IMAGENET we used excluded from the In-Distribution train and test dataset 10 classes (Carbonara, Castle, Flamingo, Lighter, Menu, Odometer, Parachute, Pineapple, Triceratops, Volcano) and use them as OoD datasets. We do not include the results for Carbonara and Menu in thr Table since no method was able to achive an AUROC≥ 0.5.Note that rotations are meaningful only for MNIST and FASHIONMNIST since other datasets will have artificial black padding at the corners which would make the OoD detection much easier.Here we present more experimental results. Specifically we perform an ablation on sketch size s in Appendix D.3.1 (on synthetic data) and on preconditioning size in Appendix D.3.2. Then in Appendix D.3.3 we perform again all the experiments in Figure 4 , where the budget is fixed to be 3p, but now with a higher budget of 10p. We score the 62 pairs ID-OoD presented in the previous section (each pair correspond to a x-position in the plots in Figure 9 ), with the exception of ImageNet because a single H100 GPU is not enough for it.Here we motivate the claim "the disadvantage of introducing noise through sketching is outweighed by a higher-rank approximation" with a synthetic experiment. For a fixed "parameter size" p = 10 6 and a given ground-truth rank R = 100 we generate an artificial Fisher matrix M . To do so, we sample R uniformly random orthonormal vectors v 1 . . . v R ∈ R p and define M = i λ i v i v ⊤ i for some λ i > 0. Doing so allows us to (1) implement x → M x without instantiating M and (2) have explicit access to the exact projection vector product so we can measure both the sketch-induded error and the lowrank-approximation-induced error, so we can sum them and observe the trade-off.For various values of k (on x-axis) and s (on y-axis), we run Skeched Lanczos on x → M x for k iteration with a sketch size s, and we obtain the sketched low rank approximation U S . To further clarify, on the x-axis we added the column "inf" which refers to the same experiments done without any sketching (thus essentially measuring Lanczos lowrankapproximation-error only) which coincides with the limit of s → ∞. The memory requirement of this "inf" setting is P k, that is the same as the second to last column where s = P . We generate a set of test Jacobians as random unit vectors conditioned on their projection onto Span(v 1 . . . v R ) having norm 1 √ 2 . We compute their score both exacltly (as in Equation 9) and sketched (as in Equation 10). In the Figure we show the difference between these two values. As expected, higher rank k leads to lower error, and higher sketch size s leads to lower error. Note that the memory requirement is proportional to the product ks, and the figure is in log-log scale. Preconditioning can be thought of as a smooth bridge between high-memory Lanczos and SKETCHED LANCZOS. We run the former for a few steps and then continue with the latter on the preconditioned matrix. Consequently, the preconditioning scores "start" from the vanilla Lanczos curve, and improve quicker memory-wise, as shown in Figure 8 .The results with memory bugdet 10p are in Table 4 , Table 5 and more extensively in Figure 9 .MNIST FASHIONMNIST Justification: The theoretical claim match the results presented in the Method Section. The experimental results match the numbers presented in the Experiment Section.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.https://github.com/IlMioFrizzantinoAmabile/uncertainty_quantificationThe original version statest that the memory footprint is O(T p), where T ≥ k is a parameter of their SVD algorithm.In this work, we are only concerned with real symmetric matrices.This can be done efficiently via power method. 


Decision: This paper proposes a memory-efficient method for producing uncertainty scores. The main idea is to use a rank-k sketch of the empirical Fisher information matrix, which can be obtained using the Lanczos iteration. The Lanczos iteration is memory-efficient, and the orthogonalization step can be done post-hoc, introducing further memory savings. 

The reviews are overall supportive of the paper. Reviewers appreciate the clear motivation, novel method, and the experimental validation of the approach. The concern around the scale of the experiments was largely resolved during the rebuttal phase with additional experiments. Given the positive feedback, I recommend acceptance.

Reviews: review: paper_summary: paper_summary: ## Summary

- This paper presents an architecture agnostic technique to compute uncertainty scores for pre-trained neural networks. The space complexity (memory usage) of their proposed technique SLU (Sketched Lanczos Uncertainty) grows logarithmically in the size of model parameters. The experiments reported in the paper use models such as LeNet and ResNet in the range 40K - 300K parameters. This is typically still smaller than pre-trained LLMs which are in millions of parameters or perhaps can be applied to LLMs in LoRA setting.
- More concretely, this paper designs an algorithm to compute the "local ensemble uncertainty estimation" score introduced in Madras et al 2019. https://arxiv.org/pdf/1910.09573
	- This earlier paper quantifies "unreliability" of a prediction on the test-example as "underspecification".
	- They state, "trained model is underspecified at a test input if many different predictions at that input are all equally consistent with the constraints posed by the training data and the learning problem specification (i.e., the model architecture and the loss function)" 
	- As a solution they introduce, "local ensembles" , a post-hoc method to measure how much a pre-trained model's prediction is underspecified for a particular test input.
- This paper provides a more memory efficient algorithm to compute this "local ensemble".
- The key observation the authors make in this paper are the following
	- Orthogonalization (used in vanilla Lanczos) approximately commutes with sketching (embedding high dim vectors to lower dim to save on memory). Using this key observation they can change the order of sketching the Lanczos vectors and orthogonalization. Doing orthogonalization post-hoc leads to significant memory-savings.
	- The other results they demonstrate is that in this low-memory budget regime, the disadvantage of introducing noise through sketching is countered by using a higher-rank approximation, leading to better performance.

- Background:
	- **Generalized Gauss Newton Matrix**
		- The local information can be characterized commonly by using the Fisher Information Matrix which is the same as Generalized Gauss Newton Matrix (GGN) Matrix. This GGN matrix is a pxp matrix where p is the number of parameters of the model. Various approximations have been proposed in the literature such as block-diagonal, diagonal etc.
		- In this work, the authors approximate the GGN matrix with a low-rank matrix. The rank-k approximation of the GGN matrix can be computed through two algorithms.
			- the Lanczos algorithm or
			- truncated singular value decomposition algorithm.
		- This is also known as Empirical Fischer matrix (Kutner et al 2019)
	- **Uncertainty Score**
		- The uncertainty at a datapoint x is measured as the variance of the prediction $f(\theta)$ with respect to distribution over parameter $\theta$.
		- For the approaches which only train a single network, various choices of covariance are proposed which are connected to the GGN matrix described above (Kunstner et al 2019)
	- Essentially, the authors in this paper show that low-rank approximation of GGN matrix is a good idea because the eigen values decay exponentially.
	- **Lanczos Algorithm**
		- The lanczos algorithm is an algorithm is an iterative method for tridiagonalizing a symmetric matrix G (pxp).
		- If you stop this algorithm at iteration k, it returns a column orthonormal matrix V (pxk) and a tridiagonal matrix T (kxk) such that $V^T.G.V = T$
		- $V^T.T.V$ approximates the projection of G onto its top k eigenvectors. Projecting G onto its top k eigenspace yields the best rank k approximation of G under any unitarily invariant norm.
		- Once the approximation $V^T.T.V$ of $G$ is available, we can retrieve an approximation of top-k eigenpairs of G by diagonalizing T into $T=W\Lambda W^T$
		- The benefit of Lanczos algorithm is that it only takes 3p space where p is the number of neural network parameters
		- However the downside of Lanczos is that its prone to numerical instability and the $V$ may not be orthogonal. Typically standard implementations re-orthogonalize $V$ after every iteration. This loses the memory benefit of Lanczos. 
		- The authors dub this version as hi-memory lanczos.
		- Instead of re-orthogonalizing at every step, an alternative strategy is to store all vectors and orthogonalize at the end. This approach still requires to store all vectors hence is not memory efficient. To make it more memory efficient, the authors propose to use sketching to store these vectors efficiently.
	- **Sketching**
		- Sketching is a technique to store high dimensional vectors to low dimensions so that expected norm of vector dot products is bounded.
		- The authors use Subsampled Random Fourier Transform (SRFT) to obtain the subspace embedding.
- Finally, the authors combine the Lanczos algorithm with Sketching to propose the Sketched Lanczos algorithm.
- The paper also proposes a slight modification of sketched lanczos which first runs hi-memory lanczos for few iterations followed by sketched lanczos iterations. The idea here is that running the high memory lanczos for few iterations results in a better conditioned matrix G'.
- Finally the low rank decomposition of G is used to compute the variance of a prediction, which also gives the uncertainty score of a datapoint.
strength_weakness: strength_weakness: ### Weakness

- As I alluded to earlier, the models used in experiments are primarily vision models ranging from 40K params to 300K params, it's not clear whether this approach would work for NLP pre-trained models typically ranging in millions of parameters.
- The datasets used are also primarily small scale vision datasets like CIFAR and MNIST.
- However, I still value in this work and novel insights from authors in optimizing the Lanczos algorithm using Sketching.
questions: questions: ### Questions

- Given that memory consumption of the algorithm is $O(k^2.\epsilon^-2)$. (Line 43) Since this does not depend on the number of parameters. Does this mean this can theoretically be applied to larger models? Also, how does this grow logarithmically in the size of model parameters?
- What is the reason for dip in the AUROC beyond a certain memory cost in Figure 3 and Figure 4?
- What is meant by "number of NNs" ? Do you mean number of neural network parameters?
-
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_summary: This paper provides a new algorithm for approximating the Fisher information matrix. Approximation of the Fisher information matrix is an important task whenever the amount of parameters is very large, such as the case for neural networks. This paper makes use of Lanczos algorithm to find an approximate spectrum of the Fisher information matrix, and particularly targets the low memory setting where such approximations are desirable. In particular, in this paper the authors show that:

-Orthogonalization approximately commutes with sketching, making it feasible to sketch Lanczos vectors on the fly and orthogonalize them afterward. 

-Empirically, they demonstrate that in low-memory-budget scenarios, the disadvantage of introducing noise through sketching is outweighed by achieving a higher-rank approximation.
strength_weakness: strength_weakness: I do not see significant weaknesses with this paper. I think it would be nice if the choice of embedding procedure was better described in the main text. I believe that this might be a problem dependent choice, and whilst the authors have provided a good default choice, it would be good to know when their choice might give poor results, albeit if this seldom occurs.
questions: questions: Can the authors provide examples of when the embedding method works/fails. I believe it would be greatly beneficial to the reader if some intuition is shed on when this compression step is a useful computational resource and when it may greatly hinder performance.
limitations: limitations: The authors have not included a limitations section. I think this must be included in a future version and include the dependence of their methods performance on the sketching algorithm used.
paper_summary: paper_summary: In this paper, the authors present a memory-efficient way of computing the uncertainty score of a model (the variance of the model's prediction w.r.t. distribution over the model's parameters). Specifically, they combined Lanczos method and sketching to obtain a low-rank approximation of the Generalized Gauss-Newton (GGN) matrix, which is critical in computing the score but has a square size of parameters.

The authors first introduce the Lanczos method since the GGN matrix is more tractable as an operator. However, Lanczos requires re-orthogonalizing the output vectors, leading to extra memory consumption; while the authors' key observation is that the score only depends on the norm of matrix-vector product (Eq.9) which is only slightly influenced by sketching, therefore they can sketch the output on the fly, and conduct orthogonalization and compute the norm in the sketched subspace to save memory. Consequently, the proposed method is called Sketched Lanczos Uncertainty (SLU).

For validation, the authors tested the performance of SLU on out of distribution (OoD) detection tasks with models including MLP, LeNet, and ResNet. Results showed that SLU outperforms existing metrics, achieving a higher AUROC mainly in the low-memory budget scheme (as shown in Fig.3, 4, and Tab.1).
strength_weakness: strength_weakness: Concerns of this paper are mainly around the experiment part:

1. For the claim that "the disadvantage of introducing noise through sketching is outweighed by a higher-rank approximation" (lines 55~57), the authors should add a simple experiment with synthetic data to verify it.

2. Experiment settings in this paper are too simple; expect to see results with larger-scale models such as Vision Transformers.
questions: questions: 1. Are there other ways of reducing the memory bottleneck in computing uncertainty scores? Since Eq.2 which authors leverage to compute the score is an approximation of the true variance, there may be better alternatives to it. For instance, is there a more computation/memory-efficient form of the covariance matrix $M$ that leads to better scores with less cost?
2. For Tab.1, why do the authors fix the memory budget as $3p$? Looking forward to justifications for this choice.
limitations: limitations: 1. The authors only considered the $M_\text{noeig}$ case where singular values are 0/1, need to improve the flexibility of SLU to fit for other $M$s.
2. The proposed method only works well for the low-memory setting. As the number of NNs increases, Sketched Lanczos deteriorates while vanilla hi-memory Lanczos work better. (Fig.3, 4)

Comments: comment: Regarding scaling to bigger models:
the overhead cost of the sketching operations is neglectable, and we can scale as long as we can perform GGN vector products. This  has essentially twice the cost of a gradient backpropagation, such that the time complexity of computing a rank $k$ approximation is upper bounded by the time complexity of running $2k$ training epochs.

As a proof of concept we computed the score on CelebA with a **43M** parameters vision transformer in the $3p$ budget setting:
| *score* | FOOD101 | Mustache | Bald  | Eyeglasses |
| ----- | ------- | -------- | ----- | ---------- |
| SLU   | 0.705   | 0.715    | 0.738 | 0.680      |
| LE    | 0.672   | 0.693    | 0.697 | 0.665      |

We want to emphasise that the model here has not been fine-tune trained and consequently the performance are not ideal. Nonetheless we see the same trend as in the other experiments: sketching improves results without increasing memory use.
Despite significant investigations, we have been unable to find failure modes beyond those predicted by the theory. Nonetheless, for a **fixed sketch matrix $S$**, one can construct an adversarial example on which the error is huge. One such example is to build a GGN matrix whose eigenvectors are **exactly** orthogonal to the span of the sketch in parameter space (which is an easier task with smaller $s$). This failure, however, depends on the specific $S$, so it will not fail when $S$ is randomized.

Finally it should be mentioned that the sketch size $s$ has to be larger than the rank $k$. If $s$ is smaller then we will attempt to orthogonalize more vectors than the dimensionality of the space they live in, which is guaranteed to fail.
I thank the authors for their response and appreciate their additional comments. 

I understand that the method is guaranteed to work with $s$ chosen sufficiently large. My question was relating to providing an example where the method may fail, for instance constructing an explicit example in the small $s$ regime. Usually understanding such examples gives insight into points on a methods efficiency. It is unclear to me when I would expect to need to select $s$ to be large, or when I may be okay with a smaller value. My question related to providing some intuition on this point.
Thanks for the great rebuttal. The two supplementary experiments in the attachment addressed my concerns (though the ViT used is still a small model with 4M params), so I increased my rating to 6.