Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures
Authors: Runa Eschenhagen
Alexander Immer
Richard E Turner
Frank Schneider
Philipp Hennig
Venue: NeurIPS 2023 spotlight (2023)
Keywords: deep learning
second-order
optimization
natural gradient
fisher
gauss-newton
k-fac
weight-sharing
TL;DR: We generalise Kronecker-Factored Approximate Curvature to modern deep neural network architectures like transformers, graph, and convolutional neural networks.

Abstract: The core components of many modern neural network architectures, such as transformers, convolutional, or graph neural networks, can be expressed as linear layers with *weight-sharing*. Kronecker-Factored Approximate Curvature (K-FAC), a second-order optimisation method, has shown promise to speed up neural network training and thereby reduce computational costs. However, there is currently no framework to apply it to generic architectures, specifically ones with linear weight-sharing layers. In this work, we identify two different settings of linear weight-sharing layers which motivate two flavours of K-FAC -- *expand* and *reduce*. We show that they are exact for deep linear networks with weight-sharing in their respective setting. Notably, K-FAC-reduce is generally faster than K-FAC-expand, which we leverage to speed up automatic hyperparameter selection via optimising the marginal likelihood for a Wide ResNet. Finally, we observe little difference between these two K-FAC variations when using them to train both a graph neural network and a vision transformer. However, both variations are able to reach a fixed validation metric target in $50$-$75$\% of the number of steps of a first-order reference run, which translates into a comparable improvement in wall-clock time. This highlights the potential of applying K-FAC to modern neural network architectures.

Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures
Runa Eschenhagen*.
Alexander Immer Richard E. Turner
Frank Schneider
Philipp Hennig
University of Cambridge
2 Department of Computer Science, ETH Zurich Max Planck Institute for Intelligent Systems
*University of Tübingen Tüübingen AI Center
Abstract
The core components of many modern neural network architectures, such as transformers, convolutional, or graph neural networks, can be expressed as linear layers with weight-sharing Cronecker-Factored Approximate Curvature (K-FAC), a second-order optimisation method, has shown promise speed up neural network training and thereby reduce computational costs. However, there is currently no framework to apply to generic architectures, specifically ones with linear weight-sharing layers. In this work, we identify two different settings of linear weight-sharing layers which motivate two flavours of K-FAC - expand and reduce. We show that they exact for deep linear networks with -hhriing in their respective setting. Notably, K- K-FAC-reduce is generally faster than K- K-FAC-expand, which leverage speed up automatic hyperparameter selection via optimising the marginal likelihood for a Wide ResNet. Finally, we observe little difference between these two K-FAC variations when using them train both graph neural network and a vision transformer. However, both variations are able to reach a fixed validation metric target 50-75% of the number of steps of a first-order reference run, which translates into a comparable improvement in wall-clock time. This highlights the potential of applying K-FAC modern neural network architectures.
Introduction
One of the key driving forces behind the success of deep learning is arguably the development and scaling of novel deep neural network (DNN) architectures like transformers (Vaswani et al., 2017) graph neural networks (Battaglia et al., 2018; Scarselli et al., 2009). While the landscape of neural network architectures seems vast, their core building blocks like attention, graph, recurrent, and convolutional layers can be expressed as simple linear layers with weight-sharing. For example, the veight-sharing can happen over a sequence of tokens. e.g. in language modelling, over image patches, by tying weights in the forward pass. The perspective of expressing the core neural network operations as linear operations has been formalised in a different context with the Tensor Programs framework (Yang, 2019), which develops a language to express arbitrary neural network computations as a composition of simple matrix multiplications and coordinate-we wise nonlinearities. One downside of relying on increasing the scale of the models to achieve better performance is the increased compute cost. To decrease these costs, we can try to improve the efficiency of the optimisation algorithms, which motivates exploring second-order methods like Newton's 's method and natural gradient descent (Amari, 1998), which uses the Fisher information matrix, or short, the Fisher.
Coorespondence to: re3930cam.co ac.uk
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Kronecker- Froneccer-Factored Approximate Curvature (Heskes, 2000: Martens & Grosse, 2015, K-FAC) is an upproximation to the Fisher, or for common loss functions equivalently, the generalised Gauss- Newton matrix (GGN). It has first been derived in the context of optimisation for linear layers and later for convolutional (Grosse & Martens, 2016) and recurrent (Martens et al., 2018) layers. In deep learning, using the Fisher/GGN instead of the Hessian has been the de-facto standard; in turn, K-FAC has arguably been one of the most popular approximations of the Fisher/GGN, probably due to its relative efficiency it approximates each layer's Fisher/GGN independently with Kronecker product. While K-FAC has also been used for transformers (Grosse et al., 2023; Osawa et al., 2022: Pauloski al., 2021; Zhang et al., 2019a) and graph neural network (Izadi al., 2020), there is no theoretical framework for these cases and for applying K-FAC to new architectures.
Contributions. We propose such framework by leveraging the idea of linear layers with weight- sharing. This concept reveals the additional structure in the Fisher/GGN due to the weight-sharing that has to be explicitly considered when applying K-FAC. We identify two settings of such layers that motivate two different flavours of the K-FAC approximation K- K-FAC-expand and K- K-FAC-reduce. The former corresponds to the Kronecker Factors for Convolution (KFC) approximation in Grosse & Martens (2016), to the simplest approximation for recurrent layers proposed in Martens et al. (2018), and has been used for transformers (Grosse et al., 2023; Osawa et al., 2022; Pauloski et al., 2021; Zhang et al., 2019a) though without discussion or motivation. Notably, K- K-FAC-reduce is exact for certain settings, like deep linear networks with convolutions and average pooling, under the same conditions that K-FAC is exact for deep linear networks without weight-sharing whereas the currently used K-FAC-expand is not. In practice, both approximations can be used in each setting and K-FAC-reduce generally has a lower computational and memory complexity than K-FAC-expand. We empirically verify this speed difference with Wide ResNet CIFAR-10. Moreover, the two K-FAC variations applied to a graph neural network and vision transformer can reach a fixed validation metric target in 50-75% of the steps of a first-order reference method, which translates into an almost equivalent decrease in wall-clock time. Although this does not demonstrate that K-FAC is superior training algorithm, it indicates the potential of extending K-FAC modern deep learning architectures. In general, K-FAC can be used as drop-in in Hessian approximation, e.g., in Laplace approximations for Bayesian deep learning (Daxberger et al., 2021; MacKay, 1992; Ritter et al., 2018) or in natural gradient variational inference (Khan et al., 2018: Zhang et al., 2018). To demonstrate this, we show that K-FAC-reduce can speed up automatic weight decay selection via approximate marginal likelihood optimisation, compared to the previously used K-FAC-expand.
show that
Background
We consider the supervised learning setup with dataset D of i.i.d. samples {an, 3 with n E RD and yn E R a neural net Se IRD - RC,, parameterised with 0 E R and loss function R: x xRC + R. which is often equivalent to a negative log likelihood, i.c. f(y, fe fo(x)) - -log p(u y )xx)). The DNN is usually trained with (stochastic) gradient-based iterative update rule, in its simplest form 0z+1 = O₂ - E. (z)). where t indicates the iteration and a is the learning rate. If we assume local quadratic approximation of the loss around the current parameter iterate Q2, we get a preconditioned gradient update step, i.e.
o.+1 = a C, VV. Ely fo, (x)),
where C, PxP is the symmetric positive definite (p.d.) preconditioner. Setting C, to the Hessian of the loss w.r.t. the model parameters H, V2 Vg f(y, fe, (æ)) yields simplified version of the classic Newton's 's method; however, since He is expensive to compute, store, and not guaranteed to be p.d., it is not commonly used in deep learning.
2.1 Second-order optimisation in deep learning
The term "second-order" is used ambiguously in the literature: can refer to using second-orde derivatives like the Hessian of the loss w.r.t. the parameters of the model, or approximations thereof, like the GGN, which however only contains second-order derivatives of the loss w.r.t. the model outputs and not the parameters. It can also refer to using some variation of the second moment of the (average) mini-batch gradient, which blurs the line between commonly used methods that are usually considered "first-order" methods like Adam (Kingma & Ba, 2015).
2

---

