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 effortless adaptation 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: The generalized 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-trivial inefficiency 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 likely to 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 the connections between LLA and neural target kernels (NTKs), we develop a Nyström approximation to NTKs to accelerate LLA.
Revised: To tackle this issue, inspired by the connections between LLA and neural tangent kernels (NTKs), we develop a Nyström approximation to NTKs to accelerate LLA.

Original: Our method benefits from the capability of popular deep learning libraries for forward mode automatic differentiation, and enjoys 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 of the 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 method can even scale up to architectures like vision transformers.
Revised: Our method can even scale up to architectures like vision transformers.

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

Original: Code is available at https://github.com/thudzj/ELLA .
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 facto solutions 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 reasoning about 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 in risk-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 to learning 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 the pretraining-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 Gaussian around its maximum, whose mean and covariance are the maximum a posteriori (MAP) and the inversion 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 the generalized 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: Immer et al.
Revised: Immer et al.

Original: [22] argue that LLA is more sensible than LA in the presence of GGN approximation;
Revised: [22] argue that LLA is more sensible than LA in the presence of GGN approximation;

Original: [Figure]
Revised: [Figure]

Original: LLA can perform on par with or better than popular alternatives on various uncertainty quantification (UQ) tasks [14, 6, 7].
Revised: 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.
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 to a 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 errors in 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 to make probabilistic predictions 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 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.
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 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).
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: The resultant accElerated LLA (ELLA) preserves the principal structures of vanilla
Revised: The resultant accElerated LLA (ELLA) preserves the principal structures of vanilla

Original: LLA yet without explicitly computing/storing the costly GGN/Jacobian matrices for the training data.
Revised: 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 regression task 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 standard image classification benchmarks to exhibit the superiority of ELLA over competing baselines in aspects 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 modern architectures 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 advance in machine learning suggests using an 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 the data, 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 methods usually 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 like variational 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, 63], particle-optimization based variational inference (POVI) [35], etc.
Revised: Monte Carlo (MCMC) [55, 4, 63], particle-optimization based variational inference (POVI) [35], etc.

Original: BNNs predict for new data x ∗ by posterior predictive p ( y | x ∗ , D ) =
Revised: BNNs predict for new data x ∗ by posterior predictive p ( y | x ∗ , D ) =

Original: E p ( θ |D ) p ( y | x ∗ , θ ) ≈
Revised: E p ( θ |D ) p ( y | x ∗ , θ ) ≈

Original: E q ( θ ) p ( y | x ∗ , θ ) ≈ 1 S (cid:80)
Revised: E q ( θ ) p ( y | x ∗ , θ ) ≈ 1 S (cid:80)

Original: Ss =1 p ( y | g θ s ( x ∗ )) where θ s ∼ q ( θ ) are i.i.d.
Revised: Ss =1 p ( y | g θ s ( x ∗ )) where θ s ∼ q ( θ ) are i.i.d.

Original: Monte Carlo (MC) samples.
Revised: Monte Carlo (MC) samples.

Original: 2.1 Laplace Approximation and Its Linearized Variant
Revised: 2.1 Laplace Approximation and Its Linearized Variant

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

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

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

Original: parameters, i.e., Σ − 1 = −∇ 2 θθ (log p ( D| θ )
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 prior p ( θ ) =
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: ˆ θ equals 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 to use the symmetric positive and semi-definite (SPSD) GGN matrix as a workaround, i.e.,
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 turns the 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 leveraged to 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: To mitigate 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 to pretrained 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 LA and 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] further advances 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 further approximations 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 are commonly 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-dimensional parameter space (e.g., the parameter space of the last layer [30]) to reduce the scale of the GGN matrix.
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 approximation errors 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: Based on 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-space approximate 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
Revised: GP ( f | g ˆ θ ( x ) , κ LLA ( x , x ′ )) with

Original: [Equation]
Revised: [Equation]

Original: By Woodbury matrix identity [60]
Revised: By Woodbury matrix identity [60]

Original: [Equation]
Revised: [Equation]

