NNGeometry: Easy and Fast Fisher Information Matrices and Neural Tangent Kernels in PyTorch
Authors: Thomas George
Venue:  (2021)

Abstract: Fisher Information Matrices (FIM) and Neural Tangent Kernels (NTK) are useful tools in a number of diverse applications related to neural networks. Yet these theoretical tools are often difficult to implement using current libraries for practical size networks, given that they require per-example gradients, and a large amount of memory since they scale as the number of parameters (for the FIM) or the number of examples x cardinality of the output space (for the NTK). NNGeometry is a PyTorch library that offers a simple interface for computing various linear algebra operations such as matrix-vector products, trace, frobenius norm, and so on, where the matrix is either the FIM or the NTK, leveraging recent advances in approximating these matrices. We here present the library and motivate our design choices, then we demonstrate it on actual deep neural networks.

Under review as a conference paper at ICLR 2021
NNGEOMETRY: EASY AND FAST FISHER INFORMA- TION MATRICES AND NEURAL TANGENT KERNELS IN PYTORCH
Anonymous authors Paper under double-bling review
ABSTRACT
Fisher Information Matrices (FIM) and (finie-width) Neural Tangent Kernels (NTK) are useful tools in a number of diverse applications related to neural net- works Pascanu & Bengio 2013); Kirkpatrick al. (2017): WW al. (2117) Liang] et al. (2019): Du et al. (2018). 8). Yet these theoretical tools are often difficult to im- plement using current libraries for practical size networks, given that they require per-example gradients, and a large amount of memory since they scale as the num- ber of parameters (for the FIM) or the number of examples x cardinality of the output space (for the NTK). NNGeometry is a PyTorch library that offers simple interface for computing various linear algebra operations such as matrix-vector products, trace, frobenius norm, and so on, where the matrix is either the FIM the NTK, leveraging recent advances in approximating these matrices. We hereby introduce the library and motivate our design choices, then demonstrate it on modern deep neural networks.
Code for this paper is available at this (anonymized) repo: https://githu: com/ Jmm /otumm///oorr////eeeeett
Practical and theoretical advances in deep learning have been accelerated by the development of an ecosystem of libraries allowing practitioners focus on developing new techniques instead of spending weeks or months re- re-implementing the wheel. In particular, automatic differentiation frameworks such as Theano (Bergstra et al.] (01T). Tensorflow (Abadi et al.] (016) or PyTorch (Paszke et al. 2019) have been the backbone for the leap in performance of last decade's increas- ingly deeper neural networks as they allow to compute average gradients efficiently, used in the stochastic gradient algorithm or variants thereof. While being versatile in neural networks that can be designed by varying the type and number of their layers, they are however specialized to the very task of computing these average gradients, more advanced techniques can be burdensome to implement.
While the popularity of neural networks has grown thanks to their always improving performance, other techniques have emerged, amongst them we highlight some involving Fisher Information Ma- trices (FIM) and Neural Tangent Kernels (NTK). Approximate 2nd order (Schraudolph) (002) or natural gradient techniques (Amari 1998) aim at accelerating training, elastic weight consolida- tion (Kirkpatrick et al. (2017) proposes to fight catastrophic forgetting in continual learning and WoodFisher (Singh & Alistarh, 2020) tackles the problem of network pruning so as to minimize its computational footprint while retaining prediction capability. These 3 methods all use the Fisher Information Matrix while formalizing the problem they aim at solving, but resort to using different approximations when going to implementation. Similarly, following the work offacot et al. (2018), a line of work study the NTK in either its limiting infinite-width regime, or during training of actual finite-size networks.
All of these papers start by formalizing the problem at hand in very concise math formula, then face the experimental challenge that computing the FIM or NTK involves performing operations for which off-the-shel automatic differentiation libraries are not well adapted. An even greater turooff comes from the fact that these matrices scale with the number of parameters (for the FIM) or the number of examples in the training set (for the empirical NTK). This is prohibitively large for modern neural networks involving millions of parameters or large datasets, a problem circumvented by series of techniques approximate the FIM (Ollivier 2015 Martens & Grosse 2015 George

---

Under review as a conference paper at ICLR 2021
using KFAC Fisher
using implicit computation
F_sfac
F 18 (model modell,
F_ F_8 (model=model),
loader=loader:
lsader=losdar,
@prrsertatioonnPAaRRPPP n_output=101
AnpresentattonnPMatttee
ALoutput=10)
V
PVactor. From srcommmodelllloeoeeeeeeeeeoooe (madel)
V. PVector,
from_mode PWector.fron_mett
(model)
VINv F. kfac.. (V)v
vIWv F. fffullvvw
Figure 1:
Computing a vector-Fisher Fisher-vectorr product v Fv, for a 10-fold classification model de- fined by model, I, can be implemented with the same piece of code for 2 representations of the FIM using NNGeometry, even if they involve very different computations under the hood.
accelerate development or analysis new techniques, allowing to spend more time on the theory and less time in fighting development bugs. NNGeometry's 's interface is designed to be as close as possible to maths formulas. summary, this paper and library contribute:
al. 2018 NNGeometry aims making use of these approximations effortless,
We introduce NNGeometry by describing and motivating design choices.
unified interface for all FIM and NTK operations, regardless of how these are ap- proximated.
Implicit operations for ability scale to large networks..
Using NNGeometry, get empirical insights on FIMs and NTKs:
We compare different approximations in different scenarios.
We scale some NTK evolution experiments to Tinylmagenet.
PRELIMINARIES
1.1 NETWORK LINEARIZATION
Neural networks are parametric functions / (ar, w) X' Rd - R where i X' covariates from an input space, and w E Rd are the network's parameters, arranged in layers composed of weight matrices and biases. The function returns value in R', such as the scores softmax classification, or è real values in c-dimensional regression. Neural networks are trained by iteratively adjusting their parameters (t++1) + w + ( using steps (((() typically computed using the stochastic gradient algorithm variants thereof, in order minimize the empirical risk of loss function.
In machine learning, understanding and being able control the properties of the solution obtained by an algorithm is of crucial interest, as it can provide generalization guarantees, or help design more efficient or accurate algorithms. Contrary expressions of the empirical risk minimizer exist, deep networks are non-linear functions, whose generalization properties and learning dynamics is not yet fully understood. Amongst the recent advances toward improving theory, is the study of the linearization (in w) of the deep network function f (ar, w): f (z. w + (I, (I, w) + (z, w) dw + ([[ów])
(kernelized) linear models, where closed-form
(I)
where J (xx w) (f(w.w is the Jacobian with respect parameters w, computed in (w, I.), map- ping changes in parameter space 5w to corresponding changes in output space using the identity öʃ (xx, w, = J (x. w) åw. For tiny steps dw, neglect the term o ([lōw[l) thus is close its linearization. happens for instance at small step sizes, or in the large-width limit with the specific parameter initialization scheme proposed by Jacot et al. 2018).
'The Landau notation (pronounced "little-o") means function whose exact value irrelevant, the property that lim o(x) 0, or in other words that negligible compared for small
2

---

Under review as a conference paper at ICLR 2021
1.2 PARAMETER SPACE METRICS AND FISHER INFORMATION MATRIX
While neural networks are trained by tuning their parameters w, the end goal of machine learning is not to find the best parameter values, but rather to find good functions, in a sense that is dependent of the task at hand. For instance different parameter values can represent the same function (Dinh) al,, 2017). On the contrary 2 parameter space steps dw1 and dw2 with same euclidean norm can provide very different changes in function (6f (zz w, w,W1) * öf (x, w, ww.).))). )). In order to quantify changes of a function, one generally defines distance on the function space. Examples of such distances are the Lgonom,, Wasserstein distances, or the KL divergence used in information geometry.
To each of these function space distances correspond a parameter space metric. We continue our exposition by focusing on the KI. divergence, which closely related to the Fisher Information Matrix, but our library can be used for other function space distances. Suppose f is interpreted as log-probability of density p. log (x, w) = f (z, w), the KL divergence gives a sense how much the probability distribution changes when adding small increment ow to the parameters of f (x, w). We can approximate it as:
[ log
(zr, w) (p w + ów)
KL (p (z, w) [pp (zz w + (w))
dp (I, (x, w)
(2)
JUEX
2 #EX (p(z, J (z. w) dw
dp (z, w) + (1ōwl(²
(3)
where we used this form (derived in appendix) in order to emphasize how steps in parameter space ow affect distances measured on the function space: equation 3 is the result of i) taking a step 5w in parameter space; ii) multiplying with J (ær, w) push the change to the function space; iii) weight this function space change using p (z, w) iv) square and sum. In particular, because of the properties of the KL divergence, there is no second derivative of f involved, even if equation B equivalent to taking the 2nd order Taylor series expansion of the KI. divergence. We can rewrite in more concise way:
KI. (I, w) (x, w + (w)) 8w Fwow + o (166w(²)
which uses the d x d FIM Fw = SeEx miz (2r, w) J (xr w) dp (ar, w). In particular, we can now define the norm [ów] F., = ów Fwow used in the natural gradient algorithm (Amari (1998 also see Martens 2020) for more thorough discussion of the FIM), in elastic weight consondation (Kirkpatrick et a 2017), or in pruning (Singh & Alistarh 2020). Other quantities also share the
same structure of a covariance of parameter space vectors, such as the covariance of loss gradients in
TONGA (Roux et al. 2008), the second moment of loss gradients (Kunstner et al. 2019 Thomas
ct al. 2020), or posterior covariances in bayesian deep learning (e.g. in Maddox et al. (2019)).
1.3 NEURAL TANGENT KERNEL
Another very active line of research around the linearization of equation T iis take inspiration from the rich literature on kernel methods by defining the neural tangent kernel (NTK):
Niw (x. v) = J (x, w) J (v. w)
(5)
In the limit of networks infinite width, Jacot et al. (2018) have shown that the tangent kernel remains constant through training using gradient descent, which allows to directly apply kernel learning theory to deep learning. While this regime is of theoretical interest, arguably does not explain what happens at finite width, where the NTK evolves during training.
While kernels are functions of the whole input space X' X, often only have access to limited number of samples in a datasets. We thus resort to using the kernel evaluated at points I; of a
We here use the notion of distance informally.
'The second moment of loss gradients sometimes called empirical Fisher.
3

---

Under review as a conference paper at ICLR 2021
Layer Collection
Describes (layen))
parameter space
* *eectors
Generator
'Mat Representations
Stores required elements in memorry
Computes jacobians
Depending representation:
required memory
matris mulris--vccor vecior opetatiots
Implements lincar algebra operations
populates
Figure 2: Schematic description of NNGeometry's main components
training or test set, called the Gram Matrix (Kw),; = kiw (æt, x2). Note that in the case where the output space is multidimensional with dimension then Kw is in fact a 4d tensor.
2 DESIGN AND IMPLEMENTATION
2.1 DIFFICULTIES
Current deep learning frameworks such as PyTorch and Tensorflow are well adapted to neural net- work training, i.e. computing average gradients over parameters, used in optimizers such as Adam and others. However, when going to more advanced algorithms or analysis techniques involving FIMs and NTKs, practitioners typically have to hack the framework's ' internal mechanisms, which time consuming, error prone, and results in each project having its own slightly different imple- mentation of the very same technique. We here list the difficulties in computing FIMs and NTKs using current frameworks:
Per-example gradient FIMs and NTKs require per-example Jacobians J (xz, w) of dataset (r:),. This can be obtained by looping through examples v, but at the cost of not using mini-batched operations, thus missing the benefit using GPUs. NNGeometry's 's Jacobian generator extensively use efficient techniques such Goodfellow 2015)
Memory usage and computational cost A FIM matrix is d x d where d is the total number of parameters. With a memory cost in O (d²). this is prohibitively costly even for moderate size networks. Typical linear algebra operations have a computational cost in either (dP) (e.g. matrix- vector product) or even O (d²) (e.g. matrix inverse). NNGeometry instead comes with recent lower memory intensive approximations.
2.2
NNGEOMETRY's S DESIGN
2.2.1 ABSTRACT OBJECTS
In section I we have worked with abstract mathematical objects 5w, 5f (sr, ww J (x, w), Fw. and Kw. We now identify these mathematical objects to Python classes NNGeometry.
We start with the parameter space, that we previously identified as d.d Closer to how they are actu- ally implemented in deep learning frameworks, vectors the parameter space w can equivalently be considered set of weight matrices and bias vectors w { bi.... Wi, br). Parameter space vectors are represented by the class PVector in NNGeometry, which is essentially a dictionary of PyTorch P Parameted cemeterss with basic algebra logic: PVectors can be readily added, substracted, and scaled by a scalar with standard python operators. As an illustration wsun wl x2 internally loops through all parameter tensors of and and returns new PVector w_avm.
4

---

Under review as a conference paper at ICLR 2021
Similarly, and more interestingly, parameter space metrics such as the FIM are represented by classes prefixed with PMat. For instance, the natural gradient dus = =n V S applies the linear opera- tor w + F- w to the parameter space vector V S, and can be implemented cleanly and concisely using delts_nntt
eta F.aol ve inabla_T,), even it internally involves different operations for different layer types, and different approximation techniques.
Function space vectors FVector define objects associated to vectors of the output space, evaluated on dataset of examples X. As an example, getting back to the linearization 8f (2,, w, dw) J (x, w) dw, we define of (X) = (ōf (Iı, w, (ww)...... 6f (In, w, (, as the IR function space vector of output changes for all examples of X. Gram matrices of the NTK are linear op- erators on this space, represented by objects prefixed with FMat. Borrowing from the vocabulary of differential geometry, we also define PushFo rd objects that are linear operator from parame- ter space to function space, and Pul 1 Back objects that are linear operator from function space to parameter space.
While the following consideration can be ignored upon first glance, the structure of the parame- ter space is internally encoded using a LayerColl eyyeeClllett.. ion object. This gives the flexibility of defining our parameter space as parameters of subset of layers, in order to treat different layers in different ways. An example use case is to use KFAC for linear layers parameters, and block diagonal for GroupNorm layers, as KFAC is not defined for the latter.
2.2.2 CONCRETE REPRESENTATIONS
These abstract objects are implemented in memory using concrete representations. NNGeometry comes with a number of representations. Amongst them, most notably, are parameter space approx- imations proposed in recent literature (Ollivier, 2015 Martens & Grosse 2015 Grosse & Martens) 2016 George et al. 2018), and an implicit representation for each abstract linear operator, that al- lows compute linear algebra operations without computing or storing the matrix in memory.
PMat (resp PMatDense) and PMatDiag represent the full dense matrix and the diagonal matrix and need no further introduction. PMat LowRank only computes and stores J (X, w) the C X n x stacked Jacobian for all examples of the given dataset.
Next come representations that do not consider neural networks as black-box functions, but instead are adapted to the layered structure of the networks: MMatBlockDial uses dense blocks of the FIM for parameters of the same layer, and puts zeros elsewhere, ignoring cross-layer covariance. PMat IitatuuasDDiee
(Ollivier 2015) uses the full diagonal and adds to each bias element the in- teraction with the corresponding row of the weight matrix. PMat uses KFAC (Martens & Grosse 2015) and its extension to convolution layers KFC (Grosse & Martens) 2016) to approxi- mate each layer blocks with the kronecker product of 2 much smaller matrices, thus saving memory and compute compared to PMat ockDi PMatEKF AC uses the EKFAC (George et al. 2018) extension of KFAC.
The last representation that comes with this first release of NNGeometry, PMat t Impl 1 i ciitt allows to compute certain linear algebra operations using the full dense matrix, but without the need to ever store it in memory, which permits scaling to large networks (see experiments in section 3) As an illustration, the vector-matrix vector product v Fvv can be computed using equation 3
Each representation comes with its advantages and drawbacks, allowing trade-off between mem- ory and approximation accuracy. For a new project, we recommend starting with a small network using the PMat representation, then gradually switching to representations with lower memory footprint while experimenting with actual modern networks.
While linear algebra operations associated to each representation internally involve very different mechanisms, NNGeometry's core contribution is to give easy access to these operations by using the same simple methods (figure
2.2.3 GENERATORS
In order compute FIMs and NTKs, need to compute Jacobians J (z,, w) for examples r coming from a dataset. NNGeometry's generator is the component that actually populates the representa- tions by computing the required elements of the matrices, depending the representation. While
5

---

Under review as a conference paper at ICLR 2021
naive idea would be to loop through examples z2. compute f (2), w) and compute gradients with respect to parameters using PyTorch's automatic differentiation, it is rather inefficient as it does not make usage of parallelism in GPUs. NNGeometry's 's generator instead allows to use minibatches of examples by intercepting PyTorch's gradients and using techniques such as those in (Goodfellow 2013) and (Rochette et al., (0T9):
Let us consider f (I. w) X x IR - RR. In order to simplify exposition, we focus on fully con- nected layers and suppose that f can be written f fa, w) = at gi (.w) o Ot 9-11 (-. w) o o 91 (Ir w) where di are activation functions and Oh are parametric affine transforma- tions that compute pre-activations 8y of layer using a weight matrix Wi and a bias vector by with the following expression:
9k (ax-1: w) = W44x-1 + bg- For each example Ta in minibatch, we denote these intermediate quantities by superscripting Sk
(6) (8) and de The back-
propagation algorithm applied to computing gradients of a sum S Σ, f (24, w) works by se- quentially computing intermediate gradients (f(r).) from top layers bottom layers. Denote (
[ffl,..) 00
@f(r..)) 88h
))
T
by DSa
the matrix obtained by stacking these gradients for a
(m)
(m) the corresponding matrix of activations of the same layer. These are already computed when performing the backpropagation algorithm, then used to obtain the average gradient w.r. the weight matrix by means of the matrix/matrix product aw; (E, J (I; w)} Ds as. The observation of Goodfellow (015)) is that we can addition
minibatch of size m, and a k ....... Qx
obtain individual gradients @f(z.,w ( an operation that can be efficiently done simultaneously az4") ag
for all examples of the minibatch using the brmn PyTorch function.
While we used this already known trick as an example of how to make profit of minibatching, NNGeometry's 's generator incorporate similar tricks in several other places, including in implicit operations.
Instead of reimplementing backpropagation as is for example done by Dangel et al. (2019), we chose to use PyTorch's 's internal automatic differentiation mechanism, as it already handles most corner cases encountered by deep learning practitioners: we do not have to reimplement backward compu- tations for every new layer, but instead we just have to compute individual gradients by intercepting gradients with respect to pre-activations DSk.
Other generators are to be added to NNGeometry in the future, either by using different ways of computing the Jacobians, or by populating representations using other matrices such as the Hessian matrix, or the KFRA approximation of the FIM (Botev et al. 2017).
3 EXPERIMENTAL SHOWCASE
Equipped with NNGeometry, we experiment with large network: We train a 24M parame- ters Resnet50 network on Tinylmagenet. We emphasize that given the size of the network, we would not have been able to compute operations involving the true F' without NNGeometry's PMat Impl icit representation, since F would require 2.3 petabytes of memory (24M x 24M x 4 bytes for float32).
3.1 QUALITY OF FIM APPROXIMATIONS
We start by comparing the accuracy of several PMat representations computing various linear algebra operations. We use Monte-Carlo estimate of the FIM, where we use samples from p (y)x) for each example æ. Here, since this Tinylmagenet is a classification task. p (y]æx) is multinoulli distribution with the event probabilities given by the softmax layer. We compare the approximate value obtained for each representation, to a "true" value, obtained using the full matrix with the PMat implicit representation. For trace and v Fv, we compare these quantities using the relative difference approx - irue
For Fv, we report the cos-angle F (Fv, Fapprox v),
truc
and for the solve operation, report the cos-angle between and (Fappoox + XT) 1 (F AT)

---

Under review as a conference paper at ICLR 2021
during training
model
PMaiDiag
PMarQuasiDiag
PMarKFAC
residual
1.0
PMaiEKFAC
0.5
10
10-7
A
y-vVll
and cos angle between v and v' = (Fapproc + +1) (F X/) for 24M parameters Resnet50 at different points during training Tinylmagenet, using different approximations Fapprox of F, for v uniformly sampled on the unit sphere (higher is better).
Figure 3: Residual Illl
during training
0.5
0.3
0.96
1,00
1.000
0.95
PhinDeod
Phiak PAC
Polat KPNC
Phlat Palet
PalakFap
Pshal PahalEEE
Palak
Deeou War
Figure 4: Cos angle between Fv and Papprox for 24M parameters Resnet50 at different points dur- ing training on Tinylmagenet, using different approximations Fappoox F, for v uniformly sampled on the unit sphere (higher is better).
Since the latter is highly dependent the Tikhonov regularization parameter X. plot the effect on the cos-angle of varying the value of A. The results can be observed in figures 3 6
From this experiment, there best representation for all linear algebra operations. Instead, this analysis suggest to use PMatKFAC when possible for operations involving the inverse FIM, and PMatEKFAC for operations involving the (forward) FIM. Other representations are less accurate, but should not be discarded as they can offer other advantages, such as lower memory footprint, and faster operations.
3.2 NEURAL TANGENT KERNEL EIGENVECTORS
In the line of Baratin et al. 2020): Paccolat et al. 2020 we observe the evolution of the NTK during training. We use the Resnet50 on the 200 classes of Tiny Imagenet, but in order be able to plot 2d matrix for analysis, we extract the function Sey,,g (I. w) = (ʃ (ar, w)) (ʃ (r. w)) namely binary classifier of class 02 vs class c1 We plot at different points during training i) the and e2 (figure 7 top row) and ii) kernel pea of
Gram matrix of examples from the 2 classes points from classes e₁ and C2 projected on the 2 first principal components (figure 7 bottom row). The Gram matrix is computed for valid set examples of classes and C2.
On this larger network, we reproduce the conclusion of Baratin et al. (2020); Paccolat et al. (2020) that the NTK evolution is not purely random during training, but instead adapts to the task in very specific way.
CONCLUSION
We introduced NNGeometry, a PyTorch library that allows to compute various linear algebra oper- ations involving Fisher Information Matrices and Neural Tangent Kernels, using efficient imple- mentation that is versatile enough given current usages of these matrices, while being easy enough save time for the user.

---

Under review as a conference paper at ICLR 2021
during training
0.25
0.4
0.4
0.20
0.3
0.15
0.10
0.1
0.0
PSteQuaaniar
PataDaa
MMteQuaDee
PlaakFa
PMatEAFaC
PaleDia
PA
Figure 5: Relative difference between Pv and Fapprox v for 24M parameters Resnet50 at different points during training on Tiny linlmmaenne, using different approximations F Fappoux F, for v uniformly sampled on the unit sphere (higher is better).
a1 inil.
Juring training
best model
0.3
0.3
0.2
0.1
0.0
0.0
PMat
PMaaIAc
PAGRBEAF
PMtat
PAtatRPAC
Figure 6: Relative difference of trace computed using Fappox and F (lower is better). As we observe, all 3 representations PMat iag, PMatQuas liag and PMаtミKF AC estimate the trace very ac- curately, since the only remaining fluctuation comes from Monte-Carke sampling of the FIM. On the other hand, the estimation provided by PMat is less accurate.
at init.
during training
best model
at Init.
during training
best model
2nd comporent
0.2
2nd component
2nd component
0.0
0.0
-0.1
-0.2
-0.2
-0.1 component
-0.15
-0.10 -0.05
component
component
Figure 7: NTK analysis for 50 examples class ej and 50 examples class C2 at various points dur- ing training. (top row) Gram matrix of the NTK. Fach row and column is normalized by /diog (G) for better visualization. We observe that the NTK encodes some information about the task later in training, since highlights intra-class examples. (bottom row) Examples are projected on the 1st principal components of the Gram Matrix at various points during training. While points are merely mixed initialization, the NTK adapts to the task and becomes good candidate for kernel PCA since examples become linearly separable as training progresses.
We hope that NNGeometry will help make progress across deep learning subfields as FIMs and NTKs are used in range of applications.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large- scale machine learning. In 12th {USENIX] symposium on operating systems design and imple- mentation ((OSDI) 16), pp. 265-283, 2016.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pascal Vin- cent, and Simon Lacoste- Julien. Implicit regularization in deep learning: A view from function space. arXiv preprini 20XXix:0008.009388 2020.
James Bergstra, Frédéric Bastien, Olivier Breuleux, Pascal Lamblin, Razvan Pascanu, Olivier De- lalleau, Guillaume Desjardins, David Warde-Farley, Ian Goodfellow, Arnaud Bergeron, et al. Theano: Deep learning on gpus with python. In NIPS 2011, BigLearning Workshop, Granada, Spain, volume 3, pp. 1-48. Citeseer, 2011.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. International Conference on Machine Learning, pp. 557-565, 2017.
Felix Dangel, Frederik Kunstner, and Philipp Hennig. Backpack: Packing more into backprop. In International Conference Learning Representations, 2019.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. arXiv preprint arXiv:1705 04933, 2017.
Simon S Du, Xiyu Zhai, Barnahas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint 18100000, 2018.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast ap- proximate natural gradient descent kronecker factored eigenbasis. In Advances in Neural Information Processing Systems, pp. 9550-9560, 2018.
lan Goodfellow. Efficient per-example gradient computations. arXiv preprint arXiv:1510.0179)
2015.
Roger Grosse and James Martens. A cromecker-factoree approximate fisher matrix for convolution layers. In International Conference on Machine Learning, 573-582, 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In Advances in neural information processing systems, pp. 8571- 8580, 2018.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska- -arnskkkkaaiikkaa al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):352-3335 2017.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approxi- mation for natural gradient descent. In Advances in Neural Information Processing Systems, pp. 4156-4167, 2019.
Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-rao metric, geom- etry, and complexity of neural networks. In Proceedings of Machine Learning Research, vol- ume 89, pp. 888-896, 2019.
Wesley J Maddox. Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems, pp. 13153-13164, 2019.
James Martens. New Insights and Perspectives on the Natural Gradient Method. Journal of Ma- chine Learning Research, 21(146):1-76,, 2020. URL http:/ jjnrr. org/papers/v211
7- 678. html

---

Under review as a conference paper at ICLR 2021
James Martens and Roger Grosse. Optimizing neural networks with stoonecker-fcctoree approximate curvature. In International conference on machine learning, pp. 2408-24111 2015.
Yann Ollivier. Riemannian metrics for neural networks i: feedforward networks. Information and
Inference: Journal of the IMA, 4(2): (()):108-553, 2015.
Jonas Paccolat, Leonardo Petrini, Mario Geiger, Kevin Tyloo, and Matthieu Wyart. Geometric compression of invariant manifolds in neural nets, 2020.
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.3584, 2013.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. In Advances in neural information processing systems, pp- 8026-8037, 2019.
Gaspar Rochette, Andre Manoel, and Eric W Tramel. Efficient per-example gradient computations in convolutional neural networks. arXiv preprint 191vv0191..00015, 2019.
Nicolas L Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gra- dient algorithm. In Advances in neural information processing systems, pp. 849-856, 2008.
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent.
Neural computation, 14(7):1723-1738, 2002.
Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximations for model
compression. In Advances in Neural Information Processing Systems, 2020.
Valentin Thomas, Fabian Pedregosa, Bart Merriënboer, Pierre-Antoins Manzagol, Yoshua Bengio, and Nicolas Le Roux. On the interplay between noise and curvature and its effect on optimization and generalization. In International Conference Artificial Intelligence and Statistics, pp. 3503- 3513. PMLR, 2020.
Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using cooneker-factoree approximation. Advances in neural information processing systems, pp. 5279-5288, 2017.
A APPENDIX
You may include other additional sections here.
10


Reviews:
[{"review_id": "uO_-1y5Me2", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper provides a PyTorch implementation for applying Fisher information matrix and neutral tangent kernel as a linear operator. This efficient implementation relies on several existing approximation methods and computational tricks. I think the work is practical as it can be used for many applications, some of which from works that are cited in this paper. The paper is reasonably well-written, but some details are missing. I'm mostly concerned with the novelty and and significance of this work, although I need to admit I'm unfamiliar with the acceptance criteria for papers on software implementation.\n\n1. I think most of the techniques in the paper are from previous work. Thus I'm not sure there's enough novelty for an acceptance. The paper did not highlight any improvement in scaling. I believe all the speedup comes from taking advantage of GPU operation using the trick from Goodfellow, and the collection of matrix approximation methods that the paper cited.\n\n2. The paper focused on implementation, but surprisingly the experiments did not include any scaling/runtime results. I think the authors should show the efficiency gain, and demonstrate the type the applications that are enabled by it.", "extra_scores": {}, "extra_text": {}}, {"review_id": "pABlBOWWy0", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper introduces a new PyTorch library for computing Fisher Information Matrices and Neural Tangent Kernel (NTK) in deep learning; with applications ranging from Frobenius norm regularization, second-order optimization, and generalization analysis. The authors begin by providing a background on Fisher matrices and NTK and then present the main components of their proposed NNGeometry library (consisting of Layer Collection, Generator, and  Concrete Representations modules). A brief experimental study is provided in the last part of the paper. \n\nAssessment: While I think that a clean and effective computational library for implementation of Fisher matrices would greatly benefit the DL/ML research community, the current work falls short of the standards of an ICLR publication in numerous ways. \n\nDetailed Comments:\n- Towards the end of the introduction, the authors claim \"NNGeometry aims at making use of these appoximations effortless...\". I really do not see how NNGeometry is capable of doing this. One of the reasons why natural gradient methods (and its approximations such as K-FAC) are rather difficult to implement is that they are \"white-box\" optimization methods and very much model-dependent, i.e., different model architectures means that the Fisher matrices (and their approximations) can look very different from each other. Please see [1], [2] for the K-FAC approximations needed for convolutional and recurrent networks respectively. As an aside, I do not believe there is even a good open-source code for K-FAC on RNNs especially given the complexity involved. I did not find anywhere in the paper how NNGeometry addresses approximations for different types of layers.\n\n- A lot of spelling mistakes and strange notations throughout the paper. Here are a few I found- \"jacobian\" is not capitalized throughout, the big \"O\" notation in Equations 3 and elsewhere in the paper are all small \"o\"'s\n\n- The discussion in Section 1.2 is a bit wordy and not precise mathematically at times. I would suggest cutting down and citing [3]. Also, somewhere in the introduction, I think the authors should make a note of the distinction between the empirical Fisher matrix and the true Fisher matrix (and perhaps cite [4]); and be clear about which one they are working with. \n\n- Given that this is a library concerned with computing FIM/NTK; there should be some comparisons with existing open-source libraries such as JAX and Neural Tangents? \n\n- Lots of issues in Sections 2.2.1 and 2.2.2. It is not exactly clear where the authors are trying to do here; and there are many imprecise/incorrect mathematical statements throughout. I believe that the purpose of Section 2.2.1 was to describe that the FIM defines a Riemannian metric on the parameter space; and that the FIM is a representation of this metric in coordinate form. This is certainly true- but I cannot see the connection of this to the NNGeometry framework. Another purpose of this subsection was the notion of duality; for example, which objects may be pushed forward/pulled back (to be more precise, which ones live on the tangent and cotangent spaces). I would encourage the authors to look at the publicly-available JAX documentation/tutorial; where it is explained nicely how all of the theory + code fits together; JVP (Jacobian-vector products) / forward-mode autodiff <--> pushforward map of tangent spaces, VJP (vector-Jacobian products) / reverse-mode autodiff <--> pullback of cotangent spaces.\n\n- It would be great if the authors were more clear and explained explicitly the tricks in the sentence \"NNGeometry's generator incorporate similar tricks in several other places, including in implicit operations\". Many of these types of tricks are known to practitioners who have had to implement FIM (and its approximations); so I am curious what is the novelty provided by NNGeometry's generator here.\n\nReferences:\n\n[1] Grosse, Roger, and James Martens. \"A kronecker-factored approximate fisher matrix for convolution layers.\" International Conference on Machine Learning. 2016.\n\n[2] Martens, James, Jimmy Ba, and Matt Johnson. \"Kronecker-factored curvature approximations for recurrent neural networks.\" International Conference on Learning Representations. 2018.\n\n[3] Martens, James. \"New insights and perspectives on the natural gradient method.\" arXiv preprint arXiv:1412.1193 (2014).\n\n[4] Kunstner, Frederik, Philipp Hennig, and Lukas Balles. \"Limitations of the empirical Fisher approximation for natural gradient descent.\" Advances in Neural Information Processing Systems. 2019.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Sl-oJ2G56Lz", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe Fisher information matrix and the neural tangent kernel matrix have been used in several recent papers to provide insight into deep neural networks, but operations involving these matrices have so far been less well supported in frameworks such as Tensorflow and PyTorch.  The current paper describes a new library (NNGeometry) for working with these matrices in PyTorch, incorporating approximate representations from prior work that are appropriate to matrix-vector product computations and evaluation of quadratic forms.\n\nThe main contribution of the paper is the software library; no new analysis or algorithms are presented.  However, this seems like a useful software artifact for the particular problem that it tackles.  I recommend accepting it if there is space.\n\n### Additional details\n\n- The cosine angle measure is a somewhat peculiar choice for comparing approximate matrix-vector products and solves with their reference versions.  It is not only the direction of these vectors that matters, but also their magnitude, and so a normwise relative error seems much more natural.  For the (regularized) linear system, the natural choice might be a normwise relative residual.\n- While I understand the space constraints, I would have liked to see more discussion of the concrete matrix representations (Section 2.2.2).  In particular, there is not even a one-sentence description of the KFAC or EKFAC representation; only the name and a reference are given.  Adding a sentence or two would probably not break the page limit, and would make the paper more accessible to readers unfamiliar with those prior papers.\n\n### Minor points\n\n- Throughout, the authors should carefully check choice of citation style and how the citation reads (e.g. \"efficient techniques such as Goodfellow (2015)\" seems to refer to the cited paper as a technique; I suggest rewording to \"... such as those in Goodfellow (2015)\" (or \"of Goodfellow (2015)\").\n- Typos: \"inn\" to \"in\" (top of page 5), \"with on a large network\" - \"with a large network\" (start of Section 3)\n\n### Update\n\nThanks to the authors for their responses, particularly to the other reviewers who were more critical.  I still believe this is worth publishing if there is room, with the caveat stated before: this is an implementation paper, and does not introduce new algorithms or analysis.  If it is accepted, it should be accepted on this basis.", "extra_scores": {}, "extra_text": {}}, {"review_id": "K91eoFt6o3K", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper describes a new PyTorch package, NNGeometry, for computing complicated neural network objects, such as the Fisher Information Matrix (FIM) and the Neural Tangent Kernel (NTK). The package uses an abstract representation to allow the user to implicitly choose between different approximations to these objects and automatically makes a bunch of efficient choices \"under the hood\" for the user. The paper concludes with a selection of experiments that compare these approximations as well as verify their validity against Monte-Carlo estimates of these matrices.\n\nThe package seems like a really fantastic contribution to PyTorch. I will definitely try it out and most likely use it for my research.\n\nOn the other hand, I am worried that this is more of an engineering contribution than a research paper. To the point, all of the methods described in this paper and implemented in the package are based on methods present in other papers. Furthermore, the authors do not use the package to make any novel research contributions. Instead, they first show how their FIM approximations compare to Monte-Carlo estimates, and then they use their NTK approximations to reproduce some NTK results after training.\n\nI feel pretty undecided about whether ICLR should publish this paper. On the one hand, I think the package is very nice; it accumulates a bunch of useful approximations and allows for an efficient interface for switching between them. However, I don't know how significant this is as a research contribution; the authors propose no new methods nor do they achieve any novel results with their package. For now I think that the paper is not quite able to be published at ICLR, but I'm interested to hear what the authors and other reviewers have to say.\n\nTo that point, I wonder if the authors might use their package to show how it allows researchers to perform certain types of experiments that couldn't have previously been performed.\n\nSome minor comments and a question:\n- I think on Figure 3, the x-axis is lambda? It would be great if a label was included.\n- In the References section, Paccolat et al. (2020) is missing its arxiv number.\n- Could the authors include some detail about the type of Monte-Carlo experiments they did to compare for Figures 3-8? \n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper provides a high-level API for working with Neural Tangent Kernels (NTK) and Fisher Information Matrices (FIM). This is an implementation paper, but such concepts are clearly useful in many tasks. However, such methods are available in many in-house code (almost every paper on FIM / NTK uses such methods) I would not say it hampers the progress. 


Pros: - The proposed methods are already widely used by the communities in the in-house codes, but no single library is available
          - The library implements state-of-the art approaches

Cons: - This is an implementation library, and no benchmarking is available. More testing is needed to showcase the library. 
           - Matrix-by-vector products are provided, but typical operations include many more, like Lanczos method for approximating solutions of linear systems and matrix function by vector product, or randomized SVD to compute low-rank approximations. I believe that this should be the part of the library in order to make it a serious competitor for existing "local implementations". Also, such kind of methods would be later or sooner part of big libraries such as PyTorch and Tensorflow. 


Decision: Reject