Original: Accelerated Linearized Laplace Approximation for Bayesian Deep Learning
Revised: Accelerated Linearized Laplace Approximation for Bayesian Deep Learning

Original: Abstract
Revised: Abstract

Original: Laplace approximation (LA) and its linearized variant (LLA) enable effortlessadaptation of pretrained deep neural networks to Bayesian neural networks.
Revised: Laplace approximation (LA) and its linearized variant (LLA) enable effortless adaptation of pretrained deep neural networks to Bayesian neural networks.

Original: Thegeneralized Gauss-Newton (GGN) approximation is typically introduced to improve their tractability.
Revised: The generalized Gauss-Newton (GGN) approximation is typically introduced to improve their tractability.

Original: However, LA and LLA are still confronted with non-trivialinefficiency issues and should rely on Kronecker-factored, diagonal, or even lastlayer approximate GGN matrices in practical use.
Revised: However, LA and LLA are still confronted with non-trivial inefficiency issues and should rely on Kronecker-factored, diagonal, or even lastlayer approximate GGN matrices in practical use.

Original: These approximations are likelyto harm the fidelity of learning outcomes.
Revised: These approximations are likely to harm the fidelity of learning outcomes.

Original: To tackle this issue, inspired by theconnections between LLA and neural target kernels (NTKs), we develop a Nyströmapproximation to NTKs to accelerate LLA.
Revised: To tackle this issue, inspired by the connections between LLA and neural target kernels (NTKs), we develop a Nyström approximation to NTKs to accelerate LLA.

Original: Our method benefits from the capabilityof popular deep learning libraries for forward mode automatic differentiation, andenjoys reassuring theoretical guarantees.
Revised: Our method benefits from the capability of popular deep learning libraries for forward mode automatic differentiation, and enjoys reassuring theoretical guarantees.

Original: Extensive studies reflect the merits ofthe proposed method in aspects of both scalability and performance.
Revised: Extensive studies reflect the merits of the proposed method in aspects of both scalability and performance.

Original: Our methodcan even scale up to architectures like vision transformers.
Revised: Our method can even scale up to architectures like vision transformers.

Original: We also offer valuableablation studies to diagnose our method.
Revised: We also offer valuable ablation studies to diagnose our method.

Revised: Code is available at https://github.com/thudzj/ELLA .

Original:  Introduction
Revised:  Introduction

Original: Deep neural networks (DNNs) excel at modeling deterministic relationships and have become de factosolutions for diverse pattern recognition problems [18, 54].
Revised: Deep neural networks (DNNs) excel at modeling deterministic relationships and have become de facto solutions for diverse pattern recognition problems [18, 54].

Original: However, DNNs fall short in reasoningabout model uncertainty [1] and suffer from poor calibration [17].
Revised: However, DNNs fall short in reasoning about model uncertainty [1] and suffer from poor calibration [17].

Original: These issues are intolerable inrisk-sensitive scenarios like self-driving [27], healthcare [34], finance [25] etc.
Revised: These issues are intolerable in risk-sensitive scenarios like self-driving [27], healthcare [34], finance [25] etc.

Original: Bayesian Neural Networks (BNNs) have emerged as effective prescriptions to these pathologies [37,21, 44, 16].
Revised: Bayesian Neural Networks (BNNs) have emerged as effective prescriptions to these pathologies [37, 21, 44, 16].

Original: They usually proceed by estimating the posteriors over high-dimensional NN parameters.
Revised: They usually proceed by estimating the posteriors over high-dimensional NN parameters.

Original: Due to some intractable integrals, diverse approximate inference methods have been applied tolearning BNNs, spanning variational inference (VI) [1, 36], Markov chain Monte Carlo (MCMC) [55,4], Laplace approximation (LA) [38, 50, 30] etc.
Revised: Due to some intractable integrals, diverse approximate inference methods have been applied to learning BNNs, spanning variational inference (VI) [1, 36], Markov chain Monte Carlo (MCMC) [55, 4], Laplace approximation (LA) [38, 50, 30], etc.

Original: LA has recently gained unprecedented attention because its post-hoc nature nicely suits with thepretraining-finetuning fashion in deep learning (DL).
Revised: LA has recently gained unprecedented attention because its post-hoc nature nicely suits with the pretraining-finetuning fashion in deep learning (DL).

Original: LA approximates the posterior with a Gaussianaround its maximum, whose mean and covariance are the maximum a posteriori (MAP) and theinversion of the Hessian respectively.
Revised: LA approximates the posterior with a Gaussian around its maximum, whose mean and covariance are the maximum a posteriori (MAP) and the inversion of the Hessian respectively.

Original: It is a common practice to approximate the Hessian with thegeneralized Gauss-Newton (GGN) matrix [40] to make the whole workflow more tractable.
Revised: It is a common practice to approximate the Hessian with the generalized Gauss-Newton (GGN) matrix [40] to make the whole workflow more tractable.

Original: Linearized LA (LLA) [14, 22] applies LA to the first-order approximation of the NN of concern.
Revised: Linearized LA (LLA) [14, 22] applies LA to the first-order approximation of the NN of concern.

Original: LLAshares with LA the same predictive mean yet possesses a more benign covariance, which mitigatesthe underfitting issue of LA [33].

Original: Immer et al.
Revised: Immer et al.

Original: [22] argue that LLA is more sensible than LA in thepresence of GGN approximation;  LLA can perform on par with or better than popular alternatives onvarious uncertainty quantification (UQ) tasks [14, 6, 7].
Revised: [22] argue that LLA is more sensible than LA in the presence of GGN approximation; [Figure] LLA can perform on par with or better than popular alternatives on various uncertainty quantification (UQ) tasks [14, 6, 7].

Original: The Laplace library [6] further substantially advances LLA’s applicability, making it a simple and competing baseline for Bayesian DL.[Figure]
Revised: The Laplace library [6] further substantially advances LLA’s applicability, making it a simple and competing baseline for Bayesian DL.

Original: The practical adoption of LA and LLA actually entails further approximations on top of GGN.
Revised: The practical adoption of LA and LLA actually entails further approximations on top of GGN.

Original: E.g.,when using Laplace to process a pretrained ResNet [18], practitioners are recommended to resort toa Kronecker-factored (KFAC) [40] or diagonal approximation of the full GGN matrix for tractability.
Revised: E.g., when using Laplace to process a pretrained ResNet [18], practitioners are recommended to resort to a Kronecker-factored (KFAC) [40] or diagonal approximation of the full GGN matrix for tractability.

Original: An orthogonal tactic is to apply LA/LLA to only NNs’ last layer [30].
Revised: An orthogonal tactic is to apply LA/LLA to only NNs’ last layer [30].

Original: Yet, the approximation errorsin these cases can hardly be identified, significantly undermining the fidelity of the learning outcomes.
Revised: Yet, the approximation errors in these cases can hardly be identified, significantly undermining the fidelity of the learning outcomes.

Original: This paper aims at scaling LLA up  in a more assurable way.
Revised: This paper aims at scaling LLA up to make probabilistic predictions in a more assurable way.

Original: We first revisit the inherent connectionsbetween Neural Tangent Kernels (NTKs) [24] and LLA [28, 22], and find that, if we can approximatethe NTKs with the inner product of some low-dimensional vector representations of the data, LLAcan be considerably accelerated.
Revised: We first revisit the inherent connections between Neural Tangent Kernels (NTKs) [24] and LLA [28, 22], and find that, if we can approximate the NTKs with the inner product of some low-dimensional vector representations of the data, LLA can be considerably accelerated.

Original: Given this finding, we propose to adapt the Nyström methodto approximate the NTKs of multi-output NNs, and advocate leveraging forward mode automaticdifferentiation ( fwAD ) to efficiently compute the involved Jacobian-vector products (JVPs).
Revised: Given this finding, we propose to adapt the Nyström method to approximate the NTKs of multi-output NNs, and advocate leveraging forward mode automatic differentiation ( fwAD ) to efficiently compute the involved Jacobian-vector products (JVPs).

Original: Theresultant accElerated LLA (ELLA) preserves the principal structures of vanilla LLA yet withoutexplicitly computing/storing the costly GGN/Jacobian matrices for the training data.
Revised: The resultant accElerated LLA (ELLA) preserves the principal structures of vanilla LLA yet without explicitly computing/storing the costly GGN/Jacobian matrices for the training data.

Original: What’s more,we theoretically analyze the approximation error between the predictive of ELLA and that of vanilla LLA, and find that it deceases rapidly as the Nyström approximation becomes accurate.
Revised: What’s more, we theoretically analyze the approximation error between the predictive of ELLA and that of vanilla LLA, and find that it deceases rapidly as the Nyström approximation becomes accurate.

Original: We perform extensive studies to show that ELLA can be a low-cost and effective baseline for Bayesian
Revised: We perform extensive studies to show that ELLA can be a low-cost and effective baseline for Bayesian

Original: DL.
Revised: DL.

Original: We first describe how to specify the hyperparameters of ELLA, and use an illustrative regressiontask to demonstrate the effectiveness of ELLA (see Figure 1).
Revised: We first describe how to specify the hyperparameters of ELLA, and use an illustrative regression task to demonstrate the effectiveness of ELLA (see Figure 1).

Original: We then experiment on standardimage classification benchmarks to exhibit the superiority of ELLA over competing baselines inaspects of both performance and scalability.
Revised: We then experiment on standard image classification benchmarks to exhibit the superiority of ELLA over competing baselines in aspects of both performance and scalability.

Original: We further show that ELLA can even scale up to modernarchitectures like vision transformers (ViTs) [11].
Revised: We further show that ELLA can even scale up to modern architectures like vision transformers (ViTs) [11].

Original:  Background
Revised:  Background

Original: Consider a learning problem on D =
Revised: Consider a learning problem on D =

Original: ( X , Y )
Revised: ( X , Y )

Original: = { ( x i , y i ) }
Revised: = { ( x i , y i ) }

Original: Ni =1 , where x i ∈ X
Revised: Ni =1 , where x i ∈ X

Original: and y i ∈ R C
Revised: and y i ∈ R C

Original: (e.g.,regression) or { 0 , 1 } C (e.g., classification) refer to observations and targets respectively.
Revised: (e.g., regression) or { 0 , 1 } C (e.g., classification) refer to observations and targets respectively.

Original: The advancein machine learning suggests using a NN g θ ( · ) :
Revised: The advance in machine learning suggests using an NN g θ ( · ) :