Original: we have:
Revised: 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 to g ˆ θ ( 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 ) of size 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 inefficiency issue.
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 dimensions to cast q ( f ) into C independent GPs following [51], and randomly subsample M ≪
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 ≪

Original: N data points to form a cheap substitute for the original gram matrix.
Revised: N data points to form a cheap substitute for the original gram matrix.

Original: Despite effective in some cases, these approximations 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 , the scalability 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 (detailed derivation in Appendix A.2; see also [9]) 
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]) [Equation] [Equation] κ ELLA (

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 each training 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-based methods [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 relies on a relatively large number of random 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 κ NTK (( x , i ) , ( x ′ , i ′ ))
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: =
Revised: =

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

Original: J ˆ θ ( 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 × Pcomputes 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 measure q 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- K eigenvalues 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 distribution over { 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.
Revised: The Nyström method draws M ( M ≥ K ) i.i.d.

Original: samples ˜ 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 orthonormal eigenvectors 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) [Equation] it is easy to see that:
Revised: Given the constraint in Equation (9), 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: Then, the mapping φ which satisfied κ NTK ( x , x ′ ) ≈ φ
Revised: Then

Original: [Equation]
Revised: [Equation]

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

Original:  φ [Equation] x ′ )
Revised: ≈ φ ( x ) φ ( 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 because the 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.

Original: Discussion
Revised: Discussion

Original: 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.
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 accomplished by 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: (17) where the model output g ˆ θ ( x ) and the JVP J ˆ θ ( x ) v k are simultaneously computed in one single forward pass .
Revised: 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 the capability 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 the eigendecomposition 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 set once 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 that perform 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.  Theoretical Analysis

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

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 ˆ
Revised: J ˆ

Original: [Equation]
Revised: [Equation]

Original: , ˜ X [
Revised: , ˜ X [

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: [ i m ] , ˆ θ ) { λ k , u k } Kk =1 = eig
Revised: [ 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]
Revised: [Equation]

Revised: return partial

Original: [Equation]
Revised: [Equation]

Original: return partial(_ φ , g ˆ θ , C , { v k } K k =1 )

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.
Revised: To reduce unnecessary complexity, this section assumes using M = K i.i.d.

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

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

Original: [Equation]
Revised: [Equation]

Original: 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 , thus with 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 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 ⊤ ˆ θ [Equation] J ˆ

Revised: [Equation]

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

Revised: ⊤

Original: X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ , X ∥ .
Revised: [Equation] ˆ [Equation] ⊤ ˆ θ , 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 results developed 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: 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: 5.1 Specification of Hyperparameters
Revised: 5.1 Specification of Hyperparameters

Original: 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.
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.

Original: We simply set σ 20 to 1 Nγ with γ as the weight decay coefficient used for pretraining according to [10].
Revised: 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 others as 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 architectures or 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 approximation error of the Nyström method by ϵ
Revised: We quantify the approximation error of the Nyström method by ϵ

Original: [Equation]
Revised: [Equation]

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

Original: 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 )
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: We vary 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 K vectors 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 choice given Figure 2.
Revised: K ∈ [16 , 32] seems to be a reasonable choice given Figure 2.

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

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

Original: Given these results, we set M = 2000 and K =in the following experiments unless otherwise stated.
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, formulated as p ( y | x ∗ , D )
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: =
Revised: =

Original: E p ( f |D ) p ( y | f ( x ∗ ))
Revised: E p ( f |D ) p ( y | f ( x ∗ ))

Original: ≈ E q ( f ) p ( y | f ( x ∗ )) =
Revised: ≈ E q ( f ) p ( y | f ( x ∗ )) =

Original: E f ∼N ( g ˆ θ ( x ∗ ) ,κ ELLA ( x ∗ , x ∗ ))
Revised: E f ∼N ( g ˆ θ ( x ∗ ) ,κ ELLA ( x ∗ , x ∗ ))

Original: p ( y | f ) .
Revised: p ( y | f ) .

Original: In classification tasks, we use 512 MC samples to approximate the last expectation and it is cheap.
Revised: 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 model misspecification [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 to seek 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 training set 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 deployed ELLA in Figure 2 (c) and Appendix C.2.
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 also provide the corresponding results of LLA ∗ 4 in Appendix C.3.
Revised: We also provide the corresponding results of LLA ∗ 4 in Appendix C.3.

Original: The V-shape NLL curves across settings reflect 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).

Original: Figurealso shows that tuning the prior precision w.r.t.
Revised: Figurealso shows that tuning the prior precision w.r.t.

Original: marginal likelihood can alleviate the overfitting of
Revised: marginal likelihood can alleviate the overfitting of

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

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

Original: To address the overfitting issue, we advocate performing early stop when fitting ELLA/LLA on 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)),
Revised: Specifically, when iterating over the training data to estimate G (see Equation (6)),

Original: we
Revised: we

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

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

Original: ELLAMAPMFVI-BFLLA *
Revised: ELLAMAPMFVI-BFLLA *

Original: LLA * -KFAC continuously record the NLL of the current ELLA posterior on some validation data, and stop when there is a trend of overfitting.
Revised: LLA * -KFAC continuously record the NLL of the current ELLA posterior on some validation data, and stop when there is a trend of overfitting.

Original: 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.
Revised: 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: Compared to tuning the prior, early stopping also helps reduce the time cost of processing big data (e.g., ImageNet [8]).
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 an MLP with 3 hidden layers and tanh activations, and we pretrain it by MAP for 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, we use 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 ∗ .

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

Original: It is clear that ELLA delivers a closer approximation to LLA than LLA-Diag and
Revised: It is clear that ELLA delivers a closer approximation to LLA than LLA-Diag and

Original: LLA ∗ .
Revised: LLA ∗ .

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

Original: LLA methods using certain metrics.
Revised: LLA methods using certain metrics.

Original: 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: 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).

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

