ViViT: Curvature Access Through The Generalized Gauss-Newton’s Low-Rank Structure
Authors: Felix Dangel
Lukas Tatzel
Philipp Hennig
Venue: Accepted by TMLR ()

Abstract: Curvature in form of the Hessian or its generalized Gauss-Newton (GGN) approximation is valuable for algorithms that rely on a local model for the loss to train, compress, or explain deep networks. Existing methods based on implicit multiplication via automatic differentiation or Kronecker-factored block diagonal approximations do not consider noise in the mini-batch. We present ViViT, a curvature model that leverages the GGN’s low-rank structure without further approximations. It allows for efficient computation of eigenvalues, eigenvectors, as well as per-sample first- and second-order directional derivatives. The representation is computed in parallel with gradients in one backward pass and offers a fine-grained cost-accuracy trade-off, which allows it to scale. We demonstrate this by conducting performance benchmarks and substantiate ViViT’s usefulness by studying the impact of noise on the GGN’s structural properties during neural network training.

Published in Transactions on Machine Learning Research (02/2023)
ViViT: Curvature Access Through The Generalized Gauss-Newton's Low-Rank Structure
Felix Dangel*
f. d.dangel@unii-tuebingn....m
University of Tübingen, Tübingen, Germany
Lukas Tatzel*
bukas.tattzelemmi-tueebingggn
University of Tübingen, Tübingen, Germany
Philipp Hennig
University of Tübingen MPI for Intelligent Systems, Tübingen, Germany
Reviewed on OpenReview: https: // openreview net/ pormm?? id= DaJTUfPXKE
Abstract
Curvature in form of the Hessian or its generalized Gauss-Newton (GGN) approximation valuable for algorithms that rely on local model for the loss to train, compress, or explain deep networks. Existing methods based implicit multiplication via automatic differentiation or Kronecker-factored block diagonal approximations do not consider noise in the mini-batch. We present VIVrT, a curvature model that leverages the GGN's low-rank structure without further approximations. It allows for efficient computation of cigenvalues, cigenvectors, as well as per-sample first- and second-order directional derivatives. The representation is computed in parallel with gradients in one backward pass and offers a fine- grained cost-accuracy trade-off, which allows it to scale. We demonstrate this by conducting performance benchmarks and substantiate VIVIT's usefulness by studying the impact of noise on the GGN's structural properties during neural network training.
1 Introduction & Motivation
The large number of trainable parameters deep neural networks imposes computational constraints on the information that can be made available to optimization algorithms. Standard machine learning libraries (Abadi et al., 2015; Paszke et al., 2019) mainly provide access to first-order information in the form of average mini-batch gradients. This is limitation that complicates the development of novel methods that may outperform the state-of-tho-art: They must use the same objects to remain casy to implement and use, and to rely on the highly optimized code of those libraries. There is evidence that this has led to stagnation in the performance of first-order optimizers (Schmidt et al., 2021). Here, we thus study how to provide efficient access to richer information, namely higher-order derivatives and their distribution across the mini-batch.
Recent advances in automatic differentiation (Bradbury al., 2020; Dangel al., 2020) have made such information more readily accessible through vectorization of algebraic structure in the differentiated loss. We leverage and extend this functionality to efficiently access curvature in form of the Hessian's 's generalized Gauss-Newton (GGN) approximation. It offers practical advantages over the Hessian and is established for training (Martens, 2010; Martens & Grosse, 2015), compressing (Singh & Alistarh, 2020), or adding uncertainty to (Ritter et al., 2018b;a; Kristiadi et al., 2020) neural networks. It is also linked theoretically to the natural gradient method (Amari, 2000) via the Fisher information matrix (Martens, 2020, Soction 9.2).
Traditional ways to access curvature fall into two categories. Firstly, repeated automatic differentiation allows for matrix-free exact multiplication with the Hessian (Pearlmutter, 1994) and GGN (Schraudolph, 2002).
"Equal contribution

---