Original: X → R C with parameters θ ∈ R P for data fitting.
Revised: X → R C with parameters θ ∈ R P for data fitting.

Original: Despite well-performing, regularly trained NNs only capture the most likely interpretation for thedata, thus miss the ability to reason about uncertainty and are prone to overfitting and overconfidence.
Revised: Despite well-performing, regularly trained NNs only capture the most likely interpretation for the data, thus miss the ability to reason about uncertainty and are prone to overfitting and overconfidence.

Original: BNNs [37, 21, 44] characterize model uncertainty by probabilistic principle and can holistically represent all likely interpretations.
Revised: BNNs [37, 21, 44] characterize model uncertainty by probabilistic principle and can holistically represent all likely interpretations.

Original: Typically, BNNs impose a prior p ( θ ) on NN parameters and chase the
Revised: Typically, BNNs impose a prior p ( θ ) on NN parameters and chase the

Original: Bayesian posterior p ( θ |D ) = p ( D| θ ) p ( θ ) /p
Revised: Bayesian posterior p ( θ |D ) = p ( D| θ ) p ( θ ) /p

Original: ( D ) where p ( D| θ ) =
Revised: ( D ) where p ( D| θ ) =

Original: (cid:81)
Revised: (cid:81)

Original: i p ( y i | x i , θ ) =
Revised: i p ( y i | x i , θ ) =

Original: (cid:81)
Revised: (cid:81)

Original: i p ( y i | g θ ( x i )) .
Revised: i p ( y i | g θ ( x i )) .

Original: Analytical estimation is usually intractable due to NNs’ high nonlinearity.
Revised: Analytical estimation is usually intractable due to NNs’ high nonlinearity.

Original: Thereby, BNN methodsusually find a surrogate of the true posterior q ( θ )
Revised: Thereby, BNN methods usually find a surrogate of the true posterior q ( θ )

Original: ≈ p ( θ |D ) via approximate inference methods likevariational inference (VI) [1, 20, 36, 62, 29], Laplace approximation (LA) [38, 50], Markov chain
Revised: ≈ p ( θ |D ) via approximate inference methods like variational inference (VI) [1, 20, 36, 62, 29], Laplace approximation (LA) [38, 50], Markov chain

