Year: 2021

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 firstand 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. As examples for VIVIT’s usefulness, we investigate the directional firstand second-order derivatives during training, and how noise information can be used to improve the stability of second-order methods.

heading: 1 INTRODUCTION & MOTIVATION
text: The large number of trainable parameters in 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 a limitation that complicates the development of novel methods that may outperform the state-of-the-art: They must use the same objects to remain easy 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 full statistics of the mini-batch loss.
Recent advances in automatic differentiation (Bradbury et al., 2020; Dangel et 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 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 nets. It is also linked theoretically to the natural gradient method (Amari, 2000) via the Fisher information matrix (Martens, 2020, Section 9.2), and has been used to investigate the generalization of neural networks (Jastrzebski et al., 2020; Thomas et al., 2020).
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). 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 generality of automatic differentiation. However, repeated matrix-vector products represent a critical factor for performance.
Secondly, Kronecker-factored approximate curvature (K-FAC) (Martens & Grosse, 2015; Grosse & Martens, 2016; Botev et al., 2017; Martens et al., 2018) constructs an 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 very 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 other applications, like uncertainty quantification with Laplace approximations (Ritter et al., 2018b;a; Kristiadi et al., 2020) that currently rely on K-FAC. Moreover, 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 a downside for applications that depend on a precise curvature proxy.
Here, we propose VIVIT (inspired by V V > in Equation (3)), a vivid 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 exact GGN to an approximation that has the cost of a single gradient computation. Our contributions are as follows:
• We highlight the GGN’s low-rank structure, and with it a structural limit for the inherent curvature information contained in a mini-batch.
• This low-rank structure allows for efficient computation of various GGN properties: The exact eigenvalue spectrum, including eigenvectors, and per-sample directional derivatives. They enable VIVIT to model curvature noise in a mini-batch, in contrast to existing methods.
• Approximations allow VIVIT to flexibly trade off computational cost and accuracy. We empirically demonstrate scalability on deep neural networks and provide a fully-featured efficient implementation in PYTORCH (Paszke et al., 2019) on top of the BACKPACK (Dangel et al., 2020) package.1
Using VIVIT, we illustrate that noise in deep learning poses a challenge for the stability of secondorder methods and give a simple example how its quantities can be used to address this problem.