Published in Transactions on Machine Learning Research (02/2023)
Iterative linear and eigensolvers can leverage such functionality to compute Newton steps (Martens, 2010; Zhang et al., 2017; Gargiani et al., 2020) and spectral properties (Sagun et al., 2017; 2018; Adams et al., 2018; Ghorbani et al., 2019; Papyan, 2019b; Yao et al., 2019; Granziol et al., 2021) on arbitrary architectures thanks to the gencrality of automatic differentiation. However, repeated matrix-vector products are potentially detrimental to performance.
Secondly, K-FAC (Kronecker-factorrs approximate curvature) (Martens & Grosse, 2015; Grosse & Martens, 2016; Botev et al., 2017; Martens et al., 2018) constructs explicit light-weight representation of the GGN based on its algebraic Kronecker structure. The computations are streamlined via gradient backpropagation and the resulting matrices are cheap to store and invert. This allows K-FAC to scale: It has been used successfully with large mini-batches (Osawa et al., 2019). One reason for this efficiency is that K-FAC only approximates the GGN's block diagonal, neglecting interactions across layers. Such terms could be useful, however, for applications like uncertainty quantification with Laplace approximations (Ritter et al., 201 2b;a;;a; Kristiadi et al., 2020; Daxberger et al., 2021) that currently rely on K-FAC. Morcover, due to its specific design for optimization, the Kronecker representation does not become more accurate with more data. It. remains a simplification, exact only under assumptions unlikely to be met in practice (Martens & Grosse, 2015). This might be downside for applications that depend on precise curvature proxy.
Here, we propose VIVIT (inspired by VVT in Equation (3)), a curvature model that leverages the GGN's low-rank structure. Like K-FAC, its representation is computed in parallel with gradients. But it allows a cost-accuracy trade-off, ranging from the eract GGN to an approximation that has the cost of a single gradient computation. Our contributions are as follows:
We present how to compute various GGN properties efficiently by exploiting its low-rank structure: The exact eigenvalues, eigenvectors, and per-sample directional derivatives (Figure 1). In contrast other methods, these quantitics allow modeling curvature noisc
We introduce approximations that allow a flexible trade-off between computational cost and accuracy, and provide fully-fcatured efficient implementation in PYToRCH (Paszke et al., 2019) on top of the BACKPACK (Dangel et al., 2020) package al https: wtwpt//wbn.coootttaaaee////iiiii
We empirically demonstrate scalability and efficiency of leveraging the GGN's low-rank structure through benchmarks on different. deep architectures. Finally, we use the previously inaccessible properties of the GGN to study how is affected by noise during training.
The main focus of this work demonstrating that many interesting curvature properties, including uncertainty, can be computed efficiently. Practical applications are discussed in Section
2 Notation & Method
X Y. The network, parameterized by 0 O, maps sample æ to a prediction ÿ f(0, x). Predictions are scored by convex loss function f Y x Y + I. (e.g. cross-entropy or square loss), which compares to the ground truth y. The training objective is the empirical risk 1 L 6:
Setting: Consider a model f O x X + Y and a dataset D of tuples (x,
[D] E(ED)e ((f y)) In the following, we consider the loss
((0) NN EN {(((0, (.)),n)).
(1)
evaluated on a mini-batch {(0,, yn) E X x Y C D with N samples. We use (0) = E(f(O, xx), ynd and (0) = f00, an) for per-sample losses and predictions. For gradients, we write gn (0) Vol, (0) and ( (O)( suppressing 0 if unambiguous. We also set θ = R and Y = g with D, C the model parameter and prediction space dimension, respectively. For classification, C is the number of classes
Hessian & GGN: Two-fold chain rule application to the split to f decomposes the Hossian of Equation (1) into two parts ())( = G(Ø) + R(0) e IR DxD, the positive semi-definite GGN
G ト N (Jofn) T (V7 (n) (Jofn) NN EN G.
(2)
2

---

Published in Transactions on Machine Learning Research (02/2023)
(a)
mb, exact
sub, exact
mb, me
(b)
C
10
density
10-6
0:
ek
Tke Ar
"ink Ank
10-3
10
10a
10
10-
10
10 10
10*
cigenvalues
Figure 1: Overview of ViVIT's quantities: (a) GGN eigenvalue distribution of DEEPOBS 3c3D architecture CIFAR-10 (Schneider et al., 2019) for settings with different costs mini-batch of size = 128. From left to right: Exact GGN, exact GGN mini-batch fraction, MC approximation of the GGN. (b) Pictorial illustration of VIVIT's quantities in the loss landscape: the contour lines visualize the loss function & (Equation (1)) and its quadratic model q around O, E I² (Equation (6)). The low-rank structure provides efficient access to the GGN's eigenvectors {ex}. The quadratic model's one-dimensional projections along the eigenvectors (black axes) are parabolas characterized by the directional derivatives Tie AE and their per-sample contributions Yaki Ank (Equation (8)). E the GGN's top-1 cigenspace.
and a residual R= I/N E (V fnl.- Here, use the Jacobian Jab that contains partial derivatives of b with respect to a, [Jably = / As the residual may alter the Hessian's definiteness undesirable in many applications - we focus on the GGN.
Low-rank structure: By basic inequalities, Equation (2) has rank(G) S NC. To make this explicit, we factorize the positive semi-definite Hessian V ln = Ee=r 8ne where 8nc E RC and denote its backpropagated version by Onc (Jofn)
TT snc D.D. Absorbing sums into matrix multiplications, we arrive at the GGN's outer product form that lies at the heart of the VIVIT concept,
G= 1 I Ee UncU = VVT with V vN (111 .1....., UNC) E DxNN
(3)
V allows for evact computations with the explicit GGN matrix, at linear rather than quadratic memory cost in D. We first formulate the extraction of relevant GGN properties from this factorization, before addressing how to further approximate V to reduce memory and computation costs.
2.1 Computing the full GGN eigenspectrum
Each GGN eigenvalue À E Rzo is a rool of the characteristic polynomial det(G AIn) with identity matrix ID E DDDD, Leveraging the factorization of Equation (3) and the matrix determinant lemma, the D-dimensional cigenproblen reduces to that the much smaller Gram matrix G = VVV NCXNC which contains pairwise scalar products of Une (see Appendix A.1),
det -
-ING- =0
(4)
With at least D NC trivial solutions, the GGN curvature is zero along most. directions in parameter space. Nontrivial solutions that give rise to curved directions are fully-contained in the Gram matrix, and hence much cheaper to compute.
Despite various Hessian spectral studies which rely on iterative eigensolvers and implicit matrix multiplication (Sagun et al., 2017; 2018; Adams et al., 2018; Ghorbani et al., 2019; Papyan, 2019b; Yao et al., 2019; Granziol et al., 2021), we are not aware works that efficiently extract the exact GGN spectrum from its Gram
We assume the overparameterized deep learning setting (NC D) and suppress the trivial rank bound D.
3

---

Published in Transactions on Machine Learning Research (02/2023)
matrix. In contrast to those techniques, this matrix can be computed in parallel with gradients in single backward pass, which results in less sequential overhead. We demonstrate in Section 3.1 that exploiting the low-rank structure for computing the leading eigenpairs is superior to power iteration based on matrix-free multiplication in terms of runtime.
Figenvalues themselves can help identify reasonable hyperparameterrs like learning rates (LeCun et al,, 1993). But we can also reconstruct the associated eigenvectors. These are directions along which curvature information is contained in the mini-batch. Let S = {(Ax: @x) ) # 0, Gēg x denote the nontrivial Gram spectrum" with orthonormal eigenvectors éj ē₁ dzk (ò represents the Kronecker delta and K = (G)) Then, the transformed vectors e₁ = 1/v Vê (k: = 1, ..... K) are orthonormal cigenvectors of G associated to eigenvalues A (see Appendix A.2), i.e.
A(),, ōk) E S
Ge Ag
(5)
The eigenspectrum also provides access to the GGN's pseudo-inverse based on V and S+, required by e.g. second-order m
2.2 Computing directional derivatives
Various algorithms rely on local quadratic approximation of the loss landscape. For instance, optimization methods adapt their parameters by stepping into the minimum of the local praxy. Curvature, in the form of the Hessian or GGN, allows to build a quadratic model given by the Taylor expansion. Let a denote the quadratic model for the loss around position O, that uses curvature represented by the GGN,
(()) const (0 0.) (0.) ( 0.) G())) 0e)
At its base point 0,, the shape of q along an arbitrary normalized direction e θ (i.e. ||e|| 1) is determined by the local gradient and curvature. Specifically, the projection of Equation (6) onto e gives rise to the (scalar) first-and second-order directional derivatives
Te Vog(((.)) e ((.) E and = Ve
As G's characteristic directions are its eigenvectors, they form a natural basis for the quadratic model. Denoting Thk Tex and A = Ae. the directional gradient and curval ure along eigenvector ék, we see from Equation (7) that the directional curvature indeed coincides with the GGN's eigenvalue.
Analogous to the gradient and GGN, the directional derivatives 7 and A inherit the loss function's sum structure (Equation (1)), iee they decompose into contributions from individual samples. Let Yuk and Ank denote these first- and second-order derivatives contributions of sample æ, in direction k
Pink = e gv.
gr
V G
gn.
and Ant = e Gnex
Veell n
(8)
Vdx
A
where V, E RxC a scaled sub-matrix of V with fixed sample index. Note that directional derivatives can be evaluated efficiently with the Gram matrix eigenvectors without explicit access to the associated directions in parameter space.
In Equation (7): gradient g and curvature G are sums over 9n and G., respectively. This implies the relationships between directional derivatives and per-sample contributions 7h I/N En= 7nk and As I/N Σ Ank Figure 1b shows a pictorial view of the quantities provided by VIVIT. Access to per-sample directional gradients Tik and curvatures Ank along G's natural directions one distinct feature of VIVIT. These quantities provide geometric information about the local loss landscape as well as about the model's directional curvature stochasticity over the mini-batch.
2in the following,
ordered cigenvalues, ie. \ >o 2... ... 2 Ax. for convenience.
3 AAppendix C.2 descrihes implicit multiplication with G-I
4

---

Published in Transactions on Machine Learning Research (02/2023)
2.3 Computational complexity
So far, we have formulated the computation of the GGN's eigenvalues (Equation (4)), including eigenvectors (Equation (5)), and per-sample directional derivatives (Equation (8)). We now analyze their computational complexity to identify critical performance factors. Those limitations can effectively be addressed with approximations that allow the costs to be decreased in a fine-grained fashion. We substantiate our theoretical analysis with empirical benchmarks in Section 3.1.
Relation to gradient computation: Machine learning libraries are optimized to backpropagate signals V/VV, f, bn and accumulate the result into the mini-batch gradient g = 1/N >x (Jofn) V,, bx- Each column Une of V also involves applying the Jacobian, but to different vector 8nc from the loss Hessian's symmetric lactorization. For popular loss functions, like square and cross-entropy loss, this factorization is analytically known and available at negligible overhead. Hence, computing V basically costs C' gradient computations it involves NC' backpropagations, while the gradient requires N. However, the practical overhead is expected to be smaller: Computations can re-use information from BACKPACK's vectorized Jacobians and enjoy additional speedup on GPUs.
The columns of V correspond to backpropagated vectors. During backpropa- gation, sub-matricos of V, associated to parameters in the current layer, become available once at a time and can be discarded immediately after their use. This allows for memory savings without any approximations.
Stage-wise discarding V:
One example is the Gram matrix G formed by pairwise scalar products of (Une) in (((NC))D) operations. The spectral decomposition S has additional cost of (((NC)*). Similarly, the terms for the directional derivatives in Equation (8) can be built up stage-wise: First-order derivatives {Tre) N.K the vectors (VT 9n E NNN |
==1,k== require m==1 that cost O(NCDD operations. Second-order derivatives are basically for
free, as (V, V R RRNNC
In=1 is available from G.
GGN eigenvectors: Transforming a Gram matrix eigenvector ē. to the GGN eigenvector es by application of V (Equation (5)) costs O(NCD) operations. However, repeated application of V can be avoided for sums of the form Σn (<//VA))e with arbitrary weights e E RR The summation can be performed in the Gram space at negligible overhead, and only the resulting vector Σx cgêR needs to be transformed. For practical example- - computing damped Newton steps see Appendix B.1.
2.4 Approximations & Implementation
Although the GGN's representation by has linear memory cost in D, it requires memory equivalent to NC model copies.* 4 Of course, this infeasible for many networks and data sets, e.g. IMAGENET (C 1000). So far, our formulation was concerned with exact computations. We now present approximations that allow N, C and D in the above cost analysis to be replaced by smaller numbers, enabling VIVIT to trade-off accuracy and performance.
MC approxination & curvature sub-sampling: To reduce the scaling in C, we can approxinate the factorization V & (0) Σe= 8ncST ne by a smaller set of vectors. One principled approach is to draw MC samples (8nm) such that Em [T]] V2, (0) as in Dangel al. (2020). This reduces the scaling of backpropagated vectors from C to the number of MC samples M (M = 1 in the following if not specified otherwise). A common independent approximation to reduce the scaling in N is computing curvature mini-batch subset (Byrd et al., 2011; Zhang et al., 2017).
Parameter groups (block-diagonal approximation): Some applications, e.g. computing Newton steps, require V to be kept in memory for performing the transformation from Gram space into the parameter
'Our implementation uses more memory-efficinn approach that avoids expanding V for linear layers by leveraging structure in their Jacobian. We describe additional optimizations in Appendix C.1, and demoustrate 50x speed-up for computing the Gram matrix over naive computation via vectorized Jacobians. The additional backpropagations are carried in 50 % of the expected time due to parallelism on the GPU.
5

---

Published in Transactions on Machine Learning Research (02/2023)
space. Still, can reduce costs by using the GGN's diagonal blocks (G 1a=1 of each layer, rather than the full matrix G. Such blocks are available during backpropagation and can thus be used and discarded step by step. In addition to the previously described approximations for reducing the costs in N and C this technique tackles scaling in D.
Implementation details: BACKPACK's functionality allows us to efficiently compute individual gradients and V a single backward pass, using either an exact or MC-factorization of the loss Hessian. To reduce memory consumption, we extend its implementation with a protocol to support mini-bater sub-sampling and parameter groups. By hooks into the package's 's extensions, we can discard buffers soon possible during backpropagation, effectively implementing all discussed approximations and optimizations.
Next, we specifically address how the above approximations affect runtime and memory requirements, and study their impact on structural properties of the GGN.
3 Experiments
For the practical use of the VIVIT concept, it is essential that (i) the computations are efficient and (ii) that we gain an understanding of how sub-sampling noise and the approximations introduced in Section 2.4 alter the structural properties of the GGN. In the following, we therefore empirically investigate VIVIT's scalability and approximation properties in the context of deep learning, where it can serve as a monitoring tool of novel quantities that have not. been explored previously to analyze training and other phenomena. The code used for the experiments available https: htte:://gtttuu github..hh..... wtt.coomooofoe con/f-daagellvivit eetttttexxrimmmntt
Experimental setting: Architectures include three deep convolutional neural networks from DEEPOBS (Schneider ct al., 2019) (2c2D FASHION-MNIST, 3C3D on CIFAR-10 and ALL-CNN-C CIFAR-100), well residual networks from He et al. (2016) CIFAR-10 based on Idelbayev (2018) all are equipped with entss-yntropy loss. Based on the approximations presented in Section 2.4, we distinguish the following cases:
mb, exact: Exact GGN with all mini-batch samples. Backpropagates NC vectors.
mb, me: MC-approximated GGN with all mini-batch samples. Backpropagates NM vectors with M the number of MC-samples.
sub, exact: Exact GGN on subset of mini-batch samples (L*/s] as in Zhang et al. (2017)). Backpropagates [//]CC vectors.
sub, me: MC-approximated GGN on a subset of mini-batch samples. Backpropagates [ vectors with M the number of MC-samples.
3.1 Scalability
We complement the theoretical computational complexity analysis from Section 2.3 with empirical studies. Results were generated on workstation with an Intel Core i7-8700K CPU (32 GB) and one NVIDIA GeForce RTX 2080 Ti GPU (11 GB). We use M = 1 in the following.
Memory performance: We consider two tasks:
1. Computing eigenvalues: The nontrivial eigenvalues (Ax (Ne, @x) S4] are obtained by forming and eigen-decomposing the Gram matrix G, allowing stage-wise discarding of V (see Sections 2.1 and 2.3).
Computing the top eigenpair: For (Λ,, ei), we compute the Gram matrix spectrum $4, extract its top eigenpair (Λ₁, ē1), and transform into parameter space by Equation (5), i.c. (A1,e1 e1 = VV This requires more memory than task V must be stored.

---

Published in Transactions on Machine Learning Research (02/2023)
As comprehensive measure for memory performance, use the largest batch size before our system runs out of memory- we call this the critical batch size Nerit-
Figure 2a tabularizes the critical batch sizes on GPU for the 3C3D architecture on CIFAR-10. As expected, computing eigenpairs requires more memory and leads to consistently smaller critical batch sizes in comparison to computing only eigenvalues. Yet, they all exceed the traditional batch size used for training (N 128, see Schneider et al. (2019)), even when using the exact GGN. With VIVrT's approximations, the memory overhead is reduced to significantly increase the applicable batch size. We report similar results for more architectures, a block-diagoone approximation (as in Zhang et al. (2017)), and on CPU in Appendix B.1, where we also benchmark third task - computing damped Newton steps.
In the large-batch regime, compuling the Gram matrix and its spectrum becomes a run time bot tleneck, see Section 2.3. As we show next, VIVIT is highly efficient in the mini-batch setting. However, for spectral analyses on the entire datasct, itcrative cigensolvers, like power itcrations, should be preferred.
Runtime performance: Here, we consider the task of computing the & leading eigenvoctors and eigenvalues of a matrix. A power iteration that computes eigempairs iteratively via matrix-vector products serves as a reference. For a fixed value of k, we repeat. both approaches 20 times and report the shortest time.
For the power iteration, we adapt the implementation from the PYHESSIAN library (Yao et al., 2019) and replace its Hessian-vector product by matrix-free GGN-vector product (Schraudolph, 2002) through PYTORCH's s automatic differentiation. We use the same default hyperparameters for the termination criterion." Similar to task 1, our method obtains the top-k eigenpairs e by computing 84, extracting its leading eigenpairs and transforming the eigenvectors θ1, E2...., ēg into parameter space by application of V.
Figure 2b shows the GPU runtime for the 3C3D architecture on CIFAR-10, using a mini-batch of size N = 128. Without any approximations to the GGN, our method already outperforms the power iteration for >> 1 and increases much slower in run time as more leading eigenpairs are requested. This means that, relative to the transformation of each eigenvector from the Gram space into the parameter space through V, the run time mainly results from computing V,G, and cigendecomposing the latter. This is consistent with the computational complexity of those operations in NC (compare Section 2.3) and allows for efficient extraction of a large number of eigenpairs. The run time curves of the approximations confirm this behavior by featuring the same flat profile. Additionally, they roquire significantly less time than the exact mini-batch computation. Appendix B.1 reports additional results for more architectures, a block-diagonal approximation, and on CPU.
3.2 Approximation quality
VIVIT is based on the Hessian's generalized Gauss-Newton approximation (see Equation (2)). In practice, the GGN is only computed on a mini-batch which yields a statistical estimator for the full-bolch GGN (i.e. the GGN evaluated on the entire training set). Additionally, we introduce curvature sub-sampling and an MC approximation (see Section 2.4), i.e. further approximations that alter the curvature's 's structural properties. In this section, we compare quantitics at different stages within this hicrarchy of approximations. We use the test problems from above and train the networks with both SGD and ADAM (details in Appendix B.2).
First, we empirically study the relationship between the GGN and the Hessian in the deep learning context. To capture solely the effect of neglecting the residual R (sce Equation (2)), we consider the noise-free case and compute H and G on the entire training set.
GGN vs. Hessian:
We characterize both curvature matrices by their top-C' eigenspace: the space spanned by the C leading cigenvectors. This is a C-dimensional subspace of the parameter space O, on which the loss function is subject to particularly strong curvature. The overlap between these spaces serves as the comparison metric. C Let _e 1-- be the set of orthonormal eigenvectors to the C' largest eigenvalues of some symmetric matrix U and :0 = span(ef ....,
ee. e.). The projection onto this subspace Eo is given by the projection matrix
We find similar results with relaxed convergence hyperparameters, see Appendix B.I. 6The power iteration computes the leading eigenpairs. Our approach allows choosing arbitrary eigenpairs.
7

---

Published in Transactions on Machine Learning Research (02/2023)
(a)
Nerit (eigenvalues)
(b)
power iteration mb, exact + sub, exact
iu,; mn
sub, me
Data
mb
sub
exact
909
4375
10"
InC
3840
6626
time
Nerin (top eigenpair)
10-1
Data
mb
sub
6GN
exact
677
3184
10
mc
3060
6029
1
2
4
8
10
top eigempairs (l))
Figure 2: GPU memory and run time performance: Performance measurements for the 3C3D architecture (D 895,210) on CIFAR-10 (C 10). (a) Critical batch sizes Nari for computing eigenvalues and the top eigenpair. (b) Run time comparison with a power iteration for extracting the k leading eigenpairs using a batch of size N 128.
pv = (e ... U (eF ... e²)T As in Gur-Ari et al. (2018), we define the overlap between two top-C' cigenspaces Er and gv the matrices U and V by
overlap(E" ,VY Tr(PUpt pV) VTr (PU) Tr (PV) E [0,
(9)
If overlap(E ,EV) then Eo and EV are orthogonal; if the overlap is 1, they are identical.
Figure 3a shows the overlap between the full-batch GGN and Hessian during training of the 3C3D network on CIFAR-10 with SGD. Except for short phase at the beginning of the optimization procedure (note the log scale for the epoch-axis), it shows a strong agreement (overlap 2 0.85) between the top-C eigenspaces. We make similar observations with the other test problems (see Appendix B.3), yet to a slightly lesser extent for CIFAR-100. Consequently, we identify the GGN as an interesting object, since it consistently shares relevant structure with the Hessian matrix.
Eigenspace under noise and approximations: VIVIT uses mini-batching to compute a statistical estimator of the full-batch GGN. This approximation alters the top-C' cigenspace, shown in Figure 3b: With decreasing mini-batch size, the approximation carries less and less structure of its full-batch counterpart, as indicated by dropping overlaps. In addition, at constant batch size, decrease in approximation quality can be observed over the course of training. This might be valuable insight for the design of second-ordes optimization methods, where this structural decay could lead to performance degradation over the course of the optimization, which has to be compensated for by growing batch-size (e.g. Martens (2010) reports that the optimal batch size grows during training).
In order to allow for fine-grained cost-accuracy trade-off, VIVIT introduces further approximations to the mini-batch GGN (see Section 2.4). Figure 3c shows the overlap between these GGN approximations and the full-batch GGN.7 The order of the approxinations is expected: With incrcasing computational effort, the opproximations improve and, despite the greatly reduced computational effort compared to the exact mini- bati- GGN, significant structure of the top-C' eigenspace is preserved. Details and results for the other test problems are reported Appendix B.4.
So far, our analysis is based on the top-C' eigenspace of the curvature matrices. We extend this analysis by studying the effect of noise and approximations the curvature magnilude along the top-C' directions in Appendix B.5.
7 comparison with the mini-batch GGN ground truth can found Appendix B.4
8

---

Published in Transactions on Machine Learning Research (02/2023)
(a) Full-batch GGN vs. full--atcc' Hessian
(b) Mini-batch GGN vs. full-batch GGN
mb 2, exact
mb exact
overlap
0.95
overlap
mb 128, exact
mb 32, exact
0.5
0.9
0.85
10a
10
10a
10
10'
10a
(c) Approximations full-batch GGN
(d)
Directional curvature SNRs
mb 128, exact
SNR[Ana]] (log scale)
sub 16, exact
10
128,
overlap
sub 16, me 1
0.5
10
10
10a
10a
10a
10
10a
10a
epoch (log scale)
epoch (log scale)
Figure 3: Curvature monitoring during training 3c3d CIFAR-10 with SGD: (a) Overlap between the top-C eigenspaces of the full-batch GGN and full-batch Hessian during training. the lop-C' eigenspaces of the mini-batch GGN and full-batch GGN during training. For each mini-batch size, 5 different
(b) Overlap between
are drawn. (c) Overlap between the top-C eigenspaces of the mini-batch GGN and
VIVIT's approximations with the full-batch GGN during training. Each approximation is evaluated on 5 mini-batches. (d) Curvature SNRs along each of the mini-batch GGN's top-C' eigenvectors during training. At fixed epoch, the SNR for the most curved direction is shown in and the SNR for the direction with the smallest curvature is shown in
3.3 Per-sample directional derivatives
A unique feature of VIVIT's quantities is that they provide notion of curoature uncerlainly through per- sample first- and second-order directional derivatives (see Equation (8)). To quantify noise in the directional derivatives, compute their signal-to-noise ratios (SNRs). For each direction ex, the SNR is given by the squared empirical mean divided by the empirical variance of the N mini-batch samples {he J. and [ respectively.
Figure shows curvature SNRs during training the 3C3D network on CIFAR-10 with SGD. The curvature signal along the top-C' eigenvectors decreases from SNR by two orders magnitude. This might be a challenging setting for second-order methods because the noise varies dramatically during different stages of training. In comparison, the directional gradients do not exhibit such pattern (see Appendix B.6). Results for the other test cases can be lound in Appendix B.6.
In this section, we gave a glimpse of the very rich quantities that are efficiently computed via VIVIT. Section 5 discusses the practical use of this information, in particular curvature uncertainty.
4 Related work
GGN spectrum & low-rank structure: Other works point out the GGN's low-rank structure. Botev et al. (2017) present the rank bound and propose alternative to K-FAC based backpropagating
9

---

Published in Transactions on Machine Learning Research (02/2023)
decomposition of the loss Hessian. Papyan (2019a) presents the factorization in Equation (3) and studies the eigenvalue spectrum's hierarchy for cross-entropy loss. In this setting, the GGN further decomposes into summands, some of which are then analyzed through similar Gram matrices. These can be obtained as contractions of G, but our approach goes beyond them as it does not neglect terms. We are not aware of works that oblain the exact spectrum and leverage a highly-efficient fully-parallel implementation. This may be because, until recently (Bradbury et al., 2020; Dangel al., 2020), vectorized Jacobians required to perform those operations efficiently were not available.
Efficient operations with large-scale matrices in deep learning: Chen et al. (2021) use Equation (3) for clement-wise evaluation of the GGN in fully-connected feod-forward neural networks. They also present a variant based on MC sampling. This element-wise evaluation is then used to construct hierarchical matrix approximations of the GGN. VIVIT instead leverages the global low-rank structure that also enjoys efficient. eigen-decomposition.
A special case of VIVIT's Gram matrix extraction is computing empirical neural tangent kernel (NTK) matrices, like Novak et al. (2022). While the NTK only requires a model (its Jacobian (Jacot et al., 2018)). the GGN also incorporates the loss function via its Hessian. For mean squared error, this Hessian in Equation (2) is proportional to the identity, and the GGN Gram matrix coincides with the scaled empirical NTK.
Another prominent low-rank matrix deep learning the un-centered gradient covariance (sometimes called empirical Fisher). Singh & Alistarh (2020) describe implicit multiplication with its inverse and apply it for neural network compression, assuming the empirical Fisher as Hessian proxy. However, this assumption has limitations, specifically for optimization (Kunstner et al., 2019). In principle though, the low-rank structure also permits the application of our methods from Scction 2.
5
Use cases
Our efficient implementation enables the community to explore deep learning through richer information that would previously have been costly. Here, we want to briefly address possible use cases their full development. and assessment, however, will amount to separate paper(s). They include:
Second-order optimization: Second-order methods use curvature to build local quadratic model of the loss. Established curvature proxies neglect the sampling-induced noise and therefore the quadratic model's reliability. VIVIT provides access to this information in the form of per-sample quantities. This offers a new dimension for improving second-order methods: Through statistics on the mini-batch distribution of directional derivatives, we might be able to adapt to the dynamics of noise (e.g. via variance-adapted step sizes).
Monitoring tool: However, to develop conceptually novel second-order optimizers, we believe that crucial intermediate step is to better understand the setting they operate in. The techniques we present primarily tackle this intermediate step. Sections 3.2 and 3.3 are examples of VIVIT's application as monitoring tool. Due to its efficiency, VIVIT could be integrated into live diagnostic tools like COCKPIT (Schneider et al., 2021) that aim at debugging optimizers or gaining insights into the inner workings of deep learning.
6 Conclusion
We have presented VIVIT, curvature model based on the low-rank structure of the Hessian's generalized Causs-Newton (GGN) approximation. This structure allows for efficient extraction of eract curvature properties, such as the GGN's full cigenvalue spectrum and directional gradients and curvatures along the associated eigenvectors. VIVIT's quantities scale by approximations that allow for a fine-grained cost-accuracy trade-off. In contrast to alternatives, these quantities itis offer a notion of curvature uncertainty across the mini-batch in the form of directional derivatives.
We empirically demonstrated the efficiency of leveraging the GGN's low-rank structure and substantiated its usefulness by studying characteristies of curvature noise on various deep learning architectures. The low-rank
10

---

Published in Transactions on Machine Learning Research (02/2023)
representation is efficiently computed in parallel with gradients during single backward pass. As it mainly relies vectorized Jacobians, it is general enough to be integrated into existing machine learning libraries in the future. For now, we provide an efficient open-source implementation in PYTORCH (Paszke et al,, 2019) by extending the existing BACKPACK (Dangel et al., 2020) library.
Acknowledgments
The authors gratefully acknowledge financial support by the European Research Council through ERC StG Action 757275 PANAMA; the DFG Cluster of Excellence "Machine Learning New Perspectives for Science", EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 011S18039A); and funds from the Ministry Science, Rescarch and Arts of the State of maden-Württemberg. Morcover, the authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Felix Dangel and Lukas Tatzel. Further, grateful to Agustinus Kristiadi, Filip de Roos, Frank Schneider, Jonathan Wenger, and Marius Hobbhahn for providing feedback to the manuscript.
11

---

Published in Transactions on Machine Learning Research (02/2023)
References
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, lly Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015.
Ryan P. Adams, Jeffrey Pennington, Matthew J. Johnson, Jamie Smith, Yaniv Ovadia, Brian Patton, and James Saunderson. Estimating the spectral density of large implicit matrices, 2018.
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 2000.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Ganss-Newton optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning, 2017.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne JAX: composable transformations Python + NumPy programs. 2020.
Richard H. Byrd, Gillian M. Chin, Will Neveitt, and Jorge Nocedal. On the use of stochastic Hessian information in optimization methods for machine learning. SIAM Journal on Optimiiaation, 2011.
Chao Chen, Severin Reiz, Chenhan D. Yu, Hans-Jouchim Bungartz, and George Biros. Fast approximnation tion of the Gauss Newton Hessian matrix for the multilayer perceplron. SIAM Journal on Matrix Analysis and Applications, 2021.
Felix Dangel, Frederik Kunstner, and Philipp Hennig, BackPACK: Packing more into backprop. In International Conference on Learning Representations, 2020.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless Bayesian deep learning. In NeurIPS, 2021.
Matilde Gargiani, Andrea Zanelli, Moritz Diehl, and Frank Hutter. On the promise of the stochastic generalized Gauss-Newton method for training DNNs, 2020.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via Hessian eigenvalue density, 2019.
Diego Granziol, Xingchen Wan, and Timur Garipov. Deep curvature suite, 2021.
Roger Grosse and James Martens. A Kronecker-factoree approximate Fisher matrix for convolution layers, 2016.
Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Yerlan Idelbayev. Proper ResNet implementation for CNFARI/CIFARIIO in Py'orrch. https://github... /tom http://pttorcc__eee_t_iiciiei
2018
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 2018.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being Bayesian, even just a bit, fixes overconfidence ReLU networks. In International Conference Machine Learning, 2020.
Frederik Kunstner, Philipp Hemmig, and Lukas Balles. Limitations of the empirical Fisher approximation for natural gradient descent. In Advances in Neural Information. Processing Systems, 2019.
Yann LeCun, Patrice Simard, and Barak Pearlmutter. Automatic learning rate maximization by on-line estimation of the hessian's eigenvectors. In S. Hanson, J. Cowan, and C. Giles (eds.), Advances in Newral Information Processing Systems, volume 5.
1993.
12

---

Published in Transactions on Machine Learning Research (02/2023)
James Martens. Deep learning via Hessian-free oplimization. In Proceedings of the 27th International Conference on Machine Learning, 2010.
James Martens, New insights and perspectives on the natural gradient method, 2020.
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In Proceedings of the 32nd International Conference on Machine Learning, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factoree curvature approximations for recurrent neural networks. In International Conference Learning Hepresenlations, 2018.
Roman Novak, Juscha Sohl-Dickstein, and Samuel S. Schoenholz. Fast finite width neural langent kernel. In International Conference on Machine Learning, 2022.
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order oplimization using Kronecker-factoreed approximate curvature for deep convolutional neural networks. In 2019 IEEE/CVF Conference on Compuler Vision and Pallern Recognition, 2019.
Vardan Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet Hessians. In Proceedings of the 36th International Conference on Machine Learning, 2019a.
Vardan Papyan. The full spectrum of deepnet Hessians at scale: Dynamics with SGD training and sample size, 2019b.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. Advonces in Neural Information Processing Systems 32. 2019.
Barak A. Pearlmutter. Fast exact multiplication by the Hessian. Newral Compulation, 1994.
Hippolyt Ritter, Aleksandar Botev, and David Barber. Online structured Laplace approximations for overcoming catastrophic forgetting. In Advances in Neural Information Processing Systems, 2018a.
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable Laplace approximation for neural networks. In International Conference on Learning Representations, 2018b.
Levent Sagun, Leon Botton, and Yann LeCun, Figenvaluus of the Hessian deep learning: Singularity and beyond, 2017.
Levent Sagun, Utku Evci, V. Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the Hessian of over-parametrized neural networks, 2018.
Robin M. Schmidt, Frank Schneider, and Philipp Hennig. Descending through a crowded valley benchmarking deep learning optimizers, 2021.
F. Schneider, L. Balles, and Hennig. DeepOBS: deep learning optimizer benchmark suite. 7th International Conference on Learning Representations, 2019.
Frank Schneider, Felix Dangel, and Philipp Hennig. Cockpit: A practical debugging tool for training deep neural networks, 2021.
Nicol Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 2002.
Sidak Pal Singh and Dan Alistarh. WoodFisher: Elficient second-order approximation for neural network compression. In Advances in Neural Information Processing Systems, 2020.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael Mahoney, PyHessian: Neural networks through the lens of the Hessian, 2019.
Huishuai Zhang, Caiming Xiong, James Bradbury, and Richard Socher. Block-diagonal Hessian-free optimization for training neural networks, 2017.
13

---

Published in Transactions on Machine Learning Research (02/2023)
ViViT: Curvature Access Through The Generalized Gauss-Newton's Low-Rank Structure (Supplementary Material)
A Mathematical details
A.1 Reducing the GGN eigenvalue problem the Gram matrix
A.2 Relation between GGN and Gram matrix eigenvectors
15
B Experimental details
15
B.1 Performance evaluation
17
B.2 Training of neural networks
32
B.3 GGN vs. Hessian..
32
B.4 Eigenspace under noise and approximations
35
B.5 Curvature under noise and approximations
39
B.6 Directional derivatives
41
C Implementation details
41
C.1 Optimized Gram matrix computation for linear layers
45
C.2 Implicit multiplication with the inverse (block-diagonal) GGN
46
A Mathematical details
A.1 Reducing the GGN eigenvalue problem the Gram matrix
For Equation (4), consider the left hand side of the GGN's characteristic polynomial det - AD)) = 0. Inserting the VIVIT factorization (Equation (3)) and using the matrix determinant lemma yields
det( A AI G) det( Beet--i +VV)
(Low-rank structure (3))
det (INC + V (-AId) 'V) det(->Ip)
(Matrix determinant lemma)
det INC - 1 - (-1)0
NC
V (-x)
(-X) det (G Alne).
(Gram matrix)
Setting the above expression zero reveals that the GGN's spectrum decomposes into NC zero eigenvalues and the Gram matrix spectrum obtained from det XNne) 0.
14

---

Published in Transactions on Machine Learning Research (02/2023)
A.2 Relation between GGN and Gram matrix eigenvectors
Assume the nontrivial Gram matrix spectrum 5 {(A,, ēx) A + 0, Gé, λςēx with orthonormal eigenvectors e 0 Oz* (6 represents the Kronecker delta) and K rank(G).. We now show that e₁ VVv are normalized eigenvectors of G and inherit orthogonality from ēμ.
To see the first, consider right-mm multipliiiii of the GGN with then expand the low-rank structure,
1 VV VV
Ge
(Equation (3) and definition of ex)
Vx
1
(Gram matrix)
Vx
A Vēk
(Eigenvector properly of ēy))
V
Orthonormality of the results from the Gram matrix eigenvector orthonormality,
e = _
(
Viv
Vē
(Definition of ej, ex)
Gē
(Gram matrix)
Vada
0/
(Eigenvector property of ox)
Vdzaa
6zk.
(Orthonormality)
B
Experimental details
Throughout this section, we the notation introduced in Section (see Table S.1). The code used for the experiments is available at https://gi wtw.ttcoo.oooc..ttttttttttt..ttt hwmmmfdddnnel/viviiee aaprra/ttaaiiiitttttererrrttttte
Table S.1: Notation for curvature approximations: The notation is introduced in Section 3. This table
recapitulates the abbreviations (referring to the approximations introduced in Section 2.4) and provides corresponding explanations.
Abbreviation
Explanation
mb, exact
Exact GGN with all mini-batch samples. Backpropagates NC vectors.
mb, mc
MC-approximated GGN with all mini-batch samples. Backpropagates NM vectors with M the number of MC-samples.
sub, exact sub, me
Exact GGN on subset of mini-batch samples (Lx//] in Zhang et al. (2017)).
Backpropagates | [N/s]C' vectors.
MC-approximated GGN on subset of mini-batch samples.
Backpropagates MM/J M vectors with M the number of MC-samples.
GGN spectra (Figure 1a): To obtain the spectra Figure la we initialize the respective architecture, then draw a mini-batch and evaluate the GGN cigenvalues under the described approximations, clipping the Gram matrix eigenvalues at 10-4, Figures S.4 and S.5 provide the spectra for all used architectures with both the full GGN and per-layer block-diagonal approximation.
15

---

Published in Transactions on Machine Learning Research (02/2023)
Fashion-MNIST 2c2d
Full network
Block-diagonal approximation
mb, exact
sub, exact
mb, me
mb, exact
sub, exact
mb, mc
10
10
density
density
10
10
10 6
10
10
10-1
10
10 10
10
10
10
10
10
3
10
10 10
10
10
10
10
10
eigenvalues
eigenvalues
CIFAR-10 3c3d
Full network
Block-diagonal approximation
mb, exact
sub, exact
mb, mc
sub, exact mu,, exact
mu,
10
10-
density
density
10
10
10
1
10
10
10 10
10t
10-
10
10a
10
3
10
10 3
10
10
10
eigenvalues
eigenvalues
CIFAR-10 ResNet-32
Full network
Block-diagonal approximation
mb, exact
sub, exact
mb, me
mb; exact
sub, exact
mb,
10
10
density
density
10
10
10
10 10-310
10
10 10 310-10-1
10'
10
10
10
10-3
10-1
10-3
10-1
10 10
eigenvalues
eigenvalues
Figure S.4: GGN spectra of different architectures under ViVIT's approxinations: Left and right columns contain results with the full network's GGN and per-layer block-diagonal approximation, respectively. The column labels are explained Table S.I.
16

---

Published in Transactions on Machine Learning Research (02/2023)
CIFAR-10 ResNet-56
Full network
Block-diagonal approximation
mb, exact
sub, exact
mb,
mb, exact
sub, exact
mb, mc
10
10
density
density
10
10
10
10-4 10-1
10º10-1 10-1
10"10-4 10-1
10a
10
10-1
I10
10-1
ID
10-1
10a
eigenvalues
eigenvalues
CIFAR-100 AII-CNN-C
Full network
Block-diagonal approximation
mb, exact
sub, exact
mb, mc
mu, exact
sub, exact
mu, Inc
10-3
10
10-4
density
density
10
10
10
10
10
10
10
10 1
10
10
10
10
10a
10
10
10
10
10 3
10
eigenvalues
eigenvalues
Figure S.5: GGN spectra of different architectures under VIVIT's approximations: Left and right columns contain results with the full network's GGN and per-layer block-diagonal approximation, respectively. The column labels explained in Table S.1.
B.1 Performance evaluation
Hardware specifications: Results in this section were generated on a workstation with an Intel Core 17-8700K CPU (32 GB) and one NVIDIA GeForce RTX 2080 GPU GBIG)..
Note: VIVIT's quantities are implemented through BACKPACK, which triggered by PYTORCH's 's gradient. computation. Consequently, they can only be computed together with PYTORCH's
gradient.
Architectures: We untrained deep convolutional and residual networks from DEEPOBS Schneider et al. (2019) and Idelbayev (2018). If net has batch normalization layers, we set them to evaluation mode. Otherwise, the loss would not obey the sum structure of Equation (1). The batch normalization layers' internal moving averages, required for evaluation mode, are initialized by performing five forward passes with the current mini-batch in training mode before.
In experiments with fixed mini-batches the batch sizes correspond to DEEPOBS' default value for training where possible (CIFAR-10: N 128, PASHION-MNIST: N 128). The residual networks batch size
17

---

Published in Transactions on Machine Learning Research (02/2023)
of N = 128. On CIFAR-100 (trained with N 256), we reduce the batch size to N = 64 to fit the exact computation on the full mini-batch, used as baseline, into memory. If the GGN approximation is evaluated on a subset of the approximation is always evaluated with a single sample (M = 1).
(sub), [*/a] of the samples are used (as in Zhang et al. (2017)). The MC
Memory performance (critical batch sizes): Two tasks are considered (see Section 3.1):
1. Computing eigenvalues: Compute the nontrivial eigenvalues {Ax (Ae, ēx) E S+}
2. Computing the top eigenpair: Compute the top eigenpair (11,
We repeat the tasks above and vary the mini-batch size until the device runs out of memory. The largest mini-batch size that can be handled by our system is denoted as Nerit, the critical batch size. We determine this number by bisection on the interval [1; 32768).
Figures S.8 to S.17a,b present the results. As described in Section 2.3, computing cigenvalues more memory-effieiett than computing eigenvectors and exhibits larger critical batch sizes. In line with the description in Section 2.4, a block-diagonal approximation is usually more memory-efficiett and results in larger critical batch size. Curvature sub-sampling and MC approximation further increase the applicable batch sizes.
In summary, we find that there always exists a combination of approximations which allows for critical batch sizes larger than the traditional size used for training (some architectures even permit exact computation). Different accuracy-cost trade-offs may be preferred, depending on the application and the computational budget. By the presented approximations, VIVIT's representation is capable to adapt over a wide range.
Runtime performance: Here, we consider the lask of computing the k leading eigenvectors and eigenvalues of matrix. VIVIT's eigenpair computation is compared with power iteration that computes eigenpairs iteratively via matrix-vector products. The power iteration bascline is based on the PYHESSIAN library Yao et al. (2019) and uses the same termination criterion (at most 100 matrix-vector products per eigenvalue; stop if the eigenvalue estimate's relative change is less than 10 10-3). In contrast to PYHESSIAN, we use a different data format and stack the computed eigenvectors. This reduces the number of for-loops in the orthonormalization step. We repeat each run time measurement 20 times and report the shortest exocution time result.
Figures S.8 to S.17c,d show the results. For most architectures, our exact method outperforms the power iteration for & > 1 and increases only marginally in runtime as the number of requested eigenvectors grows. The proposed approximations share this property, and further reduce run time.
Power iteration with relaxed hyperparameters: When VIVIT's approximations are used for computing eigenvalues, a power iteration with relaxed hyperparameters might be an alternative. We thus extend the runtime experiment from Soction 3.1 in the following way: We consider the task of computing the top-10 eigenpairs of the mini-batch GGN on CIFAR-10 10 3C3D, with the same batch size N = 128, on GPU using both VIVIT and the power iteration. The resulting approximations of the eigenvalues dx k I,..., .., 10 are compared with the exact cigenvalues d k = 1, 10 using
A Anl
10
A
as mcasure of accuracy. We run the power iteration with 20 different convergence tolerances varying on a logarithmic grid from 10-5 to 10-1, and disable termination due to exceeded iterations. Each setting is repeated 20 times, and we report the best run time and corresponding accuracy in Figure S.6.
In this setting, computing the exact top-10 eigenvalues with VIVIT is faster than approximating them with a power iteration, even when using the largest tested tolerance. Also, when using VIVIT with approximations through sub-sampling and MC-approximation, these runs require less time than the power iteration
18

---

Published in Transactions on Machine Learning Research (02/2023)
mean accoracy
0.9
mb, exact
mb, me
sub, exact
0.8
sub, mc
power iteration
0
0.5
1.5
2.5
3
time [s]
Figure S.6: Accuracy per run time performance: Comparison of VIVIT and its approximations to the power iteration for computing the top-10 eigenvalues for the 3C3D architecture (D 895,210) on CIFAR-10 (C 10). For the power ileration, different tolerances ranging from 10-5 to 10-1 are used. The red star shows the result for the default tolerance.
at similar accuracy. In the plot, we also highlighted the power iteration run with PYHESSIAN's default convergence parameters (red star marker). seems to yield a relatively good trade-off between time and accuracy for the power iteration.
These results show that, also in terms of "accuracy per run time", VIVIT is superior to power iteration in the mini-batch setting due to its increased parallelism.
Power iteration on the GGN versus power ileration on the GGN Grom matrin: In the run time evaluation of our method (Section 3.1) we compute the k leading GGN eigenpairs by computing the full GGN Gram matrix spectrum and discarding all but the leading cigenpairs. Here, we present additional results where we exchange the full diagonalization by a power iteration with identical convergence hyperparameters as the baseline; a power iteration on the GGN.
Figure S.7 visualizes the comparison for the same sctting as Figure 2b in the main text. In case of no approximations (mb, exact) where the Gram matrix dimension is largest, the power iteration can further reduce the time shown in Figure 2b. However, for the GGN approximations through sub-sampling MC approximation, the power iteration the (rather small) Gram matrix, deteriorates performance in comparison to the results reported in Figure 2b as the number of leading eigenvalues increases. In this regime (small Gram matrix, many requested eigenvalues), the simplistic power iteration can require more matrix-vector products than a sophisticated cigensolver that computes the full spectrum. As in the main text, these findings show that VIVIT (even in the exact version) outperforms the power iteration for k > 2.
Note on CIFAR-100 (large C): For data sets with a large number of classes, like CIFAR-100 (C 100), computations with the exact GGN are costly. In particular, constructing the Gram matrix G has quadratic memory cost in C, and its eigendecomposition has cubic cost time with C (see Section 2.3).
As result, the exact compulation only works with batch sizes smaller than DEEPOBS' default (N 256 for CIFAR-100, see Figures S.16 and S.17a.b). For the GGN block-diagonal approximation, which fits into CPU memory for N = 64, the exact computation of top eigenpairs is slower than a power itcration and only becomes comparable if large nunber of eigempairs is requested, see Figure S.17d.
For such data sets, the approximations proposed in Section 2.4 are essential to reduce costs. The most effective approximation to climinate the scaling with is using an MC approximation. Figures S.16 and S.17 confirm that the approximate computations scale batch sizes used for training and that computing eigenpairs takes less time than power iteration.
19

---

Published in Transactions on Machine Learning Research (02/2023)
GGN power iteration
(iub, exact) + power
(sub, exact)
power
(mb, mc)
power
(sub, mc)
power
10
time
10-1
10
2
3
4
6
7
9
10
top eigenpairs (k)
Figure S.7: Power iteration on the GGN versus power iteration on the Gram matrix: The figure considers the same setting as Figure 2b (3C3D on CIFAR-10 and GPU). However, we exchange the full diagonalization of the GGN's Gram matrix by a power iteration with identical convergence hyperparameters as the baseline power iteration on the GGN (at most 100 matrix-vector producst per eigenvalue; stop if the eigenvalue estimate's relative change less than 10-3).
Computing damped Newton steps: A Newton step -GG + 81)-19 with damping 0 can be decomposed into updates along the cigenvectors of the GGN G,
K
-GG + 01) -1)
EN
Σ k=K-1
(S.10)
+
k=1
It corresponds to Newton update along nontrivial eigendirections that uses the first- and second-order directional derivatives described in Section 2.2 and gradient descent step with learning rate 1/8 along trivial directions (with X 0). In the following, we refer to the first summand of Equation (S.10) as Newton step. As described in Section 2.3, we can perform the weighted sum in the Gram matrix space, rather than the parameter space, by computing
K
K Th
K
Σ
EK
Σ k=1
Vē = V
Σ
7kk
Ae +
Vxx
(h+ σ)
This way, only a single vector needs to be transformed from Gram space into parameter space.
Table S.2 shows the critical batch sizes for the Newton step computation (first term on the right-hand side of Equation (S.10)), using Gram matrix eigenvalues larger than 10-4 and constant damping directional derivatives d evaluated the samples as the GGN eigenvectors, but we always all mini-batch samples to compute the directional gradients "n. Using approximations, the Newton step computation scales to batch sizes beyond the traditional sizes used for training.
Second-order
20

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a) Next (eigenvalues)
(b) Nent (eigenvalues)
Nexit (top eigenpair)
Naxi (top eigenpair)
GGN
Data
sub
GGV
Data
BGN
Data
sub
GGN
Data
exact
1753
8235
exact
872
6439
exact
1012
8317
exact
993
7179
7585
12434
6708
12162
8448
12455
8336
12418
(c)
(d)
power iteration
mb, exact sub,
sub, exact
power iteration
mb, exact
sub,
mb,
mb,
sub,
100
100
time [8]
10
time
10-1
10
10-2
2
4
8
10
4
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.8: GPU memory and run time performance for the 2c2d architecture on Fashion-MNIST:: Left and right columns show results with the full network's GGN (D = 3,274,634, C= 10) and a per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Nerit for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the i leading cigenpairs using a mini-batch of size N = 128.
21

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(b)
Next (eigenvalues)
Nert (Lop eigenpair)
Nen (eigenvalues)
Nari (lop eigenpair)
GGy
mub
sub
mb
sub
GGN
mb
sub
ub
sub
4224
21161
3276
21295
4416
21453
3276
20378
24064
> 32768
me
24064
>> 32768
23723
>> 32768
me
23457
>> 32768
(c)
(d)
power iteration mb, exact +- sub, exact sub,
power iteration mb, exact +- sub, exact sub,
mb, me
mb, mc
10a
10
time [s]
time [s]
10
10
10-1
10-1
2
1
8
10
4
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.9: CPU memory and run time performance for the 2c2d architecture on Fashion-MNIST: Left and right columns show results with the full network's GGN (D = 3,274,634, C= 10) and a per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Narit for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the & leading eigenpairs using a mini-batch of size N = 128.
22

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a) Next (eigenvalues)
(b) Nent (eigenvalues)
Neci (top eigenpair)
Nexi (top eigenpair)
BON
Data
BGN
Data
mb
sub
Gax
Data
GDN
Data
exact
909
1375
exact
677
3181
exact
1379
4464
exact
1085
4415
3810
6626
3060
6029
3854
6626
3854
6624
(c)
(d)
power iteration
mb, exact sub,
sub, exact
power iteration
mb, exact sub,
sub, exact
mb,
mb,
10
100
time [8]
time
10-1
10-1
2
4
8
10
4
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.10: GPU memory and run time performance for the 3c3d architecture on CIFAR-10 Left and right columns show results with the full network's GGN (D 895,210, C 10) and per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Neri for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the leading eigenpairs using a mini-batch of size N = 128.
23

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a) Nezt (eigenvalues)
(b) Nent (eigenvalues)
Neci (top eigenpair)
Nexi (top eigenpair)
Gox
Dain
mib
BGN
Data
mb
sub
6GN
Data
mb
GBV
Data
mb
exact
2688
12708
exact
2179
11138
exact
2978
13312 20732
exact
2911
13186
14330
20828
11499
19703
14848
14656
20757
(c)
(d)
power iteration
mb, exact sub,
sub, exact
power iteration
mb, exact
sub, exact
mb,
mb,
sub,
10a
10a
time [8]
time
10a
100
10
10-1
2
4
8
10
4
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.11: CPU memory and run time performance for the 3c3d architecture on CIFAR-10 10:
Left and right columns show results with the full network's GGN (D = 895,210, C = 10) and per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Neri for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the i leading cigenpairs using a mini-batch of size N = 128.
24

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a) Next (eigenvalues)
(b) Nent (eigenvalues)
Nexi (top eigenpair)
Nexi (top eigenpair)
BON
Data
BGN
Data
mb
sub
Gax
Data
GDN Data
exact
1054
2052 2271
exact
361
1360
exact
1061
2048
exact
1040
2048
2064
1536
2176
2064
2271
2064
2271
(c)
(d)
power iteration
mb, exact sub,
sub, exact
power iteration
mb, exact sub,
sub,
mb,
mb,
10a
10a
time [8]
time is
10
10
10a
10
10
2
4
8
10
4
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.12: GPU memory and run time performance for the ResNet-32 architecture on CIFAR- 10: Left and right columms show results with the full network's GGN (D 464,154, C 10) and a per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Nerit for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the i leading cigenpairs using a mini-batch of size N = 128.
25

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(c)
(d)
power ileration
mb, exact +- sub, exact sub,
power iteration -b, exact + sub, exact
mb,
mb, sub, mc
10a
time [s]
10a
104
time
10
10
2
6
10
2
6
8
10
lop eigenpairs (k)
top eigenpairs (A)
Figure S.13: CPU memory and run time performance for the ResNet-32 architecture on CIFAR- 10: Left and right columns show results with the full network's GGN (D 464,154, C' 10) and a per-layer block-diagonal approximation, respectively. (c, d) Run time comparison with a power iteration for extracting the lcading cigenpairs using a mini-batch of size N 128.
26

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a) Nexi (eigenvalues)
(b)
Nea (top eigenpair)
Nen (eigenvalues)
Nari (lop eigenpair)
mb
6Gy
mb
sub
mb
sub
mb
sub
1247
1247
me
1262
1312
896
1240
me
1232
1259
me
1232
1255
(c)
(d)
power iteration aub, exact sub, exact
power iteration mb, exact
sub, exact
in, mc
d sub, in
mb, mn
d sub, me
10a
10
10a
time 3
time [
10
10
10a
2
6
8 10
2
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.14: GPU memory and run time performance for the ResNet-56 architecture on CIFAR- 10: Left and right columns show results with the full network's GGN (D = 853,018, C' 10) and a per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Nari for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the & leading eigenpairs using mini-batch of size N = 128.
27

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(c)
(d)
power ileration
mb, exact +- sub, exact sub,
power iteration -b, exact + sub, exact mb, sub, mc
mb;
10a
time [s]
time
10
10
10
10
2
6
10
2
8
10
lop eigenpairs (k)
top eigenpairs (A)
Figure S.15: CPU memory and rum time performance for the ResNet-56 architecture on CIFAR- 10: Left and right columns show results with the full network's GGN (D 853,018, C' 10) and a per-layer block-diagonal approximation, respectively. (c, d) Run time comparison with a power iteration for extracting the lcading cigenpairs using a mini-batch of size N 128.
28

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a) Next (eigenvalues)
(b) Nent (eigenvalues)
Nexi (top eigenpair)
Nexi (top eigenpair)
BON
Data
G6v Data
Data Gax
GDN
Data
exact
255
exact
exact
36
256
exact
253
1119
1536
745
1402
1119
1536
1119
1536
(c)
(d)
power ileration sub, exact
mb, mc
power iteralion sub, exact
mb,
sub,
sub,
10a
10
time [8]
10
time
10a
10
10
2
4
8
10
4
6
8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.16: GPU memory and run time performance for the AlI-CNN-C architecture on CIFAR-100: Left and right columns show results with the full network's GGN (D 1,387,108, C = 100)
and a per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Nerit for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the k
leading eigenpairs using mini-batch of size N = 64.
29

---

Published in Transactions on Machine Learning Research (02/2023)
Full network
Block-diagonal approximation
(a)
(b)
Nexi (eigenvalues)
Nea (Lop eigenpair)
Nem (eigenvalues)
Nari (lop eigenpair)
mb
sub
GGN
mb
sub
sub
sub
mub
sub
351
97
me
3771
4941
me
2558
4454
me
3783
4929
me
3759
4918
(c)
(d)
power iteration sub, exact
mb,
power iteration mb, exact
sub, exact
sub, me
mb,
sub, me
10a
10
time I
time
10
10'
10
10"
2
6
8 10
2
6 8
10
top eigenpairs (k)
top eigenpairs (k)
Figure S.17: CPU memory and run time performance for the All-CN-C architecture on CIFAR-100: Left and right columns show results with the full network's GGN (D 1,388,108. C = 100) and a per-layer block-diagonal approximation, respectively. (a, b) Critical batch sizes Nen for computing eigenvalues and the top eigenpair. (c, d) Run time comparison with power iteration for extracting the k leading eigenpairs using mini-batch of size N = 64.
30

---

Published in Transactions on Machine Learning Research (02/2023)
Table S.2: Memory performance for computing damped Newton steps: Left and right columns show the critical batch sizes with the full network's GGN and per-layer block-diagonal approximation, respectively.
Fashion-MNIST 2c2d
Full network
Block-diagonal approximation
Ner (GPU)
Merit (CPU)
Nert (GPU)
Neri (CPU)
Data
mb
sub
Data
mb
sub
Data
mb
sub
Data
mb
sub
exact
66
159
exact
202
487
exact
68
159
exact
210
487
me
362
528
me
1107
1639
In
528
Inc
1137
1643
CIFAR-10 3c3d
Full network
Block-diagonal approximation
Nen (GPU)
Nen (CPU)
Nert (GPU)
Nen (CPU)
G0X
mb
sub
GGX
mb
sub
GGN
mb
sub
GGN
mb
sub
exact
208
727
667
2215
349
795
1046
2123
mc
1055
1816
mc
3473
5632
mc
1659
2112
mc
4997
6838
CIFAR-10 ResNet-32
Full network
Block-diagonal approximation
Nen (GPU)
Nam (GPU)
GGX
mb
sub
GON
mb
sub
exact
344
1119
exact
1051
1851
1205
1535
2048
2208
CIFAR-10 ResNet-56
Full network
Block-diagonal approximation
Nen (GPU)
New (GPU)
Data
mb
sub
Data
mb
sub
exact
209
640
exact
767
1165
me
687
890
me
1232
1255
CIFAR-100 AII-CNN-C
Full network
Block Alack-diagonal approximation
Nea (GPU)
Nem (CPU)
Nom (GPU)
Nen (CPU)
GGN
mb
sub
6GX
mb
sub
GGN
mb
sub
GGN
mb
sub
exact
13
exact
309
exact
135
exact
504
mc
640
959
me
2015
2865
me
1079
1536
mc
3360
3920
31

---

Published in Transactions on Machine Learning Research (02/2023)
B.2 Training of neural networks
Procedure: We train the following DEEPOBS (Schneider et al., 2019) architectures with SGD and ADAM: 3C3D on CIFAR-10, 2C2D FASHION-MNIST and ALI-CNN-C on CIFAR-100 - all are equipped with cross-entropy loss. To ensure successful training, we use the hyperparameters from Dangel et al. (2020) (see Table S.3).
We also train a residual network RESNET-32 He al. (2016) with cross-entropy loss on CIFAR-10 with both SGD and ADAM. For this, we use a batch size of 128 and train for 180 epochs. Momentum for SGD was fixed to 0.9, and ADAM uses the default parameters (8) 0.9, 82 = 0.999, < 10- For both optimizers, log-equidistant
the learning rate
determined via grid search. Following (Schneider et al., 2019),
grid from 10-5 to 102 and 36 grid points. As performance metric, the best test accuracy during training (evaluated once every epoch) is used.
Results: The results for the hyperparameter grid search are reported in Table S.3. The training metrics training/tern loss/ loss/accuracy for all eight test problems are shown in Figure S.18 and 8.19.
Table S.3: Hyperparameter settings for training runs: For both SGD and ADAM, we report their learning rates a (taken from the baselines in Dangel et al. (2020) or, for RESNET-32, determined via grid search). Momentum for SGD is fixed to 0.9. ADAM uses the default parameters 31 0.9, 82 0.999,
10-8. We also report the batch size used for training and the number of training epochs.
Problem
SGD
ADAM
Batch size
Train epochs
FASHION-MNIST 2C2D
2.07
10-2
1.27
10-4
N 128
100
CIFAR-10 3C3D
= 3.79 10-3
= 2.98 10
N 128
100
CIPAR-10 RESNET-32
= 6.31 10-2
== 2.51 10-3
N -128
180
CIFAR-100 ALL-CNN-C
a 2 4.83-10-1
a 6.95 10-4
N = =256
350
B.3
GGN vs. Hessian
During training of the neural networks (see Appendix B.2), we store copy of the model (i.e. the network's current parameters) at specific checkpoints. This grid defines the temporal resolution for all subsequent computations. Since training progresses much faster in the early training stages, we use a log-grid with 100 grid points between and the number of training epochs and shift this grid by -1.
Checkpoints:
DC Overlap: Recall from Section 3.2: For the set orthonormal eigenvectors {e ]e to the C largest eigenvalues of some symmetric matrix U, let pU (ef ...., e the overlap between two subspaces E = span (el ..., eg) and ε3 span (ey ...., ex) of the matrices and V is defined by
...,, eg)_ As in Gur-Ari al. (2018),
Tr(PUPV
overlap(E" (8V) = Tr (PU Tr (PV) (PV)
E [0,1].
The overlap can be computed efficiently by using the trace's cyclic property: It holds Tr (P" = Tr (WTWW W) with W = (ef, e (eY, ..., e E RCx. Note that this is a small C' matrix, whereas Pú pv E DDD Since
Tr
= Tr ((e) ..., (e),,.. ((.,,,))))
Tr ((e) ec) T (ee)...))))
U Tr (Ic)
(Cyelic property of trace)
(Orthonormality of the eigenvectors)
= C
(and analogous Tr (PV) = C), the denominator simplifies C.
32