Original: Monte Carlo (MCMC) [55, 4, of the negative log posterior w.r.t.
Revised: Monte Carlo (MCMC) [55, 4, 63], particle-optimization based variational inference (POVI) [35], etc.

Original: the parameters, i.e., Σ − 1 = −∇ 2 θθ (log p ( D| θ ) 
Revised: BNNs predict for new data x ∗ by posterior predictive p ( y | x ∗ , D ) = E p ( θ |D ) p ( y | x ∗ , θ ) ≈ E q ( θ ) p ( y | x ∗ , θ ) ≈ 1 S (cid:80) Ss =1 p ( y | g θ s ( x ∗ )) where θ s ∼ q ( θ ) are i.i.d.

Revised: Monte Carlo (MC) samples.

Revised: 2.1 Laplace Approximation and Its Linearized Variant

Revised: Typically, LA builds a Gaussian approximate posterior in the form of q ( θ ) =

Revised: N ( θ ; ˆ θ , Σ ) , where ˆ θ denotes the MAP solution, i.e., ˆ θ = arg max θ log p ( D| θ )

Revised: + log p ( θ ) , and Σ is the inversion of the Hessian of the negative log posterior w.r.t.

Revised: parameters, i.e., Σ − 1 = −∇ 2 θθ (log p ( D| θ )

Original: + log p ( θ ))
Revised: + log p ( θ ))

Original: | θ =ˆ θ .
Revised: | θ =ˆ θ .

Original: Without loss of generality, we base the following discussion on the isotropic Gaussian priorp ( θ ) =
Revised: Without loss of generality, we base the following discussion on the isotropic Gaussian prior p ( θ ) =

Original: N ( θ ; 0 , σ 20 I P ) 1 , so −∇ 2 θθ log p ( θ ) | θ =
Revised: N ( θ ; 0 , σ 20 I P ) 1 , so −∇ 2 θθ log p ( θ ) | θ =

Original: ˆ θ amounts to I P /σ 20 .
Revised: ˆ θ equals to I P /σ 20 .

Original: Due to the intractability of the Hessian for NNs with massive parameters, it is a common practice touse the symmetric positive and semi-definite (SPSD) GGN matrix as a workaround, i.e.,[Equation] 1 =
Revised: Due to the intractability of the Hessian for NNs with massive parameters, it is a common practice to use the symmetric positive and semi-definite (SPSD) GGN matrix as a workaround, i.e.,

Original: [Equation]
Revised: [Equation]

Original: where J ˆ θ ( x ) ≜ ∇ θ g θ ( x ) | θ =ˆ θ and Λ( x , y ) ≜ −∇ 2 gg log p ( y | g ) | g = g ˆ θ ( x ) .
Revised: where J ˆ θ ( x ) ≜ ∇ θ g θ ( x ) | θ =ˆ θ and Λ( x , y ) ≜ −∇ 2 gg log p ( y | g ) | g = g ˆ θ ( x ) .

Original: When concatenating { J ˆ θ ( x i ) ∈ R C × P } Ni =1 as a big matrix J ˆ θ , X ∈ R NC × P and organizing{ Λ( x i , y i ) ∈ R C × C } Ni =1 as a block-diagonal matrix Λ X , Y ∈ R NC × NC , we have:
Revised: When concatenating { J ˆ θ ( x i ) ∈ R C × P } Ni =1 as a big matrix J ˆ θ , X ∈ R NC × P and organizing { Λ( x i , y i ) ∈ R C × C } Ni =1 as a block-diagonal matrix Λ X , Y ∈ R NC × NC , we have:

Original: [Equation]
Revised: [Equation]

Original: Yet, LA suffers from underfitting [33].
Revised: Yet, LA suffers from underfitting [33].

Original: This is probably because GGN approximation implicitly turnsthe original model g θ ( x ) into a linear one g lin θ ( x ) =
Revised: This is probably because GGN approximation implicitly turns the original model g θ ( x ) into a linear one g lin θ ( x ) =

Original: g ˆ θ ( x )+ J ˆ θ ( x )( θ − ˆ θ ) but g θ ( x ) is still leveragedto make prediction.
Revised: g ˆ θ ( x )+ J ˆ θ ( x )( θ − ˆ θ ) but g θ ( x ) is still leveraged to make prediction.

Original: I.e., there exists a shift between posterior inference and prediction [22].
Revised: I.e., there exists a shift between posterior inference and prediction [22].

Original: Tomitigate this issue, a proposal is to predict with g lin θ ( x ) , giving rise to linearized LA (LLA) [14, 28, 22].
Revised: To mitigate this issue, a proposal is to predict with g lin θ ( x ) , giving rise to linearized LA (LLA) [14, 28, 22].

Original: LA and LLA nicely fit the pretraining-finetuning fashion in DL – they can be post-hoc applied topretrained models where only the GGN matrices require to be estimated and further inverted.
Revised: LA and LLA nicely fit the pretraining-finetuning fashion in DL – they can be post-hoc applied to pretrained models where only the GGN matrices require to be estimated and further inverted.

Original: By LAand LLA, practitioners can adapt off-the-shelf high performing DNNs to BNNs easily.
Revised: By LA and LLA, practitioners can adapt off-the-shelf high performing DNNs to BNNs easily.

Original: LLA has revealed strong results on diverse UQ problems [14, 6, 7].
Revised: LLA has revealed strong results on diverse UQ problems [14, 6, 7].

Original: The Laplace library [6] furtheradvances LLA’s applicability, and evidences LLA is competitive to popular alternatives [63, 32, 39].
Revised: The Laplace library [6] further advances LLA’s applicability, and evidences LLA is competitive to popular alternatives [63, 32, 39].

Original: Scalability issue The GGN matrix of size P × P is still unamenable in modern DL scenarios, so furtherapproximations sparsifying it are always introduced.
Revised: Scalability issue The GGN matrix of size P × P is still unamenable in modern DL scenarios, so further approximations sparsifying it are always introduced.

Original: The diagonal and KFAC [40] approximations arecommonly adopted ones [50, 62], where only a diagonal or block-diagonal structure of the original
Revised: The diagonal and KFAC [40] approximations are commonly adopted ones [50, 62], where only a diagonal or block-diagonal structure of the original

Original: GNN matrix is preserved.
Revised: GNN matrix is preserved.

Original: An orthogonal tactic is to concern only a subspace of the high-dimensionalparameter space (e.g., the parameter space of the last layer [30]) to reduce the scale of the GGNmatrix.
Revised: An orthogonal tactic is to concern only a subspace of the high-dimensional parameter space (e.g., the parameter space of the last layer [30]) to reduce the scale of the GGN matrix.

Original: However, these strategies sacrifice the fidelity of the learning outcomes as the approximationerrors in these cases can hardly be theoretically measured.
Revised: However, these strategies sacrifice the fidelity of the learning outcomes as the approximation errors in these cases can hardly be theoretically measured.

Original: To this end, we develop accElerated
Revised: To this end, we develop accElerated

Original: Linearized Laplace Approximation (ELLA) to push the limit of LLA in a more assurable way.
Revised: Linearized Laplace Approximation (ELLA) to push the limit of LLA in a more assurable way.

Original:  Methodology
Revised:  Methodology

Original: In this section, we first revisit the relation of LLA to Gaussian processes (GPs) and Neural Tangent
Revised: In this section, we first revisit the relation of LLA to Gaussian processes (GPs) and Neural Tangent

Original: Kernels (NTKs) [24].
Revised: Kernels (NTKs) [24].

Original: After that, we reveal how to accelerate LLA by kernel approximation.
Revised: After that, we reveal how to accelerate LLA by kernel approximation.

Original: Basedon these findings, we develop an efficient implementation of ELLA using the Nyström method [58].
Revised: Based on these findings, we develop an efficient implementation of ELLA using the Nyström method [58].

Original: 3.1 The Gaussian Process View of LLA
Revised: 3.1 The Gaussian Process View of LLA

Original: Integrating q ( θ ) =
Revised: Integrating q ( θ ) =

Original: N ( θ ; ˆ θ , Σ ) with the linear model g lin θ ( x ) actually gives rise to a function-spaceapproximate posterior [14, 28, 22] in the form of q ( f )
Revised: N ( θ ; ˆ θ , Σ ) with the linear model g lin θ ( x ) actually gives rise to a function-space approximate posterior [14, 28, 22] in the form of q ( f )

Original: =
Revised: =

Original: GP ( f | g ˆ θ ( x ) , κ LLA ( x , x ′ )) withκ LLA ( x
Revised: GP ( f | g ˆ θ ( x ) , κ LLA ( x , x ′ )) with

Original: [Equation]
Revised: [Equation]

Original: By Woodbury matrix identity [60], we have:
Revised: By Woodbury matrix identity [60] [Equation] we have:

Original: [Equation]
Revised: [Equation]

Original: It follows that
Revised: It follows that

Original: [Equation]
Revised: [Equation]

Original: where κ NTK ( x , x ′ ) ≜
Revised: where κ NTK ( x , x ′ ) ≜

Original: J ˆ θ ( x ) J ˆ θ ( x ′ )
Revised: J ˆ θ ( x ) J ˆ θ ( x ′ )

Original: ⊤ denotes the neural tangent kernel (NTK)
Revised: ⊤ denotes the neural tangent kernel (NTK)

Original: [24] corresponding tog ˆ θ ( x ) .
Revised: [24] corresponding to g ˆ θ ( x ) .

Original: Note that κ NTK ( x , x ′ ) is a matrix-valued kernel, with values in the space of C × C matrices.
Revised: Note that κ NTK ( x , x ′ ) is a matrix-valued kernel, with values in the space of C × C matrices.

Original: The main challenge then turns into the computation and inversion of the gram matrix κ NTK ( X , X ) ofsize NC × NC .
Revised: The main challenge then turns into the computation and inversion of the gram matrix κ NTK ( X , X ) of size NC × NC .

Original: When either N or C is large, the estimation of κ LLA still suffers from inefficiencyissue.
Revised: When either N or C is large, the estimation of κ LLA still suffers from inefficiency issue.

Original: To address this, existing work [22] assumes independence among the C output dimensionsto cast q ( f ) into C independent GPs following [51], and randomly subsample M ≪ N data pointsto form a cheap substitute for the original gram matrix.
Revised: To address this, existing work [22] assumes independence among the C output dimensions to cast q ( f ) into C independent GPs following [51], and randomly subsample M ≪ N data points to form a cheap substitute for the original gram matrix.

Original: Despite effective in some cases, theseapproximations are heuristic, lacking a clear theoretical foundation.
Revised: Despite effective in some cases, these approximations are heuristic, lacking a clear theoretical foundation.

Original: 3.2 Scale Up LLA by Kernel Approximation
Revised: 3.2 Scale Up LLA by Kernel Approximation

Original: We show that if we can approximate κ NTK ( x , x ′ ) with the inner product of some explicit C × Kdimensional representations of the data, i.e., κ NTK ( x , x ′ )
Revised: We show that if we can approximate κ NTK ( x , x ′ ) with the inner product of some explicit C × Kdimensional representations of the data, i.e., κ NTK ( x , x ′ )

Original: ≈ φ ( x ) φ ( x ′ )
Revised: ≈ φ ( x ) φ ( x ′ )

Original: ⊤ with φ : X → R C × K , thescalability of the whole workflow can be unleashed.
Revised: ⊤ with φ : X → R C × K , the scalability of the whole workflow can be unleashed.

Original: Concretely, letting φ X ∈ R NC × K be the concatenation of { φ ( x i ) ∈ R C × K } Ni =1 , we have 
Revised: Concretely, letting φ X ∈ R NC × K be the concatenation of { φ ( x i ) ∈ R C × K } Ni =1 , we have (detailed derivation in Appendix A.2; see also [9])

Original: [Equation]
Revised: [Equation]

Original: The matrix G ∈ R K × K possesses a similar formula to the Σ in the original κ LLA in Equation (3),yet much smaller.
Revised: The matrix G ∈ R K × K possesses a similar formula to the Σ in the original κ LLA in Equation (3), yet much smaller.

Original: Once having φ , it is only required to perform one forward pass of g and φ for eachtraining data to estimate G .
Revised: Once having φ , it is only required to perform one forward pass of g and φ for each training data to estimate G .

Original: When K is reasonably small, e.g., < 100 , it is cheap to invert G .
Revised: When K is reasonably small, e.g., < 100 , it is cheap to invert G .

Original: 3.3 Approximate NTKs via the Nyström Method
Revised: 3.3 Approximate NTKs via the Nyström Method

Original: Typical kernel approximation means include Nyström method [45, 58] and random features-basedmethods [48, 49, 61, 43, 15, 9].
Revised: Typical kernel approximation means include Nyström method [45, 58] and random features-based methods [48, 49, 61, 43, 15, 9].

Original: Given that the latter routinely rely on a relatively large number ofrandom features to gain a faithful approximation (which means a large K ), we suggest leveraging the
Revised: Given that the latter routinely relies on a relatively large number of random features to gain a faithful approximation (which means a large K ), we suggest leveraging the

Original: Nyström method, which captures only several principal components of the kernel, to build φ .
Revised: Nyström method, which captures only several principal components of the kernel, to build φ .

Original: To comfortably apply the Nyström method, we first rewrite κ NTK ( x , x ′ ) as a scalar-valued kernelκ
Revised: To comfortably apply the Nyström method, we first rewrite κ NTK ( x , x ′ ) as a scalar-valued kernel κ NTK (( x , i ) , ( x ′ , i ′ ))

Original: [Equation]
Revised: =

Original:  x , i [Equation]
Revised: J ˆ θ ( x , i )

Original:  ( x ′ , i ′ )
Revised: J ˆ θ ( x ′ , i ′ )

Original: ⊤ where J ˆ θ ( x ,i ) ≜ ∇ θ [ g θ ( x )]
Revised: ⊤ where J ˆ θ ( x ,i ) ≜ ∇ θ [ g θ ( x )]

Original: ( i )
Revised: ( i )

Original: | θ
Revised: | θ

Original: =ˆ θ : X × [ C ] → R 1 × P 2computes the gradient of i -th output w.r.t.
Revised: =ˆ θ : X × [ C ] → R 1 × Pcomputes the gradient of i -th output w.r.t.

Original: parameters.
Revised: parameters.

Original: By Mercer’s theorem [42],
Revised: By Mercer’s theorem [42],

Original: [Equation]
Revised: [Equation]

Original: where ψ k ∈ L 2 ( X × [ C ] , q ) 3 refer to the eigenfunctions of the NTK w.r.t.
Revised: where ψ k ∈ L 2 ( X × [ C ] , q ) 3 refer to the eigenfunctions of the NTK w.r.t.

Original: some probability measureq with µ k ≥ 0 as the associated eigenvalues.
Revised: some probability measure q with µ k ≥ 0 as the associated eigenvalues.

Original: Based on this, the Nyström method discovers the top- Keigenvalues as well as the corresponding eigenfunctions for kernel approximation.
Revised: Based on this, the Nyström method discovers the top- K eigenvalues as well as the corresponding eigenfunctions for kernel approximation.

Original: By definition, the eigenfunctions can represent the spectral information of the kernel:
Revised: By definition, the eigenfunctions can represent the spectral information of the kernel:

Original: [Equation]
Revised: [Equation]

Original: while being orthonormal under q :
Revised: while being orthonormal under q :

Original: [Equation]
Revised: [Equation]

Original: In our case, q ( x , i ) can be factorized as the product of the data distribution and a uniform distributionover { 1 , ..., C } , which can be trivially sampled from.
Revised: In our case, q ( x , i ) can be factorized as the product of the data distribution and a uniform distribution over { 1 , ..., C } , which can be trivially sampled from.

Original: The Nyström method draws M ( M ≥ K ) i.i.d.samples ˜
Revised: The Nyström method draws M ( M ≥ K ) i.i.d.

Original:  X
Revised: samples ˜ X

Original: = { ( x 1 , i 1 ) , ..., ( x M , i M ) } from q to approximate the integration in Equation (8):
Revised: = { ( x 1 , i 1 ) , ..., ( x M , i M ) } from q to approximate the integration in Equation (8):

Original: [Equation]
Revised: [Equation]

Original: Applying this equation to these samples gives rise to:
Revised: Applying this equation to these samples gives rise to:

Original: [Equation]
Revised: [Equation]

Original: then we arrive at
Revised: then we arrive at

Original: [Equation]
Revised: [Equation]

Original: where K =
Revised: where K =

Original: J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X , with J ˆ θ , ˜ X ∈ R M × P as the concatenation of { J ˆ θ ( x m , i m ) ∈ R 1 × P } Mm =1 ,and ψ k = [ ψ k ( x 1 , i 1 ) , ..., ψ k ( x M , i M )] ⊤ .
Revised: J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X , with J ˆ θ , ˜ X ∈ R M × P as the concatenation of { J ˆ θ ( x m , i m ) ∈ R 1 × P } Mm =1 , and ψ k = [ ψ k ( x 1 , i 1 ) , ..., ψ k ( x M , i M )] ⊤ .

Original: This implies a scaled eigendecomposition problem of K .
Revised: This implies a scaled eigendecomposition problem of K .

Original: We compute the top- K eigenvalues λ 1 , ..., λ K of matrix K and record the corresponding orthonormaleigenvectors u 1 , ..., u K .
Revised: We compute the top- K eigenvalues λ 1 , ..., λ K of matrix K and record the corresponding orthonormal eigenvectors u 1 , ..., u K .

Original: Given the constraint in Equation (9), it is easy to see that:
Revised: Given the constraint in Equation (9) [Equation] it is easy to see that:

Original: [Equation]
Revised: [Equation]

Original: Combining Equation (13) and (10) yields the Nyström approximation of the top- K eigenfunctions:
Revised: Combining Equation (13) and (10) yields the Nyström approximation of the top- K eigenfunctions:

Original: [Equation]
Revised: [Equation]

Original: [Equation] Then, the mapping φ which satisfied κ NTK ( x , x ′ ) 
Revised: Then, the mapping φ which satisfied κ NTK ( x , x ′ ) ≈ φ [Equation] x

Original: ≈ φ ( x ) φ ( x ′ )
Revised: [Equation] φ [Equation] x ′ )

Original: ⊤ can be realized as:
Revised: ⊤ can be realized as:

Original: [Equation]
Revised: [Equation]

Original: Connection to sparse approximations of GPs A recent work [57] has shown that sparse variational
Revised: Connection to sparse approximations of GPs A recent work [57] has shown that sparse variational

Original: Gaussian processes (SVGP) [53, 3] is algebraically equivalent to the Nyström approximation becausethe Evidence Lower BOund (ELBO) of the former encompasses the learning objective of the latter.
Revised: Gaussian processes (SVGP) [53, 3] is algebraically equivalent to the Nyström approximation because the Evidence Lower BOund (ELBO) of the former encompasses the learning objective of the latter.

Original: That’s to say, ELLA implicitly deploys a sparse approximation to the original GP predictive of LLA. 
Revised: That’s to say, ELLA implicitly deploys a sparse approximation to the original GP predictive of LLA. Discussion

Revised: In this work, the developed Nyström approximation to NTKs is mainly used for accelerating the computation of the predictive covariance of LLA, while it can also be easily applied to future work on the practical application of NTKs.

Original: 3.4 Implementation
Revised: 3.4 Implementation

Original: As shown, the estimation of φ on a data point x degenerates as K JVPs, which can be accomplishedby invoking fwAD for K times:
Revised: As shown, the estimation of φ on a data point x degenerates as K JVPs, which can be accomplished by invoking fwAD for K times:

Original: [Equation]
Revised: [Equation]

Original:  where the model output g ˆ θ ( x ) and the JVP J ˆ θ ( x ) v k are simultaneously computed in one singleforward pass .
Revised: (17) where the model output g ˆ θ ( x ) and the JVP J ˆ θ ( x ) v k are simultaneously computed in one single forward pass .

Original: Prevalent DL libraries like PyTorch [47] and Jax [2] have already been armed with thecapability for fwAD .
Revised: Prevalent DL libraries like PyTorch [47] and Jax [2] have already been armed with the capability for fwAD .

Original: Algorithm 2 details the procedure of building φ in a PyTorch-like style.
Revised: Algorithm 2 details the procedure of building φ in a PyTorch-like style.

Original: With that, we can trivially compute the posterior q ( f ) according to Equation (6), as detailed in
Revised: With that, we can trivially compute the posterior q ( f ) according to Equation (6), as detailed in

Original: Algorithm 1.
Revised: Algorithm 1.

Original: The procedures for estimating G and the posterior can both be batched for acceleration.
Revised: The procedures for estimating G and the posterior can both be batched for acceleration.

Original: Computation overhead The estimation of φ involves M forward and backward passes of g ˆ θ and theeigendecomposition of a matrix of size M × M .
Revised: Computation overhead The estimation of φ involves M forward and backward passes of g ˆ θ and the eigendecomposition of a matrix of size M × M .

Original: After that, we only need to make { v k } Kk =1 persistent,which amounts to storing K more NN copies .
Revised: After that, we only need to make { v k } Kk =1 persistent, which amounts to storing K more NN copies .

Original: Estimating G − 1 requires scanning the training setonce and inverting a matrix of K × K , which is cheap.
Revised: Estimating G − 1 requires scanning the training set once and inverting a matrix of K × K , which is cheap.

Original: The evaluation of q ( f ) embodies that ofφ , i.e., performing K forward passes under the scope of fwAD .
Revised: The evaluation of q ( f ) embodies that of φ , i.e., performing K forward passes under the scope of fwAD .

Original: This is similar to other BNNs thatperform S forward passes with different MC parameter samples to estimate the posterior predictive.
Revised: This is similar to other BNNs that perform S forward passes with different MC parameter samples to estimate the posterior predictive.

Original: Theoretical Analysis

Original: In this section, we theoretically analyze the approximation error of κ ELLA ( x , x ′ ) to κ LLA ( x , x ′ ) .

Original: To reduce unnecessary complexity, this section assumes using M = K i.i.d.

Original: MC samples when performing the Nyström method.

Original: Then, κ ELLA can be reformulated as (details deferred to Appendix A.3)

Original: # g ˆ θ : NN pre-trained by MAP; ( X , Y ) :
Revised: # g ˆ θ : NN pre-trained by MAP; ( X , Y ) :

Original: # training set; C : number of classes # M , K , σ 20 : hyper-parameters def estimate_ G ( φ , X , Y , K , σ 20 ): G = zeros( K , K ) for ( x , y ) in ( X , Y ) : g x , φ x = φ ( x ) Λ x ,
Revised: # training set; C : number of classes # M , K , σ 20 : hyper-parameters def estimate_ G ( φ , X , Y , K , σ 20 ): G = zeros( K , K ) for ( x , y ) in ( X , Y ) : g x , φ x = φ ( x ) Λ x ,

Original: y = hessian(nll
Revised: y = hessian(nll

Original: ( g x , y ), g x )
Revised: ( g x , y ), g x )

Original: G += φ ⊤ x Λ x , y φ x return G + eye( K ) /σdef
Revised: G += φ ⊤ x Λ x , y φ x return G + eye( K ) /σdef

Original: _q_f( φ , G − 1 , x ) g x , φ x = φ ( x ) κ x , x = φ x G − 1 φ ⊤ x return N ( g x , κ x , x ) φ = build_ φ ( g ˆ θ , X , C , M , K )
Revised: _q_f( φ , G − 1 , x ) g x , φ x = φ ( x ) κ x , x = φ x G − 1 φ ⊤ x return N ( g x , κ x , x ) φ = build_ φ ( g ˆ θ , X , C , M , K )

Original: G − 1 = inv(estimate_ G ( φ , X , Y , K , σ 20 ))
Revised: G − 1 = inv(estimate_ G ( φ , X , Y , K , σ 20 ))

Original: q_f
Revised: q_f

Original: = partial(_q_f, φ , G − 1 ) def build_ φ ( g ˆ θ , X , C , M , K ): def _ φ ( g ˆ θ , C , { v k } Kk =1 , x ): φ x = zeros( C , K ) for k in range( K ): with fwAD .enable(): g x , jvp = g
Revised: = partial(_q_f, φ , G − 1 ) def build_ φ ( g ˆ θ , X , C , M , K ): def _ φ ( g ˆ θ , C , { v k } Kk =1 , x ): φ x = zeros( C , K ) for k in range( K ): with fwAD .enable(): g x , jvp = g

Original: (ˆ θ , v k ) ( x ) φ x [: , k ]
Revised: (ˆ θ , v k ) ( x ) φ x [: , k ]

Original: = jvp return g x , φ x
Revised: = jvp return g x , φ x

Original: J ˆ θ , ˜ X
Revised: J ˆ θ , ˜ X

Original: = zeros( M ,dim( ˆ θ ))
Revised: = zeros( M ,dim( ˆ θ ))

Original: for m in range( M ): x m =
Revised: for m in range( M ): x m =

Original: uniform_sample( X ) i m =
Revised: uniform_sample( X ) i m =

Original: uniform_sample( [ C ] )
Revised: uniform_sample( [ C ] )

Original: J ˆ θ , ˜ X [m ] =
Revised: J ˆ [Equation] , ˜ X [

Revised: [Equation]

Original:  grad( g ˆ θ ( xm )
Revised: ] = grad( g ˆ θ ( x

Original:  [ i m ] , ˆ θ ) { λ k , u k } Kk =1 = eig
Revised: [Equation] [ i m ] , ˆ θ ) { λ k , u k } Kk =1 = eig

Original: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X ,top = K ) for k in range( K ): v
Revised: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X ,top = K ) for k in range( K ): v

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation][Equation]
Revised: [Equation]