Original: As shown, LLA-KFAC comes pretty close to LLA.
Revised: As shown, LLA-KFAC comes pretty close to LLA.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 5.4 CIFAR-10 Classification
Revised: 5.4 CIFAR-10 Classification

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

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

Original: Apart from MAP, LLA ∗ , LLA-Diag,
Revised: Apart from MAP, LLA ∗ , LLA-Diag,

Original: 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.
Revised: 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 GGN matrices 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 are all post-hoc applied to the pretrained models, so the comparisons will be fair.
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, we use M = 2000 and K = 20 for ELLA; 6 we use 20 MC samples to estimate
Revised: Regarding the setups, we use M = 2000 and K = 20 for ELLA; 6 we use 20 MC samples to estimate

Original: 0.2 Confidence Threshold 0.1 0.
Revised: 0.2 Confidence Threshold 0.1 0.

Original: E rr o r on e x a m p l e s w it
Revised: E rr o r on e x a m p l e s w it

Original: h c on f i d e n ce
Revised: h c on f i d e n ce

Original: ELLAMAPMFVI-BFLLA *
Revised: ELLAMAPMFVI-BFLLA *

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

Original: T e s t N LL (b)
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

Original: T
Revised: T

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

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

Original: (c)
Revised: (c)

Original: [Table]
Revised: [Table]

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

Original: We have enabled the tuning of the prior precision for all LLA baselines, but not for ELLA.
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 exhibits superior 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 ∗ -KFAC can 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]), we thus exclude them 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/robustness benchmark CIFAR-10 corruptions [19] and report the results in Figure 3 and Appendix C.5.
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 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.
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 SVHN test 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 average error 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 makes less 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 K impacts 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 the efficiency of ELLA.
Revised: We see that K ∈ [20 , 30] can already lead to good performance, reflecting the efficiency of ELLA.

Original: Another benefit of ELLA is that with it, we can actively control the performance vs. cost trade-off.
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 CIFARtest data.
Revised: Figure 4 (c) shows the comparison on the time used for predicting all CIFARtest data.

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

Original: LLA methods.
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 settings are 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 Laplace would cause out-of-memory (OOM) errors or suffer from very long fitting procedures, thus take only
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

Original: MAP and MFVI-BF as baselines.
Revised: MAP and MFVI-BF as baselines.