heading: 2 NOTATION & METHOD
text: Consider a model f : Θ × X → Y and a dataset {(xn,yn) ∈ X × Y}Nn=1. For simplicity we use N for both the mini-batch and training set size. The network, parameterized by θ ∈ Θ, maps a sample xn to a prediction ŷn. Predictions are scored by a convex loss function ` : Y × Y → R (e.g. cross-entropy or square loss), which compares to the ground truth yn. The training objective L : Θ→ R is the empirical risk
L(θ) = 1N ∑N n=1 `(f(θ,xn),yn) . (1)
We use `n(θ) = `(f(θ,xn),yn) and fn(θ) = f(θ,xn) for per-sample losses and predictions. For gradients, we write gn(θ) = ∇θ`n(θ) and g(θ) = ∇θL(θ), suppressing θ if unambiguous. We also set Θ = RD and Y = RC with D,C the model parameter and prediction space dimensions, respectively. For classification, C is the number of classes.
Hessian & GGN: Two-fold chain rule application to the split ` ◦ f decomposes the Hessian of Equation (1) into two parts∇2θL(θ) = G(θ) +R(θ) ∈ RD×D; the positive semi-definite GGN
G = 1N ∑N n=1 (Jθfn) > ( ∇2fn`n ) (Jθfn) = 1 N ∑N n=1Gn (2)
and a residual R = 1/N ∑N n=1 ∑C c=1 ( ∇2θ[fn]c ) [∇fn`n]c. Here, we use the Jacobian Jab that contains partial derivatives of b with respect to a, [Jab]ij = ∂[b]i/∂[a]j . As the residual may alter the Hessian’s definiteness – an undesirable property for many applications – we focus on the GGN.
Low-rank structure: By basic inequalities, Equation (2) has rank(G) ≤ NC.2 To make this explicit, we factorize the positive semi-definite Hessian ∇2fn`n = ∑C c=1 sncs > nc, where snc ∈ RC and denote its backpropagated version by vnc = [Jθfn]>snc ∈ RD. Absorbing sums into matrix multiplications, we arrive at the GGN’s outer product representation that lies at the heart of VIVIT,
G = 1N ∑N n=1 ∑C c=1 vncv > nc = V V > with V = 1√ N (v11 v12 . . . vNC) ∈ RD×NC . (3)
1Code available at https://github.com/PwLo3K46/vivit. 2We assume the overparameterized deep learning setting (NC < D) and suppress the trivial rank bound D.
V allows for exact 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.

heading: 2.1 COMPUTING THE FULL GGN EIGENSPECTRUM
text: Each GGN eigenvalue λ ∈ R satisfies the characteristic polynomial det(G − λID) = 0 with identity matrix ID ∈ RD×D. Leveraging the VIVIT factorization of Equation (3) and the matrix determinant lemma, the D-dimensional eigenproblem reduces to that of the much smaller Gram matrix G̃ = V >V ∈ RNC×NC which contains pairwise scalar products of vnc (see Appendix A.1),
det(G− λID) = 0 ⇔ (−λ)D−NC det(G̃− λINC) = 0 . (4) With at least D −NC trivial solutions that represent vanishing eigenvalues, the GGN curvature is flat 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. For example, the left panel of Figure 1a visualizes the full, exact GGN’s empirical spectral density on a mini-batch for a deep convolutional neural net on CIFAR-10. It reproduces the characteristics that have been reported by numerous works, e.g. Sagun et al. (2018): An extensive amount of vanishing or small eigenvalues and a small number of larger outliers.
Despite these 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 of works that extract the exact GGN spectrum from its Gram matrix. In contrast to those techniques, this matrix can be computed in parallel with gradients in a single backward pass, which results in less sequential overhead. In fact, our approach allows for plots like Figure 1 to be efficiently live-monitored during training, which may be interesting for practitioners that seek to better understand their model (Schneider et al., 2021).
Eigenvalues themselves can help identify reasonable hyperparameters, like learning rates. But we can also reconstruct the associated eigenvectors in parameter space. These are directions along which curvature information is contained in the mini-batch. Let S̃+ = {(λk, ẽk)|λk 6= 0, G̃ẽk = λkẽk}Kk=1 denote the nontrivial Gram spectrum with orthonormal eigenvectors ẽ>j ẽk = δjk (δ represents the Kronecker delta and K = rank(G)). Then, the transformed set of vectors ek = 1/ √ λkV ẽk are orthonormal eigenvectors ofG associated to eigenvalues λk (see Appendix A.2),
∀ (λk, ẽk) ∈ S̃+ : G̃ẽk = λkẽk =⇒ GV ẽk = λkV ẽk . (5)
The eigenspectrum provides access to the GGN’s pseudo-inverse based on V and S̃+, required by e.g. second-order methods (see Section 4.2). As we show next, quadratic models defined through the GGN naturally decompose along the eigenvectors from S+ = {(λk, ek) |λk 6= 0,Gek = λkek}Kk=1.

heading: 2.2 COMPUTING DIRECTIONAL DERIVATIVES
text: Various algorithms update their solution by constructing a local quadratic approximation of the loss landscape. For instance, optimization methods adapt their parameters by stepping to the local proxy’s minimum. Such quadratic models are based on a second-order Taylor expansion and hence require some form of curvature. Let q(θ) denote a quadratic model for the loss around position θt ∈ Θ that uses curvature represented by the GGN,
q(θ) = const + (θ − θt)>g(θt) + 1
2 (θ − θt)>G(θt)(θ − θt) . (6)
At its base point θt, 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 first-and second-order directional derivatives
γe = e >∇θq(θt) = e>g(θt) ∈ R , (7a)
λe = e >∇2θq(θt)e = e>G(θt)e ∈ R . (7b)
AsG’s characteristic directions are its eigenvectors, they form a natural basis for the quadratic model. Denoting γk = γek and λk = λek the directional gradients and curvatures along eigenvector ek, we see from Equation (7b) that the directional curvature indeed coincides with the GGN’s eigenvalue.
In addition to the mean gradient and curvature along an eigenvector ek, we can expand the sum over samples from Equation (1) to obtain the per-sample contributions to each derivative. Let γnk and λnk denote these first- and second-order derivatives contributions of sample xn in direction k, i.e.
γnk = e > k gn = ẽ>k V >gn√ λk , (8a)
λnk = e > kGnek =
ẽ>k V >VnV > n V ẽk λk = ‖V >n V ẽk‖2 λk , (8b)
where Vn ∈ RD×C is the VIVIT factor ofGn corresponding to 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 curvatureG are sums over gn andGn, respectively. This structure also carries over to the directional derivatives, i.e. γk = 1/N ∑N n=1 γnk and λk = 1/N ∑N n=1 λnk.
Access to per-sample directional gradients γnk and curvatures λnk along G’s natural directions is a distinct feature of VIVIT. Not only do these quantities provide geometric information about the local loss landscape but also about its directional stochasticity over the mini-batch. Incorporating such knowledge about the noise into algorithms that rely on quadratic models provides a promising way to increase their performance and stability. In Section 4.2 we show how to use this information to make second-order optimization methods more robust against noise.

heading: 2.3 COMPUTATIONAL COMPLEXITY
text: So far, we have formulated the computation of GGN eigenvalues (Equation (4)) including eigenvectors (Equation (5)) and per-sample directional derivatives (Equation (8)). Now, we analyze their computational complexity in more detail 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 evaluate their effectiveness on deep neural networks to demonstrate that VIVIT scales.
Relation to gradient computation: Machine learning libraries are optimized to backpropagate signals 1/N∇fn`n and accumulate the result into the mini-batch gradient g = 1/N ∑N n=1[Jθfn]
>∇fn`n. Each column vnc of V also involves applying the Jacobian, but to a different vector snc from the loss Hessian’s symmetric factorization. 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 as it involves NC backpropagations, while the gradient requires N . However, the practical overhead is expected to be smaller: Computations can re-use information from the Jacobians and enjoy significant additional speedup on parallel processors like GPUs. Because our implementation relies on BACKPACK’s vectorized Jacobians, we expect similar run time performance as its second-order extensions which have the same backpropagation.
Stage-wise discarding V (GGN eigenvalues & directional derivatives): The columns of V correspond to backpropagated vectors. During backpropagation, sub-matrices of V , associated to parameters in the current layer, become available once at a time. As many of the above GGN properties can be computed by contraction and accumulation of these stage-wise VIVIT factors, they can be discarded immediately. This allows for memory savings without any approximations.
One example is the Gram matrix G̃ formed by pairwise scalar products of {vnc}N,Cn=1,c=1 in O((NC)2D) operations. The spectral decomposition S̃+ has additional cost of O((NC)3). Similarly, the terms for the directional derivatives in Equation (8) can be built up stage-wise: First-order derivatives {γnk}N,Kn=1,k=1 require the vectors {V >gn ∈ RNC}Nn=1 that cost O(N2CD) operations. Second-order derivatives are basically for free, as {V >n V ∈ RC×NC}Nn=1 is available from G̃.
GGN eigenvectors: Raising one Gram matrix eigenvector ẽk to the GGN eigenvector ek through application of V (Equation (5)) costs O(NCD) operations. However, repeated application of V can be avoided for raising weighted sums of the form ∑ k( ck/ √ λk)ek with arbitrary weights ck ∈ R.
The summation can be performed in the Gram space at negligible overhead, and only a single vector∑ k ckẽk needs to be transformed. For instance, this allows for efficient aggregation of Newton steps along directions in the Gram space before transforming them to parameter space (see Section 4.2).

heading: 2.4 APPROXIMATIONS & IMPLEMENTATION
text: Although the GGN’s representation by V has linear memory cost inD, it requires memory equivalent to NC model copies. Of course, this is 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 and C in the above cost analysis to be replaced by smaller numbers, enabling VIVIT to trade-off accuracy and performance.
Curvature sub-sampling & MC approximation: To reduce the scaling inC, we can approximate the factorization ∇2fn`n(θ) = ∑C c=1 sncs > nc by a smaller set of vectors. One principled approach is to draw MC samples {s̃nm} such that Em [̃snms̃>nm] = ∇2fn`n(θ) as in Dangel et al. (2020). This reduces the scaling of backpropagated vectors from C to the number of MC samples (1 in the following). Another commonly used independent approximation to reduce the scaling in N is to compute curvature only on a subset of mini-batch samples (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 into the parameter space. Still, we can reduce costs by using the GGN’s diagonal blocks {G(i)}Li=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.
Concrete example & implementation details: To assess the quality of the above approximations, Figure 1a shows the exact GGN eigenvalue spectrum (left) in comparison to its approximation through curvature sub-sampling (center) and MC (right). Even though the amount of backpropagated vectors is reduced by a factor of 8 and 10, respectively, the approximated spectra capture essential features. We also tabularize the critical batch sizes Ncrit at which their computations experience out-of-memory errors in Figure 1b. On a standard GPU, they exceed the traditional batch size used for training, even for the exact scheme. Approximations further increase the applicable batch size.3
3The critical batch sizes in Figure 1b and c differ strongly for similar reductions of 8 and 10 in the number of backpropagated vectors. This is because most neural networks have final layers with many parameters. During initial stages of backpropagation, expanding V for these weights critically affects peak memory, and thus Ncrit. This can be improved by leveraging structure in the Jacobian (see Appendix C.1).
As a concrete example for block-diagonal approximations, we group weights and biases layerwise and compute the exact Newton step implied by the GGN’s block-diagonal approximation (similar to Zhang et al. (2017)). Figure 1c shows the critical batch size and supports our approach’s scalability.
BACKPACK’s functionality allows us to efficiently compute individual gradients and V in 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-batch sub-sampling and parameter groups. By hooks into the package’s extensions, we can discard buffers as soon as possible during backpropagation, effectively implementing all discussed approximations and optimizations.

heading: 3 RELATED WORK
text: 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 an alternative to K-FAC based on backpropagating a 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 obtain the exact spectrum and leverage a highly-efficient fully-parallel implementation. This may be because, until recently (Bradbury et al., 2020; Dangel et al., 2020), vectorized Jacobians required to perform those operations efficiently were not available.
Efficient operations with low-rank matrices in deep learning: Chen et al. (2021) use Equation (3) for element-wise evaluation of the GGN in fully-connected feed-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 and linear solves.
Another prominent low-rank matrix in deep learning is the un-centered gradient covariance (sometimes called empirical Fisher). Singh & Alistarh (2020) describe implicit multiplication with its inverse4 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 Section 2.

heading: 4 EXPERIMENTS
text: Different phases are encountered in the course of neural network training (Frankle et al., 2020). We view those regimes in light of our newly accessible quantities and identify challenges for second-order optimizers. Then, on a simple example, we use those quantities to stabilize such methods.

heading: 4.1 NOISE DURING TRAINING
text: The interaction between gradient, curvature, and their stochasticity is crucial for the behavior of optimization methods (Thomas et al., 2020). Here, we aim to identify characteristic features of gradient noise, curvature noise, and their interaction to gain insights into the conditions under which deep learning optimizers are operating. Training different architectures from the DEEPOBS problem set using the baselines from Dangel et al. (2020), we evaluate the GGN’s nontrivial eigenvectors and per-sample directional derivatives on a fixed held-out mini-batch with VIVIT. Figure 2 shows results for the 3C3D architecture trained on CIFAR-10 using SGD and cross-entropy loss. A broader evaluation that includes more problems, training with ADAM, and a description of procedural details, is given in Appendix B.2. We make the following observations.
Directional gradient-curvature correlations: Figure 2a shows pairs of GGN eigenvalues λk and the associated directional gradient magnitude |γk| at different training stages.5 Like for a quadratic function, gradients and curvatures are positively correlated.
4For completeness, we describe implicit multiplication with the inverse GGN in Appendix C.2. 5We show |γk| because the directional gradient varies in sign, depending on the eigenvector’s orientiation.
High gradient overlap with top eigenspace & shrinking of non-trivial eigenspace: Similar to Gur-Ari et al. (2018), we observe high alignment between the mini-batch gradient and the GGN’s top eigenspace. Specifically, we compute the normalized gradient overlap γ2k/‖g‖2 with direction k. Due to their additivity, gradient overlaps from multiple directions can be grouped by summation. The histograms in Figure 2b show such overlaps: Each bin summarizes directions of similar curvature. The accumulated normalized overlap of the gradient with these directions is shown as the height of the histogram bar. The gradient aligns mostly with high-curvature directions. In particular, its overlap with flat directions, 1 − ∑K k=1
γ2k/‖g‖2, is smaller than the axis limits. We observe this alignment throughout training even though the GGN’s active sub-space dimension K decreases.
Curvature signal vanishes during training & gradient signal is consistently small: To quantify the noise in both the directional gradients and curvatures, we compute their signal-to-noise ratios (SNR). It is given by the squared empirical mean divided by the empirical variance of the mini-batch samples {λnk}Nn=1 and {γnk}Nn=1 for each non-trivial direction k. Figures 2c,d show the evolution of both ratios. At early stages, the curvature signal along some of the high-curvature directions and parts of the bulk dominates over the noise (SNR(λk) > 1), i.e. curvatures of per-sample GGN s are similar. As training proceeds, the signal decreases until all directions are dominated by noise (SNR(λk) < 1). In comparison, the directional gradients do not exhibit such a pattern. They are always strongly corrupted by noise (SNR(γk) < 1), irrespective of the curvature, even though their mean γk correlates with λk (Figure 2a).
On the one hand, the high overlap between gradient and the top GGN eigenspace encourages gradient pre-conditioning with the GGN pseudo-inverse for second-order optimization, as negligible gradient
information will be projected out. On the other hand, however, noise in both the gradients and curvatures, which can vary among directions and during training over several orders of magnitude, represents a challenge for such methods. A naive Newton step based on the mean gradient and curvature along each direction may be unstable due to large noise, eliminating all previously made progress. We now turn to those methods and investigate how to improve their stability by equipping them with a noise-aware adaptation strategy to act individually along directions.

heading: 4.2 USE-CASE: IMPROVING DAMPING IN SECOND-ORDER METHODS
text: Spectral decomposition of Newton’s method into one-dimensional problems: Second-order optimizers are based on a local quadratic approximation q (cf. Equation (6)) of the objective function. A common approach to make such methods work in practice is to regularize the curvature matrix with a damping term δID, 0 < δ ∈ R. The resulting quadratic is minimized by the update ∆θ = −(G + δID)−1g. In the GGN eigenvector basis from S+ and the trivial spectrum S0 = {(λk, ek) | λk = 0,Gek = 0}Dk=K+1 we observe different updates in the respective sub-spaces,
∆θ = −(G+ δID)−1g = − ∑K k=1 γk λk+δ ek − ∑D k=K+1 γk δ ek . (9)
In the non-trivial eigenspace, ∆θ takes damped Newton-type steps, while flat directions are updated with SGD at a learning rate δ−1. Motivated by Section 4.1 and Gur-Ari et al. (2018), we omit the SGD update along flat directions in the following due to their negligible overlap with the gradient.
Global damping versus directional damping: Due to mini-batching, the gradient and curvature of the quadratic model are corrupted by sub-sampling noise. Hence, the mini-batch averages γk = 1/N ∑N n=1 γnk and λk = 1/N ∑N n=1 λnk in Equation (9) may deviate considerably from the true values on the noise-free training loss. This noise can lead to overly large steps and thus cause the optimization procedure to become unstable. Damping helps to increase stability as it decreases the step size. However, its relative effect on the step length in a direction ek depends sensitively on the scale of the curvature λk. Also, even at comparable curvature scale, stronger damping may be required in the presence of high uncertainty over γk and λk. It is thus expected that treating all directions identically yields diminishing returns along certain directions while it is required to keep uncertain ones stable. This suggests a directional damping, i.e. an update of the form ∆θ = − ∑K k=1 γk λk+δk ek.
Noise-aware directional damping via bootstrap: We use the quadratic mini-batch model’s loss reductionR(δk) to assess the step sk = − γkλk+δk ∈ R in direction ek,
R(δk) = q(θ)− q(θ + skek) = −sk ( 1 N ∑N n=1 γnk ) − 12s 2 k ( 1 N ∑N n=1 λnk ) . (10)
From observations on a mini-batch, our goal is to minimize the training loss. Hence, we want to choose a damping such that its corresponding update not only reduces q, but consistently decreases the loss over all other mini-batch models as well. Ideally, one would compute directional derivatives on additional samples for this purpose. Instead, we use a resampling technique based on the nonparametric bootstrap (Efron, 1979) to simulate samples for R(δk) on other batches, only using information contained in γnk, λnk (details in Appendix B). For a given damping δk, this provides an indicator of what reduction in training loss to expect with the respective update. Taking the 5% percentile of the bootstrap-generated samples, we obtain a confident lower bound toR(δk). We then choose the δk that maximizes this lower bound from candidates on a discrete grid.
Evaluation on a noisy quadratic: We consider a quadratic loss function L(θ) = θ>Gθ with Gii = i
2 for i ∈ {1, ..., D = 20} and initialize θ at 100 · 1D. At each step, the optimizer observes unit vectors as directions {ek}Kk=1, and noisy directional derivatives {γnk} N,K n=1,k=1, {λnk} N,K n=1,k=1. Specifically, we design the noise to provide unbiased samples with constant variance (see Appendix B for details). We compare our directional damping based on Equation (10) with constant damping and SGD. The results reported in Figure 3 were obtained from multiple runs with different random seeds.
With constant damping, there is a trade-off between large steps that may cause instabilities (small damping) and slower progress due to smaller steps (large damping). Both extremes are observed in our experiment. For δ = 10−4, there are unstable runs, as indicated by the loss mean and the maximum final distance to the minimum. For large constant damping, the behavior becomes more
stable and increasingly resembles SGD. Compared to constant damping, the directional noise-aware damping δk provides the smallest median final distance to the minimum with relatively low variance.
The right panel of Figure 3 shows the directional damping δk in relation to the underlying true curvature k2 for a single run. We make two observations: Firstly, the relative damping increases over time to compensate for the vanishing gradient signal as the optimization approaches the minimum. Secondly, since constant noise is applied, the SNR in low curvature directions (yellow) is smaller than in high-curvature directions (red). This is also reflected in the directional damping, which tends to assign larger dampings to the noisier (yellow) directions. The directional bootstrap damping works as expected and yields stable runs with consistently good final performance.
We provided evidence that VIVIT’s quantities may be required for the improvement of second-order methods. As the main contribution of this paper is in delivering these quantities, we deliberately designed simplistic experiments — a full derivation and empirical evaluation would amount to a separate paper of its own right (and would be incommensurate with the space limitations).

heading: 5 CONCLUSION
text: We have presented VIVIT, a curvature model based on the low-rank structure of the Hessian’s generalized Gauss-Newton approximation. This structure allows for efficient extraction of curvature properties, such as the full eigenvalue spectrum and directional gradients and curvatures along the associated eigenvectors. In contrast to alternatives, VIVIT offers statistics of these directional derivatives across the mini-batch, and thus a rich noise model.
We demonstrated the utility of these new quantities by studying noise characteristics of representative deep learning tasks. We find that they pose challenges to the stability of second-order methods, and showed, in a simplistic toy model, how VIVIT can provide quantities to improve their stability.
VIVIT’s representation is efficiently computed in parallel with gradients during a single backward pass. As it mainly relies on vectorized Jacobians, it is even general enough to be integrated into existing machine learning libraries in the future. For the moment, we provide an efficient open-source implementation in PYTORCH by extending the existing BACKPACK library.

Review summaries:
This work proposes a curvature model VIVIT based on generalized Gauss-Newton (GGN) approximation for the training of neural networks with a convex loss function. The low-rank structure of VIVIT allows for efficient eigen-value decomposition, which also gives per-sample directional derivatives and curvatures. To further improve the efficiency, sampling within the mini-batch and among the coordinates of the prediction can be applied to trade off computational cost with accuracy. As an application example, VIVIT is used to provide noise-aware directional damping which improves the stability of second-order methods.

This paper highlights how the low-rank structure of the generalized Gauss-Newton (GGN) approximation of the Hessian can be used as a computationally efficient tool to study the loss landscape of deep neural networks. In particular, authors discuss methods to compute the full spectrum of the GGN and thus providing access to per-sample directional gradients and curvature approximation. Through the lens of the GGN spectrum, authors make observations on the geometry of the loss landscape and its evolution during training, and propose an adaptive damping technique for second-order optimizers that utilizes the GGN curvature information.