Original: u

Original: [Equation]
Revised: [Equation]

Original: / √ λ k return partial(_ φ , g ˆ θ , C , { v k } Kk =1 )κ ELLA (
Revised: return partial(_ φ , g ˆ θ , C , { v k } K k =1 )

Original:  [Equation] [Equation]
Revised:  Theoretical Analysis

Original: [Equation] = [Equation] θ , ˜ X (cid:104)
Revised: In this section, we theoretically analyze the approximation error of κ ELLA ( x , x ′ ) to κ LLA ( x , x ′ ) .

Original: [Equation] , ˜ X
Revised: To reduce unnecessary complexity, this section assumes using M = K i.i.d.

Original: [Equation] ⊤ ˆ θ , X Λ X , Y [Equation]
Revised: MC samples when performing the Nyström method.

Original: ˆ θ , X
Revised: Then, κ ELLA can be reformulated as (details deferred to Appendix A.3)

Original: [Equation]
Revised: [Equation]

Original: ˆ θ , ˜ X

Original: + J ˆ θ , ˜

Original: X J ⊤ ˆ θ , ˜ X /σ 20 (cid:105) − 1 J ˆ θ , ˜ X (cid:124)

Original: (cid:123)(cid:122) (cid:125) Σ ′ J ˆ θ ( x ) ⊤ . Thereby, the gap between κ ELLA ( x , x ′ ) and κ LLA ( x , x ′ ) can be upper bounded by E = ∥
Revised: Thereby, the gap between κ ELLA ( x , x ′ ) and κ LLA ( x , x ′ ) can be upper bounded by E = ∥

Original: Σ ′ − Σ ∥ ,where ∥ · ∥ represents the matrix 2-norm (i.e., the spectral norm).
Revised: Σ ′ − Σ ∥ , where ∥ · ∥ represents the matrix 2-norm (i.e., the spectral norm).

Original: In typical cases, P ≫ M , thuswith high probability K =
Revised: In typical cases, P ≫ M , thus with high probability K =

Original: J ˆ θ , ˜ X
Revised: J ˆ θ , ˜ X

Original: J ⊤ ˆ θ , ˜ X ≻ 0 .
Revised: J ⊤ ˆ θ , ˜ X ≻ 0 .

Original: We present an upper bound of E as follows.
Revised: We present an upper bound of E as follows.

Original: Theorem 1 (Proof in Appendix A.4) .
Revised: Theorem 1 (Proof in Appendix A.4) .

Original: Let c Λ be a finite constant associated with Λ , and E ′ the errorof Nyström approximation ∥
Revised: Let c Λ be a finite constant associated with Λ , and E ′ the error of Nyström approximation ∥

Original: J ˆ θ , X J ⊤ ˆ θ , ˜ X
Revised: J ˆ θ , X J ⊤ ˆ θ , ˜ X

Original: ( J ˆ θ , ˜ X J ⊤ ˆ θ [Equation] ˜ X ) − 1 J ˆ θ , ˜X
Revised: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X ) − 1 J ˆ θ , ˜

Original:  J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ∥ .
Revised: X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ∥ .

Original: It holds that
Revised: It holds that

Original: [Equation]
Revised: [Equation]