Original: Table 2 presents 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 posterior so it is necessary to properly adjust the mean of the Gaussian approximate posterior.
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.

Original: 5 Skew intensity
Revised: 5 Skew intensity

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

Original: ELLAMAP    5 Skew intensity 0.
Revised: ELLAMAP    5 Skew intensity 0.

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

Original: ELLAMAP
Revised: 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 MAP in Table 3 as all other baselines incur OOM errors or crushingly long 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 ImageNet corruptions [19].
Revised: Figure 5 shows the results of ELLA and MAP on ImageNet corruptions [19].

Original: They are consistent with those for CIFARcorruptions.
Revised: They are consistent with those for CIFARcorruptions.

Original: We reveal by this experiment that ELLA can be a more 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 Bayesian posterior with a Gaussian distribution, 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 pretrained models 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: 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.
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 heuristically construct 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 explicitly explore 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 has been shown that LA/LLA can perform on par with or better than deep ensembles and cyclical MCMC on 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 relationships between 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 embrace deep ensembles to capture multiple posterior modes [13].
Revised: Also of note that LA can embrace deep ensembles to capture multiple posterior modes [13].

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

Original: By contrast, we focus on leveraging kernel approximation to accelerate Laplace approximation.
Revised: By contrast, we focus on leveraging kernel approximation to accelerate Laplace approximation.

Original: Thus, the focus of the two works is distinct.
Revised: Thus, the focus of the two works is distinct.

Original: Indeed, this work shares a similar idea with the Sec 4.3 in [9] that LLA can be accelerated by kernel approximation.
Revised: Indeed, this work shares a similar idea with the Sec 4.3 in [9] that LLA can be accelerated by kernel approximation.

Original: But, except for such an idea, this work differentiates from [9] in aspects like motivations, techniques, implementations, theoretical backgrounds, and applications.
Revised: But, except for such an idea, this work differentiates from [9] in aspects like motivations, techniques, implementations, theoretical backgrounds, and applications.

Original: Our implementation, theoretical analysis, and some empirical findings are all novel.
Revised: Our implementation, theoretical analysis, and some empirical findings are all novel.

Original:  Conclustion
Revised:  Conclustion

Original: This paper proposes ELLA, a simple, effective, and reliable approach for Bayesian deep learning.
Revised: This paper proposes ELLA, a simple, effective, and reliable approach for Bayesian deep learning.

Original: ELLA addresses the unreliability issues of existing approximations to LLA and is implemented based on Nyström method.
Revised: ELLA addresses the unreliability issues of existing approximations to LLA and is implemented based on Nyström method.

Original: We offer theoretical guarantees for ELLA and perform extensive studies to testify its efficacy and scalability.
Revised: We offer theoretical guarantees for ELLA and perform extensive studies to testify its efficacy and scalability.

Original: ELLA currently accounts for only the predictive, and extending it to estimate model evidence for model selection [38] deserves future investigation.
Revised: ELLA currently accounts for only the predictive, and extending it to estimate model evidence for model selection [38] deserves future investigation.

Original: Using model evidence to select the number of remaining eigenpairs K for ELLA is also viable.
Revised: Using model evidence to select the number of remaining eigenpairs K for ELLA is also viable.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: This work was supported by  NSF of China Projects (Nos. 62061136001, 62076145, 62106121, 
Revised: This work was supported by the National Key Research and Development Program of China (No. 2017YFA0700904), NSF of China Projects (Nos. 62061136001, 62076145, 62106121, U19B2034,

Original: U19B2034, U1811461, U19A2081, 6197222); Beijing NSF Project (No. JQ19016); a grant from 
Revised: U1811461, U19A2081, 6197222); Beijing NSF Project (No. JQ19016); a grant from Tsinghua

Original: Tsinghua Institute for Guo Qiang; and the High Performance Computing Center, Tsinghua University.
Revised: Institute for Guo Qiang; and the High Performance Computing Center, Tsinghua University.

Original: J.Z was also supported by the XPlorer Prize
Revised: J.Z was also supported by the XPlorer Prize

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

Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] (b) Did you describe the limitations of your work? [Yes] See Section Related Work.