There have been numerous works on "second-order" methods in deep learning. For example, Martens (2010) proposes to use Hessian-free optimisation, Gupta et al. (2018) introduce an architecture- agnostic structured version of AdaGrad (Duchi et al., 2011) called Shampoo, Goldfarb et al. (2020) develop K-BFGS, Kronecker-facton introduce Tensor Normal Training, which is similar to Shampoo, but approximates the Fisher; just to name a few among many more methods (Yang et al., 2022, 2023; Yao et al., 2021). While these methods have been shown to require fewer steps than commonly used first-order methods to reach similar performance in many settings, the mechanism behind their benefits is not fully understood. It could be orthogonal to unique properties of the true GGN Fisher but might be more related to being able to adapt to the stochastic gradient noise (Kunstner et al., 2019) or to the fact that K-FAC with a specific but commonly used form of damping can be viewed as approximation to gradient descent on neurons instead of weights (Benzing, 2022). In any case, generalising K-FAC to modern architectures is promising direction because it is used in many training algorithms that can potentially speed up training and train a wider range of architectures. Moreover, it can be used to study the mechanisms behind the benefits of second-order methods in deep learning.
stochastic quasi-Newtor method, and Ren & Goldfarb (2021)
2.2 Linear weight-sharing layers
Here, we present three examples of how the core building blocks of modern DNN architectures can be expressed as linear layers with weight-sharing. DNNs have a layered structure, i.c. they can be written as fo fer o...o fo. o.... o for. with 0- = concat (0,,... Oz,... O2) and L the number of layers; concal(......" concatenates vector inputs to a larger vector. For a linear layer, we have Sa, (2) <(W,z + be), where I E I2 M,,i,, We W Pr. Pr.in beR R.,,ou, 0, = concat (vec W2), be) E ER?, and P: = =Pt,out + an element-wise nonlinearity. For simplicity, subsume the bias into the weights. We interested in linear layers with weights shared across an additional input dimension of size R. i.e. with inputs X ERRx The linear layer is applied this input as X WT and the weight matrix W is shared across the additional first dimension, i.e. every element in W contributes to each row of the output.
( vectorises a matrix by concatenating its column vectors and @ is
Example 1: Attention. The transformer architecture is popular in vision (Dosovitskiy et al., 2021), language (Brown et al., 2020), and ppatio-temporal modelling (Bi et al., 2022). The inputs and outputs described
of the attention operation (Bahdanau et al., 2015) in transformers are exactly shaped above, where R is the sequence length, e.g. the number of tokens in language data or the number of image patches in vision. For the i-th row of the output matrix O and the r-th row of the inputs X, the attention operation is generally defined as
R
Arr
where A E RR is the attention matrix, which is normalised such that Σ A.,- 1: alternatively. we can write it using matrix operations as O= the inputs, where each weighting factor indicates the importance of the r-th sequence element to the element with index i. The most common choice for A is a variation of so-called self-attention,
=0=A AX. Intuitively, the outputs are a weighted sum of
defined as A- softmax,ow W Wx X where Wq E R Poee x and WK E
Po Pone D are weight
matrices and softmax,ow (-) applies the softmax function softmax(x)) (x); cxp(x:)/ /xp(a:)/E/= exp(2;) for an input æ R to each row of the input matrix. This corresponds to a dot-produce similarity measure of two independent linear projections of the input, allowing for asymmetric relationships between sequence elements. The two weight matrices We. WK. and any linear projection of the attention output applied to all R sequence elements they are linear weight-sharing layers.
Example 2: Convolution. typical input to a 2d convolution is an image X RR x x where Cin is the number of input channels, H is the height, and W the width of the image. We choose kernel of size K x K, such that we have a rank-4 weight tensor of shape Cout X x K K, where Cout is the number of output channels. For simplicity, following Grosse & Martens (2016), we assume that the convolution is performed with a stride of and that the padding is chosen such that the height and width of the inputs maintained, i.c. the output image will be O' e BC X WWW
'Our framework also applies to weight-shand sharing where the weights are used multiple times during the forward pass; however, we focus the presentation input with explicit weight-sharing dimension.
3

---

The elements of the output tensor of the convolution are then defined as
Cin K
Ceith Lo -ie----Wddeeee Wd.kk.
(3)
the input elements Xc,4+k-154k'- are defined as 0 outside of the image's boundaries. Now, we can unfold the input tensor to matrix X E ER HW where we have extracted and vectorised the K x K sized patches for each of the H-W spatial locations. The weight tensor can be reshaped to the matrix W E TRCouxCCa.t
Using these quantities, we can write the convolution as O = XWT E RWWxxoom Again, we can recognise the canonical form of linear weight-sharing layer here, with R = HW, i.c. the weights are shared across all spatial locations.
Example 3: Graph neural network layer. We follow Battaglia et al. (2018) since their formulation encompasses many graph neural network (GNN) layers. For the full example and notation, please refer to Appendix A. Once again, the core operations in these kinds of models can be expressed as linear weight-sharing layers. Here, the weights are shared across all R = Ny nodes or R = Ne edges of graph, depending on whether node or edge features are being updated.
2.3 The generalised Gauss-Newton and the Fisher matrix
To avoid computing the Hessian, approximations like the generalised Gauss Neustn matrix (GGN) are usually used in deep learning (Botev et al., 2017; Martens, 2010; Schraudolph, 2002). Defining A (f = 4% (yn. f=(2z)) ERCxC
we have
N Σ Jololaaa))((((am))oooooo.)) n=1
GGN
(4)
which is equivalent the Hessian when fe is linear model or the residual 9 (ff yr is zero for all data points, where g-1 is the inverse link function, c.g. the softmax function for classification. For all likelihoods of the exponential family with natural parameterisation the GGN coincides with the Fisher information matrix (Martens, 2014; 4; Wang, 2010). This is, for example, the case for the categorical likelihood and the mean-square error loss corresponding to a Gaussian likelihood. Setting C, in Equation (1) to the Fisher leads to natural gradient descent (Amari, 1998). We will only state our derivations for the GGN explicitly, but they follow analogously for the Fisher.
common cross-entropy loss corresponding
2.4 Kronecker-Factoree Approximate Curvature
The GGN is semi-p.d. approximation to the Hessian, of size Px P, which is still prohibitively large for DNNs. K-FAC was proposed as an efficient approximation a neural network's Fisher (Heskes, 2000; Martens & Grosse, 2015) and GGN (Botev et al., 2017) matrix. 2 First, we only approximate the GGN for each layer, ignoring cross-layer interactions, which results block-diagonal matrix. We define Ja, =o, (x,) = Ja s,.. To(a,))Ja, St, E Cxxx: as the Jacobian of the model outputs
w.r.t. the parameters of the f-th layer. Now we can write st,n Weag, = (a & It. vec We) and with this we have Jo, 8t,n aT,, Ip Additionally, by defining br,n. := -a So (Ix)² R Pt.o XO as the transposed Jacobian of the model outputs w.r.t. the pre-activations of the f-th layer, note that Jo, (x.,) (a. . I T aggn & b K-FAC approximates the GGN of the l-th layer
I N
GG ~
ag an
s
Σ
b A (Jo(
(5)
F3 where we have replaced the (transposed) Jacobians the definition of the GGN in Equation (4) by ag, & be, and approximated the sum of Kronecker products with Kronecker product of sums.3
In the literature, the name "K-FAC" is used refer the approximation of the Fisher/GNN and to specific algorithm using this approximation (Martens Grosse, 2015); here, we use to refer to just the approximation. "Note that the Kronecker product A @ B is in reverse order compared to in Figure 1, because we concatenate the columns in the derivation and the rows in the implementation when flattening matrix with vec(-).

---

ae in_a
by out_.rads
/
b,bb
GÔNG; B8A [GGN(0.) G
K-FAC-roduce K-FAC-expand
N-
N-
Prin
Pt. Pt.uu
Pr
Pe
Priin
Pioou
Peçn
Pr.out
Prán
Pe.out
Pr
Figure 1: Visualisation of K-FAC-expar and K-FAC-reduce in the expand setting. The shown quantities are for single layer within a deep linear network. We have N = 4. = 2, Pt.i 8, and Pt Pt,in Pt,out 64. As we have seen in Section 3.2, K-FAC-expand exact for the expand case in this setting and K-FAC-reduce is not. For better visibility, the color scale is not the same for all quantities, except for the approximation error (right) where black represents zero.
K-FAC for linear weight-sharing layers
In this section, we propose to distinguish between two different linear shagiigg settings when applying K-FAC. The two settings motivate two corresponding approximations, K- K-FAC-expand and K-FAC-reduce. However, both approximations can be applied each setting, see Listing for illustration with code. More details on the derivations and the proofs for Proposition 1 and Proposition 2 can be found in Appendix B.
3.1 The expand and reduce settings
Within network, linear layers with weight-sharing can be classified based on the point where the weight-sharing dimension aggregated. Specifically, there are three possible aggregation points, which will define the setting for the f-th layer:
(i) Before the f-th layer, i.e. Ar, E F Rxx Pe.in is reduced to āg.n E R Pe.iu before being multiplied with the weight matrix of the layer W.. We are in the setting of regular linear layer and no additional considerations are necessary when using K-FAC.
(ii) After the per-example loss, i.c. there will be N R labels and outputs the model. The per-example loss is applied to each of the N R output-label pairs. - This is the expand setting, as the loss for each of the N data points expanded with R terms.
(iii) between the pre-activation of the f-th layer St, E Rt, and the model output fo(x,,) i.c. before the final aggregation over the per-example losses. - This is the reduce setting, as all seight-sharing dimensions have been reduced during the forward pass.
If assume single weight-sharing dimension which will be reduced once, the number of loss terms determines which setting applies to all linear weight-sharing layers within the model. So for our purposes, can identify the setting we are in simply by looking the form the loss.
3.2 K-FAC-expand
The expand setting can be characterised by a loss with N R terms, similar to N R i.i.d. examples, Lexpand (fe. D):= -I
Σ -N log lf (x)))) where fe(x,) is the r-th row of the model output fe (x,,) IR RxC and y is the r-th row of the label Y, E Ctx A typical example of this type of loss function is language translation, where N the dataset size and the sequence length.
We can express the Jacobian of the r-th row the model output fo(x,) E RRCC w.r.t. O, as Jo, (xr)) - 2m=1 Jar,..m Sa..oo(oee (o J.J, Jo, 8t.n,m- Since the weights O, are shared across the weight- sharing dimension of size R, we can write the r-th row St., as 8t, Waat,n,r and we have Jo, St, a & I as for regular linear layer. We denote bi := Jein fo( (..
5

---

at,x,m &......m Deriving the approximation the same way as we would with N examples is not possible, since we cannot directly write each of the N-R loss terms as Kronecker product without any approximation. The Jacobians (m.)r could be approximated with Kronecker product of sums, but this requires access > br...r.m and would not be exact in the simple settings we consider later. In contrast, what can be implemented in practice without additional backward passes and what has been used for convolutional neural networks (Grosse & Martens, 2016) and language transformers (Grosse et al., 2023; Osawa et al., 2022; Pauloski et al., 2021: Zhang et al., 2019a)
Hence, we have Je, (zx)) =
it
N
n Σ
G
ar @ be
((Jo(en)))
T ainmm & by b
pm
mmm
m1''=1
(6)
Σ ar @ (e, ((fo(2)))) (al @ by mmmm))
where all terms with m mm' are ignored.* 4 Consequently, we can apply the regular K-FAC approxi- mation over N-R R terms instead of the usual N terms, resulting in what we call the K-FAC-expand approximation:
K-FAC- K-FAC-expand
R b b A (ff(()))) n, k,n,r,
G
expand
a a mm
NR
A2
B,
In the context of convolutions, this approximation has been derived by explicitly stating the assump- tions on the activations and pre-activation derivatives (Grosse & Martens, 2016).
For deep linear networks without any weight-sharing layers, K-FAC known to exact assuming a Gaussian likelihood (Bernacchia et al., 2018). While this holds for the full GGN/Fisher, we only focus on the block-diagonal case here. To motivate K-FAC-expand, we want to show that the same also holds for deep linear networks with weight-sharing in the expand setting. However, it does not hold for simplistic transformer models since the dot-pne product attention mechanism in transformers directly correlates elements across the weight-sharing dimension; we discuss this in more detail in Appendix B. A deep linear network is defined as model of the form
fx = -WL... W...WWe... We
Wi = Wxx
(8)
where R and WT Cx We e Pt. Pr,t (with Pt, = Pt-11 and W1 E Pl Decomposing single weight matrix W into separate ones is a common framework for theoretical analysis since it creates nonlinear training dynamics for gradient-based training algorithms, while still having analytical solutions (Bernacchia et al., 2018; Saxe et al., 2014). For an input X Rxx to this type of network, we can show:
Proposition (Exactness of K-FAC-expand for deep linear network in the expand setting). For layer E of deep linear network defined as in Equation (8) and a Gaussian likelihood with p.d. covariance matrix CxCC K-FAC-expand exact in the expand setting.
3.3 K-FAC-reduce
The reduce setting is characterised by loss with just N loss terms, i.e. Croduce D) = >>NI log ffy,.ff(x.. (xx)). and thus the weight-sharing dimension must have been reduced some- where in the forward pass of the neural network fe. A typical instance where this type of loss is used in model with linear weight-sharing layers image classification with convolutional neural network or vision transformer.
Although the work transformers does not scale by 1/R R in Equation (7), Appendix B.2.1 for details.

---

de data
by out_.grads
/
bbbb
GÔN;; B&A -GGN(0))GG&rr) GÔNo,
K-FAC-expand
Ptán
Pt. Pt.uu
N
Pr.
Pe
Priin
Prout
Peçn
Pr.out
K-FAC-roduce
N-R
N-R
Prán
Ptouu
Pr
Pe
Figure 2: Visualisation of K-FAC-expard and K-FAC-reduce in the reduce setting. This is similar to Figure 1, but for the reduce setting, where K-FAC-reduce is exact and K-FAC-expand is (Section 3.3).
Since Ae, RR Rx Pt.in now matrix, we have St,, = Ag, W Rx Hence, Jo, St., and Js, Js.... (x are now both multi-dimensionn arrays. Luckily, we can avoid dealing with this directly by writing fo(x,) Σx= fo (a. where 8t,,r E Rr is the r-th row of St, and St,n, W n,r, Using this equivalence, we approximate the (transposed) Jacobians the GGN for layer as
n Σ ar 2 b
R
Jo (o
ai
b
Here, we have approximated the sum of Kronecker products with Kronecker product of sums over R terms of each of the N per-input Jacobians. This approximation has been proposed in Tang et al. 12021) to improve the efficiency of their proposed K-FAC variation for convolutions and in a different context for invariance learning with deep neural networks via differentiable Laplace approximations (Immer et al., 2022). We can apply the same approximation usual to the sum over the N data points and call the resulting final approximation K-FAC-reduce and can write explicitly as
K-FAC-reduce
G
reduce
Oz
R
R Σ
N
R Σ
NR
arp
a
Σ
Σ by. A(fo(X,))
b
(10)
n=l
r=
p=l
A
B
To ensure that K-FAC-reduce is exact in the simple setting of Proposition
have choose appropriate aggregation function Rxxxx,.om Pi. - Pr.out A simple and relevant case for which this holds is a scaled sum, i.e. 2(St,n)) = E r- St,,r with e R. Both vision transformers and convolutional neural networks with average pooling use scaled sums as the aggregation function (with 1/R). With this, we can state the corresponding statement to Proposition 1.
Rx
Proposition 2 (Exactness K-FAC-reduce for deep linear network in the reduce setting). For layer & of deep linear network (Equation (8)), Gaussian likelihood with p.d. covariance matrix EE EEROX CXC and scaled sum aggregation function, K-FAC-reduce is exact in the reduce setting.
Notably, for deep linear model with convolution layers and average pooling, K-FAC reduce will be exact under the assumptions of Proposition 2, whereas the approximation proposed in the literature (Grosse Martens, 2016) (which corresponds K-FAC-expand) is not. Moreover, the calculation of A costs ONN (N... for K-FAC-expand and only O Pe, (Pt,in R)) for Å, and K-FAC- reduce. The same holds for B, and Вr, with complexities of O(NRCP, (C + Pt.our)) and ONCPP. (Ct (C+tt.outPR)) The space complexity of K-FAC's overhead also differs by a factor of R. with O(N R(Pt,in in + Pt.ou)) for K-FAC-expand and (N Pr,ou))) for K-FAC-reduce.
7

---

4 Experiments
In Section 4.1, we empirically validate the speed up of K-FAC-reducee compared to K-FAC-expand in line with their computational complexities with Wide ResNet CIFAR-I0 and show how this scales with the batch size. Additionally, in Sections 4.2 and 4.3, we compare the two approximations based their downstream optimisation performance. We validate that both K-FAC variations can reduce the number of required steps to achieve a fixed validation target compared to a well-tuned first-order reference optimiser. For this, use GNN and vision transformer (ViT), two of the examples in Section 2.2. We leverage the codebase and validation metric targets of the MI Commmnnss al algori thmic-eff wweefffciiene ficeecy (AlgoPerf) repository' (Dahl et al., 2023), and use their target-setring runs as reference runs. Both K-FAC variations use the same learning rate schedule as those target- setting runs but the warmup and expected number of steps are multiplied by 0.75 to account for the fact that second-order methods tend to require fewer steps. Additionally, we tune the learning rate and damping for K-FAC via random search and set all other hyperparameters the same values as the reference target-setting run. To exclude the adjusted learning rate schedule and the additional tuning of the learning rate as confounding factors, we additionally applied the same tuning to the reference runs; however, none of these runs hit the target validation performance. Despite these considerations, the experiments in Sections 4.2 and 4.3 are not meant to demonstrate that K-FAC a superior training algorithm (see Appendix C. However, the reference runs allow us to put the training performances of both K-FAC flavours into perspective and also indicate the potential of extending second-order methods like K-FAC to modern neural network architectures Lastly, to demonstrate use case of K-FAC other than as a preconditioner, we automatically learn the weight decay hyperparameter by optimising a Laplace approximation of the marginal likelihood in Section 4.4. Our K-FAC implementation leverages the ASDL package (Osawa et al., 2023). See Appendix C for more details on the experiments and additional results.
4.1 Update step speed with K-FAC-expand and K- K-FAC-reduce
To empirically validate the smaller computational com- plexity of K-FAC-reduce compared K-FAC-expand, we time a single preconditioned gradient update step for Wide ResNet on CIFAR-10 with both approximations and five different batch sizes on an NVIDIA V100 GPU (Ta- ble 1). As expected based on the complexities of the two approximations, we see increasing gains in speed with increasing batch size. For the largest tested batch size, 2048, K-FAC-reduce is faster than K-FAC-expand is for half the batch size. Moreover, K- K-FAC-expand runs out-of-memory (OOM) at a batch size of 2048.
Table 1: Timing [s] of an update step with the two K-FAC variants for a Wide ResNet on CIFAR-10.
Batch size
K-FAC
128
256
512
1024
2048
expand reduce
0.24
0.38
0.75
1.36
0OM
0.17
0.24
0.43
0.63
1.17
Notably, we implement both approximations with the naive unfold operation (c.f. Section 2.2), but K-FAC-reduce could be implemented even more efficiently by never explicitly allocating memory for the full unfolded layer input, which has been shown lead to a further speed-up of about 4.5x (Dangel, 2023). To highlight the speed difference of the two K-FAC variations for language models, compute K-FAC GGN approximation for nanoGPT (Karpathy, 2023), a popular implementation of GPT-2 (Radford et al., 2019), on the full DART dataset (Nan et al., 2021). K-FAC-reduce is significantly faster than K-FAC-expand, which is currently used for transformers, only taking 70% of the time of K-FAC-expand. We expect this discrepancy to become even more pronounced for larger batch sizes and sequence lengths.
4.2 Graph neural network on ogbg-molpeha
We use basic GraphNetwork (Battaglia et al., 2018) with multi-layer perceptrons the update functions (c.f. Appendix A) and train it on the The target validation metric is a mean average precision (mAP) of 0.28098. The reference algorithm is SGD with Nesterov momentum, with the tuned hyperparameters from the AlgoPerf target-setting runs. For both K-FAC variations, the statistics and preconditioner are updated every 10 iterations. We report the mean and standard error for five runs with different random seeds. The reference run reaches the target in 31,887= + 2,070 steps, whereas K-FAC-expand takes 106 + 863 and K-FAC-
dataset (Hu et al., 2020) (Figure 3).
$ "attps://bit..1//aaggpaa
8

---

0.10
Reference (Nesteniv)
K. FAC-exxand
0.08
Training
Loss.
Validation
0.06
Validation target
0.04
0.02
0.5
0.4
mAP
0.3
0.2
10
20
30
40
0.5
1.0
1.5
2.0
2.5
3.0
Step
x 10
Wall-clock time [hours]
Figure 3: Training results for graph neural network on ogbg-molpcba, Both K-FAC variations require 50 % the steps which almost directly translates into reduced wall-clock time. The K-FAC statistics are updated every 10 steps. We show runs with five different random seeds for each method.
reduce 15,566 + 429 steps, about 51 % and 49 % of the reference run, respectively. In wall-clock time, the reference run takes about 2.33 + 0.22 hours, whereas K-FAC-expand takes 1.32 + 0.11 and K-FAC-reduce 1.23 + 0.03 hours, about 57 % and 53 % of the reference run, respectively. The reduced number of steps almost directly translates into reduced wall-clock time since the update step time is dominated by the data pipeline. K-FAC-reduce appears to be slightly faster in steps and wall-clock time than K-FAC-expand, but not significantly. Moreover, the runs with K-FAC-reduce have lower variance.
4.3 Vision transformer on ImageNet
We train ViT (Dosovitskiy et al., 2021) on LSVRC-2012 ImageNet (Russakovsky et al., 2015) and use NAdamW (Loshchilov & Hutter, 2019) with the hyperparameters from the AlgoPerf target- setting runs as the reference run (Figure 4); each run repeated with three different random seeds. For K-FAC-expand and K-FAC-reduce, iterations. The reference run reaches the target validation accuracy 0.77309 in 122,139 + 1119 steps about 26. 10 + 0.24 hours. K-FAC-expam only takes 47,464 = 245 steps 19.57 + 0.13 hours, and K-FAC-reduce 92,550 + 731 steps or 20.28 0.00 hours. This corresponds to about 72 % of the steps and 75 % of the time of the reference for K-FA K-FAC-expand and 76 % of the steps and 78 % of the reference run't time for K-FAC-reduce. While the difference in speed between both K-FAC variations does become less relevant due to the infrequent K-FAC updates, step with K-FAC-expand takes about 1.36 as much as with K-FAC-reduce, which has significant impact the runtime when the K-FAC statistics run where the update is performed every step. While K-FAC-expand requires slightly fewer steps and wall-clock time than K-FAC-reduce here, the difference does not appear significant.
update the K-FAC statistics and preconditioner every 50
update
computed more frequently; Figure 5 for
4.4 K-FAC for automatic hyperparameter selection via marginal likelihood optimisation
K-FAC has also been used for marginal likelihood maximisation Bayesian neural net- works using Laplace approximations (Daxberger et al., 2021; Immer et al., 2021). In par- ticular, K-FAC can be used to obtain a scalable Laplace approximation and optimise its marginal likelihood with respect regularisation parameters, e.g., weight decay, during training.

---

Reference (NAdamW)
FAC-expand
2.5
F-C--reduce
Training
Loss
2.0
Validation
Validation target
1.5
0.5
Accuracy 1
0.7
0.6
0.5
20
40
60
80
100
120
5
15
20
Step
100
Wall-clock time [hours]
Figure 4: Training results for vision transformer ImageNet. The K-FAC statistics updated every 50 steps. Due to amortising the costs, the reduced number of steps to the target translates into reduced wall-clock time. We show runs with three different random seeds for each method.
Therefore, we compare K-FAC-reduce and K-FAC-expand in this setting following the setup of Daxberger al. (2021) on CIFAR-I0 with a Wide ResNet (Zagoruyko & Komodakis, 2016) and optimise weight decay parameter per layer every five epochs during 100 epochs of training. Table 2 shows that K-FAC-reduce significantly faster than K-FAC-expand in this setting, but K-FAC-expand provides better test negative log likelihoods (NLL) in both settings.
Table 2: Performance of K-FAC variants for marginal likelihood maximisation on CIFAR-10 with Wide ResNet with and without data augmentation (DA).
K-FAC
DA
NLL 1
Acc [%]
Time [%)
X
0.42
88.9
expand
100
0.24
92.5
reduce
x
0.70
86.7
0.35
93.5
5 Discussion and conclusion
We have leveraged the idea of linear weight-sharing layers to generalisc K-FAC to a broader range of architectures. This resulted in two distinct flavours of the approximation, which exact for simple cases with deep linear networks in their respective setting. Maybe surprisingly, we see little difference between the two K-FAC variations in terms of their downstream optimisation performance. Moreover, in some settings, the overhead of K-FAC can be amortised or is negligible. This is, for example, the case when 1) the data loading dominates the algorithm's 's update step time (c.f. the GNN experiment), 2) when the K-FAC update frequency can be reduced sufficiently (c.f. the ViT experiment), 3) when underutilised resources be leveraged compute K-FAC updates, for example when using pipeline parallelism for language models (Osawa et al., 2022). In these settings, the improvement wall-clock time of K-FAC-reduce compared K-FAC-expand insignificant. However, in applications where the K-FAC overhead still significantly impacts the runtime, e.g. in the online weight decay selection experiment in Section 4.4, or when there are memory constraints (c.f. Section 4.1), K-FAC-reduce can greatly reduce the runtime compared to the previously used K-FAC-expand. Some challenges remain when applying K-FAC in modern deep learning settings. For example, the memory overhead can be prohibitive, even for K-FAC-reduce, which could be addressed by sparse structured Kronecker factors (Grosse al., 2023; Zhang et al., 2018). Also, numerical instabilities might arise in low fow-precision settings due to the need for matrix inversions decompositions; this could be addressed by inverse-free methods relying on the K-FAC approximation (Lin al., 2023).
rendered
10