Original: E ′ has been extensively analyzed by pioneering work [12, 5, 26], and we simply adapt the resultsdeveloped by Drineas and Mahoney [12] to our case.
Revised: E ′ has been extensively analyzed by pioneering work [12, 5, 26], and we simply adapt the results developed by Drineas and Mahoney [12] to our case.

Original: We denote the maximum diagonal entry and the
Revised: We denote the maximum diagonal entry and the

Original: M + 1 -th largest eigenvalue of J ˆ θ , X J ⊤ ˆ θ , X by c κ and ˜ λ
Revised: M + 1 -th largest eigenvalue of J ˆ θ , X J ⊤ ˆ θ , X by c κ and ˜ λ

Original: M +1 respectively.
Revised: M +1 respectively.

Original: Theorem 2 (Error bound of Nyström approximation) .
Revised: Theorem 2 (Error bound of Nyström approximation) .

Original: With probability at least 1 − δ , it holds that:
Revised: With probability at least 1 − δ , it holds that:

Original: [Equation]
Revised: [Equation]

Original: Plugging this back to E , we arrive at the corollary below.
Revised: Plugging this back to E , we arrive at the corollary below.

Original: Corollary 1.
Revised: Corollary 1.

Original: With probability at least 1 − δ , the following bound exists:
Revised: With probability at least 1 − δ , the following bound exists:

Original: [Equation]
Revised: [Equation]

Original: As desired, the upper bound of E decreases along with the growing of the number of MC samples M .
Revised: As desired, the upper bound of E decreases along with the growing of the number of MC samples M .

Original:  Experiments
Revised:  Experiments

Original: We first discuss how to specify the hyperparameters of ELLA, then expose an interesting finding.
Revised: We first discuss how to specify the hyperparameters of ELLA, then expose an interesting finding.

Original: After that, we compare ELLA to competing baselines to evidence its merits in efficacy and scalability.
Revised: After that, we compare ELLA to competing baselines to evidence its merits in efficacy and scalability.

Original: 5.1 Specification of Hyperparameters

Original: Given an NN g ˆ θ pretrained by MAP, we need to specify the prior variance σ 20 , the number of MCsamples in Nyström method M , and the number of remaining eigenpairs K before applying ELLA.

Original: We simply set σ 20 to 1 Nγ with γ as the weight decay coefficient used for pretraining according to [10].

Original: 4 8 16 32 64 128 256 512 1024 2000 K 

Revised: 4 8 16 32 64 128 256 512 1024 2000 K

Original: M 0.
Revised: M 0.

Original: (a) 4 8 16 32 64 128 256 512 1024 2000 K 

Revised: (a) 4 8 16 32 64 128 256 512 1024 2000 K

Original: M 0.
Revised: M 0.

Original: (b)    40000 N 0.
Revised: (b)    40000 N 0.

Original: T e s t N
Revised: T e s t N

Original: LL
Revised: LL

Original: ResNet20 ResNet
Revised: ResNet20 ResNet

Original: ResNet44 ResNet (c)
Revised: ResNet44 ResNet (c)

Original: [Table]
Revised: Given an NN g ˆ θ pretrained by MAP, we need to specify the prior variance σ 20 , the number of MC samples in Nyström method M , and the number of remaining eigenpairs K before applying ELLA. We simply set σ 20 to 1 Nγ with γ as the weight decay coefficient used for pretraining according to [10].

Original: We perform an empirical study to inspect how M and K affect the quality of the Nyström approximation and the ELLA GP covariance.
Revised: We perform an empirical study to inspect how M and K affect the quality of the Nyström approximation and the ELLA GP covariance.

Original: We take 2000 MNIST images as training set X , and 256 othersas validation set X val .
Revised: We take 2000 MNIST images as training set X , and 256 others as validation set X val .

Original: The architecture contains 2 convolutions and a linear head.
Revised: The architecture contains 2 convolutions and a linear head.

Original: Batch normalization [23] and ReLU activation are used.
Revised: Batch normalization [23] and ReLU activation are used.

Original: The number of parameters P is 29 , 034 .
Revised: The number of parameters P is 29 , 034 .

Original: Larger architecturesor larger X will render the exact evaluation of κ LLA unapproachable.
Revised: Larger architectures or larger X will render the exact evaluation of κ LLA unapproachable.

Original: We quantify the approximationerror of the Nyström method by 
Revised: We quantify the approximation error of the Nyström method by ϵ

Original: [Equation]
Revised: [Equation]

Original:  X ∥ ,and that from κ ELLA to κ LLA by ϵ ELLA ≜ 1 | X val | [Equation] x
Revised: J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ∥ / ∥

Original: [Equation] X val ∥ κ ELLA ( x , x ) − κ LLA ( x , x )∥ /
Revised: J ˆ θ , X J ⊤ ˆ θ , X ∥ , and that from κ ELLA to κ LLA by ϵ ELLA ≜ 1 | X val | (cid:80) x ∈ X val ∥ κ ELLA ( x , x ) − κ LLA ( x , x )

Original: ∥  κ LLA ( x , x ) ∥ .
Revised: ∥ / ∥ κ LLA ( x , x ) ∥ .

Original: Wevary M from 4 to 2000 and K from 4 to M , and report the approximation errors in Figure 2 (a)-(b).
Revised: We vary M from 4 to 2000 and K from 4 to M , and report the approximation errors in Figure 2 (a)-(b).

Original: We notice that 1) the larger K the better; 2) when fixing K , ϵ Nyström and ϵ ELLA decay as M grows; 3) ϵ Nyström decays more rapidly than ϵ ELLA , echoing Theorem 1.
Revised: We notice that 1) the larger K the better; 2) when fixing K , ϵ Nyström and ϵ ELLA decay as M grows; 3) ϵ Nyström decays more rapidly than ϵ ELLA , echoing Theorem 1.

Original: Given that ELLA needs to store Kvectors of size P , small K is desired for efficiency.
Revised: Given that ELLA needs to store K vectors of size P , small K is desired for efficiency.

Original: K ∈ [16 , 32] seems to be a reasonable choicegiven Figure 2.
Revised: K ∈ [16 , 32] seems to be a reasonable choice given Figure 2.

Revised: Besides, Appendix C.1 includes a direct study on how the test NLL of ELLA varies w.r.t.

Revised: M and K on CIFAR-10 [31] benchmark.

Original: Unless otherwise stated, we set M = 2000 and K = 20 in the following experiments.
Revised: Given these results, we set M = 2000 and K =in the following experiments unless otherwise stated.

Original: ELLA (or LLA) finds a GP posterior, so predicts with a variant of the aforementioned posterior predictive:
Revised: ELLA (or LLA) finds a GP posterior, so predicts with a variant of the aforementioned posterior predictive, formulated as p ( y | x ∗ , D )

Original: [Equation]
Revised: = E p ( f |D ) p ( y | f ( x ∗ )) ≈ E q ( f ) p ( y | f ( x ∗ )) = E f ∼N ( g ˆ θ ( x ∗ ) ,κ ELLA ( x ∗ , x ∗ ))

Original:  In classification tasks, we use 512 MC samples to approximate the last expectation and it is cheap.
Revised: p ( y | f ) . In classification tasks, we use 512 MC samples to approximate the last expectation and it is cheap.

Original: 5.2 The Overfitting Issue of LLA
Revised: 5.2 The Overfitting Issue of LLA

Original: Reinspecting Equation (1) and (6), we see, with more training data involved, the covariance in
Revised: Reinspecting Equation (1) and (6), we see, with more training data involved, the covariance in

Original: LA, LLA, and ELLA shrinks and the uncertainty dissipates.
Revised: LA, LLA, and ELLA shrinks and the uncertainty dissipates.

Original: However, under ubiquitous modelmisspecification [41], should the uncertainty vanish that fast?
Revised: However, under ubiquitous model misspecification [41], should the uncertainty vanish that fast?

Original: We perform a study with ResNets [18]on CIFAR-10[31] to find an answer.
Revised: We perform a study with ResNets [18] on CIFAR-10 to seek an answer.

Original: Concretely, we randomly subsample N data from the trainingset of CIFAR-10, and fit ELLA on them.
Revised: Concretely, we randomly subsample N data from the training set of CIFAR-10, and fit ELLA on them.

Original: We depict the test negative log-likelihood (NLL), accuracy,and expected calibration error (ECE) [17] of the  ELLA in Figure 2 (c) and Appendix B.1.
Revised: We depict the test negative log-likelihood (NLL), accuracy, and expected calibration error (ECE) [17] of the deployed ELLA in Figure 2 (c) and Appendix C.2.

Original: We alsoprovide the corresponding results of LLA ∗  in Appendix B.1.
Revised: We also provide the corresponding results of LLA ∗ 4 in Appendix C.3.

Original: The V-shape NLL curves across settingsreflect the overfitting issue of ELLA and LLA ∗ (or more generally, LLA). 
Revised: The V-shape NLL curves across settings reflect the overfitting issue of ELLA and LLA ∗ (or more generally, LLA). Figurealso shows that tuning the prior precision w.r.t.

Revised: marginal likelihood can alleviate the overfitting of

Revised: LLA ∗ to some extent, which may be the reason why such an issue has not been reported by previous works.

Revised: But also of note that tuning the prior precision cannot fully eliminate overfitting.

Original: To address this issue, we advocate performing early stop when applying ELLA/LLA to big data .
Revised: To address the overfitting issue, we advocate performing early stop when fitting ELLA/LLA on big data .

Original: Specifically, when iterating over the training data to estimate G (see Equation (6)), we continuouslyrecord the NLL of the current ELLA posterior on some validation data, and stop when there is a trendof overfitting.
Revised: Specifically, when iterating over the training data to estimate G (see Equation (6)), we [Table] [Table]

Revised: [Figure]

Revised: E xp ec t e d C a li b r a ti on E rr o r

Original:  If we cannot access a validation set easily, we can apply strong data augmentation tosome training data to form a substitute of the validation set.Cross validation is also required by [7],but in the sake of finding a proper Gaussian prior variance σ 20 .
Revised: ELLAMAPMFVI-BFLLA * LLA * -KFAC continuously record the NLL of the current ELLA posterior on some validation data, and stop when there is a trend of overfitting. If we cannot access a validation set easily, we can apply strong data augmentation to some training data to form a substitute of the validation set.

Original: An additional benefit of early stoppingis that it significantly reduces the time cost for handling large data (e.g., ImageNet [8]).[Figure] [Figure]
Revised: Compared to tuning the prior, early stopping also helps reduce the time cost of processing big data (e.g., ImageNet [8]).