Revised: (c) Did you discuss any potential negative societal impacts of your work?

Revised: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Revised: [Yes] 2.

Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Revised: [Yes] See Appendix. 3.

Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Revised: [Yes] See supplemental material.

Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Revised: [Yes] See Section Experiments.

Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

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

Revised: [No] 4.

Revised: 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: [Yes] (b) Did you mention the license of the assets?

Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Revised: [No] (d) Did you discuss whether

Revised: and how consent was obtained from people whose data you’re using/curating?

Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Revised: [N/A] 5.

Revised: 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 applicable?

Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Revised: [N/A]

Revised: A Proof 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]
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 as square 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 ⊤
Revised: UU ⊤

Original: [Equation]
Revised: [Equation]

Original: ⊤
Revised: ⊤

Original: [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]

Original: =
Revised: =

Original: [Equation] [Equation]
Revised: [Equation] [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: [Equation]
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: ≻ 0 , let P =
Revised: ≻ 0 , let P =

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

Original: ˜ X
Revised: ˜ X

Original: K − 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]
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 the eigenvalues 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 equal to 1 /σ 2 0 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 and n × 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]

Original: , X ) ∥
Revised: , X ) ∥

Original: [Equation]
Revised: [Equation]

Original: X , Y
Revised: X , Y

Original: [Equation]
Revised: [Equation]

Original: θ , X
Revised: θ , X

Original: [Equation]
Revised: [Equation]

Original: ⊤
Revised: ⊤

Original: [Equation]
Revised: [Equation]

Original: θ , ˜ X
Revised: θ , ˜ X

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

Original: [Equation]
Revised: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: ˆ θ , ˜ X J ⊤ ˆ θ , X − J ˆ θ , X J ⊤ ˆ θ
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 handling regression 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 to the 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]

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

Original: 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: 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.

Original: 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: On the one hand, the cost of eigendecomposing the M × M matrix in our case is low considering that we set M < 10 4 .

Original: 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: 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.

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

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

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

Original: (i.e., evaluate the
Revised: (i.e., evaluate the

Original: 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: 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.

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

Original: Jacobian-vector products (JVPs).
Revised: Jacobian-vector products (JVPs).

Original: But this work proposes to first compute v k
Revised: But this work proposes to first compute v k

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

Original: 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: 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 .

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

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

Original: Secondly, we leverage forward-mode autodiff (fwAD) to efficiently compute the JVPs where we do not need to explicitly compute and
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]

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

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

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

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

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

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

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

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

Original: [Figure]
Revised: [Figure]

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

Original: 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.
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.

Original: These techniques make the evaluation of ELLA equal to performing K forward passes under the scope of fwAD.
Revised: These techniques make the evaluation of ELLA equal to performing K forward passes under the scope of fwAD.

Original: This is similar to other BNNs that perform 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: C More Experimental Results C.1 Test NLL of ELLA Varies w.r.t. K and M
Revised: C More Experimental Results C.1 Test NLL of ELLA Varies w.r.t. K and M

Original: Table 4 presents how the test NLL of ELLA varies w.r.t.
Revised: Table 4 presents how the test NLL of ELLA varies w.r.t.

Original: M and K on CIFAR-10 with ResNetarchitecture.
Revised: M and K on CIFAR-10 with ResNetarchitecture.

Original: 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.
Revised: 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.

Original: C.2 Test Accuracy and ECE of ELLA Vary w.r.t. N
Revised: C.2 Test Accuracy and ECE of ELLA Vary w.r.t. N

Original: Figure 6 presents how the test accuracy and ECE of ELLA vary w.r.t.
Revised: Figure 6 presents how the test accuracy and ECE of ELLA vary w.r.t.

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

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

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

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

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

Original: We can also see that properly tuning the prior precision can alleviate the overfitting of LLA ∗ to some extent but cannot fully resolve it.
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] [Figure]
Revised: [Figure] [Figure]

Original: ResNet
Revised: ResNet

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

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

Original: 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: 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.

Original: 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: 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).

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

Original: 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.
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 ELLA surpasses 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]
Revised: [Figure]

Original: g a ti v
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]