---

Published in Transactions on Machine Learning Research (02/2023)
Fashion-MNIST 2c2d SGD
100
accuracy in %
loss
train
train
test
test
20
40
60
80
100
0
20
40
60
80
100
epoch
epoch
Fashion-MNIST 2c2d Adam
100
accuracy in %
80
loss
train
60
train
test
test
40
20
20
60
100
0
40
80
100
epoch
epoch
CIFAR-10 3c3d SGD
80
accuracy in
60
loss
train
train
test
test
40
20
20
40
60
80
100
0
20
40
60
80
100
epoch
epoch
CIFAR-10 3c3d Adam
80
accuracy in %
loss
train
60
train
test
test
40
20
20
60
100
20
60
100
epoch
epoch
Figure S.18: Training metrics (1): Training /tst loss/accuracy for all test problems.
33

---

Published in Transactions on Machine Learning Research (02/2023)
CIFAR-10 ResNet-32 SGD
100
10
loss (log scale)
accuracy in %
10
train test
60
train
test
40
10
20
My
10 1
50
100
150
50
100
150
epoch
epoch
CIFAR-10 ResNet-32 Adam
100
10a
80
loss (log scale)
srcuuracy in %
10
train
60
train
lest
Lest
100
20
10
0
50
100
150
0
50
100
150
epoch
epoch
CIFAR-100 AlI-CNN-C SGD
60
4
accuracy in
40
loss
3
train
train
test
test
20
100
200
300
100
200
300
epoch
epoch
CIFAR-100 AlI-CNN-C Adam
80
accuracy in %
60
loss
train
train
40
test
2
20
1
100
200
300
100
200
300
epoch
epoch
Figure S.19: Training metrics (2): Fraining/tess test loss/accuracy for all test problems.
34