Original: 5.3 Illustrative Regression
Revised: 5.3 Illustrative Regression

Original: We build a regression problem with N = 16 samples from y = sin 2 x + ϵ, ϵ ∼ N (0 , 0 . 2) as shown in
Revised: We build a regression problem with N = 16 samples from y = sin 2 x + ϵ, ϵ ∼ N (0 , 0 . 2) as shown in

Original: Figure 1.
Revised: Figure 1.

Original: The model is a MLP with 3 hidden layers and tanh activations, and we pretrain it by MAPfor 1000 iterations.
Revised: The model is an MLP with 3 hidden layers and tanh activations, and we pretrain it by MAP for 1000 iterations.

Original: For ELLA, we set M = 16 and K = 5 for efficiency.
Revised: For ELLA, we set M = 16 and K = 5 for efficiency.

Original: Unless stated otherwise, weuse the interfaces of Laplace [6] to implement LLA, LLA-KFAC, LLA-Diag, and LLA ∗ .
Revised: Unless stated otherwise, we use the interfaces of Laplace [6] to implement LLA, LLA-KFAC, LLA-Diag, and LLA ∗ .

Revised: The hyperparameters of the competitors are equivalent to those of ELLA except for some dedicated ones like M and K .

Original: Comparedto LLA-KFAC, LLA-Diag, and LLA ∗ , ELLA delivers a closer approximation to LLA.
Revised: It is clear that ELLA delivers a closer approximation to LLA than LLA-Diag and LLA ∗ .

Revised: We further quantify the quality of the predictive distributions produced by these approximate

Revised: LLA methods using certain metrics.

Revised: Considering that in this case, the predictive distribution for one test datum is a Gaussian distribution, we use the KL divergence between the Gaussians yielded by the approximate LLA method and vanilla LLA as a proxy of the approximation error (averaged over a set of test points).

Revised: The results are reported in Appendix C.4.

Original:  5.4 CIFAR-10 Classification
Revised: As shown, LLA-KFAC comes pretty close to LLA. Yet, LLA seems to underestimate in-between uncertainty in this setting, so ELLA seems to be a more reliable (instead of more accurate) approximation than LLA-KFAC. We also highlight the higher scalability of ELLA than LLA-KFAC (see Figure 4(c)), which reflects that ELLA can strike a good trade-off between efficacy and efficiency. 5.4 CIFAR-10 Classification

Original: Then, we assess ELLA on CIFAR-10 [31] with ResNet architectures [18].
Revised: Then, we evaluate ELLA on CIFAR-10 benchmark using ResNet architectures [18].

Original: We obtain pretrainedmodels from the open source community.
Revised: We obtain pretrained MAP models from the open source community.

Original: Apart from MAP, LLA ∗ , LLA-Diag, LLA-KFAC, wefurther introduce last-layer LLA with KFAC approximation (LLA ∗ -KFAC) and mean-field VI via Bayesian finetuning [10] (MFVI-BF) 4 as baselines.
Revised: Apart from MAP, LLA ∗ , LLA-Diag, LLA-KFAC, we further introduce last-layer LLA with KFAC approximation (LLA ∗ -KFAC) and mean-field VI via Bayesian finetuning [10] (MFVI-BF) 5 as baselines.

Original: LLA cannot be directly applied as the GGNmatrices are huge.
Revised: LLA cannot be directly applied as the GGN matrices are huge.

Original: These methods all locate in the family of Gaussian approximate posteriors and areall post-hoc applied to the pretrained models to establish fair comparison.
Revised: These methods all locate in the family of Gaussian approximate posteriors and are all post-hoc applied to the pretrained models, so the comparisons will be fair.

Original: Regarding the setups, weuse M = 2000 and K = 20 for ELLA; 5 we use 20 MC samples to estimate the posterior predictiveof MFVI-BF (as it incurs 20 NN forward passes), and use 512 ones for the other methods as stated.
Revised: Regarding the setups, we use M = 2000 and K = 20 for ELLA; 6 we use 20 MC samples to estimate 0.2 Confidence Threshold 0.1 0. E rr o r on e x a m p l e s w it h c on f i d e n ce

Revised: ELLAMAPMFVI-BFLLA *

Revised: LLA * -KFAC (a)      30 K 0.

Revised: T e s t N LL (b)

Original: {+E L L A M A P M F V I- B F L L A * L L A * - K F A C L L A - D i a g L L A - K F A C 
+}
Revised: {+E L L A M A P M F V I- B F L L A * L L A * - K F A C L L A - D i a g L L A - K F A C 
+}

Revised: T

Revised: i m e s p e n t p r

Revised: e d i c ti ng a ll t e s t d a t a ( s ) 14.

Revised: (c)

Revised: [Table]

Revised: the posterior predictive of MFVI-BF (as it incurs 20 NN forward passes), and use 512 ones for the other methods as stated.

Revised: We have enabled the tuning of the prior precision for all LLA baselines, but not for ELLA.

Original: We present the comparison on test accuracy, NLL, and ECE in Table 1.
Revised: We present the comparison on test accuracy, NLL, and ECE in Table 1.

Original: As shown, ELLA exhibitssuperior NLL and ECE across settings.
Revised: As shown, ELLA exhibits superior NLL and ECE across settings.

Original: MFVI-BF also gives good NLL.
Revised: MFVI-BF also gives good NLL.

Original: LLA ∗ and LLA ∗ -KFACcan improve the uncertainty and calibration of MAP, yet underperforming ELLA.
Revised: LLA ∗ and LLA ∗ -KFAC can improve the uncertainty and calibration of MAP, yet underperforming ELLA.

Original: LLA-Diag and
Revised: LLA-Diag and

Original: LLA-KFAC fail for unclear reasons (also reported by [9]), so are excluded from the following studies.
Revised: LLA-KFAC fail for unclear reasons (also reported by [9]), we thus exclude them from the following studies.

Original: We then examine the models on the widely used out-of-distribution (OOD) generalization/robustnessbenchmark CIFAR-10 corruptions [19] and report the results in Figure 3 and Appendix B.1.
Revised: We then examine the models on the widely used out-of-distribution (OOD) generalization/robustness benchmark CIFAR-10 corruptions [19] and report the results in Figure 3 and Appendix C.5.

Original: It isprominent that ELLA surpasses the baselines in aspects of NLL and ECE at various levels of skew,implying its ability to make conservative predictions for OOD inputs.
Revised: It is prominent that ELLA surpasses the baselines in aspects of NLL and ECE at various levels of skew, implying its ability to make conservative predictions for OOD inputs.

Original: We further evaluate the models on the combination of the CIFAR-10 test data and the OOD SVHNtest data.
Revised: We further evaluate the models on the combination of the CIFAR-10 test data and the OOD SVHN test data.

Original: The predictions on SVHN are all regarded as wrong due to label shift.
Revised: The predictions on SVHN are all regarded as wrong due to label shift.

Original: We plot the averageerror rate of the samples with ≤ τ ( τ ∈ [0 , 1] ) confidence in Figure 4 (a).
Revised: We plot the average error rate of the samples with ≤ τ ( τ ∈ [0 , 1] ) confidence in Figure 4 (a).

Original: As shown, ELLA makesless mistakes than the baselines under various confidence thresholds.
Revised: As shown, ELLA makes less mistakes than the baselines under various confidence thresholds.

Original: Figure 4 (b) displays how Kimpacts the test NLL.
Revised: Figure 4 (b) displays how K impacts the test NLL.

Original: We see that K ∈ [20 , 30] can already lead to good performance, reflecting theefficiency of ELLA.
Revised: We see that K ∈ [20 , 30] can already lead to good performance, reflecting the efficiency of ELLA.

Revised: Another benefit of ELLA is that with it, we can actively control the performance vs. cost trade-off.

Original: Figure 4 (c) shows the comparison on the time used for predicting all CIFAR-10test data.
Revised: Figure 4 (c) shows the comparison on the time used for predicting all CIFARtest data.

Original:  ELLA is slightly slower than MFVI and substantially faster than non-last-layerLLA.
Revised: We note that ELLA is slightly slower than MFVI and substantially faster than non-last-layer

Original: [Table] N e g a ti v e L og - li k e li hood ELLAMAP [Table]
Revised: LLA methods.

Original: 5.5 ImageNet Classification
Revised: 5.5 ImageNet Classification

Original: We apply ELLA to ImageNet classification [8] to demonstrate its scalability.
Revised: We apply ELLA to ImageNet classification [8] to demonstrate its scalability.

Original: The experiment settingsare identical with those for CIFAR-10.
Revised: The experiment settings are identical with those for CIFAR-10.

Original: We observe that all LLA methods implemented with Laplacewould cause out-of-memory (OOM) errors, thus take only MAP and MFVI-BF as baselines.
Revised: We observe that all LLA methods implemented with Laplace would cause out-of-memory (OOM) errors or suffer from very long fitting procedures, thus take only MAP and MFVI-BF as baselines.

Original: Table 2presents the comparison on test accuracy, NLL, and ECE with ResNet architectures.
Revised: Table 2 presents the comparison on test accuracy, NLL, and ECE with ResNet architectures.

Original: We see that ELLA maintains its superiority in ECE while MFVI-BF can induce higher accuracy and lower NLL.
Revised: We see that ELLA maintains its superiority in ECE while MFVI-BF can induce higher accuracy and lower NLL.

Original: This may be attributed to that the pretrained MAP model lies at a sharp maxima of the true posteriorso it is necessary to properly adjust the mean of the Gaussian approximate 
Revised: This may be attributed to that the pretrained MAP model lies at a sharp maxima of the true posterior so it is necessary to properly adjust the mean of the Gaussian approximate posterior. 5 Skew intensity

Original: [Table]
Revised: N e g a ti v e L og - li k e li hood ELLAMAP    5 Skew intensity 0. E xp ec t e d C a li b r a ti on E rr o r ELLAMAP

Original: [Table]
Revised: [Table]

Original: We lastly apply ELLA to ViT-B [11].
Revised: We lastly apply ELLA to ViT-B [11].

Original: We compare ELLA to MAPin Table 3 as all other baselines incur OOM errors or crushinglylong running time.
Revised: We compare ELLA to MAP in Table 3 as all other baselines incur OOM errors or crushingly long running time.

Original: As shown, ELLA beats MAP in multiple aspects.
Revised: As shown, ELLA beats MAP in multiple aspects.