---

Acknowledgements
The authors thank Kazuki Osawa for fruitful discussions and support with using and extending the ASDL library. The authors also thank Wu Lin for pointing out a mistake in the original equation of the K-FAC-expand approximation. Finally, the authors thank Felix Dangel for helpful comments on this manuscript. Runa Eschenhagen is supported by ARM and the Cambridge Trust. Alexander Immer is funded by the Max Planck ETH Center for Learning Systems (CLS). Richard Turner is supported by Google, Amazon, ARM, Improbable and EPSRC grant EP/T005386/1 Frank Schneider and Philipp Hennig are funded by the Deutsche Forschungsgemeinschaan (DFG, German Research Foundation) in the frame of the priority programme SPP 2298 "Theoretical Foundations of Deep Learning" HE 71 7114/5-1, as well as 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: (IIS1803AA) and funds from the Ministry of Science, Research and Arts of the State of Baden-Wi Wiii-Wittttti Württemberg.
References
Amari, S. Natural gradient works efficiently in learning. Neural computation, 10(2), 1998.
Ba, L. J., Kiros, J. R., and Hinton, G. E. Layer normalization, 2016.
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. In JCLR, 2015.
Battaglia, P. W., Hamrick, B., Bapst, V., Sanchez-Gonzalez, A., Zambaldi, V. F., Malinowski, M., Tacchetti, A., Raposo, D., Santoro, A., Faulkner, R., Gülçehre, Ç.,, Song, H. F,, Ballard, A. J., Gilmer, J., Dahl, G. E., Vaswani, A., Allen, K. R., Nash, C., Langston, V., Dyer, C., Heess, N., Wierstra, D., Kohli, P., Botvinick, M. M., Vinyals, O.,, Li, Y., and Pascanu, R. Relational inductive biases, deep learning, and graph networks. arXiv 1806.01261, 2018.
Benzing, F. Gradient descent on neurons and its link to approximate second-order optimization. In ICML, 2022.
Bernacchia, A., Lengyel, M., and Hennequin, G. Exact natural gradient in deep linear networks and its application to the nonlinear case. NeurlPS, 2018.
Bi, K., Xie, L., Zhang, H., Chen, X., Gu, X., and Tian, Q. Pangu-weather: A 3d high-resolution model for fast and accurate global weather forecast. arXiv 2211.02556, 2022.
Botev, A. and Martens, J. KFAC-JAX, 2022. URL http http::////g........o///////////p weoon//eepmidd///e ktac-jax...a....--a jax.
Botev, A., Ritter, H., and Barber, D. Practical Gauss-Newton optimisation for deep learning. In ICML,, 2017.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-MMlne, S., and Zhang. Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: tigb.b........ com/ google/jaaaa
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P. Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, and Amodei, D.
Language models are few-shot learners. In NeurlPS, 2020.
Chaudhari, P., Choromanska, A., Soatto, S., LeCun, Y., Baldassi, C., Borgs, C., Chayes. J. T., Sagun, L., and Zecchina, R. Entropy-SGD: Biasing gradient descent into wide valleys. In ICLR, 2017.
Dahl, F., Schneider, F., Nado, Z.,, Agarwal, N., Sastry, S., Hennig, P., Medapati, S., Eschenhagen, R., Kasimbeg, P., Suo, D., Bae, J., Gilmer, J., Peirson, A. L.., Khan, B., Anil, R., Rabbat, M., Krishnan, S., Snider, D., Amid, E., Chen, K., Maddison, C. J., Vasudev, R., Badura, M., Garg, A., and Mattson, P. Benchmarking neural network training algorithms. arXiv 2306.07 179, 2023.
Dangel, F. Convolutions through the lens of tensor networks. arXiv 2307 02275.022,.0... 2023.

---

Dangel, F., Kunstner, F., and Hennig, P. BackPACK: Packing more into backprop. In ICLR, 2020.
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace redux-effortless Bayesian deep learning. NeurIPS, 2021.
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. JMLR, 12(61), 2011.
Goldfarb, D., Ren, Y., and Bahamou, A. Practical Quasi-Newton methods for training deep neural networks. In NeurlPS, 2020.
Graves, A. Practical variational inference for neural networks. In NIPS, 2011.
Grosse, R., Bae, J., Anil, C.,, Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E., Hubinger, E., Lukošiūté, K., Nguyen, K., Joseph, N., McCandlish, S., Kaplan, J., and Bowman, S. R. Studying large language model generalization with influence functions. arXiv 2308.03296, 2023.
Grosse, R. B. and Martens, J. A frateceer-factored approximate Fisher matrix for convolution layers. In ICML, 2016.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned stochastic tensor optimization. In ICML, 2018.
Heskes, T. On "natural" learning and pruning in multilayered perceptrons. Neural Computation, 12 (4), 2000.
Hinton, G. E., Srivastava, N., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Improving neural networks by preventing co-adaptation feature detectors. arXiv 1207.0580, 2012.
Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovee, J. Open graph benchmark: Datasets for machine learning graphs. In NeurIPS, 2020.
Immer, A., Bauer, M., Fortuin, V., Rätsch, G., and Khan, M. E. Scalable marginal likelihood estimation for model selection in deep learning. ICML, 2021.
Immer, A., van der Ouderaa, T. F., Rätsch, G., Fortuin, V., and van der Wilk, M. Invariance learning in deep neural networks with differentiable Laplace approximations. In NeurIPS, 2022.
Izadi, M. R., Fang, Y., Stevenson, R., and Lin, L. Optimization of graph neural networks with natural gradient descent. In IEEE BigData, 2020.
Karpathy, A. nanoGPT, 2023. URL http://////gccuuu.o.oaarraaraeee
Khan, M., Nielsen, D., Tangkaratt, V., Lin, W., Gal, Y., and Srivastava, A. Fast and scalable Bayesian deep learning by weight-perturbatiin in Adam. In ICML, 2018.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In ICLR, 2015.
Kipf, T. N. and Welling, M. Semi-supervised classification with graph convolutional networks. In ICLR, 2016.
Kunstner, F., Balles, L., and Hennig, Limitations of the empirical Fisher approximation for natural gradient descent. In NeurlPS, 2019.
Lin. W. Duruisseaux, V., Leok. M., Nielsen, F., Khan, M. E., and Schmidt, M. Simplifying momentum- based positive-definite submanifold optimization with applications to deep learning. In ICML, 2023.
Liu, K., Ding, R., Zou, 7.,, Wang, I.,, and Tang. W. A comprehensive study weight sharing in graph networks for 3d human pose estimation. In Computer Vision ECCV, 2020.
12

---

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In ICLR. 2019.
MacKay, D. J. Bayesian interpolation. Neural computation, 4(3), 1992.
Martens, J. Deep learning via Hessian-free optimization. In ICML, 2010.
Martens, New insights and perspectives on the natural gradient method. JMLR, 21(146), 2014.
Martens, J. and Grosse, R. Optimizing neural networks with Kronecker-factored approximate curvature. In ICML, 2015.
Martens, J., Ba, J., and Johnson, M. Kronecker-factored curvature approximations for recurrent neural networks. ICLR, 2018.
MI.Commons. Algorithms Working Group, 2022.
https://mlccomoons org/ en/ atggrrps///////////
research- algorithms/ Last accessed: 20. 12.2022.
Nan, L., Radev, D. R., Zhang, R., Rau, A., Sivaprasad. A., Hsieh, C., Tang, X., Vyas, A., Verma, N., Krishna, P., Liu, Y., Irwanto, N., Pan, J., Rahman. F., Zaidi, A., Mutuma, M., Tarabar, Y., Gupta, A., Yu, T., Tan, Y. C., Lin, X. V., Xiong, C.,, Socher, R., and Rajani, N. F. DART: open-domain structured data record to text generation. In NAACI-HLLT, 2021.
Osawa, K., Li, S., and Hoefler, T. Pipefisher: Efficient training of large language models using pipelining and Fisher information matrices. arXiv 221 1.1133, 2022.
Osawa, K., Ishikawa, S., Yokota, R., Li, S., and Hoefler, T. ASDL: A unified interface for gradient preconditioning in PyTorch. arXiv
2023.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, Lin, Z., Gimelshein, N., Antiga, L., et al. PyTorch: An imperative style, high-performance deep learning library. In NeurlPS, 2019.
Pauloski, G., Huang, Q., Huang, I.., Venkataraman, S., Chard, K., Foster, T., and Zhang, KAISA: adaptive second-order optimizer framework for deep neural networks. In International Conference for High Performance Computing, Networking, Storage and Analysis (SC21), 2021.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. 2019.
Ren, and Goldfarb, D. Tensor normal training for deep learning models. In NeurlPS, 2021.
Ritter, H.,, Botev, A., and Barber, D. A scalable Laplace approximation for neural networks. In ICLR, 2018.
Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. ImageNet large scale visual recognition challenge. IJCV, 115, 2015.
Saxe, A. M., McClelland, L., and Ganguli, S. Exact solutions the nonlinear dynamics of learning in deep linear neural networks. In ICLR, 2014.
Scarselli, F., Gori, M., Tsoi, A. C., Hagenbuchner, M. and Monfardini, G. The graph neural network model. IEEE Transactions on Neural Networks, 20, 2009.
Schraudolph, Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7), 2002.
Sen, P., Namata, G., Bilgic, M., Getoor, L., Galligher, B., and Eliassi-Rad, T. Collective classification in network data. Al Magazine, 29(3), 2008
Tang, Z., Jiang, F., Gong, M. Li, H., Wu, Y., Yu, F., Wang, Z., and Wang, M. SKFAC: Training neural networks with faster cronecker-factoree approximate curvature. CVPR, 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, I.., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In NIPS, 2017.
13

---

Wang. Y. Fisher scoring: An interpolation family and its Monte Carlo implementations. Comput. Stat. Data Anal., 54(7), 2010.
Yang, G. Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are gaussian processes. In NeurIPS, 2019.
Yang, M., Xu, D., Wen, Z., Chen, M., and Xu, Sketch-based empirical natural gradient methods for deep learning. Journal of Scientific Computing, 92(3), 2022.
Yang, M., Xu, D., Cui, Q- Wen, Z., and Xu, P. An efficient Fisher matrix approximation method for large-scale neural network optimization. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(5), 2023.
Yao, Z., Gholami, A.. Shen, S., Mustafa, M., Keutzer, K., and Mahoney, M. W. ADAHESSIAN: an adaptive second order optimizer for machine learning. In AAAI, 2021.
Zagoruyko, and Komodakis, N. Wide residual networks. In BMVC, 2016.
Zhang, G., Sun, S., Duvenaud, D., and Grosse, R. Noisy natural gradient as variational inference. In ICML, 2018.
Zhang. G., Li, I.., Nado, Z.,, Martens, J., Sachdeva, S., Dahl, G. E., Shallue, C. J., and Grosse, R. B. Which algorithmic choices matter at which batch sizes? Insights from a noisy quadratic model. In NeurIPS, 2019a.
Zhang, H., Dauphin, Y. N., and Ma. T. Fixup initialization: Residual learning without normalization. In ICLR, 2019b.
14

---