---

Published in Transactions on Machine Learning Research (02/2023)
Procedure: For each checkpoint, compute the top-C' cigenvalues and associated eigenvectors of the full-batch GGN and Hessian (i.e. GGN and Hessian are both evaluated on the entire training set) using an iterative matrix-free approach. We then compute the overlap between the top-C' eigenspaces as described above. The cigspaces (i.c. the top-C' eigenvalues and associated cigenvectors) are stored on disk such that they can be used as reference by subsequent experiments.
The results for all test problems are presented in Figure S.20. Except for a short phase at the beginning of the optimization procedure (note the log scale for the epoch-axis), strong agreement (note the different limits for the overlap-axis) between the top-C' eigenspaces is observed. We make similar observations for all test problems, yet to slightly lesser extent for CIFAR-100. A possible explanation for this would be that the 100-dimensional eigenspaces differ in the eigenvectors associated with relatively small curvalure. The corresponding eigenvalues already transition into the bulk of the spectrum, where the "sharpness of separation' decreases. However, since all directions are equally weighted in the overlap, overall slightly lower values are obtained.
Results:
B.4 Eigenspace under noise and approximations
Procedure (1): We use the checkpoints and the definition of overlaps between eigenspaces from Ap- pendix B.3. For the approximation of the GGN, we consider the cases listed in Table S.4.
Table S.4: Considered cases for approximation of the eigenspace: We use a different of cases for the approximation of the GGN's full-batch eigenspace depending the test problem. For the test problems with C = 10, we use M = 1 M MC-sample, for the CIFAR-100 AL-CNN-C lest problem (C 100), we use M = 10 MC-samples in order to reduce the computational costs by the same factor.
Problem
Cases
FASHION-MNTST 2c2D
mb, exact with mini-batch sizes N E [2, 8, 32, 128) mb, me with N = 128 and M = MC-sample sub, exact using 16 samples from the mini-batch sub, me using 16 samples from the mini- batch and M MC-sample
CIFAR-10 3C3D and CIFAR-10 10 RESNET-32
CIFAR-100 ALL-CNN-C
mb, exact with mini-batch sizes N E {2, 8, 32, 128} mb, me with N = 128 and M = 10 M MC-samples sub, exact using samples from the mini-batch sub, me using 16 samples from the mini-batch and M 10 MC-samples
For every checkpoint and case, compute the top-C eigenvectors of the respective approximation to the GGN. The eigenvectors are either computed directly using VIVIT (by transforming eigenvectors of the Gram matrix into parameter space, see Section 2.1) or, if not applicable (because memory requirements exceed Norit, see Section 3.1), using an iterative matrix free approach. The overlap is computed in reference to the GGN's full-batch top-C cigenspace (see Appendix B.3). We extract 5 mini-batches from the training data and repeal the above procedure for each mini-batch (i.e. we obtain overlap measurements for every checkpoint and case). The same 5 mini-batches are used over all checkpoints and cases.
Results (1): The results be found in Figure S.21 and S.22. All test problems show the same characteristics: With decreasing computational effort, the approximation carries less and less structure of its full-batch counterpart, as indicated by dropping overlaps. In addition, for fixed approximation method, a decrease in approximation quality can be observed over the course of training.
Procedure (2): Since VIVIT's GGN approximations using curvature sub-sampling and/or the MC approximation (the cases mb, mc as well as sub, exact and sub, me in Table S.4) are based on the mini-batch GGN, we cannot expect them to perform better than this baseline. We thus repeat the analysis
35

---

Published in Transactions on Machine Learning Research (02/2023)
Pashion-MNIIST 2c2d SGD
Pashion-MNNIST 2c2d Adam
0.95
overlap
0.9
overlap
0.9
0.8
0.85
10
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d SGD
CIFAR-10 3c3d Adam
overlap
0.95
overlap
0.95
0.9
0.9
0.85
0.85
10
10a
10a
10"
10a
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 ResNet-32 32 SGD
CIFAR-10 ResNet-32 Adam
0.8
0.8
overlap
overlap
0.6
0.6
0.4
0.4
0.2
0.2
100
10a
10a
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-100 AII-CN--C SGD
CIFAR-100 AlI-CNN-C Adam
0.8
0.8
overlap
overlap
0.6
0.6
0.4
0.4
0.2
10
10a
10a
10a
10
10a
epoch (log scale)
epoch (log scale)
Figure S.20: Full-batch GGN vs. full-batch Hessian: Overlap between the top-C' eigenspaces of the full-batch GGN and full-batch Hessian during training for all test problems.
36

---

Published in Transactions on Machine Learning Research (02/2023)
Fashion-MNIST 2c2d SGD
mmb exact
128, exact
mb 8, exact
sub 16, exact
mb 32, exact
amb 128, me
overiap
128, exact
overlap
rsub 16, mo
0.5
0
10
10
10
10
10
10
epoch (log scale)
epoch (log scale)
MaShion-MNIST 2c2d Adam
mb 2, exact
amb 128, exact
mb exact
euu 16,
mb 32, exact
128, me
overlap
mb 128, exact
overlap
sub 16, mc
0.5
0.5
10
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d SGD
exact mb 8, exact
128,
sub 16, exact
mb 32, exact
mb 128, me
overiap
ubb exact
overlap
subb 16,
0.5
10
10
10
10"
10
10
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d Adam
mb 2, exact
zub 128, exact
mb exact
sub 16,
mb 32, exact mb 128, exact
128, me
overlap
overlap
sub 16, me
0.5
0.5
10a
10
10a
10a
10
10a
epoch (log scale)
epoch (log scale)
Figure S.21: VIVIT vs. full-batch GGN (1): Overlap between the top-C' eigenspaces of different GGN approximations and the full-batch GGN during training for all test problems. Each approximation is evaluated on 5 different mini- batches.
37

---

Published in Transactions on Machine Learning Research (02/2023)
CIFAR-10 ResNet-32 SGD
wmb exact
mb 128, exect sub
mb
mb 32, exact
mub 128, me
overlap
mb 128, exact
overlap
sub 16, me
0.5
0.5
0
10
10
102
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 ResNet-32 Adam
128, exact
mb 8, exact
sub 16, exact
mb 32, exact
mb 128, me
overlap
288 exact
overlap
0.5
0.5
0
100
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-100 AIL-CNN-C SGD
mb exact
mb 128, exact
mb
sub 128, me 10
mb 32, exact
overlap
mb 128, exact
overlap
rsub 16, me 10
0.5
0.5
0
10
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR- 100 AII-CNN-C Adam
128,
mb 8, exact
sub 16, exact
mb 32, exact
ub 128, 10
overlap
128,
overlap
0.5
0.5
0
100
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
Overlap between the top-C' eigenspaces of different GGN approximations and the full-batch GGN during training for all test. problems. Each approximation is evaluated 5 different mini- batches.
Figure S.22: VIVIT vs. full-batch GGN (2):
38

---

Published in Transactions on Machine Learning Research (02/2023)
from above but use the mini-batch GGN with batch-size N GGN. Of course, the mini-batch reference top-C eigenspace is always evaluated on the same mini-batch as the approximation.
128 ground truth instead of the full-batch
Results (2): The results be found in Figure S.23. Over large parts of the optimization (note the log scale for the epoch-axis), the MC approximation seems to be better suited than curvature sub-sampling (which has comparable computational cost). For the CIFAR-100 ALL-CNN-C test problem, the MC approximation stands out particularly early from the other approximations and consistently yields higher overlaps with the mini-batch GGN.
B.5 Curvature under noise and approximations
GGN and Hessian are predominantly used to locally approximate the loss by quadratic model q (see Equation (6)). Even if the curvature's eigenspace is completely preserved in spite of the approximations, they can still alter the curvature magnilude along the eigenvectors.
Procedure: Table S.5 gives an overview over the cases considered in this experiment.
Table S.5: Considered cases for approximation of curvature: We use a different set of cases for the approximation of the GGN depending on the test problem. For the test problems with C = 10, we use MC-sample, for the CIFAR-100 ALL-CNN-C test problem (C = 100), we use M = 10 MC-samples in order reduce the computational costs by the same factor.
Problem
Cases
FASHION-MNIST 2c2D CIFAR-10 3c3D and CIFAR-10 RESNET-32
mb, exact with mini-batch size N 128
mb, me with N 128 and M MC-sample sub, exact using 16 samples from the mini-batch sub, me using 16 samples from the mini-bateh and M MC-sample
CIFAR-100 ALL-CNN-C
mb, exact with mini-batch size N 128
mb, with N 128 and M 10 MC-samples sub, exact using 16 samples from the mini-batch
sub, me using 16 samples from the mini-batch and M 10 MC-samples
Due to the large computational effort needed for the evaluation of full-batch directional derivatives, subset of the checkpoints from Appendix B.3 is used for two test problems: We use every second checkpoint for CIFAR-10 RESNET-32 and every forth checkpoint for CTFAR-100 ALI-CNN-C.
For each checkpoint and case, we compute the top-C' eigenvectors {ex]e of the GGN approximation ((ap) either with VIVIT or using an iterative matrix-free approach (as in Appendix B.4). The second- order directional derivative of the corresponding quadratic model along direction eş is then given by
(ap) = e, G(M))e (see Equation (7)). As reference, we compute the full-batch GGN (ib)) and the
resulting directional derivatives along the same eigenvectors {ex]e ie. ( e ((fb)e The average (over all C directions) relative error is given by
(ap) {
(Ab) X
(bb)
k=
The procodure above repeated on mini-batches from the training data. (i.c. we obtain average relative errors for every checkpoint and case) except for the CIFAR-100 ALL-CNN-C test problem, where we perform only single run to keep the computational effort. manageable.
39

---

Published in Transactions on Machine Learning Research (02/2023)
Fashion-MNIST 2c2d SGD
Fashion-MNISTT 2c2d Adam
16, exact
euub 16, exact
ub 128, me 1
mb 128, me
sub 16,
sub
overlap
overlap
0.5
0.5
100
10
10
10
10'
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d SGD
CIFAR-10 3c3d Adam
suh 16, exact
smb 16, exart amb 128,
mb 128, me 1
vuu 16, me
veub
overlap
overlap
0.5
0.5
10
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 ResNet-32 SGD
CIFAR-10 ResNet-32 Adam
sub 16, exact
sub 16, exact
mb 128,
mmb 128, I
16, me
overlap
overlap
0.5
0.5
0
10"
10a
10a
10
10a
10a
epoch (log scale)
epoch (log scale)
CIFAR-100 AII-CN-C SGD
CIFAR-100 AlI-CNN-C Adam
sub 16, exact
sub 16, exact mb
mb 128,
sub 16, mc 10
sub 16, me 10
overlap
overlap
0.5
0.5
0
10
10
102
10
10'
10a
epoch (log scale)
epoch (log scale)
Figure S.23: VIVIT vs. mini-batch GGN: Overlap between the top-C' eigenspaces of different GGN approximations and the mini-batch GGN during training for all test problems. Each approximation is evaluated on different mini-batches.
40

---

Published in Transactions on Machine Learning Research (02/2023)
Results:
The results can be found in Figure S.24. We observe similar results as in Appendix B.4: With increasing computational effort, the approximated directional derivatives become more precise and the overall approximation quality decreases over the course of the optimization. For the RESNET-32 architecture, the average errors are particularly large.
B.6 Directional derivatives
Procedure: We use the chockpoints from Appendix B.3. For every checkpoint, compute the top-C' eigenvectors of the mini-batch GGN (with a mini-batch size of N 128) using an ilerative matrix-free method. We also compute the mini-batch gradient. The first- and second-order directional derivatives of the resulting quadratic model (see Equation (6)) are given by Equation (8).
We use these directional derivatives (rink |= == {Ank) ) I==1,==1 to compute signal-to-noise ratios (SNRs) along the top-C' eigenvectors. The curvature SNR along direction ex is given by the squared sample mean divided by the empirical variance of the samples {Ank A i.c.
N.C.
X m= (Ank - A₂)²
N
SNR
where
Ap-
Ank
1/N-1 Em
N
n=1
(and similarly for (Tnn))_m))
The results can be found in Figure S.25 and S.26. These plots show the SNRs in C distinct colors that are generated by linearly interpolating in the RGB color space from black to light red (o). At each checkpoint, the colors are assigned based on the order of the respective directional curvature d The SNR that belongs the direction with the smallest curvature is shown in black and the SNR that belongs to the direction with the largest curvature is shown in light red. The color thus encodes only the order of the top-C' directional curvatures not their magnitude. We use this color encoding to reveal potential correlations between SNR and curvature.
Results:
We find that the gradient SNR along the top-C' eigenvectors is consistently small (in comparison to the curvature SNR) and remains roughly on the same level during the optimization. The curvature signal decreases as training procoeds. The SNRs along the top-C cigendirections do not appear to show any significant correlation with the corresponding curvatures. Only for the CIFAR-100 test problems we can suspect a correlation between strong curvature and small curvature SNR.
C Implementation details
Layer view of backpropagation: Consider a single layer 1ots that transforms inputs z (i-1) gh into outputs (() I84(9 by means of a parameter ((1) _dd( During backpropagation for V, the layer receives vectors Snc = (J ( fn) sne from the previous stage (recall Vz6n = EE TT) Parameter contributions
Une to V are obtained by application of its Jacobian,
(1) = (Jocn 8nc ne
= (Joc) T
(Ja(OS)) 8nc
(Chain rule)
= (Joc (1)) T 8
(Definition of 8 nc)
(S.11)
Consequently, the contribution of (0)) V, denoted by V()) xNC
vi
(8)
(
with fifl = (Joc)) fn) 8nc-
(S.12)
Uii
12
NC
41

---

Published in Transactions on Machine Learning Research (02/2023)
Pashion-MNIST 2c2d SGD
Fashion-MNISTT 2c2d Adam
rel. error (log scale)
av. rel. error (log scale)
10a
10
10
mb 128, exact
mb 128,
sub 16, exnnt
sub 16, exact
mh 128, me 1
mb 128, me
sub 16,
10
10"
10
10a
10
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d SGD
CIFAR-10 3c3d Adam
av. rel. error (log scale)
av. rel. error (log scale)
10a
10a
10
10
mb 128, exact
mbb 128, excact
sub 16, exact
10
sub 16, exact
mb 128,
mb
sub 16, me
10-2
sub 16, me
10
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 ResNct-32 SGD
CIFAR-10 ResNet-32 Adam
av. rel. error (log scale)
BXY rel. error (log scale)
10
10a
10a
10a
mb 128, exact
mbb 128, exact sub exact mb 128, me sub 16, me
10a
sub 16, exart
10a
128, me
sub 16, mc
10
10
10
10
10a
10"
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-100 AII-CNN-C SGD
CIFAR-100 ALI-CNN-C Adam
av: rel. error (log scale)
sealey Bol) eoou yorr
10
10a
10a
emb 128, exact
mb 128, exact
10-0.5 0.5
mb 128, me 10
mb 128, me 10
10-1
sub 16, 10
XX
suub 16, 10
10"
10
10a
10"
10
10a
epoch (log scale)
epoch (log scale)
Figure S.24: ViVIT's vs. full-batch quadratic model: Comparison between approximations to the quadratic model and the full-batch model terms of the average relative error for the directional curvature during training for all test problems.
42

---

Published in Transactions on Machine Learning Research (02/2023)
Fashion-MNIST 2c2d SGD
SNR SNR(pat) [Tak] (log seale)
10-3
SNH[Ank] (log scale)
10
100
10
10-1
10
10-8
10"
10
10a
10c
10
10a
epoch (log scale)
epoch (log scale)
Fashion-MNIST 2c2d Adam
SNR( (log stale)
10
SNR[A (log scale)
10
10
10
10
10-8
10
-2
10
10a
102
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d SGD
SNR(1mll) (log scale)
SNR(A (log scale)
10
10
10
10
10"
10a
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 3c3d Adam
10
SNR (ra)] (log scale)
[An[ (log scale)
10a
10-5
10-1
10-9
10a
10
10a
10
10'
10a
epoch (log scale)
epoch (log scale)
Directional SNRs (1): SNR along each of the mini-batch GGN's top-C eigenvectors during training for all test problems. At fixed epoch, the SNR for the most curved direction is shown in and for the least curved direction in
Figure 8.25:
43

---

Published in Transactions on Machine Learning Research (02/2023)
CIFAR-10 ResNet-32 SGD
SNR[1na] (log scale)
10
SNR/Aoot (log scale)
100
10
10
10-22
10
10
10
10a
10
10
10a
epoch (log scale)
epoch (log scale)
CIFAR-10 ResNet-32 Adam
(NR( (log scale)
10
Ana/ SaR (log scale)
10
10-5
SNR
10-2
10"
10
10a
10"
10a
10a
epoch (log scale)
epoch (log scale)
CIFAR-100 AII-CNN-C SGD
10
10a
SNR[Tnk] (log scale)
SNR[Ane] (log scale)
10
10a
10
10-7
10a
10
10
2
10
50
100
200
350
10
50
100
200
350
epoch (log scale)
epoch (log scale)
CIFAR- 100 AlI-CNN-C Adam
10a
10
SNR[Tnk] (log scale)
SNR Ans (log scale)
10a
10-4
10a
10
10
10-8
10a
10-10
10
10
epoch (log scale)
epoch (log scale)
Figure S.26: Directional SNRs (2): SNR along each the mini-batch GGN's top-C eigenvectors during training for all test problems. At fixed epoch, the SNR for the most curved direction is shown in and for the least curved direction in
44

---

Published in Transactions on Machine Learning Research (02/2023)
C.1 Optimized Gram matrix computation for linear layers
Our goal is to efficiently extract 060's contribution to the Gram matrix G, given by
i (ii v E NCNCC
(S.13)
Gram matrix via expanding v One way to construct G is first expand ( (Equation (S.12)) via the Jacobian Jgn then contract it (Equation (S.13)). This can be a memory bottleneck for large linear layers which are common in many architectures close to the network output. However if only the Gram matrix rather than is required, structure in the Jacobian can used to construct G without expanding V and thus reduce this overhead.
Optimization for linear layers: Now, let T be linear layer with weights W R4 1) i.e. g(i) vec
vec(W (W)) = with column stacking convention for vectorization,
Tr( 0
(
w( 2 (1-1))
The Jacobian is
J0:0 1 = 2
(i-1)
(S.14)
Its structure can be used directly compute entries of the Gram matrix without expanding V
(1)
T (6) Une on'e
(Equation (S.13))
(me)(w'e))
= 0/c
(t) (J (Jow)
T
(1)
T snd
T
(T)
2 In(O)
((-1)
Inco
(Equation (S.14))
T
(i-1) (i-1)
6 In(
(2)
(Equation (S.11))
= Anc
(1-1) T (i-1) (1)
(2-11 T (i-1) E R)
2n
Zx¹
I T
2n
(2(1111)
(() ")
n'e)
We see that the Gram matrix is built from two Gram matrices based {21 (i-1) Jn=1 and {s { I==1,== that
N.C
(1) N.C
require O(N²) and (((CC)) memory, respectively. In comparison, the naïve approach via v e wdd fx XNC scales with the number of weights, which is often comparable to D. For instance, the 3C3D architecture on CIFAR-10 has D = 895,210 and the largest weight matrix has ( = 589,824, whereas NC = 1,280 during training (Schneider et al., 2019).
To evaluate the efficiency of our optimization for linear layers, we consider a multi-layer perceptron with 1024 512 256 128 + 64 32 C 10 units, each activated by ReLU except for the last, and mean squared error loss function. We set the batch size N = 128 and randomly generate synthetic inputs and labels. The following tasks are considered (results for GPU, reporting the smallest run time over 20 repetitions):
Run time comparison:
T1: Mini-batch gradient computation: /grad Ri 0.774 ms (1.0x)
T2: Naive computation of explicit V: ty s 9.25 ms (11.9x)
T3: Naive computation of the Gram matrix & via explicit V: tv. G Po 233 ms (301x)
T4: Optimized computation of G via implicit V: by G.opt 4.28 ms (5.5x)
45

---

Published in Transactions on Machine Learning Research (02/2023)
The main finding of this comparison is that, for this setting, optimized approach (T4) computes the Gram matrix >50 faster than the naive approach (T3), and requires less than C' gradient computations (T1). Detailed description:
Naive computations of V (T2) and & (T3) are slow: From the complexity analysis in Section 2.3, we expect tv s Cigrad- Empirically, we observe a factor of 11.9. Despite the use of a GPU, performance does not improve. We believe this may be because the operation requires allocating and writing large amounts of data into memory for V. Naively expanding this erplicil V, then computing G (T3), would cost fv.o/tgrad 301 gradients.
Our optimized approach to implicitly work with V in combination with exploiting structure in the Jacobian (T4) fast: With our optimizations, we observe tv optt &grad 5.5. This is > 50x speed-up over the naive approach (T3)! It also requires less than C gradient computations, i.e. almost 50% of the expected overhead is compensated by parallelism on the GPU.
In our work, we first extended the BACKPACK package by the naive operations T2 and T3 described above. These are already superior to for-loop based approaches with PYTORCH's built-in automatic differentiation because BACKPACK implements vectorized vector-Jacobian products. On top of that, we then further significanlly improved performance through structural tricks (T4), as shown above.
C.2 Implicit multiplication with the inverse (block-diagonal) GGN
Inverse GGN-vector products: A damped Newton step requires multiplication by (G + -Ip)=I By means of Equation (3) and the matrix inversion lemma,
(8In G (8In + VVT)-
(Equation (3))
1
Ip + vvT)
8
1
I
1
v
INC INC+VTIL ) VT
(Matrix inversion lemma)
8
8
1
Ip V (8ING VONNC+VVV)" VTV)' 'vi]
(Gram matrix)
8
1 Io V (ŌING + G) G)-1v'.
(S.15)
8
Inverse GGN-vector products require inversion of the damped Gram matrix well as applications of V, V for the transformations between Gram and parameter space.
Inverse block-diagonal GGN-vector products: Next, we replace the full GGN by its block diagonal approximation G z GBDA = diag(G
G ((.).... with
G T xa x
and (i) as Equation (S.12). Then, inverse multiplication reduces to each block,
G G
diag (1))) _((2)-1 ......
BDA
If agaain a damped Newton step is considered, we can reuse Equation (S.15) with the substitutions
(G,DVV. ( + (G)) (i), vi G())
to apply the inverse and immediately discard the ViViT factors: At backpropagation of layer ())
STT be replaced by other easy-to-invert matrices.
46

---

Published in Transactions on Machine Learning Research (02/2023)
1. Compute v using Equation (S.12).
Compute (()) using Equation (S.13).
3. Compute (ÓING + (()
4. Apply the inverse in Equation (S.15) with the above substitutions to the target vector.
Discard V V(O G and (6Ixc + G( Proceed to layer i
Note that the above scheme should only be used for parameters that satisfy di) NC, i.e. dim(G tim(G()) Low-dimensional parameters can be grouped with others to increase their joint dimension, and to control the block structure of GBDA.
47


Reviews:
[{"review_id": "YobsT6UIpw", "reviewer": "Reviewer_Wsxx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- I think it's far from obvious that this method of computing the top eigenvalues and eignvectors of the minibatch GN would be much faster, in wall clock time, than running power method using a GNVP oracle, as the authors show in Figure 2.  Leveraging Backpack seems to be crucial here for the performance.\n- It's nice to see an examination of several approximations which further bring down the time complexity of ViViT.\n- It's nice to see a study comparing the top eigenspace of the minibatch GGN to that of the population GGN (Figure 3b).  Ap\n\nWeaknesses\n- The method cannot scale to large minibatch sizes, and it's unclear whether the Gauss-Newton spectrum computed over small minibatches is useful, since this matrix can differ substantially from the population version (e.g. see Figure 3(b)).\n- Any applications to second-order optimization seem quite speculative. \n- The authors say that a benefit of ViViT is that it provides access to the \"example-wise\" directional curvatures.  That is, for some direction $\\mathbf{e}$, by using the baseline approach of Gauss-Newton vector product on the minibatch we could compute the overall minibatch directional curvature $\\mathbf{e}^T \\mathbf{G} \\mathbf{e}$, but by using ViViT we could get not only that but also the \"example-wise\" curvatures $\\mathbf{e}^T \\mathbf{G}_i \\mathbf{e}$ for each example $i$ in the minibatch, where $\\mathbf{G} = \\sum_i \\mathbf{G}_i$.  (I think this is also what the authors are getting at in the abstract when they write that \"existing works ... do not consider noise in the mini-batch.\"). However, it's not clear to me what is the benefit of knowing these example-wise curvatures.\n- It's not clear to me that the \"MC approximation\" approach of retaining $M$ (out of $C$) classes from each datapoint is better than just naively subsampling an equivalent number of terms from the $\\sum_n \\sum_c$ double sum.  Could this authors comment on this?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper describes an efficient way to compute the eigenvalue decomposition of the minibatch Gauss-Newton matrix -- information which could potentially be useful for scientific or algorithmic (second order optimization) purposes.  The Gauss-Newton matrix is a positive semidefinite approximation to the Hessian which consists of (a) gradients of the network, and (b) second derivatives of the loss function.  One might be interested in computing the eigenvectors and eigenvalues of this matrix over a minibatch.\n\nOne way to compute the top eigenvalues and eigenvectors of the minibatch GNVP would be to run the power method with an oracle that computes matrix-vector products between the Gauss-Newton matrix and arbitrary vectors (easy to write in PyTorch).  In contrast, this paper presents an _alternative_ method which (a) can deliver _all_ of the eigenvectors of the minibatch GN matrix, rather than just the top ones, and (b) is apparently faster at even computing just the top ones.\n\nThe basic idea is to leverage the fact that the minibatch Gauss-Newton matrix, which has dimension $D \\times D$, where $D$ is the number of network parameters, is related to a Gram matrix of dimension $NC \\times NC$, where $N$ is the minibatch size and $C$ is the number of classes.  (If the loss function is mean squared error, this Gram matrix is the neural tangent kernel.) In particular, the GN matrix is $\\mathbf{V} \\mathbf{V}^T$ and the Gram matrix is $\\mathbf{V}^T \\mathbf{V}$, for a certain matrix $\\mathbf{V}$ of shape $D \\times NC$.  The GN matrix and the Gram matrix share non-zero eigenvalues, and given an eigenvector of one of these matrices, we can recover the corresponding eigenvector of the other matrix by multiplying by $\\mathbf{V}$ or $\\mathbf{V}^T$.  Since $NC \\ll D$ in settings of interest, if we want to compute the eigenvectors of the GN matrix $\\mathbf{V} \\mathbf{V}^T$, it is computationally more efficient to compute the eigenvectors of the Gram matrix $\\mathbf{V}^T \\mathbf{V}$, and then multiply by $\\mathbf{V}$ to get the eigenvectors of the GN matrix.\n\nIf one were using PyTorch plain's autograd engine, computing  $\\mathbf{V}$  would cost $NC $ times the cost of the minibatch gradient.  But helpfully, by leveraging the \"Backpack\" pytorch library (which provides access to the gradients for individual examples in the minibatch), the matrix $\\mathbf{V}$ can be computed using only $C$ times the cost of computing the minibatch gradient.  (So, if the network had just one output unit, then $\\mathbf{V}$ could be computed at no additional cost over computing the minibatch gradient; for $C > 1$, the authors propose an approximation that retains this time complexity.)\n\nTo compute the eigenvalue spectrum of the minibatch Gauss-Newton matrix, the authors therefore propose: (a) compute the minibatch $\\mathbf{V}$ using Backpack, at $C$ times the cost of computing the minibatch gradient; form the $NC \\times NC$ Gram matrix $\\mathbf{V}^T \\mathbf{V}$; compute the eigendecomposition of this matrix; recover the eigenvectors of the Gauss-Newton matrix by multiplying the Gram matrix eigenvectors by $\\mathbf{V}$.\n\nThe authors show that this method (called \"ViViT\") is faster in wall clock time than the baseline power method approach that I described above, they discuss and test several approximations to speed up ViViT further, and they study the extent to which the top eigenspace of the minibatch GN matrix is similar to that of the population GN matrix.\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- I found it confusing that the letter $N$ was used to denote both the size of the training dataset and the size of the minibatch.  The paper should be explicit about when $N$ is intended to denote the training dataset size, and when $N$ is intended to denote the minibatch size.  Notably, Section 2 seems to use $N$ to mean the size of the training dataset, yet Footnote 1 in Section 2 asserts that modern NNs are always in the 'overparameterized' setting, which is only true if $N$ means the minibatch size.  (e.g. the ResNet-32 studied in this work has 464k parameters, which is 'underparameterized' since the dataset size is 50k and the num classes is 10.)\n\n- In Appendix B.1, what does the column called \"mc, full\" mean?  I don't believe that this label was explained.\n\nOptional, might strengthen the work:\n-  In section 3.1, when the authors assess the wall clock performance of ViViT on the task of computing the top eigenpairs of the minibatch Gauss-Newton matrix, it might be interesting to also test the approach of (1) explicitly computing the gram matrix $\\tilde{\\mathbf{G}}$ and (2) running power method on that gram matrix.  You could then compare that approach to the baseline (already in the paper) of running power method directly on $\\mathbf{G}$ using GNVP products; these approaches are directly comparable since both are based on the power method, and hence, for a fair comparison, you just need to make sure that the power method tolerance is set to the same value for both methods.  By contrast, in the current draft, making a comparison between ViViT and the baseline is a bit of a headache because by changing the power method tolerance, you can speed up or slow down the baseline arbitrarily (as noted on p. 17).  This is just a suggestion.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "a02heTLngA", "reviewer": "Reviewer_bWHi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* well-written paper\n* time complexity and wall time are better than the power iteration method in the case of mini-batches\n\nWeaknesses:\n* while the paper allows access to higher-order derivatives on a single sample basis, it does not specify any use cases apart from broad descriptions (monitoring tool, second-order optimization, etc.). The lack of concrete examples showing the benefit of this approach is a major weakness as. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose ViViT: a method for calculating the generalized Gauss-Newton (GGN) approximation of the Hessian. \nTheir method is exact, computationally efficient, and gives access to first and second-order directional derivatives of single samples. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Suggestions for rephrasing: \n* Page 1: *There is evidence that this has led to stagnation in the performance of the first-order optimizers.* Can you elaborate on the causes of stagnation? How is this also relevant for second-order optimizers?\n* Page 2: *Finally, we use the previously inaccessible properties of the GGN to study how it is affected by noise during training.* As the evolution of the GGN during training is not central to the paper, --during training-- here adds another layer of noise that is not addressed. \n* Page 6: *These insights substantiate ViViT's value as a monitoring tool* Insights are not rigorous to validate the proposed method, this has to be shown through concrete examples. ", "broader_impact_concerns": "--", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hgSQ4M3yOd", "reviewer": "Reviewer_kVKm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The efficiency of ViViT is better than previous methods in computing top-$k$ eigenvectors when $k$ is larger than 1, and the computation time almost doesn't grow with $k$. This makes ViViT useful for computing the top eigenspaces of neural network Hessian.\n\n- This paper is clearly written. The problem background is introduced, the notations and settings are well-defined, and the sections are well-organized and form a smooth logic flow. The related works appear to be properly cited.\n\n- The authors empirically studied not only the efficiency and approximation accuracy of ViViT, but also the factors influencing its performance, including the variance in first and second-order directional derivatives during training. This could help with the understanding of neural network Hessians.\n\nWeaknesses:\n\n- There might not be enough theoretical or empirical guarantees for the performance of the approximated versions of ViViT. The accuracy of the approximation depends on the distribution of the vectors $s_{nc}$, which is related to the data distribution and network structure and weights. The experiments measure the performance by top-$C$ eigenspace overlap on standard convolutional networks, and it could be hard to justify whether it can be generalized to more network structures or beyond top-$C$ eigenspaces.\n\n- The wallclock running time of ViViT compared to that of backpropagation is not provided. Although ViViT can utilize the intermediate results from the computation of gradients, it might be unclear how much overhead it introduces in practice.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a method ViViT to compute the eigenvalues and eigenvectors of the generalized Gauss-Newton (GGN) approximation of neural network Hessian matrices. ViViT is more efficient than previous methods by utilizing the low-rank nature of GGN, and it has approximation versions that are more efficient than ViViT but are less accurate. The authors also did experiments for convolutional networks on FASHION-MNIST and CIFAR to verify the efficiency of ViViT and analyze the factors that influence the performance of ViViT.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- It could be better if the authors could add theoretical or empirical justifications for the approximated versions of ViViT.\n\n- It might be better to provide the running time comparison to gradient computation.\n\n- Figure 1b seems a bit confusing. It looks like an overlap of two figures, and the curves for $\\gamma_k$ and $\\lambda_k$ are partially hidden. Could the authors add more explanation on the figure or in the text description?", "broader_impact_concerns": "This paper focuses on a general problem of computing and approximating the second-order information of neural network training, and I do not see any immediate concerns about the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This manuscript is clearly written, addresses a topic of wide interest in the machine learning community (second order information in large-scale optimization), and delivers interesting (empirical) computational contributions.  Among the weaknesses raised by reviewers are lack of theoretical guarantees and of experimental evaluations in large-scale settings.  While I agree with these critiques, I do not think they deem the manuscript unsuitable for publication, particularly since the text is transparent about said limitations.  A more significant shortcoming, raised by multiple reviewers, is that the paper does not provide sufficient use cases for the quantities it allows computing.  The authors provide certain (rather speculative) possibilities for use cases, which I believe suffice for publication.  However, the work will be much stronger if it were to include more concrete suggestions.  I encourage the authors to address this point in the camera-ready version of the text.
Decision: Accept (poster)