Original: Figure 5 shows the results of ELLA and MAP on ImageNetcorruptions [19].
Revised: Figure 5 shows the results of ELLA and MAP on ImageNet corruptions [19].

Original: They are consistent with those for CIFAR-10corruptions.
Revised: They are consistent with those for CIFARcorruptions.

Original: We reveal by this experiment that ELLA can be amore applicable and scalable method than most of existing BNNs.
Revised: We reveal by this experiment that ELLA can be a more applicable and scalable method than most of existing BNNs.

Original:  Related Work
Revised:  Related Work

Original: LA [38, 50, 14, 30, 22, 6, 7] locally approximates the true posterior with a Gaussian, analogous to VI with Gaussian variationals [16, 1, 36, 52, 46] and SWAG [39].
Revised: LA [38, 50, 14, 30, 22, 6, 7] locally approximates the Bayesian posterior with a Gaussian distribution, analogous to VI with Gaussian variationals [16, 1, 36, 52, 46] and SWAG [39].

Original: LA can be applied to pretrainedmodels effortlessly while the acquired posteriors are potentially restrictive (see Section 5.5).
Revised: LA can be applied to pretrained models effortlessly while the acquired posteriors are potentially restrictive (see Section 5.5).

Original: VIenjoys higher flexibility yet relies on costly training; BayesAdapter [10] seems to be a remedy to theissue but its accessibility is still lower than LA.
Revised: VI enjoys higher flexibility yet relies on costly training; BayesAdapter [10] seems to be a remedy to the issue but its accessibility is still lower than LA.

Original: SWAG stores a series of SGD iterates to heuristicallyconstruct an approximate posterior and is empirically weaker than LA/LLA [6].
Revised: SWAG stores a series of SGD iterates to heuristically construct an approximate posterior and is empirically weaker than LA/LLA [6].

Original: Though more expressive approaches like deep ensembles [32] and MCMC [55, 4, 63] can explicitlyexplore diverse posterior modes, they face limitations in efficiency and scalability.
Revised: Though more expressive approaches like deep ensembles [32] and MCMC [55, 4, 63] can explicitly explore diverse posterior modes, they face limitations in efficiency and scalability.

Original: What’s more, it hasbeen shown that LA/LLA can perform on par with or better than deep ensembles and cyclical MCMCon multiple benchmarks [6].
Revised: What’s more, it has been shown that LA/LLA can perform on par with or better than deep ensembles and cyclical MCMC on multiple benchmarks [6].

Original: This may be attributed to the un-identified, complicated relationshipsbetween the parameters space and the function space of DNNs [59].
Revised: This may be attributed to the un-identified, complicated relationships between the parameters space and the function space of DNNs [59].

Original: Also of note that LA can embracedeep ensembles to capture multiple posterior modes [13].
Revised: Also of note that LA can embrace deep ensembles to capture multiple posterior modes [13].

Original: Conclustion
Revised: [9] introduces a general kernel approximation technique using neural networks.

Original: This paper proposes ELLA, a simple, effective, and reliable approach for Bayesian deep learning. ELLA addresses the unreliability issues of existing approximations to LLA and is implemented basedon Nyström approximation.
Revised: By contrast, we focus on leveraging kernel approximation to accelerate Laplace approximation. Thus, the focus of the two works is distinct. Indeed, this work shares a similar idea with the Sec 4.3 in [9] that LLA can be accelerated by kernel approximation.

Original: We offer theoretical guarantees for ELLA and perform extensive studiesto testify its efficacy and scalability. ELLA currently accounts for only the predictive, and extendingit to estimate model evidence for model selection [38] deserves future investigation.
Revised: But, except for such an idea, this work differentiates from [9] in aspects like motivations, techniques, implementations, theoretical backgrounds, and applications. Our implementation, theoretical analysis, and some empirical findings are all novel.

Original:  Checklist
Revised:  Conclustion

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]
Revised: This paper proposes ELLA, a simple, effective, and reliable approach for Bayesian deep learning.

Original: (b) Did you describe the limitations of your work?

Original: [Yes] See Section Related Work.

Original: (c) Did you discuss any potential negative societal impacts of your work? [N/A](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]
Revised: ELLA addresses the unreliability issues of existing approximations to LLA and is implemented based on Nyström method.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [N/A] (b) Did you include complete proofs of all theoretical results?
Revised: We offer theoretical guarantees for ELLA and perform extensive studies to testify its efficacy and scalability.

Original: [Yes] See Appendix.3.

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)? [Yes] See supplemen- tal material. (b) Did you specify all the training details (e
Revised: ELLA currently accounts for only the predictive, and extending it to estimate model evidence for model selection [38] deserves future investigation.

Original: g., data splits, hyperparameters,

Original: how they

Original: were chosen)? [Yes] See Section Experiments. (c) Did you report error bars (e g., with respect to  the random seed after running experiments multiple times)?
Revised: Using model evidence to select the number of remaining eigenpairs K for ELLA is also viable.

Original: [No] (d) Did you include the total amount of compute and the type of resources used (e g., typeof GPUs, internal cluster, or cloud provider)? [No]
Revised: Acknowledgments and Disclosure of Funding

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: This work was supported by NSF of China Projects (Nos. 62061136001, 62076145, 62106121,

Original: [Yes](b) Did you mention the license of the assets? [N/A] (c) Did you include any new assets either in the supplemental material or as a URL? [No]
Revised: U19B2034, U1811461, U19A2081, 6197222); Beijing NSF Project (No. JQ19016); a grant from

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content? [N/A] If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: Tsinghua Institute for Guo Qiang; and the High Performance Computing Center, Tsinghua University.

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: J.Z was also supported by the XPlorer Prize

Original: [N/A]

Original: A Proof 
Revised: A Proof A.1 Derivation of Equation (5)

Original: A.1 Derivation of Equation (5)

Original: Given
Revised: Given

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: A.2 Derivation of Equation (6)
Revised: A.2 Derivation of Equation (6)

Original: [Equation]

Original: (Woodbury matrix

Original: [Equation]
Revised: [Equation]

Original: A.3 Derivation of Equation (18)
Revised: A.3 Derivation of Equation (18)

Original: When M = K , the eigenvalues and eigenvectors of K =
Revised: When M = K , the eigenvalues and eigenvectors of K =

Original: J ˆ θ , ˜ X
Revised: J ˆ θ , ˜ X

Original: J ⊤ ˆ θ , ˜ X can be organized assquare matrices diag ( λ ) =
Revised: J ⊤ ˆ θ , ˜ X can be organized as square matrices diag ( λ ) =

Original: diag ( λ 1 , ..., λ K ) and U =
Revised: diag ( λ 1 , ..., λ K ) and U =

Original: [ u 1 , ..., u K ] respectively.
Revised: [ u 1 , ..., u K ] respectively.

Original: And φ ( x ) =[
Revised: And φ ( x ) =

Original:  J ˆ θ ( x ) v 1 , ..., J ˆ θ ( x ) v K ]
Revised: [ J ˆ θ ( x ) v 1 , ..., J ˆ θ ( x ) v K ]

Original: =
Revised: =

Original: J ˆ θ ( x ) J ⊤ ˆ θ , ˜ X
Revised: J ˆ θ ( x ) J ⊤ ˆ θ , ˜ X

Original: U diag ( λ ) − 12 .
Revised: U diag ( λ ) − 12 .

Original: Of note that U is a orthogonal matrix where
Revised: Of note that U is a orthogonal matrix where

Original: UU ⊤ = U ⊤U
Revised: UU ⊤ [Equation] ⊤

Revised: [Equation]

Original: = I K ⇒
Revised: = I K ⇒

Original: U − 1 = U ⊤ , and by definition, J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X
Revised: U − 1 = U ⊤ , and by definition, J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X

Original: = U diag ( λ ) U ⊤ .
Revised: = U diag ( λ ) U ⊤ .

Original: Thenκ ELLA ( x , x ′ ) =
Revised: Then κ ELLA ( x , x ′ ) =

Original: φ ( x ) (cid:104) (cid:88)
Revised: φ ( x ) (cid:104) (cid:88)

Original: [Equation]
Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: [Equation]

Original: A.4 Proof of Theorem
Revised: A.4 Proof of Theorem

Original: Theorem 1.
Revised: Theorem 1.

Original: Let c Λ be a finite constant associated with Λ , and E ′ the error of Nyström approximation∥
Revised: Let c Λ be a finite constant associated with Λ , and E ′ the error of Nyström approximation ∥

Original: J ˆ θ , X J ⊤ ˆ θ , ˜ X
Revised: J ˆ θ , X J ⊤ ˆ θ , ˜ X

Original: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X ) − 1 J ˆ θ , ˜
Revised: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X ) − 1 J ˆ θ , ˜

Original: X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ∥ .
Revised: X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ∥ .

Original: It holds that
Revised: It holds that

Original: E ≤ σ 40 c Λ E ′ + σ 20 .
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: With K
Revised: With K

Original: =
Revised: =

Original: J ˆ θ , ˜ X J ⊤ˆ
Revised: J ˆ θ , ˜ X J ⊤

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: X ≻ 0 , let P =
Revised: ≻ 0 , let P =

Original: J ⊤ ˆ θ , ˜ XK
Revised: J ⊤ ˆ θ , ˜ X

Original: [Equation]

Original:  1 J ˆ θ , ˜ X .
Revised: K − 1 J ˆ θ , ˜ X .

Original: By Woodbury matrix identity,
Revised: By Woodbury matrix identity,

Original: [Equation]
Revised: [Equation]

Original: It is interesting to see that P is a projector: P 2 = P and PJ
Revised: It is interesting to see that P is a projector: P 2 = P and PJ

Original: ⊤ ˆ θ , ˜ X
Revised: ⊤ ˆ θ , ˜ X

Original: = J ⊤ ˆ θ , ˜ X .
Revised: = J ⊤ ˆ θ , ˜ X .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: By Woodbury matrix identity again,
Revised: By Woodbury matrix identity again,

Original: [Equation]
Revised: [Equation]

Original: Let T =
Revised: Let T =

Original: ( PJ ⊤ ˆ θ , X Λ X , Y J ˆ θ , X P + I P /σ 20 ) − 1 ,
Revised: ( PJ ⊤ ˆ θ , X Λ X , Y J ˆ θ , X P + I P /σ 20 ) − 1 ,