A Weight-sharing in graph neural networks
In this section, we expand on the third example in Section 2.2 on graph neural networks (GNNs) and show how they use linear layers with weight-sharing.
Graph convolutional network for node classification. A popular type of GNN is called graph convolutional network (Kipf& & Welling. 2016, GCN). A GCN defines a convolution operation on graph structures, by repeatedly aggregating feature information over the neighbourhood of a node. As regular convolutional neural network, it also uses weight-sharing; see Liu al. (2020) for a comprehensive discussion on weight-sharing in GCNs. However, in contrast to the other models presented here, they utilise a slightly different type of weight-sharing, which will become apparent in Equation (12). Nevertheless, we briefly mention this case here, since the only work on K-FAC for GNNs has been on this model architecture and we will explicitly show how K-FAC was applied in this case in Appendix B.3.2; this relies on the notation introduced here.
A graph is defined as E), where V is the set of N nodes and & the set of edges. The edges can be encoded relative to the nodes in an adjacency matrix C' I NVN with Ci3y = 0 if there is no edge and Ct2- Here, we focus on node classification, e.g. classifying scientific publications which are represented as nodes in a citation network into topics (Sen et al., 2008).
if there is one. Typically, they are used for node and graph classification tasks.
The L-th GCN layer is defined as
(x) o(CXW,)
(1I)
which is identical to a regular dense linear layer from Section 2.2, but the input matrix X E R Nx Peii, which has the N node features Σn of size Ptii stacked in the rows, is first transformed by the normalised adjacency matrix Ć == (D + IN) ) (C + (D) + In)- where D is the diagonal node degree matrix of the graph and I is the N x N identity matrix.
Defining
N
En =
ŌnjØz = Σ Enjajj
(12)
j=1
JEN(m)
we can express the forward pass for a single node and layer as
(āx)) - o (V..),
(13)
where (ooreeN(n)) N(n) = {j fL....,,NHC,, * is the neighbourhood of the node with index n. Notably, the forward pass for a single node H,, depends its neighbourhood, i.c. we cannot express the forward pass for the node without access to the feature information of the nodes in its neighbourhood ((n). Moreover, we can now see that the forward pass through the linear layer, i.e. the matrix multiplication of the weight matrix W. with the transformed input En. does not need the notion of weight-sharing anymore, in the sense, that we do not need batched matrix-vector product over weight-sharing dimension. This is because we aggregate over each node's neighbourhood, over which the weights are shared, before the matrix-vector product. Hence, in contrast to the Getpor introduced in the next paragraph, this model does not require special consideration when applying K-FAC (c.f. setting (i) in Section 3.1).
GraphNetwork for graph classification. One more general formulation of a GNN is an instance of the GraphNetwork introduced in Battaglia al. (2018). The GraphNetwork in its general form takes a graph g = (u. V. V)E) where u E RD= are the global features of the graph, and V and E the sets of nodes and edges, respectively, just before. We can also write the i-th graph of dataset of N graphs as a 5-tuple x := (æ X X rn, 8n), with global features xx E RDD, I...., N. The two vectors r. E R.N. and 8n E N contain the indices of the receiving and sending nodes of each edge, respectively. Using these indices, we define XV n,P. E N Dv and XV n.s. E N. D which contain the node features X at indices 8, and rn. respectively. Note, that these graph inputs unfortunately cannot trivially be batched by stacking them, since the number of nodes N.V or edges N are not necessarily the same for all n NJ..
node features X E N.NY DDv and edge features XE E N N xDg for all
15

---

A GraphNetwork block updates the 3-tuple (. X ) by using three update functions Ø.
X - (X X n,rx' X
n.e. xx
x - 0 (X V
(14)
x - x [
and three permutation-innariiat aggregation functions
X p P (xx)
t p p (x X E n
(15)
X V t p" (X.)).
Examples of these aggregation functions include element-wise summation, mean, or maximum.
One forward pass through GraphNetwork block corresponds the following steps, where each step is executed for all n E NN]
1. Update edges _ with of (X X n.r. X X 2.t)
2. Aggregate updated edges over all nodes in X E NR xDg using p E
3. Update nodes X using o² (XX X xxi).
4. Aggregate updated edges over all graphs in X E RR using p E X D)
5. Aggregate updated nodes over all graphs in X E Dv using P Vuu (X,) )
6. Update global features x with 6" X x
In this work, we consider graph classification; for example, molecules can be represented as graphs and could classify them according to some chemical property (e.g. the moolpbaa dataset used in Section 4.2). We specifically consider a GraphNetwork instance with simple MLPs for all update functions and an element-wise sum for the aggregation functions p. Moreover, multiple of these GraphNerwork blocks can be stacked on top of each other. To classify the input graphs, an MLP is applied to the global features x; after they are updated by the last GraphNetwork block.
To be more precise, the update functions are in this case specified as
EXE XV n,r." n,s," x()) concat (XE XV n.r." XY repeat Ng ( (
XV
V (XX F x") concat X X repeat Ny (x")) 0² (m V E W''cocat XV x)
(16)
with W E E +D.
w E (DDx and W" E R D. X3D.
Note, that this is a simplification since in reality, the update functions 6 are MLPs with ReLU activations, layer normalisation (Ba et al., 2016), and dropout (Hinton al., 2012). Also, we omit the potential bias vectors. However, these components are not relevant for deriving K-FAC for the linear layers within these networks, which why we can omit them here for simplicity.
Most importantly, we can now observe that this type of GNN shares its weights over each graph's edges and nodes: just for the attention convolution operations described in Section 2.2, we apply the transposed weight matrices from the right side of the input of the layers of type df and Φi i.e. for updating the edge and node features. However, since the number of edges N and the number of nodes NY is not necessarily the same for all N graphs, we now have a weight-sharing dimension of size R,, which depends on the n-th input. Also, note that the update function d² used to update the global features x; does not use any weight-sharing, as there is just one feature vector per data point. We have specifically introduced this notation of the inputs to show that the edge and node feature update functions are exactly examples of the concept of linear weight-sharing layer, introduced in Section 2.2. This might not be immediately obvious when only looking at the original notation used in Battaglia et al. (2018) 8); therefore, this can be seen as instructive example for expressing neural network architecture in terms of linear weight-sharing layers. Consequently, our framework for K-FAC can directly be applied to this architecture, as we show in Appendix B.3.2.
16

---

B Extended derivation and discussion of K-FAC-expand and K-FAC-reduce
B.1 Background: K-FAC for regular linear layers
Cronecker-Factoree Approximate Curvature (Heskes, 2000; Martens & Grosse, 2015, K-FAC) was proposed an efficient approximation neural network's Fisher information matrix. The Fisher information matrix is defined as
F())
Σ
E By-piulfa(r... () Vg (og)()
(17)
By-ppylla(m..) Ve (oo.:lee))ellooooooooollllllll.d bo (z())))
Notably, the labels y are samples from the predictive distribution of the model and are not the empirical labels from the data. Replacing y with yn leads to the empirical Fisher (EF), which simply the uncentered covariance of the empirical gradient. While replacement for the Fisher (Chaudhari et al., 2017; Graves, 2011; 1; Kingma & Ba, 2015), it can give rise to very different downstream behaviour when used for optimisation (Kunstner et al., 2019).
commonly used as
First, in all of this work, we focus on layer-wise K-FAC approximation of the Fisher, i.c. it is approximated by block- matrix
F()) Rs diag( F(O))..... F(O))....FF().)) ERPXP
(18)
where F(O:) TRPxxP, and diag(.,..... build block-diagonal matrix with the input matrices as blocks.
To derive K-FAC, we first note that the pre-activation for layer and the 1-th data point æ,, can be expressed as 8 = Wiat,n with W, E FR oont and az E Pr, the input to the b-th layer (or equivalently, the activation of the & 1-th layer). We have omitted explicit bias parameter b;, since it can always be subsumed in W.. Hence, by applying the chain rule, the gradient the loss w.r.t. the weights of the 8-th layer can be written V (o. V Cly. fe (z,))a?,n gr,n a
R on it
Using these insights, K-FAC then replaces the sum of expectations over Kronecker products with Kronecker product of two sums of expectations, i.c.
N
F(O,)=
E
(19a)
n=l
I Wywptullscc...looooe...... vee
(19b)
E Ey-rlolfte..)laan (a @ n
(19c)
N
E Py-plullloo.a)aa...... @ gn
(19d)
N
an aen
x
B
(19c)
n=1
n=1
A
=
where A, IR x and Ge E Pr. omt oPu. For this derivation, we have used three convenient properties of the Kronecker product (using matrices A, B, C, D with appropriate dimensions):
vec( ABC) (CT A)vec(B) and (A& B) A BT for Equation (19c), and (A BC D) = AC & BD for Equation (19d).
Moore generally, is defined with expectation p(æ) well.
17

---