Original: so PJ ⊤ ˆ θ , X Λ X , Y J ˆ θ , X P = T − 1 − I P /σ 20 .
Revised: so PJ ⊤ ˆ θ , X Λ X , Y J ˆ θ , X P = T − 1 − I P /σ 20 .

Original: It follows that
Revised: It follows that

Original: [Equation]
Revised: [Equation]

Original: As a result
Revised: As a result

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where we leverage the fact that the eigenvalue of the projector P is either 1 or 0 .
Revised: where we leverage the fact that the eigenvalue of the projector P is either 1 or 0 .

Original: And
Revised: And

Original: [Equation]
Revised: [Equation]

Original: where λ min ( · ) denotes the smallest eigenvalue of a matrix.
Revised: where λ min ( · ) denotes the smallest eigenvalue of a matrix.

Original: The last inequality holds due to that theeigenvalues of PJ ⊤ ˆ θ , X Λ X , Y J ˆ θ , X P + I P /σ 20 and J ⊤ ˆ θ , X Λ X , Y J ˆ θ , X
Revised: The last inequality holds due to that the eigenvalues of PJ ⊤ ˆ θ , X Λ X , Y J ˆ θ , X P + I P /σ 20 and J ⊤ ˆ θ , X Λ X , Y J ˆ θ , X

Original: + I P /σ 20 are larger than or equalto 1 /σ 20 as PJ ⊤ ˆ θ , X Λ X , Y J
Revised: + I P /σ 20 are larger than or equal to 1 /σ 2 0 as PJ ⊤ ˆ θ , X Λ X , Y J

Original: ˆ θ , X P and J ⊤ ˆ θ , X Λ X , Y J ˆ θ , X are SPSD.
Revised: ˆ θ , X P and J ⊤ ˆ θ , X Λ X , Y J ˆ θ , X are SPSD.

Original: To estimate the upper bound of ∥ T − Σ ∥ , we lay out the following Lemma.
Revised: To estimate the upper bound of ∥ T − Σ ∥ , we lay out the following Lemma.

Original: Lemma 1 (Proposition of Weinstein–Aronszajn identity) .
Revised: Lemma 1 (Proposition of Weinstein–Aronszajn identity) .

Original: If A and B are matrices of size m × n andn × m respectively, the non-zero eigenvalues of AB and BA are the same.
Revised: If A and B are matrices of size m × n and n × m respectively, the non-zero eigenvalues of AB and BA are the same.

Original: It follows that
Revised: It follows that

Original: [Equation]
Revised: [Equation]

Revised: , X ) ∥

Original: [Equation]
Revised: [Equation]

Original: ∥ Λ X , Y( J ˆ θ , X J ⊤ ˆ θ , ˜ X
Revised: X , Y

Original: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X ) − 1 J ˆ θ , ˜ X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ) ∥ ≤ σ 40 ∥ Λ X , Y ∥∥ J ˆ θ , X J ⊤ ˆ θ , ˜ X
Revised: [Equation]

Original: ( J ˆ θ , ˜ X J ⊤ ˆ θ ,˜ X
Revised: θ , X

Original: [Equation]
Revised: [Equation]

Original: −

Original: [Equation] ˆ θ [Equation] ˜ X
Revised: ⊤

Original: [Equation]
Revised: [Equation]

Original: ⊤ ˆ θ 
Revised: θ , ˜ X

Original: [Equation] X [Equation] ˆ θ 
Revised: ( J ˆ θ , ˜ X J ⊤ ˆ θ , ˜ X

Original: [Equation]
Revised: [Equation]

Original: X
Revised: −

Original: [Equation]
Revised: [Equation]

Original:  ⊤ ˆ θ
Revised: ˆ θ , ˜ X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ

Original: [Equation]
Revised: [Equation]

Original: X ∥ .
Revised: X ∥ .

Original: It is easy to show that ∥ Λ X , Y ∥ ≤ c Λ with c Λ as a finite constant.
Revised: It is easy to show that ∥ Λ X , Y ∥ ≤ c Λ with c Λ as a finite constant.

Original: For example, when handlingregression tasks, − log p ( y | g ) typically boils down to 1 2 σ 2 noise ( g − y )
Revised: For example, when handling regression tasks, − log p ( y | g ) typically boils down to 1 2 σ 2 noise ( g − y )

Original: ⊤ ( g − y ) where σ 2 noise refers tothe variance of data noise, so −∇ 2 gg log p ( y | g ) =
Revised: ⊤ ( g − y ) where σ 2 noise refers to the variance of data noise, so −∇ 2 gg log p ( y | g ) =

Original: 1 σ 2 noise I C and c Λ = 1 σ 2 noise .
Revised: 1 σ 2 noise I C and c Λ = 1 σ 2 noise .

Original: When facing classification cases, − log p ( y | g ) becomes the softmax cross-entropy loss.
Revised: When facing classification cases, − log p ( y | g ) becomes the softmax cross-entropy loss.

Original: It is guaranteed that c Λ ≤ 2 because
Revised: It is guaranteed that c Λ ≤ 2 because

Original: [Equation]
Revised: [Equation]

Original: To summarize, we have the following bound for the approximation error E :
Revised: To summarize, we have the following bound for the approximation error E :

Original: [Equation]
Revised: [Equation]

Revised: B How Does ELLA Address the Scalability Issues of Nyström Approximation

Revised: Existing works like NeuralEF [9] highlight the scalability issues of the Nyström approximation when applied to NTKs, but they have been successfully addressed by ELLA, which reflects our technical novelty.

Revised: On the one hand, the cost of eigendecomposing the M × M matrix in our case is low considering that we set M < 10 4 .

Revised: As shown in Figure 2(a)(b), M = 2000 can lead to reasonably small approximation errors for both the Nyström approximation and the resulting ELLA covariance.

Revised: On the other hand, the prediction cost at test time is reduced by two innovations in this work.

Revised: Firstly, looking at Equation (14), typical Nyström method first computes J ˆ θ ( x , i )

Revised: J ⊤ ˆ θ , ˜ X

Revised: (i.e., evaluate the

Revised: NTK similarities between ( x , i ) and all training data ˜ X ) and then multiplies the result with u k ∈ R M to get the evaluation of k -th eigenfunction.

Revised: So it needs to store J ˆ θ , ˜ X ∈ R M × P and compute M

Revised: Jacobian-vector products (JVPs).

Revised: But this work proposes to first compute v k

Revised: = J ⊤ ˆ θ , ˜ X

Revised: u k , k = 1 , ..., K (we omit the scalar multiplier) and then obtain the evaluation of k -th eigenfunction by J ˆ θ ( x , i ) v k .

Revised: Namely, we only need to store K vectors of size P and compute K JVPs.

Revised: In most cases, M =and K = 20 , so the benefits of our technique are obvious.

Revised: Secondly, we leverage forward-mode autodiff (fwAD) to efficiently compute the JVPs where we do not need to explicitly compute and

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Original: B More experiments
Revised: [Table] [Figure] [Table] [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: [Figure]
Revised: [Figure]

Revised: store the Jacobian J ˆ θ ( x , i ) but to invoke once fwAD.

Original: B.1 More  Results on CIFAR-
Revised: Further, fwAD enables the parallel evaluation of JVPs over all output dimensions, i.e., we can concurrently compute J ˆ θ ( x , i ) v k , i = 1 , ..., C in one forward pass. These techniques make the evaluation of ELLA equal to performing K forward passes under the scope of fwAD. This is similar to other BNNs that perform S forward passes with different MC parameter samples to estimate the posterior predictive. C More Experimental Results C.1 Test NLL of ELLA Varies w.r.t. K and M

Original: [Figure]
Revised: Table 4 presents how the test NLL of ELLA varies w.r.t. M and K on CIFAR-10 with ResNetarchitecture. Considering the results in Figure 2, we conclude that under the same K , a larger M can bring closer approximation but it does not necessarily lead to a better test NLL. C.2 Test Accuracy and ECE of ELLA Vary w.r.t. N

Original: Figure 6 presents how the test accuracy and ECE of ELLAon CIFAR-10 vary w.r.t. 
Revised: Figure 6 presents how the test accuracy and ECE of ELLA vary w.r.t. the number of training data N on CIFAR-10.

Revised: C.3 Test NLL of LLA* Varies w.r.t. N

Revised: Figure 7 shows how the test NLL of LLA ∗ varies w.r.t.

Original: the number oftraining data N .
Revised: the number of training data N on CIFAR-10.

Revised: These results, as well as those in Section 5.2, reflect that the overfitting issue of LLA generally exists.

Revised: We can also see that properly tuning the prior precision can alleviate the overfitting of LLA ∗ to some extent but cannot fully resolve it.

Original: Figure 7, Figure 8, and Figure 9 show the results of the considered methods on CIFAR-10 corruptionsusing ResNet-20, ResNet-32, and ResNet-44 architectures respectively.
Revised: [Figure] [Figure] ResNet [Figure]

Revised: [Table]

Revised: C.4 Comparison on the Approximation Error to Vanilla LLA

Revised: Table 5 lists the discrepancies between the predictive distributions of the approximate LLA method and vanilla LLA in the illustrative regression case detailed in Section 5.3.

Revised: Here, considering the predictive distribution for one test datum is a Gaussian distribution, we use the KL divergence between the Gaussians yielded by the method of concern and LLA as a proxy of the approximation error (averaged over a set of test points).

Revised: C.5 More Results on CIFAR-10 Corruptions

Revised: Figure 8, Figure 9, and Figure 10 show the results of the considered methods on CIFAR-10 corruptions using ResNet-20, ResNet-32, and ResNet-44 architectures respectively.

Original: We can see that ELLAsurpasses the baselines in aspects of NLL and ECE at various levels of skew.
Revised: We can see that ELLA surpasses the baselines in aspects of NLL and ECE at various levels of skew.

Original: These results signify
Revised: These results signify

Original: ELLA’s ability to make conservative predictions for OOD inputs.
Revised: ELLA’s ability to make conservative predictions for OOD inputs.

Original: Figure 10 shows how the test NLL of LLA ∗ on CIFAR-10 varies w.r.t.
Revised: [Figure]

Original: the number of training data N . These results, as well as those in Section 5.2, reflect that the overfitting issue of LLA generally exists. We can also see that properly tuning the prior precision can alleviate the overfitting of LLA ∗ to someextent but cannot fully resolve it.
Revised: g a ti v

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: N e g a ti v e L og - li k e li
Revised: N e g a ti v e L og - li k e li

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: L

Original: [Figure]

Original: [Figure]