We can see that the approximation is exact in the trivial case of single data point, i.c. N = 1. Moreover, is also exact the case of a single linear layer or deep linear network and a Gaussian likelihood (Bernacchia et al., 2018).
K-FAC is more efficient than naive block-wise approximation because we only have store and invert two Kronecker factors instead of a larger dense matrix for each layer, which reduces the memory complexity from (i P2 mout) ( to O12)) foot)) and the computational complexity of the preconditioning of the gradient with the approximate Fisher from ( P3 Pi to O E,in P.out) E,in
since
F(0))) g(())) = (A @ G Gz)) g(0z)
(20)
vec (G7 V w; L(y. So (z.))) (m))))7))
with g(0+) vec (Vw, (xx)) and the property (A @ B) & B-1
Alternatively, we can derive K-FAC for the GGN (Botev et al., 2017), which will recover the same result for the Fisher in Equation (5) for many common loss functions, as we have learned Section 2.3. We define Jo, (z:) == Jo, (x,,) = J. fo St,n E RCx as the Jacobian of the model outputs w.r.t. the parameters of the f-th layer and (z.() == H Cyyn (a(w.). CCCC as the Hessian of the loss w.r.t. the model outputs. Now we can write se,n Waag, (a & or vec( We) and with this we have Jo, S,, at,,, @ Pr Additionally, by defining ben fo
E IR Pt. XC as the transposed Jacobian of the model outputs w.r.t. to the
pre-activee re-aativvaiion of the ê-th layer, we note that Jo, (xx,)) (a I b ag, b
Replacing the (transposed) Jacobians in the definition of the GGN by this expression, we have
N
GGN((:))
Σ Jo, (((o
atnn @ bbb ((f())(((( & be,")
(21)
ai a a 2 (ff..((((
n=
N
N Σ
ag
a
be
N
=:A,
=:B,
This derivation is a bit more convenient for our purposes, as it does not require us to keep track of the expectation over the labels y, while still being equivalent the Fisher for common loss functions (c.f. Section 2.3). Moreover, in our context, it will be useful to have the Jacobians Jo, (æ,, separate from the loss; therefore, we will only explicitly write our results for the GGN.
B.2 K-FAC for linear weight-sharing layers
B.2.1
The expand setting and K-FAC-expand
The expand setting can identified by a loss with terms, which corresponds to assuming N R i..dd examples,
The Expand Setting
N
Lexpand (..D) D) =
Σ Σ log (l
(22)
where fa(w,), is the r-th row of the model output (xx,) E RRCC and yo,r the r-th row of the label Y, E fxxC typical example of this type of loss function is language translation, where N is the number of training examples and R is the sequence length.
18

---

Note that we are not assuming our inputs to have an additional veight-sharing dimension since we only require that the input to the ε-th layer has this additional dimension, Ar, R Rx 1) This obviously does not exclude the case where already has this weight-sharing dimension, e.g. sequence of tokens in translation tasks.
We can express the Jacobian of the r-th row of the model output fa(x..) E IR RxC w.r.t the parameters Og
Peont Ofo(2n)) Dri nS
(Jo, (an)))is
(23)
aSt i.mmmm 20
or in matrix form
Jo. (3n)r
J m ( St,n, Σ m=1
(24)
Since the weights Og are shared across the weight-sharing dimension of size R. can write the r-th row of Se, as 8g,n, = Weag,n. and we have Jo. 8E,nr = ar,n, @ Pe. as for regular K-FAC
(c.f. Appendix B. We denote be,nrr == x xx))) Hence, we have
12
Jo, (2x)²
Σ
0 m (aF & IREwar)
=
mml
(25)
ae & bgn.r.m.
On a high level, applying K-FAC model trained with this type of loss just requires treating the problem as we had N It independent examples and deriving the approximation in the same way as we would with N examples (c.f. Appendix B.1),
GG ΣΣ
Jo, (z, ((.)))((fo(. (x.,r
m=1m=
(26)
R
R Σ mm/=1
ΣΣ
Σ at e be A(fo(Tn)))
a a e bj
However, cannot directly write each of the N R loss terms as Kronecker product without any approximation. One approach could be to approximate each Je, (In)r with a Kronecker product of sums, as the K-FAC approximation does for the sum of the N data points, but then would have be able access Σm be,n.rm- Moreover, this would not even be exact in the simple settings we consider later. In contrast, what can be implemented in practice without additional backward passes (c.f. Appendix B.4) and what has been used for convolutional neural networks (Grosse Martens, 2016) and language transformers (Grosse et al., 2023; Osawa et al., 2022; Pauloski et al., 2021: Zhang et al., 2019a) is
GGN((.)
ΣΣΣ
(armm be,.r.m) ((o(())))) (atm m.....m
m=1
(27)
R &
Σ
Σ
(at af In
Σ
b (Gn
where we ignore all the with o m' which allows to express each the N R terms as Kronecker product.? Consequently, we can apply the regular K-FAC approximation over N R terms instead of just N terms usual. We call the resulting approximation K-FAC-expand:
'However. the authors of the work on transformers do not discuss this extension of K K-FAC at all; although Zhang al. (2018) do explicitly discuss the tying of the embedding and linear output layer weights.
19

---

K-FAC-expand
N
N
R
R
G expand
a a Enmm
ΣΣΣ
bea A ((o(( n.m.
NR
n= m=1
m=
B2
(28)
There is one simple where the exact expression in Equation (26) is identical to the approximation in Equation (27). When bg,n.r.m 0 for all * m. both expression are equivalent to
N
aag,.r @ be,,,rr) (() (a bl.r)))
(29)
n=1r=1
With other words, when (z,) is independent of all pre ane-activations st,n,m with + the two expressions coincide. This does not even hold for simplistic transformer models since the self-attention mechanism (Section 2.2) in transformers directly correlates elements across the weight- sharing dimension - we discuss this in more detail in Appendix B.3.1. Alternatively, we could also scale Equation (27) by R, which leads to a scaling of 1/N instead of 1 1/N R in Equation (28). This might be a better approximation when be...mm * for r * m and also what has been used by previous work on transformers (Grosse al., 2023; Osawa al., 2022: Pauloski al., 2021; Zhang et al., 2019a). However, we choose to use the scaling in Equation (28) since the condition above holds for networks that simply stack multiple linear weight-sharing layers; this allows us show that the approximation in Equation (28) is exact in the same simple cases as regular K-FAC.
For typical neural network with nonlinear activation functions, K-FAC is only approximation. However, for regular individual linear layers and deep linear networks, K-FAC is known to be exact assuming Gaussian likelihood (Bernacchia al., 2018). While this holds for the full GGN/Fisher, we only focus on the block-diagonal case here. To motivate K-FAC-expand, we want to show that similar statements hold for single linear weight-sharing layer and deep linear networks with weight-sharing in the expand setting. First, we state a simple condition for which the approximation indeed exact; this line of reasoning could also be applied K-FAC for regular linear layers since only the effective number of data points changes from N to N. Note, that we could also state more trivial sufficient conditions for the exactness of the approximation, i.e. N R and when all inputs to a layer dg,n. are the same for all n E [1....,N] N) and E {1..... R}. We do not state these types of conditions explicitly from now on.
Lemma 3 (Sufficient condition for exactness of K-FAC-expand in the expand setting). Let C₁ or Pr. be a constant matrix for layer f. If be 0 for all T + m and be ((f )m)))))))) ( fm.m. = C2 for all n E (1....,N) and m E {1...., R), then the K- FAC approximation in Equation (7) is equal to the exact GGN/Fisher of the f-th layer in the expand setting.
Proof. As mentioned before, when br.n.r.m 0 for all r + m the last line of Equation (26) and the first line of Equation (27) both simplify Equation (29). Hence, can directly show that the second and third approximation in Equation (27) equal the exact expression for the GGN of layer e from there. We have
N
R
R ΣV
N
ar
2
bi b ((fo(e)))))
NR
R a.m.ra
RC))
NR
n=l
(30)
N
ag,,r alar
&C
R
(arm,
20

---

where we have used the assumption that ((f.m..mAfffaaa 1 m.n.m.m C2 is the same for all
[I....,N]] N. and m E (1...., RJ.
L
Leveraging this simple insight, we can provide an example of single layer where the assumptions of Lemma 3 are fulfilled.
Proposition 4 (Exactness of K-FAC-expand for single linear layer in the expand setting). For single linear weight-sharing layer and Gaussian likelihood with p.d. covariance malrix E RCxc K-FAC-expand is exact in the expand setting.
Proof. We can write fo( Wr and hence bi, n,,,. for + m. Moreover, have A(fo(zn),) E-1 and b mem,m,m Ic (Pe = C for single layer). Hence,
for all n E (I....,NN .N and m (1..... RJ. Therefore, the desired result follows from Lemma 3.
A natural question might be if the same result also holds for deep linear networks. A deep linear network is here defined as model of the form
fo(x) WT. W.... Wi Wæ,
(31)
where E R.D and W. E Ix We E Re PP, (with Pt,nn Pt-t, and W1 P..uut D Decomposing single weight matrix W into L separate ones a common framework for theoretical analysis since it creates nonlinear training dynamics for gradient-based training algorithms, while still having analytical solutions (Bernacchia et al., 2018; Saxe al., 2014). We adopt the notation of Bernacchia et al. (2018) and define
W = WT. W4+1
(32)
as the product of the weight matrices ahead of We and
w = Wi-1... W1
(33)
as the product of the weight matrices behind of We. Hence, we can write fe(Ø) = W W, W
Note, that now
at W/ En, E Pr.in
(34)
and
ben,rr W ERProuxC XC
(35)
Using these insights, we can now easily state the result for deep linear networks.
Proposition (Exactness of K-FAC-expand for deep linear network in the expand setting).
For layer & of a deep linear network defined as in Equation (8) and a Gaussian likelihood with p.d. covariance matrix EE CCCC K-FAC-expante is exact in the expand setting.
Proof. We can write fo( = W Wnn, and hence be. for r A m. We have (f)())))) E-1 and bgnurr = We' Hence, be,n,, A (fo(z,)) ) an.,m, = W w for all n (I...., N} and m [I...., RJ. Therefore, the desired result follows from Lemma 3.
B.2.2 The reduce setting and K-FAC-reduce
The reduce setting is characterized by a loss with just N loss terms, i.c.
The Reduce Setting
N
CroducelJo. D) =
log log.(unlfo(ww.))
(36)
where the crucial observation is that the weight-sharing dimension must have been reduced somewhere in the forward pass of the neural network fe. A typical instance where this type of loss is used together
21

---

with model with linear weight-sharing layers is image classification with vision transformer or convolutional neural network. Note, that the inputs on and labels y,, do not have weight-sharing dimension here; in general, it is also possible for the inputs to have this additional dimension of size R already.
Since Aen E R is now a matrix, we have St. = At.n W E RRxP Hence, Ja.St. St. and fol.)) are now both multi-dimensional arrays. Luckily, can simplify this by writing
P Poui 0fo (In)) 0S
(Jonf(Zz))
(37)
iStmrr a0r.j
or in matrix form
R
Jo.Jo( = J Ja So 8t,
Σ
(38)
where St,n,r Pt.tt the r-th row of St,, and 8g,,, W.. Using this equivalence we can approximate the GGN for layer l as
N Jo ( (A(J((z))))
G Σ
n=1
N
ft
R
J oL. se..rr
A(folEn))
Σ
Jo Join st,,rr
n=1
N
R
R
Ot,mr @ bear
A (Jo(Z.))
Σ
at,n @ bejer
7=
N
R
R
R
x
Σ be
((o((an))
T
Σ
b
R
R
ag aT a.)
A(S
((o
A
B
(39)
where we have approximated the sum over the R Kronecker products with a Kronecker product of sums for each of the N per-input Jacobians, before applying the same type of approximation usual to the sum over the N data points. This approximation has been proposed in Tang et al. (2021) to improve the efficiency of their proposed K K-FAC variation (SKFAC) for convolutions by reducing the spatial dimension, purely based on the empirical observation that it works well in practice. The idea to approximate the Jacobians within the GGN with context of invariance learning with deep neural networks via differentiable Laplace approximations in Immer al. (2022). We call the approximation Equation (39) K-FAC- C-rducce and highlight the difference to K- K-FAC-expand, can rewrite it as
connecker--roducce has also been proposed the
22

---

K-FAC-reduce
G
reduce
R Σ
R
E a
Σ bennr (X))(Xn))
bl
iin,
NR
n=1
r=ll
=B.
(40)
As for K-FAC-expand, we want show that this approximation can exact in the case a single layer or a deep linear network and Gaussian likelihood. First, we state an analogous condition to Lemma
Lemma 5 (Sufffcient condition for exactness of K-FAC-reduce in the reduce setting). Let Din R. XC be conslant matrix for layer E and data point En. Further, let C2 E R Pr o be a constant matrix for layer E. If holds for each a that br,n.r = Din for all E (l....,11 R] and (ff((((( Equation (10) is equal to the exact GGN of the E-th layer in the reduce setting.
Ce for all E {l....,N}. then the K-FAC-reduce approximation in
Proof. We start with the first approximation and derive the exactness of this step under assump- tions. We have
agnar
bennr
A(fo(zn))
1
Σ ai
bearr
R
R
N
R
1
R ar,nr @ RDDm
A(fo(znn)
ag, & RRDe
(41)
R
R
R
ain.r & benrr A(fo(zn))
at,rr @ benr
where we have used the assumption that for each n, we have b = Den, for all r E {1...., R). Now we consider the second approximation in Equation (39). Analogously, we have
N Σ
āg
Be, A(F((())())))
N
2 NC
(42)
1=
(ār a..n)
where we have used that (f,,,((f(x = Ct for all e (I....,N..
Until now, we not have to explicitly take the aggregation function z :R RxxPoo + P Pro into account, since its Jacobian is simply subsumed in b Since want to verify that the approximation in the reduce case also exact in the simple scenarios from Proposition and Proposition 2, we now have to also check if the Jacobian Jer.. zt,n with zt,n 2(S(,n) R Prou is the same for all E {l....,RR to make sure the first condition in Lemma 5 is fulfilled. Maybe
23

---

the simplest case where this holds is scaled sum, i..
x(S,,n)
Σ Stn,r
cS 1R
(43)
= (1h & cI (e,(
=(1h cr
K Peo (5,n)
with c E R and the commutation matrix
Pront
K (R,,Po
Σ
Σ
(e om..)) @ (e pe
(44)
where Gi.j is the j-th canonical vector of dimension i. This is linear function in vec(Se,,. and we have Jes,n., Rt,n = clpt.ou for all r E {I...., R}. In particular, when c = the aggregation function is a simple sum and when c 1/ is the mean. Notably, it is not sufficient for to be linear in vec(St,n). because as soon as we have a weighted sum with weights c, E IR and they are the same for all r E fI....: R). the Jacobians Ja.... Zt,n will also not be the same anymore. Both vision transformers and convolutional neural networks with average pooling use scaled sums as the aggregation function (with 1/n)
.......,.
After clarifying the role of the aggregation function in the exactness of K- K-FAC-reduce, we can now state a similar statement Proposition 4.
Proposition 6 (Exactness of K-FAC-reduce for single linear layer in the reduce setting). For single linear layer, a Gaussian likelihood with p.d. covariance matrix Σ E R RCCC and a scaled sum as defined Equation (43) as the aggregation function applied to the output of the linear function, T-FAC-reduce is exact in the reduce setting.
Proof. We have (f))(zn)) = -111 and bi,nr (o z1,)) {1...., R} and n E (1....,N] N} (Pt,out = =C' for a single layer). Hence, be, I -II = = for all n E NJ. Therefore, the desired result follows from Lemma 5.
elc for all r
Just for K-FAC-expand, can extend this result deep linear networks.
Proposition 2 (Exactness K-FAC-reduce for deep linear network in the reduce setting). For layer E of a deep linear network (Equation (8)), a Gaussian likelihood with p.d. covariance matrix EEROXC and scaled sum aggregation function, K-FAC-reduce is exact the reduce setting.
Proof. We have A(fo(zn)r) E-1 and D Se =mm) cW." for all r {1...., .......,) R} and (I.....NY N}. Hence,
b (f c2 BewWaa---ii w-'We
for all n E (1,....N, Therefore, the desired result follows from Lemma 5.
To summarise, the difference between the expand and the reduce setting is at what point the aggre- gation over the additional weight-sharing dimension happens. If this dimension is not aggregated before the per-example loss, i.e. if the loss can be expanded to N R instead of N terms, we call it the expand setting. If the aggregation happens inside the model, we call the reduce setting, Both settings motivate an approximation each, K-FAC-expand and K K-FAC-reduce. Moreover, we presented simple cases where the approximations are exact. In Figure and Figure we verify this numerically, and also show that using the inappropriate approximation results in an inexact computation.
Remark 1. In practice, however, both approximations can be applied in each of the two settings.

---

B.3 Examples of K-FAC for linear weight-sharing layers in the wild
B.3.1 K-FAC for self-attention
While we have mentioned (vision) transformers for translation and image classification as prototypical examples for the expand and the reduce setting, we have mostly ignored how linear veight-sharing layers are used within the architecture and how this affects the approximation quality of K-FAC- expand and K-FAC-reduce. Linear weight- weight-sharing layers are crucial for the self-attention mechanism in Section 2.2. To gain some intuition for models using this type of attention mechanism, we look at network that only consists of one simplified variation of the self-attention mechanism used in transformers (we ignore the softmax function, but also consider linear projection of the output with weight matrix W ), i.e.
= = W X X, W.
(45)
SSQq.
=: Kin
SS
We can observe that is no longer a linear function in the input X, E IR ReD and that we have three linear weight-sharing layers involved in this operation. First, consider the expand setting, i.e. the output (..)
, is not reduced before the loss is applied.
Simplified self-attention in the expand setting. Since we want to understand if K-FAC-expand can be exact this case, we first derive the Jacobians appearing in the derivation of K-FAC-expand in Equation (27) for all three involved layers, i.e. Jaq Xx),, , for the layer with weights WQ,
Ja Us fo xr for the layer with weights WK and J Jay.,, fo (X.,,, for the layer with weights WV.
We can simply write the r-th row of the output of the layer with the weight matrix W as a function of sQ
so fo(Xn); = 8 n. Sk.n Sy
(46)
Therefore, we have
J Jaq. So (x())), SV Sk, 1Q CCCxxx
(47)
with Pvt and bQ.n.r.m 0 for all # r. which is the first assumption necessary for Lemma 3 to hold. While bQ... is not the same for all E {1....,.N N). it is the same for all r E (1...., R) and hence, under the same assumptions as Proposition 4, K-FAC-expand is exact for the layer with weights Wº in the special case of a single data point, N ==
For the other two involved linear layers, we cannot express the r-th row of fe( as function of the r-th row of SK/V.n. i.c. elements from all rows of SK/V.a contribute to the r-th row of the output matrix. We can also see this by directly deriving JaK/v..... (x(X,,, which will be generally non-zero and dependent and directly state the results. For the second layer with weight matrix WK we have
and m. We omit the explicit derivation by taking the partial derivatives
J Jak.n. So (X,)r &v Q. E CCXPa... PQ.
(48)
Moreover, for the third layer with weight matrix WV, we have
J J (X,))) = 8 =
E CCC
(49)
This means that the assumption of Lemma 3 that the R. elements along the weight-sharing dimension are independent does not hold, since the Jacobians depend on r and m. The approximation leads to an inexact computation, even though only linear layers are involved and Gaussian likelihood is used. Similarly, we can inspect the corresponding reduce case.
Simplified self-attention in the reduce setting. Assuming we use a scaled sum 2 with factor as the aggregation function, we can further rewrite the Jacobians occurring in Equation (39) as
R
Jo. (e.e(f((Xn)
Σ Jer... x(fo(Xn)) St.n,
r=l
atn.r b
(50)
R
R c Σ
So(
ag,nr
r=1
m=1
25

---

where Jos (e are the same Jacobians we have derived for the expand case. Since according to Lemma 5 we need all bg, to be the same for all r E (1,..., RJ for the first approximation to he exact under the assumptions of Proposition 6, K-FAC-reduce is only exact when N and only for the layer with weights W - just as K-FAC-expand in the expand setting.
When extend this scenario above statements regarding the special case where K-FAC-expand and K-FAC-reduce are exact for the layer with weights W only hold for the last block. While we omit explicit derivation, intuitively. this can be seen by the fact that we cannot rewrite the r-th row of this model's output as a function of only the r-th row of the layer's 's output St.q. of all layers with weights W besides for the layer in the last block, i.c. the layer in the L-th block with weights W.-
network consisting of L blocks defined in Equation (45), the
This shows that even without explicit nonlinear activation functions, the self-attention mechanism in transformer models breaks the two approximations, Hence, it is not inherently clear how useful it is to consider the corresponding approximation in the expand and reduce setting. This is especially relevant given that we know that the computational complexity of K-FAC-reduce is generally smaller than of K-FAC-expand and given the similar downstream optimisation performance we report in Section 4.
B.3.2 K-FAC for GNNs
Beyond transformers, we have introduced GNNs as a class of models that also use linear weight- sharing layers. There are many types of GNNs and we will only explicitly cover two of them here.
node classification with GCN. We first consider a GCN layer, described in Ap- pendix A. This specification K-FAC for GNNs has been previously derived for semi-supervised node classification in Izadi et al. (2020) and we include it for completeness, since it is, to the best of our knowledge, the only case where K-FAC has been applied to GNN. The only difference to the normal derivation of K-FAC is that the inputs āg, to the L-th layer for the node with index n now depend on its neighbourhood A (n)., since
Related work:
den
Σ jEN(n)
Enjac. E PRP....
(51)
Using this notation, the definition of the K-FAC GGN for node classification is simply
N
GGN(O:) = Σ Jo, (x)
(52)
d A (Jo(X) bb,
ag agn
B
where X E NND, be, Jar.. fe, (X), and 8e, Wrā Again, it is important to note that we need to have access to the whole neighbourhood of x,, to be able to write the n-th term of the GGN, which is why the input the model is the matrix X containing all nodes for each loss term. Also, depending the sparsity of Ć, i.e. the size of neighbourhoods, we might have multiple identical terms. In the extreme case of all neighbourhoods being the same, e.g. in the case of a fully connected graph, i.e. are values of Ć are the same, all terms of the GGN will be the same.
According to the first of the three cases in Section 3. 1. we do not need to think in terms of the expand and reduce settings here - as opposed to the case of the GraphNetwork we consider next, because we aggregate over each node's neighbourhood before the forward pass through a linear layer.
Graph classification with GraphNetwork. Now, we want to look at a more general architecture, an instance of the GraphNetwork introduced in Battaglia et al. (2018) and described in Appendix A. It is important to note that while the inputs and the GraphNetwork block structure look different from our standard input and linear layer, this case can be treated the same. This architecture is therefore good didactic example of how to apply the presented framework of thinking about K-FAC for linear weight-sharing layers to new model architectures. In contrast to the original description in Battaglia
26

---

et al. (2018), we already defined the inputs to a GraphNetwork block according to our definition of an input that leads to weight-sharing, i.e. with an additional weight-sharingg dimension of size R. This is in fact the crucial step to be able to apply our framework in this setting. As noted in Appendix A, the inputs cannot be trivially batched in this formulation. This is not an issue for our derivation, but requires special consideration in the implementation, which we will consider in Appendix B.4.
First, we note that we consider the task of graph classification. Hence, our loss has the same form as Equation (36), which means that the weight-sharing dimensions have to be reduced at some point during the forward pass and we are in the reduce setting. Notably, this is the setting of the ogbg-molpcba workload of the AlgoPerf benchmark from the experiment in Section 4.2, Following this line of reasoning, we would simply have to apply the corresponding K-FAC approximation. Since the imputs take more complex form than in our description of the reduce case, we still have to adopt the notation from Appendix A to concretely write down the approximation.
To recap, the weight-sharing dimension of size Rs of graph x depends on the input graph itself (indicated by the index n) and which update function within a GraphNetwork block we want derive K-FAC-reduce for. For of this dimension is going to be R, N whereas it will be R2 = N for σy. In the case of d² this layer is applied, and we can simply apply the regular K-FAC approximation. We can define the inputs to layer of type o² as
do not have weight-sharing dimension, has been reduced before
X A = concat(X:, XV n. n,s," repeat N.
N (xxi)) E R No (Dg+2Dy+DD)
(53)
and as
men = concat(XX) X repeat (xx)) E R.N. x(Dv+
(54)
for oY. Correspondingly, we have b = fo(x)) E R N. for oF and b
fe (XA) E R DDxCC with Si, W. RNNED and St, A Wy
R.N DDv respectively.
Using this notation, we can approximate the GGN for layer l, assuming its type is either @ or oY as
GG = Σ Jo (xf) (x
Rm
Rm
R
1
ag
b be,no
A(F((()))
Og
&
be
RS
Ry
=tdenn
=ibe,n
Σ b A (fo (fo (1)8))
ag,n a
A
B2
(55)
analogously to Equation (39). However, there is one important difference: since R, now depends on the n-th data point, it makes a difference where we insert the scaling by 1/R,,, since Å, and/or B, now weighted sums. To avoid having one potentially non-uniformly weighted and one unweighted sum, we choose to not just include the scaling by 1/Rt R,, in â as in Equation (39), or in 6 Instead, to try to keep the overall scale and the weighting of the N terms balanced, we simply include a scaling by 1/ RR,, in both, āzn and b
B.4 Practical Considerations
While we have discussed theoretically how to apply K-FAC to linear weight-sharing layers, now turn to implementation details and computational considerations that are crucial for the practical application of the approximations.
Implementation details. There are multiple libraries that implement K-FAC for popular deep learning frameworks like Jax (Bradbury et al., 2018) and PyTorch (Paszke et al., 2019), c.g. KFAC- JAX
27

---

Check if there Jactore can directly be calculated.
veight-sharing disension; not, the Kronecker
2
if data ndin y Hint- size M, seight-sharing dimension R, feature dimension Pe Pe,in/out N, R, P_in in_data.shape
P_out out_.gadd shape[2]
if spproxination -- "expand';
Flatten the ueight-sharing dimenssin into the mins batch dinensson. in_data in_data. viev P_in) / nath, (R)
9
10
out_grade out_grrds, Vieu vie P_out)
11
alif approximation == 'raduce":
12
Redsce the weight-sharing dimension with meas and sue. in_data in_data deddaaaaaennnaeee (din-1)
14
out_ out out_grads. sun(dim=1)
15
Calculate Kronecker factors Az/Âc and B./B,.
16
torch. natmal (in (in_data.T, in_data) /
17
B
torch natmal (out._grads. out_grada)
Illustration of K-FACdexpand and K-FAC-reduce with code. This piece of code calculates the approximations one mini-batch and for one layer. We receive the inputs to the layer, in_ in_data, from a forward hook and the gradients of the loss w.r.t. the outputs of the layer, out tat_grads, from a backward hook. Here assume that only one additional weight-sharing dimension exists and that the first dimension always the mini-batch dimension. Note that the dimensions of in_ data and out- are the same, independently of the setting we are in (c.f. Section 3.1). This illustrates why we can choose to use either approximation, K-FAC-expand or K-FAC-reduce, in each setting, The actual implementation in ASDL very similar for torch. -n. Linear modules, the logic is just separated into multiple functions and allows for multiple weight-sharing dimensions. However, the adjustments for the GraphNetwork are more involved.
Listing 1:
(Botev & Martens, 2022) for Jax and BackPACK (Dangel et al., 2020), ASDL (Osawa et al., 2023),
and KFAC-PyTorch (Pauloski et al., 2021) for PyTorch. We focus the implementation of K- FAC-expand and K-FAC-reduce within ASDL, which we also use for the experiments in Section 4. K-FAC is implemented using forward and backward hooks, which allow specific layer and the gradients of the loss w.r.t. the layer outputs which are the ingredients we need for all K-FAC approximations. Notably, this requires that linear layers are implemented with torch. mn. Linear instances, since otherwise, the implementation with hooks does not work. The default implementation of multi-head attention in PyTorch does indeed not use the required linear modules, so the implementation has to be adjusted to work with common K-FAC implementations like ASDL. In contrast, other methods like Shampoo (Gupta al., 2018) and Tensor Normal Training (Ren & Goldfarb, 2021) are agnostic to the architecture. Assuming the implementation of the model is appropriate, K- X-FAC-expand and K-FAC-reduce in their simplest form only require a minor adjustment in the code base for regular K- K-FAC, which is presented in Listing 1.
get the inputs
However, if we wanted to use K-FAC-expand in the expand and K-FAC-reduce in the reduce setting, we would need to find a way of automatically determining the setting we are in. For all models considered here, i.e. the (vision) transformer and GNN, only one of the settings applies to all linear layers with weight-sharing. Hence, using single additional forward pass, we could check if any linear weight-sharing layers are used and what the shape of the model output is. From this, we can deduce if the expand or the reduce case applies. As we mentioned before, this might not even be desirable, as it is unclear if we should always use the approximation theoretically motivated by the setting. Alternatively, single flag set by the user can determine if K- K-FFCC-eppad or K K-FAC-reduce applied to all linear weight-sharing layers.
This implementation obviously assumes that we even have an explicit weight-sharing dimension. In the case of K-FAC-reduce for the Getworr introduced in Appendix A, we have to adopt our implementation due to the batching technique that is employed for graph inputs in practice. Since each graph in mini-batch M of size M might have different weight-sharing dimension Rm. i.c. the number of nodes and the number of edges of each graph, we cannot batch them trivially. As a solution, the inputs for each graph stated in Equation (53) and Equation (54) are simply concatenated in the first dimension, which results in a dimension of size R. == -MM_ Rm, To apply K-FAC-expand here, we do not have to modify anything. besides scaling the approximation for
28

---

each mini- by 1/RM instead of 1/M.* 8 To apply K-FAC-reduce, we can use scatter mean, which aggregates tensor elements according to indices, to implement the mean (with the square root scaling from Equation (55)) operation without having an explicit weight-sharing dimension. Unfortunately, this creates two issues. First, we have to know that this adjustment to K-FAC even required for a specific layer since we cannot deduce it from the shape of the layer inputs. Second, the scatter mean requires additional information, since we need to know to which graphs the nodes/edges in the input belong. One approach to resolve these issues is to define custom layer type for this type of linear layer, which has an attribute containing the indices of all nodes/edges for each graph in the batch. However, this requires changes to the model architecture implementation, because additional attributes have to be set for the regular linear modules.
Besides the changes necessary for K- K-FAC-expand and K-FAC-reduce, we can use the same additional algorithmic tools often used for optimisation with K-FAC. Typically, damping is used (Martens & Grosse, 2015), i.e. scalar is added to the diagonal of the two Kronecker factors A and B or the diagonal of their product - the latter corresponds to adding the Hessian of an isotropic Gaussian prior over the weights. Also, since we usually operate in the stochastic setting and only compute the K-FAC approximation on mini-batches, sometimes an exponential moving average over the Kronecker factors is used. However, in our experiments, we do not use such moving average and only compute K-FAC for a single mini-bated batch (besides for Figure 5, as mentioned in Appendix C.4) and still reduce the update frequency of the K-FAC statistics and the preconditioner.
Computational considerations. Besides the implementation details, we also have to consider the computational cost when deploying K-FAC approximations in practice. Here, we have to respect the same constraints as with regular K-FAC. When we have a large output dimension C, it is expensive or even unfeasible to propagate the C' x C' loss Hessian H fo (Bn, fo (xx)) for each of the N data points through the computation graph. Instead, we use the fact that we have
E Py--piullo )} V log logp(v)lf(<z..))
Ho (yn, fo(or))
(56)
and take S Monte Carlo (MC) samples from the model's predictive distribution y. ~ ffylff(z,)) Taking a single sample results in rank-1 I MC approximation of the true loss Hessian and only requires the propagation of a single vector through the computation graph for each data point.
C Additional experimental details and results
C.1 MI.Commons' AlgoPerf benchmark for training algorithms
The goal of the AlgoPerf benchmark for training algorithms (Dahl et al., 2023) is to measure "train- ing speedups due to algorithmic improvements" (MLCommons, 2022). Specifically, the benchmark defines multiple workloads, where one workload is defined by the combination of a dataset, a neural network model, loss function, and target performance defined by some evaluation metric. For example, training a ViT on ImageNet using the cross- cross-entropy loss until a target validation accuracy of 77.309% has been reached constitutes workload. In the AlgoPerf benchmark, training algorithms are compared on fixed hardware in terms of the wall-clock time they require to reach the fixed validation target performance.
For our experiments, we use two of the workloads of the AlgoPerf benehmark as realistic deep learning problems, to showcase the optimisation behaviour of our two K-FAC variations (Sections 4.2 and 4.3). Similar to the benchmark, we measure both the number of steps and the wall-clock time necessary reach the target validation metric. To put the training behaviour of K-FAC-expand and K-FAC-reduce into perspective, we compare them to the target-setting run of the AlgoPerf benchmark, denoted reference run in the experiments (Sections 4.2 and 4.3). To determine the fixed validation target for each workload, four standard algorithms (AdamW, NAdamW, SGD with Nesterov momentum (Nesterov), and SGD with heavy ball momentum) were each tuned with budget of 200 trials for every workload. The combination of algorithm and hyperparameters that reached the highest validation performance within fixed time was then run 20 times and the median of the best achieved validation metric was set as the validation target. As our reference run, we use the workload-specific, best-performing algorithm and hyperparameter combination. Note that this should
As explained below Equation (55), we could also choose a different way of scaling here. We choose this it enables simple implementation in contrast to K-FAC-reduce.
29

---

Table 3: Results from Table I with standard errors over 10 epochs.
Batch size
K-FAC
128
256
512
1024
2048
expand
0.24 0.01
0.38 + 0.02
0.75 = 0.03
1.36 + 0.05
OOM
reduce 0.17 = 0.01
0.24 = 0.02
0.43 = 0.03
0.63 = 0.04
1.17 = 0.08
not be confused with the baseline algorithms in the AlgoPerf benchmark, since the target-setting runs are tuned for the best possible validation performance each workload.
K-FAC is run with most of the hyperparameters of the reference run and a tuning budget of 15 runs for tuning the learning rate and damping. This means that most hyperparameters are not directly tuned for K-FAC. To enable a fairer comparison, we also use the same additional budget for tuning the learning rate of the reference run, the same learning rate schedule, and tuning goal as K-FAC. However, none of these runs reaches the validation metric target, which is why the reference runs correspond to the original target-setting runs. Despite this, the experiments presented in Sections 4.2 and 4.3 are not meant to demonstrate that K-FAC is superior training algorithm. To provide actual evidence for this claim we would have run valid and optimised submission the entire AlgoPerf benchmark or a benchmark of comparable rigour. Nevertheless, using the well-tuned reference run allows us to put K-FAC's s training performance into perspective.
The two AlgoPerf workloads used in this paper are:
Graph neural network on oghg-molpeba. The workload consists of a (binary) cross-entropy loss, the GraphNetwork instance described in Appendix A, and the gggg-molpeba molecular property prediction dataset (Hu et al., 2020). Each molecule is represented as a graph, where atoms are nodes and edges are chemical bonds. The task is predict whether not molecule has certain chemical properties; there are 128 different properties. For training, we have about 350k examples and almost 44k for validation and testing. The validation target mean average precision (mAP) determined by the target-setting runs is 0.28098 and the best performing reference run algorithm is Nesterov.
Vision transformer ImageNet. This workload also uses the cross-entropy loss, vision trans- former architecture, and the LSVRC-2012 ImageNet (short: ImageNet) image dataset (Russakovsky et al., 2015). The goal is to classify images into one of 1000 classes. There are about 1.3 million training, 50k validation, and 10k test examples. The validation target accuracy determined by the target-setting runs 0.77309 and the best performing reference run algorithm NAdamW.
C.2 Update step speed with K-FAC-expand and K-FAC-reduce
We provide the full results of the update step timing exper- iment with standard errors over 10 epochs in Table 3. The model architecture, optimiser, and data setup are exactly the same as described in Appendix C.5. The full results of the GPT-2 (nanoGPT) timing experiment are presented in Table 4; the mean and standard error are computed based on three runs.
Table 4: Timing of K-FAC GGN approx- imation for GPT-2 on full DART dataset.
K-PAC
Absolute time [min]
Relative time [%]]
expand roduce
9.55 0.13
100
6.688 0.04
70
C.3
Graph neural network on ogbg-molpeba
For this workload, we use a training batch size of 512 and a single NVIDIA V100 32GB GPU for each run. The reference run algorithm (Nesterov) uses a learning rate of 2.491 11.9977228669184422 3 equal 0.9449369031 191439991117744, weight decay set 1.285 964 054 102 592 7 and linear warmup for 3,000 steps and polynomial schedule with decay steps factor of 1861509027889633 and an end factor of 1c-3. The two K-FAC variations use the exact same hyperparameters, but the warmup and expected number of steps (60,000) are multiplied by 0.75 and the learning rate and the damping are tuned via random search. The search space for the learning rate is log uniform values in 10.1, 10) and for the damping in [1e-3, 1j. We choose the hyperparameters of the run that first reaches the validation target. This setting a learning rate of 0.9687190219 0. 1.78819653391908 for K-FAC-expann and learning rate of 0. 0.55885766511 14111300633 and damping
and damping of
30

---

of 0.0579230199994011 for K-FACCreeuce. We then repeated each run five times. The Kronecker factors and the preconditioner are computed every 10 iterations and we use single sample MC approximation of the Fisher, as explained in Appendix B.4.
C.4 Vision transformer on ImageNet
We use training batch size of 1,024, e 1e-8 for NAdamW, and 4x NVIDIA V100 32 GPUs for all runs on this workload. The reference run algorithm (NAdamW) is using learning rate of 31 equal to 0.8895758153482811 Bg to 0.9975504782314613,, weight decay set to 0.088354002596553223 linear warmup for 6999 steps, and a cosine decay schedule. The two K-FAC variations use the exact same hyperparameters. but the warmup and expected number of steps (140,000) is multiplied by 0.75 and the learning rate and the damping are tuned via random search. The search space for the learning rate and the damping is log uniform values in [1c-4, 1e-2]. We choose the hyperparameters of the run that first reaches the validation target. For both K-FAC variations, the best setting is learning rate 0. 0001266293334000453 of 0.000165240111 0111166511154228860 19254426572. The Kronecker factors and the preconditioner are computed every 50 iterations and we use a single sample MC approximation the Fisher, as explained Appendix B.4.
0.000844507456619757777
and damping
We also conduct a run where we update the Kronecker factors every step, use an exponential moving average with a factor lema_decay ASDL) equal 32, update the preconditioner every 10 steps, set the damping to 1e-5, and use all the other hyperparameters of a reference run setting, including the learning rate. The reference run setting corresponds to an earlier target-setting run result from the AlgoPerf repository and uses a learning rate of about 2e-3, 31 0.7132, 32 0.9982, and the weight decay is set to 0.026595. Moreover, it clips the gradients to keep their norm below 1. As before, we also multiply the number of warmup and expected number of steps by 0.75. Due to the high update frequency the Kronecker factors, we can see the significant wall-clock time difference between K-FAC-expand and K-FAC-reduce in Figure 5. Both variations are similar in terms of steps to the target, K-FAC-expand takes about 92.6k and K-FAC-reduce takes about 93.7k steps, but whereas K- K-FAC-expant takes about 50 hours, K-FAC-reduce reaches the target after about 37 hours. Note, that both variations are still significantly slower than the NAdam W reference run which only takes about 25 hours, but 117.4k steps.
C.5 K-FAC for automatic hyperparameter selection via marginal likelihood optimisation
K-FAC is used for the Hessian approximation within the Laplace approximation by first accumulating it over the entire data set and using the eigendecomposition of Kronecker factors as in Immer et al. (2021). This is done every five epochs during training to update the weight decay parameters per layer of the neural network. Therefore, the computation time of K- makes significant contribution to the overall training time: reduce takes on average 75% of the overall time of expand, a result of reducing the wall-clock time of single marginal likelihood update step by about 50%, shown in Table 2. Here, we provide additional experimental details and the full table with standard errors.
For the marginal likelihood optimisation experiment, we use the same settings as Daxberger et al. (2021), i.e., Wide ResNet 16-4 with Fixup initialisation and parameters instead of batch normalisa- tion (Zagoruyko & Komodakis, 2016; Zhang et al., 2019b). We use a batch size of 128, an initial learning rate of 1 cosine decayed to le-6 and weight decay of 5e-3 training for 100 epochs. We use standard data augmentation with horizontal flips and random crops. Results are averaged over three seeds and the timings are obtained running all runs on a single NVIDIA A100 GPU and locally measuring the hyperparameter update time. The full results with standard errors are presented in Table
31

---

Reference (NAdamW)
3.0
X-FAC-cxpand
K-FAC-rdduce
Training
Loss
2.0
Validation
Validation target
1.5
0.5
0.8
Accuracy f
0.7
0.6
0.5
40
60
100
120
10
20
30
Step
x10
Wall-clock time [hours]
Figure 5: Training results for vision transformer ImageNet. This is similar to Figure 4, but K-FAC statistics are updated every step and different hyperparameters are used. Due to K-FAC's overhead, the wall-clock time is not reduced in this setting. Moreover, the discrepancy in speed between K-FAC-expand and K-FAC-reduce becomes apparent.
Table 5: Results from Table 2 with standard errors over three random seeds. The update time is the average time per full K-FAC Laplace approximation to the marginal likelihood, which requires full data set pass with computation and eigendeecomposition of the Kronecker factors.
K-FAC
Data Augmentation
NLL
Accuracy [%]
Update Time [s] J
X
0.422 + 0.013
88.89 + 0.24
expand
196 + 22
/
0.244 + 0.004
92.52 + 0.12
X
0.703 + 0.012
86.71 + 0.13
reduce
99 +
0.352 +'t 0.008
93.50 + 0.06
32


Reviews:
[{"review_id": "lUKHVA8RgU", "reviewer": "Reviewer_t9fV", "summary": "This paper considers linear layers with weight-sharing, which is used in many neural network architectures, such as transformers, convolutional and graph neural networks. Two variants of K-FAC, K-FAC –expand and K-FAC-reduce are considered for two different settings of linear weight sharing layers. The authors also show that they are exact for deep linear networks with weight-sharing in their respective setting. Numerical results show the efficiency of the proposed K-FAC variants.", "questions": "To improve the clarity of the paper, it is suggested to provide a detailed description of the current implementations mentioned in section 3.2 [13, 29, 31, 43] in section 2, which will help readers follow the paper more easily and understand the differences more clearly. Besides, some mathematical derivations in appendix can be put in the main text.\n\nIt would be better to mention som related works. 1.) the NG+ method [1], which is similar to Shampoo and uses a direct approach to approximate the Fisher matrix.  2.) AdaHessian method [2], which dynamically incorporates the curvature of the loss function via ADAptive estimates of the Hessian matrix.  3.) the SENG method [3], which employs a sketching method to approximate the empirical Fisher matrix and is implemented in the ASDL package.\n\nThe computational complexity of using regular K-FAC approximation over N R terms instead of the usual N terms in the KFAC-expand method is higher, and when the batch size is large, as reported in the numerical experiments, it also increases memory usage.\n\nIn the KFAC-reduce setting, where $\\widehat{A_l}$ (also $\\widehat{B_l}$) is a symmetric matrix with rank at most N, this rank is small when a smaller batch size is used. However, it is unclear whether this reduction preserves sufficient curvature information. Do the authors consider use moving averaging in the implementation?\n\n\n[1] Yang, Minghan, et al. \"An efficient fisher matrix approximation method for large-scale neural network optimization.\" IEEE Transactions on Pattern Analysis and Machine Intelligence 45.5 (2022): 5391-5403.\n\n[2] Yao, Zhewei, et al. \"Adahessian: An adaptive second order optimizer for machine learning.\" proceedings of the AAAI conference on artificial intelligence. Vol. 35. No. 12. 2021.\n\n[3] Yang, Minghan, et al. \"Sketch-based empirical natural gradient methods for deep learning.\" Journal of Scientific Computing 92.3 (2022): 94.", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is well-written and organized. The use of second-order methods is typically influenced by the specific network architecture. In this paper, we focus on linear weight-sharing layers, which are crucial and commonly used in Transformer, convolutional, and graph networks, but have not received adequate attention in the realm of second-order methods. In my opinion, a more thoughtful analysis of the architecture could enhance the performance of second-order methods like KFAC and others. Undoubtedly, this paper will expand the applicability of second-order methods and their potential benefits for machine learning tasks. \n\nThis paper provides a correct mathematical representation of gradient for the weight-sharing linear layer structure, which is not trivial and can facilitate the development of related optimization methods.", "weaknesses": "One of my concerns is the applicability of second-order methods, particularly for modern architectures such as the LLM model. As noted in the AlgoPerf paper[1], the use of KFAC requires detailed architectural information, which may limit its use in practice. In contrast, methods like Shampoo and NG+ appear to offer more flexibility. Additionally, while the direct approach is effective in practice, there is a lack of theoretical understanding regarding the rationale for using the KFAC-reduce approximation.\n\nThe computational complexity of the proposed method appears to be higher than previous approaches, which could pose challenges for implementation. However, there is no direct comparison with the current implementation of the KFAC method in this setting, and it would be beneficial to include such a comparison. Additionally, it is recommended to compare the proposed method with other second-order approaches to gauge its effectiveness.\n\n[1] Dahl G E, Schneider F, Nado Z, et al. Benchmarking Neural Network Training Algorithms[J]. arXiv preprint arXiv:2306.07179, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and the detailed comments.\n\n>One of my concerns is the applicability of second-order methods, particularly for modern architectures such as the LLM model. As noted in the AlgoPerf paper [1], the use of KFAC requires detailed architectural information, which may limit its use in practice. (...)\n\nLLMs specifically can be easily implemented in a way that is compatible with K-FAC (in fact, many popular implementations already are, e.g. nanoGPT, which implements GPT-2).\nTo demonstrate this, **we compute a K-FAC GGN approximation for nanoGPT on the full DART dataset** [1]. In line with the complexities and our other results K-FAC-reduce is significantly faster, only taking 70% of the time of K-FAC-expand. Please find the table with the results in our reply to Reviewer hvpV.\n\nMore generally, we agree that second-order methods, and K-FAC in particular, are often not easy to apply to modern neural network architectures. However, we think that our work contributes to extending its applicability, since we describe how K-FAC can be applied to basically arbitrary neural network layers (as long as they can be expressed as linear weight-sharing layers). Despite this theoretical extension, we fully agree that K-FAC is still not agnostic to the model implementation, whereas other methods like Shampoo are. Since K-FAC has a more direct connection to the Fisher/GGN than the preconditioner used in Shampoo and other similar methods, we believe it is still a valuable tool to study the behaviour of second-order methods in deep learning. But, as mentioned before, it might also be able to significantly speed up training in specific use cases like LLM training.\n\n>Additionally, while the direct approach is effective in practice, there is a lack of theoretical understanding regarding the rationale for using the KFAC-reduce approximation.\n\nK-FAC-reduce is just as well motivated as K-FAC-expand, as they are derived by following the same principles, but in two different settings.\n\n>The computational complexity of the proposed method appears to be higher than previous approaches, which could pose challenges for implementation. However, there is no direct comparison with the current implementation of the KFAC method in this setting, and it would be beneficial to include such a comparison.\n\nThere seems to be a misunderstanding: K-FAC-reduce has actually lower computational complexity than K-FAC-expand, which has previously been used for linear weight-sharing layers. There is no third way of applying K-FAC to linear weight-sharing layers in the literature that we are aware of.\n\n>Additionally, it is recommended to compare the proposed method with other second-order approaches to gauge its effectiveness.\n\nWhile we fully agree that it will be interesting to see how K-FAC compares to other second-order approaches, we decided to not focus on this question in our paper. We do not intend to make any general claims about K-FAC as an optimiser, but want to generalise K-FAC to modern neural network architectures.\n\nTo make more general statements about K-FAC’s performance as a general purpose optimiser we would have to use a different experimental setup, e.g. run a valid submission to the AlgoPerf benchmark. \n\n>To improve the clarity of the paper, it is suggested to provide a detailed description of the current implementations mentioned in section 3.2 [13, 29, 31, 43] in section 2, which will help readers follow the paper more easily and understand the differences more clearly. (...)\n\nWe will highlight the differences between the different implementations more clearly and will use the additional page to move some parts of the derivations to the main text. To clarify, K-FAC has only been applied to the architectures we consider here in the form of K-FAC-expand (up to the scaling factor, see the footnote on page 5). While we mention this in lines 50-53 of the paper and in section 3.2, we will highlight this point in the revision of the paper. Also, note that we provide a Python code snippet in Listing 1 in the appendix, illustrating the actual implementation.\n\n>It would be better to mention some related works. 1.) the NG+ method [1], which is similar to Shampoo and uses a direct approach to approximate the Fisher matrix. 2.) AdaHessian method [2], which dynamically incorporates the curvature of the loss function via ADAptive estimates of the Hessian matrix. 3.) the SENG method [3], which employs a sketching method to approximate the empirical Fisher matrix and is implemented in the ASDL package.\n\nThank you for the suggestions, we will include these references in section 2.1.\n\n>The computational complexity of using regular K-FAC approximation over N R terms instead of the usual N terms in the KFAC-expand method is higher, and when the batch size is large, as reported in the numerical experiments, it also increases memory usage.\n\nYes, this is correct. In fact, K-FAC-expand is the currently used K-FAC variation for linear weight-sharing layers and the here derived K-FAC-reduce offers an alternative with lower computational complexity by a factor of $R$.\n\n>In the KFAC-reduce setting, where (also ) is a symmetric matrix with rank at most N, this rank is small when a smaller batch size is used. However, it is unclear whether this reduction preserves sufficient curvature information. Do the authors consider use moving averaging in the implementation?\n\nWe agree that it will be interesting to see how moving averages over the Kronecker factors will influence the approximation quality and downstream optimisation performance, but only use local estimates in our experiments, which seems to be enough to achieve the reported performance.\n\nWe hope that we have sufficiently addressed your concerns. Please let us know if there still is anything unclear.\n\n[1] L. Nan et al. DART: Open-Domain Structured Data Record to Text Generation. NAACL-HLT, 2021."}, {"review_id": "bwBjLoqkDl", "reviewer": "Reviewer_1aFR", "summary": "The paper formalizes two variants of the K-FAC optimisation method for a family of neural network architectures that can be expressed as linear layers with weight sharing. Experiments are conducted to validate that the proposed method can be applied to different kinds of neural network architectured speeding-up training.", "questions": "See above, in the weaknesses section.  Basically, contextualization of the proposed method with other similar approaches. \n\n\n-------------------------------\n\n\nI have read the rebuttal, other reviewers comments, feedback and discussion and I got a better understanding of the contribution of the paper. I increase my rating to accept.", "limitations": "Limitations of the method are partially addressed.", "rating": 7, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is a bit out of the scope of my specific field of expertise and it is possible that I failed to fully understand all technical details and mathematical formulation, but the general mathematical formulation of the method seems sound and complete. The family or achitectures where the two variants of the proposed method can be applied is formally defined, and the conditions for which exact approximations can be obtained are also defined.", "weaknesses": "I miss a better contextualitzation, both from a theoretical and experimental point of view,  of the proposed method with other existing approaches, cited in the related work, that also make use of K-FAC and that have also been applied to similar architectures using convolutional networks, transformers or GNNs. From a theoretical point of view, a better justification of the contribution of the proposed method with respect to other applications of K-FAC. From an experimental point of view, a comparison of the proposed method with other works also using K-FAC.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and their feedback.\n\n>I miss a better contextualitzation, both from a theoretical and experimental point of view, of the proposed method with other existing approaches, cited in the related work, that also make use of K-FAC and that have also been applied to similar architectures using convolutional networks, transformers or GNNs. From a theoretical point of view, a better justification of the contribution of the proposed method with respect to other applications of K-FAC. From an experimental point of view, a comparison of the proposed method with other works also using K-FAC.\n\nThanks for pointing this out. To clarify, K-FAC has only been applied to the architectures we consider here in the form of K-FAC-expand (up to the scaling factor, see the footnote on page 5). While we mention this in lines 50-53 of the paper and in section 3.2, we will emphasise this point more in the next revision of the paper. Since there is no known way to apply K-FAC to the considered architectures besides K-FAC-expand and K-FAC-reduce, we also can’t compare to other variations.\n\nSince you mentioned that the paper is a bit out of the scope of your specific expertise, please don’t hesitate to ask for further clarifications, we understand that it can be hard to place the work into the correct context without familiarity with this specific field of research."}, {"review_id": "O8vwU1LWgi", "reviewer": "Reviewer_M1Nr", "summary": "The researchers propose a framework for applying the Kronecker-Factored Approximate Curvature (K-FAC) method to neural networks with linear weight-sharing layers. They identify two different settings of linear weight-sharing layers, leading to two variations of K-FAC: expand and reduce. They demonstrate that both variations are exact for deep linear networks with weight-sharing in their respective settings. They find that K-FAC-reduce is generally faster than K-FAC-expand and leverage this to accelerate automatic hyperparameter selection for a Wide ResNet. They also show that both K-FAC variations can achieve a target validation metric in a significantly reduced number of steps compared to a first-order baseline, resulting in a comparable improvement in training time.", "questions": "1. what does Exactness of K-FAC-reduce and K-FAC-expand? why is this property important?\n2. How does it fare in lower precision setting, is this still numerically stable? ", "limitations": "providing more intuition in the theory section will be highly appreciated. From the reader points of view, who are not familiar with the notation setup, it looks like a \"wall of math\", everyone complains about.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Very interesting perspective by looking at modern NN building blocks as weight shared linear layers. \nSolid results and thorough experiments. It's great that the authors measure actual wall-clock time speedup instead of simply recording the steps. This shows that it's a method with practical merits.", "weaknesses": "Minor issues in presentations, such as the plots are not really hard to read, and organization of the propositions are too clustered to be read comfortably.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and appreciate that they also found the concept of linear weight-sharing layers to describe modern neural network layers useful.\n\n>Minor issues in presentations, such as the plots are (...) really hard to read, and organization of the propositions are too clustered to be read comfortably.\n\nWe will improve the readability of the plots by summarising the lines for runs with different seeds and changing the colour scheme. Moreover, we will use the additional space to improve the readability of the propositions. In case you have any other concrete suggestions, we would highly appreciate it.\n\n>what does Exactness of K-FAC-reduce and K-FAC-expand? why is this property important?\n\nWe use the exactness of the two K-FAC variations in their respective settings purely as motivation for their derivation, since this is exactly the original motivation used for deriving K-FAC for regular linear layers [1][2]. **The proposed variations should really be understood as a generalisation of K-FAC and not an alternative method.** The statements do not say anything about the downstream optimisation performance or even just the approximation quality in deep nonlinear networks.\n\n>How does it fare in lower precision setting, is this still numerically stable?\n\nWe have not tested this experimentally, but we expect that K-FAC might encounter numerical issues to a similar extent as all other optimisation methods that have to invert dense matrices. However, there is an inverse-free natural gradient method that also relies on K-FAC, which we expect to do much better in lower precision settings [3]. Crucially, since a K-FAC approximation to the Fisher/GGN is also needed for this method, our work by extension enables this method for the here covered modern neural network architectures.\n\n>providing more intuition in the theory section will be highly appreciated.\n\nWe are happy to use the additional page for more details and intuitions on the derivations. Specifically, we will add an intuitive visualisation of the the different settings K-FAC can be applied in and how this relates to the different K-FAC variations.\n\n[1] J. Martens, R. Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. ICML, 2015.\n\n[2] A. Bernacchia, M. Lengyel, G. Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. NeurIPS, 2018.\n\n[3] W. Lin, V. Duruisseaux, M. Leok, F. Nielsen, M. E. Khan, M. Schmidt. Simplifying Momentum-based Positive-definite Submanifold Optimization with Applications to Deep Learning. ICML, 2023."}, {"review_id": "LJ5l3alpv4", "reviewer": "Reviewer_T2nD", "summary": "This paper propose two methods called KFAC-reduce and KFAC-expand to conduct the secon-order based training for neural networks. The key idea is to approximate the Gauss–Newton matrix with low-rank matrix decomposition. The experiments on several tasks show that the propsed method can achieves faster convergence rate compared to NAdamW optimizer.", "questions": "- Fig 4., the loss is not flat and still tend to decrease. Could the author extend the training steps until the loss curve becomes flat?", "limitations": "Please check the weakness sections.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed method has well-grounded theoretical support.\n- The proposed method has shown better results than the NAdamW/SGD optimizer.", "weaknesses": "My biggest concern comes from the evaluation sections.\n- The proposed method is only verified on limited number of tasks and networks (4 configurations). The author should show more results with various configurations.\n- The method is only compared with two type of optimizers. For example, SGD is the baseline for GNN experiments while NAdamW is the baseline for the ViT experiment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and comments.\n>The proposed method is only verified on limited number of tasks and networks (4 configurations). The author should show more results with various configurations.\n>The method is only compared with two type of optimizers. For example, SGD is the baseline for GNN experiments while NAdamW is the baseline for the ViT experiment.\n\nWe believe that these points are the result of a misunderstanding. The goal of our experimental evaluation is not to make any general statement about the algorithmic superiority of the K-FAC variations as an optimiser.\n\nIn contrast, the goal of the paper is to generalise K-FAC to modern neural network architectures and **the proposed variations should really be understood as a generalisation of K-FAC and not an alternative method.** We use our experiments as a “proof of concept”, to show that the K-FAC variations have the potential to speed up training, to empirically validate the difference in the computational complexity of the two variations, and present an example of another use case, online approximate marginal likelihood optimisation for automatic weight decay tuning. But we fully agree that our results are very promising and hope that our work can pave the way for a rigorous comparison, e.g. in form of a submission to the AlgoPerf benchmark [1].\n\n\n>Fig 4., the loss is not flat and still tend to decrease. Could the author extend the training steps until the loss curve becomes flat?\n\nThis is due to our experimental setting: as described in the experiments section, we follow the setup of the AlgoPerf benchmark and only train until a validation target performance is reached. While we think that it is interesting to investigate the behaviour of K-FAC at convergence, this is an orthogonal question and requires a different experimental setup. Since we did not tune for the best validation performance at convergence, running the experiments for more steps will not allow us to make any statement about the performance at convergence. E.g., the learning rate schedule will stop adopting the learning rate soon after the target is reached.\n\n[1] G. E. Dahl, F. Schneider, Z. Nado, N. Agarwal et al. Benchmarking Neural Network Training Algorithms. Preprint (arXiv: 2306.07179), 2023."}, {"review_id": "OODUhwKTWw", "reviewer": "Reviewer_NJ6c", "summary": "This paper introduces two different settings of linear weight-sharing layers to fit for the second-order optimization method Kronecker-Factored Approximate Curvature (K-FAC). Experiments on GNNs and ViTs achieve less number of steps compared with the first-order baseline.", "questions": "How about the comparison with other efficient optimization methods other than weight sharing?", "limitations": "The authors discussed the runtime of the two proposed schemes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Experiments include GNNs and ViTs, also the results are quite inspiring in the acceleration of steps.", "weaknesses": "Fig. 4 doesn’t show the loss and accuracy variations after convergence.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">How about the comparison with other efficient optimization methods other than weight sharing?\n\nWe believe there is a fundamental misunderstanding. **Linear weight-sharing is a property of the layers of the model architecture and not of the optimisation method.** In our paper, we show how we can apply K-FAC to model architectures with linear weight-sharing layers by leveraging their structure.\n\n>Fig. 4 doesn’t show the loss and accuracy variations after convergence.\n\nThis is due to our experimental setting: as described in the experiments section, we follow the setup of the AlgoPerf benchmark and only train until a validation target performance is reached. While we think that it is interesting to investigate the behaviour of K-FAC at convergence, this is an orthogonal question and requires a different experimental setup. Since we did not tune for the best validation performance at convergence, running the experiments for more steps will not allow us to make any statement about the performance at convergence. E.g., the learning rate schedule will stop adopting the learning rate soon after the target is reached."}, {"review_id": "KDtc4G76jB", "reviewer": "Reviewer_hvpV", "summary": "This paper extends Kronecker-Factored Approximate Curvature (K-FAC) into training generic neural networks, especially for transformers and graph neural networks. This is done by taking advantage of the weight sharing mechanism. The authors propose two flavors of K-FAC (expand and reduce), depending on the dimension of the output tensor. The authors prove that they are exact for deep linear networks. The proposed methods reach a validation metric target of a well-tuned first-order baseline with smaller steps and less wall-clock time.", "questions": "In figure 4, the K-FAC statistics are updated every 50 steps.\n\n1. Is the same setting used in other figures?\n2. What if we change the number of steps?\n", "limitations": "This paper has no negative societal impact.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper extends K-FAC into transformers and graph neural networks. Although some previous works try to extend K-FAC into transformers, their approximation is inexact, as pointed by the authors.\n\n2. The proposed methods have mathematical foundations. Because the authors prove the exactness under deep linear settings.\n\n3. The proposed methods reach a validation metric target of a well-tuned first-order baseline with smaller steps and less wall-clock time. So it is possible to replace first order methods.\n\n4. This paper provides sufficient supplementary material, including mathematical derivations, additional experiments, and reproduction codes.\n", "weaknesses": "1. The performance of transformers is evaluated on ImageNet. Since transformers are mostly used for language modeling, it is better to evaluate your optimizers on some language tasks (a very big dataset is unnecessary).\n\n2. The authors prove the exactness under deep linear settings. If possible, It is better to provide some error bounds for non-linear settings.\n\n3. It is helpful to provide the time complexity (not wall-clock time) and space complexity for both K-FAC-expand and K-FAC-reduce. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and their comments.\n\n>The performance of transformers is evaluated on ImageNet. Since transformers are mostly used for language modeling, it is better to evaluate your optimizers on some language tasks (a very big dataset is unnecessary).\n\nWe believe that transformers are interesting in many settings and vision transformers in particular are certainly a widely used model class -- according to the paperswithcode website on vision transformers they are currently even more popular than ResNets and the original paper proposing them has over 19k citations. Moreover, K-FAC(-expand) has been applied to transformers on language tasks before [1][2][3], but never to a vision transformer.\nHowever, we fully agree that it would be great to have results for a language task. To demonstrate the applicability of K-FAC to language models, **we compute a K-FAC GGN approximation with nanoGPT, a popular implementation of GPT-2, on the full DART dataset** [4]. In the following table we present the timing results of the whole computation for K-FAC-expand (currently used for transformers) and K-FAC-reduce. In line with the complexities and our other results K-FAC-reduce is significantly faster, only taking 70% of the time of K-FAC-expand. We show the mean and standard error of the timing of three runs.\n| K-FAC  | Absolute time (minutes) | Relative time (%) |\n|--------|-------------------------|-------------------|\n| expand | 9.55 +/- 0.13           | 100               |\n| reduce | 6.68 +/- 0.04           | 70                |\n\n\n>The authors prove the exactness under deep linear settings. If possible, It is better to provide some error bounds for non-linear settings.\n\nWe use the exactness in the deep linear setting as motivation for our derivation, analogous to the original derivation of K-FAC for regular linear layers [5][6].\nWe believe it is hard to derive useful error bounds for the deep non-linear setting, without making strong assumptions. Moreover, it is unclear if statements about the approximation quality translate in downstream optimisation performance in practice.\n\n>It is helpful to provide the time complexity (not wall-clock time) and space complexity for both K-FAC-expand and K-FAC-reduce.\n\nWe state the time complexity of both variations in lines 242-244 in the main text of the paper. We will add the space complexity, which is $\\mathcal{O}(NR(P_{\\ell, \\mathrm{in}} + P_{\\ell, \\mathrm{out}}))$ for K-FAC-expand and $\\mathcal{O}(N(P_{\\ell, \\mathrm{in}} + P_{\\ell, \\mathrm{out}}))$ for K-FAC-reduce, where $N$ is the dataset size, $R$ is the size of the weight-sharing dimension (e.g. the number of tokens for language data), and $P_{\\ell, \\mathrm{in}}$ is the input and $P_{\\ell, \\mathrm{out}}$ the output dimension of the layer.\n\n>In figure 4, the K-FAC statistics are updated every 50 steps.\n>1. Is the same setting used in other figures?\n>2. What if we change the number of steps?\n\n1. This setting is only used for Figure 4. In Figure 3 the K-FAC statistics are updated every ten steps, as described in section 4.2 of the paper. We will add this information to the caption of the figure.\n2. In Figure 5 in the appendix you can see a run where the statistics are updated every step. There is little to no difference in the optimisation performance. For practical applications it is an interesting question how sparsely K-FAC can be updated. Also, note that sparse updates are already commonly used for second-order methods, see e.g. [7][8][9].\n\nWe hope that we have sufficiently addressed your concerns and questions. Please let us know if there still is anything unclear.\n\n[1] G. Zhang, L. Li, Z. Nado, J. Martens, S. Sachdeva, G. E. Dahl, C. J. Shallue, R. Grosse. Which algorithmic choices matter at which batch sizes? Insights from a noisy quadratic model. NeurIPS, 2019.\n\n[2] J. G. Pauloski, Q. Huang, L. Huang, S. Venkataraman, K. Chard, I. T. Foster, Z. Zhang. KAISA: an adaptive second-order optimizer framework for deep neural networks. International Conference for High Performance Computing, Networking, Storage and Analysis (SC21), 2021.\n\n[3] K. Osawa, S. Li, T. Hoefler. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. Preprint (arXiv: 2211.14133), 2022.\n\n[4] L. Nan et al. DART: Open-Domain Structured Data Record to Text Generation. NAACL-HLT, 2021.\n\n[5] J. Martens, R. Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. ICML, 2015.\n\n[6] A. Bernacchia, M. Lengyel, G. Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. NeurIPS, 2018.\n\n[7] V. Gupta, T. Koren, Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. ICML, 2018.\n\n[8] Y. Ren, D. Goldfarb. Tensor normal training for deep learning models. NeurIPS, 2021.\n\n[9] K. Osawa, S. Ishikawa, R. Yokota, S. Li, T. Hoefler. ASDL: A unified interface for gradient preconditioning in PyTorch. Preprint (arXiv: 2305.04684), 2023."}]

Author rebuttal:

Meta-review: The paper generalizes the second-order KFAC optimization method to linear weight-sharing layers. The authors sufficiently addressed the reviewers' concerns regarding the applicability of the method to large LLMs and comparison to other second-order methods. The technique is useful for making second-order methods applicable to large networks and encourages further research in this important area.
Decision: Accept (spotlight)