Title: Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures
Venue: NeurIPS 2023
Accepted: yes
Keywords: deep learning

second-order

optimization

natural gradient

fisher

gauss-newton

k-fac

weight-sharing
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

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.

1 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) or 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 weight-sharing can happen over a sequence of tokens, e.g. in language modelling, over image patches, or 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-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 method and natural gradient descent (Amari, 1998), which uses the Fisher information matrix, or short, the Fisher.

Kronecker-Factored Approximate Curvature (Heskes, 2000; Martens & Grosse, 2015, K-FAC) is an approximation to the Fisher, or for common loss functions equivalently, the generalised GaussNewton 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 as it approximates each layer's Fisher/GGN independently with a Kronecker product. While K-FAC has also been used for transformers (Grosse et al., 2023; Osawa et al., 2022; Pauloski et al., 2021; Zhang et al., 2019a) and a graph neural network (Izadi et al., 2020), there is no theoretical framework for these cases and for applying K-FAC to new architectures.

Contributions. We propose such a framework by leveraging the idea of linear layers with weightsharing. 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-FAC-expand and 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-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 a Wide ResNet on CIFAR-10. Moreover, we show that 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 a superior training algorithm, it indicates the potential of extending K-FAC to modern deep learning architectures. In general, K-FAC can be used as a drop-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.

2 Background

We consider the supervised learning setup with a dataset D of N i.i.d. samples { x n , y n } N n =1 , with x n ∈ R D and y n ∈ R C , a neural net f θ : R D → R C , parameterised with θ ∈ R P , and a loss function glyph[lscript] : R C × R C → R , which is often equivalent to a negative log likelihood, i.e. glyph[lscript] ( y , f θ ( x ))= -log p ( y | f θ ( x )) . The DNN is usually trained with a (stochastic) gradient-based iterative update rule, in its simplest form θ t +1 = θ t -α ∇ θ t glyph[lscript] ( y , f θ t ( x )) , where t indicates the iteration and α is the learning rate. If we assume a local quadratic approximation of the loss around the current parameter iterate θ t , we get a preconditioned gradient update step, i.e.

where C t ∈ R P × P is the symmetric positive definite (p.d.) preconditioner. Setting C t to the Hessian of the loss w.r.t. the model parameters H t := ∇ 2 θ t glyph[lscript] ( y , f θ t ( x )) yields a simplified version of the classic Newton's method; however, since H t 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: it can refer to using second-order 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).

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 architectureagnostic structured version of AdaGrad (Duchi et al., 2011) called Shampoo, Goldfarb et al. (2020) develop K-BFGS, a Kronecker-factored stochastic quasi-Newton method, and Ren & Goldfarb (2021) 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 or 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 an approximation to gradient descent on neurons instead of weights (Benzing, 2022). In any case, generalising K-FAC to modern architectures is a 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.

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.e. they can be written as f θ = f θ L ◦ . . . ◦ f θ glyph[lscript] ◦ . . . ◦ f θ 1 , with θ = concat ( θ 1 , . . . , θ glyph[lscript] , . . . , θ L ) and L the number of layers; concat( · , . . . , · ) concatenates vector inputs to a larger vector. For a linear layer, we have f θ glyph[lscript] ( x ) = φ ( W glyph[lscript] x + b glyph[lscript] ) , where x ∈ R P glyph[lscript], in , W glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], in , b ∈ R P glyph[lscript], out , θ glyph[lscript] = concat ( vec( W glyph[lscript] ) , b glyph[lscript] ) ∈ R P glyph[lscript] , and P glyph[lscript] = P glyph[lscript], out P glyph[lscript], in + P glyph[lscript], out ; vec( · ) vectorises a matrix by concatenating its column vectors and φ is an element-wise nonlinearity. For simplicity, we subsume the bias into the weights. We are interested in linear layers with weights shared across an additional input dimension of size R , i.e. with inputs X ∈ R R × D . 1 The linear layer is applied to this input as XW ᵀ 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.

Example 1: Attention. The transformer architecture is popular in vision (Dosovitskiy et al., 2021), language (Brown et al., 2020), and spatio-temporal modelling (Bi et al., 2022). The inputs and outputs of the attention operation (Bahdanau et al., 2015) in transformers are exactly shaped as described 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

where A ∈ R R × R is the attention matrix , which is normalised such that ∑ R r =1 A i,r = 1 ; alternatively, we can write it using matrix operations as O = AX . Intuitively, the outputs are a weighted sum of 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, defined as A = softmax row ( XW ᵀ Q W K X ᵀ ) , where W Q ∈ R P out × D and W K ∈ R P out × D are weight matrices and softmax row ( · ) applies the softmax function softmax( x ) i := exp( x i ) / ∑ J j =1 exp( x j ) for an input x ∈ R J to each row of the input matrix. This corresponds to a dot-product similarity measure of two independent linear projections of the input, allowing for asymmetric relationships between sequence elements. The two weight matrices W Q , W K , and any linear projection of the attention output are applied to all R sequence elements - they are linear weight-sharing layers.

Example 2: Convolution. A typical input to a 2 d convolution is an image X I ∈ R C in × H × W , where C in is the number of input channels, H is the height, and W the width of the image. We choose a kernel of size K × K , such that we have a rank4 weight tensor of shape C out × C in × K × K , where C out is the number of output channels. For simplicity, following Grosse & Martens (2016), we assume that the convolution is performed with a stride of 1 and that the padding is chosen such that the height and width of the inputs are maintained, i.e. the output image will be O I ∈ R C out × H × W .

The elements of the output tensor of the convolution are then defined as

the input elements X c,i + k -1 ,j + k ′ -1 are defined as 0 outside of the image's boundaries. Now, we can unfold the input tensor to a matrix X ∈ R HW × C in K 2 , where we have extracted and vectorised the K × K sized patches for each of the H · W spatial locations. The weight tensor can be reshaped to the matrix W ∈ R C out × C in K 2 . Using these quantities, we can write the convolution as O = XW ᵀ ∈ R HW × C out . Again, we can recognise the canonical form of a linear weight-sharing layer - here, with R = HW , i.e. 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 = N V nodes or R = N E edges of a 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-Newton matrix (GGN) are usually used in deep learning (Botev et al., 2017; Martens, 2010; Schraudolph, 2002). Defining Λ ( f θ ( x n )):= ∇ 2 f θ glyph[lscript] ( y n , f θ ( x n )) ∈ R C × C , we have

which is equivalent to the Hessian when f θ is a linear model or the residual g -1 ( f θ ( x n )) -y n is zero for all data points, where g -1 is the inverse link function, e.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; Wang, 2010). This is, for example, the case for the common cross-entropy loss corresponding to a categorical likelihood and the mean-square error loss corresponding to a Gaussian likelihood. Setting C t 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.

2.4 Kronecker-Factored Approximate Curvature

The GGN is a semi-p.d. approximation to the Hessian, of size P × P , which is still prohibitively large for DNNs. K-FAC was proposed as an efficient approximation to 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 in a block-diagonal matrix. We define J θ glyph[lscript] ( x n ):= J θ glyph[lscript] f θ ( x n )= J s glyph[lscript],n f θ ( x n ) J θ glyph[lscript] s glyph[lscript],n ∈ R C × P glyph[lscript] as the Jacobian of the model outputs w.r.t. the parameters of the glyph[lscript] -th layer. Now we can write s glyph[lscript],n = W glyph[lscript] a glyph[lscript],n =( a ᵀ glyph[lscript],n ⊗ I P glyph[lscript], out ) vec( W glyph[lscript] ) and with this we have J θ glyph[lscript] s glyph[lscript],n = a ᵀ glyph[lscript],n ⊗ I P glyph[lscript], out . Additionally, by defining b glyph[lscript],n := J s glyph[lscript],n f θ ( x n ) ᵀ ∈ R P glyph[lscript], out × C as the transposed Jacobian of the model outputs w.r.t. the pre-activations of the glyph[lscript] -th layer, we note that J θ glyph[lscript] ( x n ) ᵀ =( a ᵀ glyph[lscript],n ⊗ I P glyph[lscript], out ) ᵀ b glyph[lscript],n = a glyph[lscript],n ⊗ b glyph[lscript],n . K-FAC approximates the GGN of the glyph[lscript] -th layer as

where we have replaced the (transposed) Jacobians in the definition of the GGN in Equation (4) by a glyph[lscript],n ⊗ b glyph[lscript],n and approximated the sum of Kronecker products with a Kronecker product of sums. 3

ᵀ

P

/lscript,

in

N

·

R

a

/lscript

=

in_data

K-FAC-

expand

P

/lscript,

out

N

·

R

b

/lscript

=

out_grads

P

/lscript,

in

P

/lscript,

in

A

=

a

/lscript

a

/lscript

/N

P

/lscript,

out

P

/lscript,

out

B

=

b

/lscript

b

/lscript

P

/lscript

P

/lscript

ˆ

GGN

θ

/lscript

=

B

⊗

A

P

/lscript

P

/lscript

|

GGN

(

θ

/lscript

)

-

ˆ

GGN

θ

/lscript

|

N

K-FAC-

reduce

N

P

/lscript,

in

P

/lscript,

out

P

/lscript

P

/lscript

ᵀ

3 K-FAC for linear weight-sharing layers

In this section, we propose to distinguish between two different linear weight-sharing settings when applying K-FAC. The two settings motivate two corresponding approximations, K-FAC-expand and K-FAC-reduce . However, both approximations can be applied in each setting , see Listing 1 for an 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 a network, linear layers with weight-sharing can be classified based on the point where the weight-sharing dimension is aggregated. Specifically, there are three possible aggregation points, which will define the setting for the glyph[lscript] -th layer:

(i) Before the glyph[lscript] -th layer, i.e. A glyph[lscript],n ∈ R R × P glyph[lscript], in is reduced to ˜ a glyph[lscript],n ∈ R P glyph[lscript], in before being multiplied with the weight matrix of the layer W glyph[lscript] . → We are in the setting of a regular linear layer and no additional considerations are necessary when using K-FAC.

(ii) After the per-example loss , i.e. there will be N · R labels and outputs of 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 is expanded with R terms.

(iii) In between the pre-activation of the glyph[lscript] -th layer S glyph[lscript],n ∈ R R × P glyph[lscript], out and the model output f θ ( x n ) , i.e. before the final aggregation over the per-example losses. → This is the reduce setting, as all weight-sharing dimensions have been reduced during the forward pass.

If we assume a 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, we can identify the setting we are in simply by looking at the form of 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, L expand ( f θ , D ):= -∑ N n =1 ∑ R r =1 log p ( y n,r | f θ ( x n ) r ) , where f θ ( x n ) r is the r -th row of the model output f θ ( x n ) ∈ R R × C and y n,r is the r -th row of the label Y n ∈ R R × C . A typical example of this type of loss function is language translation, where N is the dataset size and R is the sequence length.

We can express the Jacobian of the r -th row of the model output f θ ( x n ) ∈ R R × C w.r.t. θ glyph[lscript] as J θ glyph[lscript] ( x n ) r = ∑ R m =1 J s glyph[lscript],n,m f θ ( x n ) r J θ glyph[lscript] s glyph[lscript],n,m . Since the weights θ glyph[lscript] are shared across the weightsharing dimension of size R , we can write the r -th row of S glyph[lscript],n as s glyph[lscript],n,r = W glyph[lscript] a glyph[lscript],n,r and we have J θ glyph[lscript] s glyph[lscript],n,r = a ᵀ glyph[lscript],n,r ⊗ I P glyph[lscript], out , as for a regular linear layer. We denote b glyph[lscript],n,r,m := J s glyph[lscript],n,m f θ ( x n ) ᵀ r .

Hence, we have J θ glyph[lscript] ( x n ) ᵀ r = ∑ R m =1 a glyph[lscript],n,m ⊗ b glyph[lscript],n,r,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 a Kronecker product without any approximation. The Jacobians J θ glyph[lscript] ( x n ) r could be approximated with a Kronecker product of sums, but this requires access to ∑ R m =1 b glyph[lscript],n,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) is

glyph[negationslash]

where all terms with m = m ′ are ignored. 4 Consequently, we can apply the regular K-FAC approximation over N · R terms instead of the usual N terms, resulting in what we call the K-FAC-expand approximation:

K-FAC-expand

ˆ

GGN

expand

θ

glyph[lscript]

:=

[

1

NR

N

∑

n

=1

R

∑

m

=1

a

glyph[lscript],n,m

a

ᵀ

glyph[lscript],n,m

]

︸

︷︷

︸

=

A

glyph[lscript]

⊗

[

N

∑

n

=1

R

∑

m

=1

R

∑

r

=1

b

glyph[lscript],n,r,m

Λ

(

f

θ

(

x

n

)

r

)

b

ᵀ

glyph[lscript],n,r,m

]

︸

︷︷

︸

=

B

glyph[lscript]

(7)

In the context of convolutions, this approximation has been derived by explicitly stating the assumptions on the activations and pre-activation derivatives (Grosse & Martens, 2016).

For deep linear networks without any weight-sharing layers, K-FAC is known to be 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 even hold for simplistic transformer models since the dot-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 a model of the form

where x ∈ R D and W L ∈ R C × P L, in , W glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], in (with P glyph[lscript], in = P glyph[lscript] -1 , out ) , and W 1 ∈ R P 1 , out × D . Decomposing a single weight matrix W into L 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 ∈ R R × D to this type of network, we can show:

Proposition 1 ( Exactness of K-FAC-expand for deep linear network in the expand setting ) . For layer glyph[lscript] of a deep linear network defined as in Equation (8) and a Gaussian likelihood with p.d. covariance matrix Σ ∈ R C × C , K-FAC-expand is exact in the expand setting.

3.3 K-FAC-reduce

The reduce setting is characterised by a loss with just N loss terms, i.e. L reduce ( f θ , D ) := -∑ N n =1 log p ( y n | f θ ( x n )) , and thus the weight-sharing dimension must have been reduced somewhere in the forward pass of the neural network f θ . A typical instance where this type of loss is used in a model with linear weight-sharing layers is image classification with a convolutional neural network or vision transformer.

ᵀ

P

/lscript,

in

N

a

/lscript

=

in_data

K-FAC-

expand

P

/lscript,

out

N

b

/lscript

=

out_grads

P

/lscript,

in

P

/lscript,

in

A

=

a

/lscript

a

/lscript

/N

P

/lscript,

out

P

/lscript,

out

B

=

b

/lscript

b

/lscript

P

/lscript

P

/lscript

ˆ

GGN

θ

/lscript

=

B

⊗

A

P

/lscript

P

/lscript

|

GGN

(

θ

/lscript

)

-

ˆ

GGN

θ

/lscript

|

N

·

R

K-FAC-

reduce

N

·

R

P

/lscript,

in

P

/lscript,

out

P

/lscript

P

/lscript

ᵀ

Since A glyph[lscript],n ∈ R R × P glyph[lscript], in is now a matrix, we have S glyph[lscript],n = A glyph[lscript],n W ᵀ glyph[lscript] ∈ R R × P glyph[lscript], out . Hence, J θ glyph[lscript] S glyph[lscript],n and J S glyph[lscript],n f θ ( x n ) are now both multi-dimensional arrays. Luckily, we can avoid dealing with this directly by writing J θ glyph[lscript] f θ ( x n ) = ∑ R r =1 J s glyph[lscript],n,r f θ ( x n ) J θ glyph[lscript] s glyph[lscript],n,r , where s glyph[lscript],n,r ∈ R P glyph[lscript], out is the r -th row of S glyph[lscript],n and s glyph[lscript],n,r = W glyph[lscript] a glyph[lscript],n,r . Using this equivalence, we can approximate the (transposed) Jacobians in the GGN for layer glyph[lscript] as

Here, we have approximated the sum of Kronecker products with a Kronecker product of sums over R terms of each of the N per-input Jacobians. This approximation has been proposed in Tang et al. (2021) 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 as usual to the sum over the N data points and call the resulting final approximation K-FAC-reduce and can write it explicitly as

To ensure that K-FAC-reduce is exact in the simple setting of Proposition 1, we have to choose an appropriate aggregation function z : R R × P glyph[lscript], out → R P glyph[lscript], out . A simple and relevant case for which this holds is a scaled sum , i.e. z ( S glyph[lscript],n ) = c ∑ R r =1 s glyph[lscript],n,r with c ∈ R . Both vision transformers and convolutional neural networks with average pooling use scaled sums as the aggregation function (with c = 1 / R ). With this, we can state the corresponding statement to Proposition 1.

Proposition 2 ( Exactness of K-FAC-reduce for deep linear network in the reduce setting ) . For layer glyph[lscript] of a deep linear network (Equation (8) ), a Gaussian likelihood with p.d. covariance matrix Σ ∈ R C × C , and a scaled sum aggregation function, K-FAC-reduce is exact in the reduce setting.

Notably, for a 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 to K-FAC-expand) is not. Moreover, the calculation of A glyph[lscript] costs O ( NRP 2 glyph[lscript], in ) for K-FAC-expand and only O ( NP glyph[lscript], in ( P glyph[lscript], in + R )) for ˆ A glyph[lscript] and K-FACreduce. The same holds for B glyph[lscript] and ˆ B glyph[lscript] , with complexities of O ( NRCP glyph[lscript], out ( C + P glyph[lscript], out )) and O ( NCP glyph[lscript], out ( C + P glyph[lscript], out + R )) . The space complexity of K-FAC's overhead also differs by a factor of R , with O ( NR ( P glyph[lscript], in + P glyph[lscript], out )) for K-FAC-expand and O ( N ( P glyph[lscript], in + P glyph[lscript], out )) for K-FAC-reduce.

4 Experiments

In Section 4.1, we empirically validate the speed up of K-FAC-reduce compared to K-FAC-expand in line with their computational complexities with a Wide ResNet on CIFAR-10 and show how this scales with the batch size. Additionally, in Sections 4.2 and 4.3, we compare the two approximations based on 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, we use a GNN and a vision transformer (ViT), two of the examples in Section 2.2. We leverage the codebase and validation metric targets of the MLCommons algorithmic-efficiency ( AlgoPerf ) repository 5 (Dahl et al., 2023), and use their target-setting runs as reference runs. Both K-FAC variations use the same learning rate schedule as those targetsetting 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 to 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 is a superior training algorithm (see Appendix C.1). 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 a 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-FAC-reduce

To empirically validate the smaller computational complexity of K-FAC-reduce compared to K-FAC-expand, we time a single preconditioned gradient update step for a Wide ResNet on CIFAR-10 with both approximations and five different batch sizes on an NVIDIA V100 GPU (Table 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-FAC-expand runs out-of-memory (OOM) at a batch size of 2048 . 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 to lead to a further speed-up of about 4 . 5 × (Dangel, 2023). To highlight the speed difference of the two K-FAC variations for language models, we compute a 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-molpcba

We use a basic GraphNetwork (Battaglia et al., 2018) with multi-layer perceptrons as the update functions (c.f. Appendix A) and train it on the ogbg-molpcba dataset (Hu et al., 2020) (Figure 3). 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 16 , 106 ± 863 and K-FAC-

0

.

02

0

.

04

0

.

06

0

.

08

0

.

10

Loss

↓

Reference (Nesterov)

K-FAC-expand

K-FAC-reduce

Training

Validation

Validation target

10

20

30

40

Step

10

3

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

mAP

↑

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

Wall-clock time [hours]

×

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 a 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 targetsetting runs as the reference run (Figure 4); each run is repeated with three different random seeds. For K-FAC-expand and K-FAC-reduce, we update the K-FAC statistics and preconditioner every 50 iterations. The reference run reaches the target validation accuracy of 0 . 77309 in 122 , 139 ± 1119 steps or about 26 . 10 ± 0 . 24 hours. K-FAC-expand only takes 87 , 464 ± 245 steps or 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 run for K-FAC-expand and 76 % of the steps and 78 % of the reference run's 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, one update step with K-FAC-expand takes about 1 . 36 × as much as with K-FAC-reduce, which has a significant impact on the runtime when the K-FAC statistics are computed more frequently; see Figure 5 for a 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.

4.4 K-FAC for automatic hyperparameter selection via marginal likelihood optimisation

K-FAC has also been used for marginal likelihood maximisation in Bayesian neural networks using Laplace approximations (Daxberger et al., 2021; Immer et al., 2021). In particular, K-FAC can be used to obtain a scalable Laplace approximation and optimise its marginal likelihood with respect to regularisation parameters, e.g., weight decay, during training.

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

Loss

↓

K-FAC-expand

K-FAC-reduce

Training

Validation

Validation target

20

40

60

80

100

120

Step

10

3

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

↑

5

10

15

20

25

Wall-clock time [hours]

×

Therefore, we compare K-FAC-reduce and K-FAC-expand in this setting following the setup of Daxberger et al. (2021) on CIFAR-10 with a Wide ResNet (Zagoruyko & Komodakis, 2016) and optimise a weight decay parameter per layer every five epochs during 100 epochs of training. Table 2 shows that K-FAC-reduce is significantly faster than K-FAC-expand in this setting, but K-FAC-expand provides better test negative log likelihoods (NLL) in both settings.

5 Discussion and conclusion

We have leveraged the idea of linear weight-sharing layers to generalise K-FAC to a broader range of architectures. This resulted in two distinct flavours of the approximation, which are 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 update step time (c.f. the GNN experiment), 2) when the K-FAC update frequency can be reduced sufficiently (c.f. the ViT experiment), or 3) when underutilised resources can be leveraged to compute K-FAC updates, for example when using pipeline parallelism for language models (Osawa et al., 2022). In these settings, the improvement in wall-clock time of K-FAC-reduce compared to K-FAC-expand is rendered 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 et al., 2023; Zhang et al., 2018). Also, numerical instabilities might arise in low-precision settings due to the need for matrix inversions or decompositions; this could be addressed by inverse-free methods relying on the K-FAC approximation (Lin et al., 2023).

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 ICLR , 2015.

Battaglia, P. W., Hamrick, J. 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. In NeurIPS , 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://github.com/deepmind/kfac-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. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

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, I., and Amodei, D. Language models are few-shot learners. In NeurIPS , 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, G. E., Schneider, F., Nado, Z., Agarwal, N., Sastry, C. 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.07179, 2023.

Dangel, F. Convolutions through the lens of tensor networks. arXiv 2307.02275, 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. In 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 NeurIPS , 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¯ ut˙ e, 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 Kronecker-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 of feature detectors. arXiv 1207.0580, 2012.

Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning on 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. In 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 https://github.com/karpathy/nanoGPT .

Khan, M., Nielsen, D., Tangkaratt, V., Lin, W., Gal, Y., and Srivastava, A. Fast and scalable Bayesian deep learning by weight-perturbation 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, P. Limitations of the empirical Fisher approximation for natural gradient descent. In NeurIPS , 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, Z., Wang, L., and Tang, W. A comprehensive study of weight sharing in graph networks for 3d human pose estimation. In Computer Vision ECCV , 2020.

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, J. 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. In ICLR , 2018.

MLCommons. Algorithms Working Group, 2022. https://mlcommons.org/en/groups/ 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 NAACL-HLT , 2021.

Osawa, K., Li, S., and Hoefler, T. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. arXiv 2211.14133, 2022.

Osawa, K., Ishikawa, S., Yokota, R., Li, S., and Hoefler, T. ASDL: A unified interface for gradient preconditioning in PyTorch. arXiv 2305.04684, 2023.

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. PyTorch: An imperative style, high-performance deep learning library. In NeurIPS , 2019.

Pauloski, J. G., Huang, Q., Huang, L., Venkataraman, S., Chard, K., Foster, I. T., and Zhang, Z. KAISA: an 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, I. Language models are unsupervised multitask learners. 2019.

Ren, Y. and Goldfarb, D. Tensor normal training for deep learning models. In NeurIPS , 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, J. L., and Ganguli, S. Exact solutions to 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, N. N. 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. AI 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 Kronecker-factored approximate curvature. In CVPR , 2021.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In NIPS , 2017.

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, P. 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, S. 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, L., 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.

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 a regular convolutional neural network, it also uses weight-sharing; see Liu et 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 G := ( V , E ) , where V is the set of N nodes and E the set of edges. The edges can be encoded relative to the nodes in an adjacency matrix C ∈ R N × N with C ij = 0 if there is no edge and C ij = 1 if there is one. Typically, they are used for node and graph classification tasks. 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).

The glyph[lscript] -th GCN layer is defined as

which is identical to a regular dense linear layer from Section 2.2, but the input matrix X ∈ R N × P glyph[lscript], in , which has the N node features x n of size P glyph[lscript], in stacked in the rows, is first transformed by the normalised adjacency matrix ˆ C := ( D + I N ) -1 2 ( C + I N )( D + I N ) -1 2 , where D is the diagonal node degree matrix of the graph and I N is the N × N identity matrix.

Defining

we can express the forward pass for a single node and layer as

glyph[negationslash]

where N ( n ) := { j ∈ { 1 , . . . , N }| ˆ C nj = 0 } is the neighbourhood of the node with index n. Notably, the forward pass for a single node x n depends on its neighbourhood, i.e. we cannot express the forward pass for the node without access to the feature information of the nodes in its neighbourhood N ( n ) . Moreover, we can now see that the forward pass through the linear layer, i.e. the matrix multiplication of the weight matrix W glyph[lscript] with the transformed input ˜ x n , does not need the notion of weight-sharing anymore, in the sense, that we do not need a batched matrix-vector product over a 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 GraphNetwork 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 et al. (2018). The GraphNetwork in its general form takes a graph G = ( u , V , E ) where u ∈ R D u are the global features of the graph, and V and E are the sets of nodes and edges, respectively, just as before. We can also write the i -th graph of a dataset of N graphs as a 5-tuple X G n := ( x u n , X V n , X E n , r n , s n ) , with global features x u n ∈ R D u , node features X V n ∈ R N V n × D V , and edge features X E n ∈ R N E n × D E for all n = 1 , . . . , N . The two vectors r n ∈ R N E n and s n ∈ R N E n contain the indices of the receiving and sending nodes of each edge, respectively. Using these indices, we define X V n, r n ∈ R N E n × D V and X V n, s n ∈ R N E n × D V which contain the node features X V n at indices s n and r n , respectively. Note, that these graph inputs unfortunately cannot trivially be batched by stacking them, since the number of nodes N V n or edges N E n are not necessarily the same for all n ∈ { 1 , . . . , N } .

A GraphNetwork block updates the 3-tuple ( x u n , X V n , X E n ) by using three update functions φ ,

and three permutation-invariant aggregation functions ρ

Examples of these aggregation functions include element-wise summation, mean, or maximum.

One forward pass through a GraphNetwork block corresponds to the following steps, where each step is executed for all n ∈ { 1 , . . . , N } :

Update edges X E n with φ E ( X E n , X V n, r n , X V n, s n , x u n ) .

Aggregate updated edges over all nodes in ˜ X E n ∈ R N V n × D E using ρ E → V ( X E n ) .

Update nodes X V n using φ V ( X V n , ˜ X E n , x u n ) .

Aggregate updated edges over all graphs in ¯ X E n ∈ R D E using ρ E → u ( X E n ) .

Aggregate updated nodes over all graphs in ¯ X V n ∈ R D V using ρ V → u ( X V n ) .

Update global features x u n with φ u ( x u n , ¯ X V n , ¯ X E n ) .

In this work, we consider graph classification; for example, molecules can be represented as graphs and we could classify them according to some chemical property (e.g. the ogbg-molpcba 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 ρ . Moreover, multiple of these GraphNetwork blocks can be stacked on top of each other. To classify the input graphs, an MLP is applied to the global features x u n after they are updated by the last GraphNetwork block.

To be more precise, the update functions are in this case specified as

Note, that this is a simplification since in reality, the update functions φ are MLPs with ReLU activations, layer normalisation (Ba et al., 2016), and dropout (Hinton et 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 is 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 as for the attention or 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 φ E and φ V , i.e. for updating the edge and node features. However, since the number of edges N E n and the number of nodes N V n is not necessarily the same for all N graphs, we now have a weight-sharing dimension of size R n , which depends on the n -th input. Also, note that the update function φ u used to update the global features x u n 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 a linear weight-sharing layer, as introduced in Section 2.2. This might not be immediately obvious when only looking at the original notation used in Battaglia et al. (2018); therefore, this can be seen as an instructive example for expressing a 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.

B Extended derivation and discussion of K-FAC-expand and K-FAC-reduce

B.1 Background: K-FAC for regular linear layers

Kronecker-Factored Approximate Curvature (Heskes, 2000; Martens & Grosse, 2015, K-FAC) was proposed as an efficient approximation to a neural network's Fisher information matrix. The Fisher information matrix is defined as 6

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 y n leads to the empirical Fisher (EF), which is simply the uncentered covariance of the empirical gradient. While it is commonly used as a replacement for the Fisher (Chaudhari et al., 2017; Graves, 2011; Kingma & Ba, 2015), it can give rise to very different downstream behaviour when used for optimisation (Kunstner et al., 2019).

First, in all of this work, we focus on a layer-wise K-FAC approximation of the Fisher, i.e. it is approximated by a block-diagonal matrix

where F ( θ glyph[lscript] ) ∈ R P glyph[lscript] × P glyph[lscript] and diag( · , . . . , · ) build a block-diagonal matrix with the input matrices as blocks.

To derive K-FAC, we first note that the pre-activation for layer glyph[lscript] and the n -th data point x n can be expressed as s glyph[lscript],n = W glyph[lscript] a glyph[lscript],n , with W glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], in and a glyph[lscript],n ∈ R P glyph[lscript], in , the input to the glyph[lscript] -th layer (or equivalently, the activation of the glyph[lscript] -1 -th layer). We have omitted an explicit bias parameter b glyph[lscript] , since it can always be subsumed in W glyph[lscript] . Hence, by applying the chain rule, the gradient of the loss w.r.t. the weights of the glyph[lscript] -th layer can be written as ∇ W glyph[lscript] L ( y , f θ ( x n )) = ∇ s glyph[lscript],n L ( y , f θ ( x n )) a ᵀ glyph[lscript],n =: g glyph[lscript],n a ᵀ glyph[lscript],n ∈ R P glyph[lscript], out × P glyph[lscript], in .

Using these insights, K-FAC then replaces the sum of expectations over Kronecker products with a Kronecker product of two sums of expectations, i.e.

where A glyph[lscript] ∈ R P glyph[lscript], in × P glyph[lscript], in and G glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], out . For this derivation, we have used three convenient properties of the Kronecker product (using matrices A , B , C , D with appropriate dimensions): vec( ABC ) = ( C ᵀ ⊗ A )vec( B ) and ( A ⊗ B ) ᵀ = A ᵀ ⊗ B ᵀ for Equation (19c), and ( A ⊗ B )( C ⊗ D ) = AC ⊗ BD for Equation (19d).

We can see that the approximation is exact in the trivial case of a single data point, i.e. N = 1 . Moreover, it is also exact in the case of a single linear layer or a deep linear network and a Gaussian likelihood (Bernacchia et al., 2018).

K-FAC is more efficient than a naive block-wise approximation because we only have to store and invert two Kronecker factors instead of a larger dense matrix for each layer, which reduces the memory complexity from O ( P 2 glyph[lscript], in P 2 glyph[lscript], out ) to O ( P 2 glyph[lscript], in + P 2 glyph[lscript], out ) and the computational complexity of the preconditioning of the gradient with the approximate Fisher from O ( P 3 glyph[lscript], in P 3 glyph[lscript], out ) to O ( P 3 glyph[lscript], in + P 3 glyph[lscript], out ) , since

Alternatively, we can derive K-FAC for the GGN (Botev et al., 2017), which will recover the same result as for the Fisher in Equation (5) for many common loss functions, as we have learned Section 2.3. We define J θ glyph[lscript] ( x n ) := J θ glyph[lscript] f θ ( x n ) = J s glyph[lscript],n f θ ( x n ) J θ glyph[lscript] s glyph[lscript],n ∈ R C × P glyph[lscript] as the Jacobian of the model outputs w.r.t. the parameters of the glyph[lscript] -th layer and Λ ( f θ ( x n )) := H f θ L ( y n , f θ ( x n )) ∈ R C × C as the Hessian of the loss w.r.t. the model outputs. Now we can write s glyph[lscript],n = W glyph[lscript] a glyph[lscript],n = ( a ᵀ glyph[lscript],n ⊗ I P glyph[lscript], out )vec( W glyph[lscript] ) and with this we have J θ glyph[lscript] s glyph[lscript],n = a ᵀ glyph[lscript],n ⊗ I P glyph[lscript], out . Additionally, by defining b glyph[lscript],n := J s glyph[lscript],n f θ ( x n ) ᵀ ∈ R P glyph[lscript], out × C as the transposed Jacobian of the model outputs w.r.t. to the pre-activations of the glyph[lscript] -th layer, we note that J θ glyph[lscript] ( x n ) ᵀ = ( a ᵀ glyph[lscript],n ⊗ I P glyph[lscript], out ) ᵀ b glyph[lscript],n = a glyph[lscript],n ⊗ b glyph[lscript],n .

Replacing the (transposed) Jacobians in the definition of the GGN by this expression, we have

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 to the Fisher for common loss functions (c.f. Section 2.3). Moreover, in our context, it will be useful to have the Jacobians J θ glyph[lscript] ( x n ) 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 be identified by a loss with N · R terms, which corresponds to assuming N · R i.i.d. examples,

The Expand Setting

where f θ ( x n ) r is the r -th row of the model output f θ ( x n ) ∈ R R × C and y n,r is the r -th row of the label Y n ∈ R R × C . A 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.

Note that we are not assuming our inputs x n to have an additional weight-sharing dimension since we only require that the input to the glyph[lscript] -th layer has this additional dimension, i.e. A glyph[lscript],n ∈ R R × D . This obviously does not exclude the case where x n already has this weight-sharing dimension, e.g. a sequence of tokens in translation tasks.

Wecan express the Jacobian of the r -th row of the model output f θ ( x n ) ∈ R R × C w.r.t. the parameters θ glyph[lscript] as

Since the weights θ glyph[lscript] are shared across the weight-sharing dimension of size R , we can write the r -th row of S glyph[lscript],n as s glyph[lscript],n,r = W glyph[lscript] a glyph[lscript],n,r and we have J θ glyph[lscript] s glyph[lscript],n,r = a ᵀ glyph[lscript],n,r ⊗ I P glyph[lscript], out , as for regular K-FAC (c.f. Appendix B.1). We denote b glyph[lscript],n,r,k := J s glyph[lscript],n,k f θ ( x n ) ᵀ r . Hence, we have

On a high level, applying K-FAC to a model trained with this type of loss just requires treating the problem as if we had N · R independent examples and deriving the approximation in the same way as we would with N examples (c.f. Appendix B.1),

However, we cannot directly write each of the N · R loss terms as a Kronecker product without any approximation. One approach could be to approximate each J θ glyph[lscript] ( x n ) r with a Kronecker product of sums, as the K-FAC approximation does for the sum of the N data points, but then we would have to be able to access ∑ R m =1 b glyph[lscript],n,r,m . 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

glyph[negationslash]

where we ignore all the terms with m = m ′ , which allows us to express each of the N · R terms as a Kronecker product. 7 Consequently, we can apply the regular K-FAC approximation over N · R terms instead of just N terms as usual. We call the resulting approximation K-FAC-expand :

or in matrix form

K-FAC-expand

glyph[negationslash]

There is one simple case where the exact expression in Equation (26) is identical to the approximation in Equation (27). When b glyph[lscript],n,r,m = 0 for all r = m , both expression are equivalent to

glyph[negationslash]

With other words, when f θ ( x n ) r is independent of all pre-activations s glyph[lscript],n,m with m = r 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 weightsharing 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 /NR in Equation (28). This might be a better approximation when b glyph[lscript],n,r,m = 0 for r = m and is also what has been used by previous work on transformers (Grosse et al., 2023; Osawa et al., 2022; Pauloski et 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 to show that the approximation in Equation (28) is exact in the same simple cases as regular K-FAC.

glyph[negationslash]

glyph[negationslash]

For a typical neural network with nonlinear activation functions, K-FAC is only an approximation. However, for regular individual linear layers and deep linear networks, K-FAC is known to be 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 similar statements hold for a 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 is indeed exact; this line of reasoning could also be applied to K-FAC for regular linear layers since only the effective number of data points changes from N · R to N . Note, that we could also state more trivial sufficient conditions for the exactness of the approximation, i.e. N = R = 1 and when all inputs to a layer a glyph[lscript],n,r are the same for all n ∈ { 1 , . . . , N } and r ∈ { 1 , . . . , R } . We do not state these types of conditions explicitly from now on.

glyph[negationslash]

Lemma 3 ( Sufficient condition for exactness of K-FAC-expand in the expand setting ) . Let C glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], out be a constant matrix for layer glyph[lscript] . If b glyph[lscript],n,r,m = 0 for all r = m and b glyph[lscript],n,m,m Λ ( f θ ( x n ) m ) b ᵀ glyph[lscript],n,m,m = C glyph[lscript] for all n ∈ { 1 , . . . , N } and m ∈ { 1 , . . . , R } , then the KFAC approximation in Equation (7) is equal to the exact GGN/Fisher of the glyph[lscript] -th layer in the expand setting.

glyph[negationslash]

Proof. As mentioned before, when b glyph[lscript],n,r,m = 0 for all r = m the last line of Equation (26) and the first line of Equation (27) both simplify to Equation (29). Hence, we can directly show that the second and third approximation in Equation (27) equal the exact expression for the GGN of layer glyph[lscript] from there. We have

where we have used the assumption that b glyph[lscript],n,m,m Λ ( f θ ( x n ) m ) b ᵀ glyph[lscript],n,m,m = C glyph[lscript] is the same for all n ∈ { 1 , . . . , N } and m ∈ { 1 , . . . , R } .

Leveraging this simple insight, we can provide an example of a 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 a single linear weight-sharing layer and a Gaussian likelihood with p.d. covariance matrix Σ ∈ R C × C , K-FAC-expand is exact in the expand setting.

glyph[negationslash]

Proof. We can write f θ ( x n ) r = W glyph[lscript] x n,r and hence b glyph[lscript],n,r,m = 0 for r = m . Moreover, we have Λ ( f θ ( x n ) r ) = Σ -1 and b glyph[lscript],n,m,m = I C ( P glyph[lscript], out = C for a single layer). Hence,

for all n ∈ { 1 , . . . , N } and m ∈ { 1 , . . . , R } . 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 a model of the form

f

θ

(

x

) =

W

L

. . .

W

glyph[lscript]

. . .

W

1

x

=

Wx

,

(31)

where x ∈ R D and W L ∈ R C × P L, in , W glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], in (with P glyph[lscript], in = P glyph[lscript] -1 , out ) , and W 1 ∈ R P 1 , out × D . Decomposing a single weight matrix W into L 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). We adopt the notation of Bernacchia et al. (2018) and define

as the product of the weight matrices ahead of W glyph[lscript] and

as the product of the weight matrices behind of W glyph[lscript] . Hence, we can write f θ ( x ) = W a glyph[lscript] W glyph[lscript] W b glyph[lscript] x . Note, that now and

Using these insights, we can now easily state the result for deep linear networks.

Proposition 1 ( Exactness of K-FAC-expand for deep linear network in the expand setting ) . For layer glyph[lscript] of a deep linear network defined as in Equation (8) and a Gaussian likelihood with p.d. covariance matrix Σ ∈ R C × C , K-FAC-expand is exact in the expand setting.

glyph[negationslash]

Proof. Wecan write f θ ( x n ) r = Wx n,r and hence b glyph[lscript],n,r,m = 0 for r = m . Wehave Λ ( f θ ( x n ) r ) = Σ -1 and b glyph[lscript],n,r,r = W a ᵀ glyph[lscript] . Hence, b glyph[lscript],n,m,m Λ ( f θ ( x n ) m ) b ᵀ glyph[lscript],n,m,m = W a ᵀ glyph[lscript] Σ -1 W a glyph[lscript] for all n ∈ { 1 , . . . , N } and m ∈ { 1 , . . . , R } . 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.e.

The Reduce Setting

where the crucial observation is that the weight-sharing dimension must have been reduced somewhere in the forward pass of the neural network f θ . Atypical instance where this type of loss is used together

with a model with linear weight-sharing layers is image classification with a vision transformer or a convolutional neural network. Note, that the inputs x n and labels y n do not have a weight-sharing dimension here; in general, it is also possible for the inputs to have this additional dimension of size R already.

Since A glyph[lscript],n ∈ R R × P glyph[lscript], in is now a matrix, we have S glyph[lscript],n = A glyph[lscript],n W ᵀ glyph[lscript] ∈ R R × P glyph[lscript], out . Hence, J θ glyph[lscript] S glyph[lscript],n and J S glyph[lscript],n f θ ( x n ) are now both multi-dimensional arrays. Luckily, we can simplify this by writing

where s glyph[lscript],n,r ∈ R P glyph[lscript], out is the r -th row of S glyph[lscript],n and s glyph[lscript],n,r = W glyph[lscript] a glyph[lscript],n,r . Using this equivalence we can approximate the GGN for layer glyph[lscript] as

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 as 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-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 a Kronecker-product has also been proposed in the context of invariance learning with deep neural networks via differentiable Laplace approximations in Immer et al. (2022). We call the approximation in Equation (39) K-FAC-reduce and to highlight the difference to K-FAC-expand, we can rewrite it as or in matrix form

K-FAC-reduce

As for K-FAC-expand, we want to show that this approximation can be exact in the case of a single layer or a deep linear network and a Gaussian likelihood. First, we state an analogous condition to Lemma 3.

Lemma 5 ( Sufficient condition for exactness of K-FAC-reduce in the reduce setting ) . Let D glyph[lscript],n ∈ R P glyph[lscript], out × C be a constant matrix for layer glyph[lscript] and data point x n . Further, let C glyph[lscript] ∈ R P glyph[lscript], out × P glyph[lscript], out be a constant matrix for layer glyph[lscript] . If it holds for each n that b glyph[lscript],n,r = D glyph[lscript],n for all r ∈ { 1 , . . . , R } and ˆ b glyph[lscript],n Λ ( f θ ( x n )) ˆ b ᵀ glyph[lscript],n = C glyph[lscript] for all n ∈ { 1 , . . . , N } , then the K-FAC-reduce approximation in Equation (10) is equal to the exact GGN of the glyph[lscript] -th layer in the reduce setting.

Proof. We start with the first approximation and derive the exactness of this step under our assumptions. We have

where we have used the assumption that for each n , we have b glyph[lscript],n,r = D glyph[lscript],n for all r ∈ { 1 , . . . , R } . Now we consider the second approximation in Equation (39). Analogously, we have

where we have used that ˆ b glyph[lscript],n Λ ( f θ ( x n )) ˆ b ᵀ glyph[lscript],n = C glyph[lscript] for all n ∈ { 1 , . . . , N } .

Until now, we did not have to explicitly take the aggregation function z : R R × P glyph[lscript], out → R P glyph[lscript], out into account, since its Jacobian is simply subsumed in b glyph[lscript],n,r . Since we want to verify that the approximation in the reduce case is also exact in the simple scenarios from Proposition 1 and Proposition 2, we now have to also check if the Jacobian J s glyph[lscript],n,r z glyph[lscript],n with z glyph[lscript],n := z ( S glyph[lscript],n ) ∈ R P glyph[lscript], out is the same for all r ∈ { 1 , . . . , R } , to make sure the first condition in Lemma 5 is fulfilled. Maybe

the simplest case where this holds is a scaled sum, i.e.

with c ∈ R and the commutation matrix

where e i,j is the j -th canonical vector of dimension i . This is a linear function in vec( S glyph[lscript],n ) and we have J s glyph[lscript],n,r z glyph[lscript],n = c I P glyph[lscript], out for all r ∈ { 1 , . . . , R } . In particular, when c = 1 the aggregation function is a simple sum and when c = 1 /R it is the mean. Notably, it is not sufficient for z to be linear in vec( S glyph[lscript],n ) , because as soon as we have a weighted sum with weights c r ∈ R and they are not the same for all r ∈ { 1 , . . . , R } , the Jacobians J s glyph[lscript],n,r z glyph[lscript],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 c = 1 / R ).

After clarifying the role of the aggregation function in the exactness of K-FAC-reduce, we can now state a similar statement to Proposition 4.

Proposition 6 ( Exactness of K-FAC-reduce for single linear layer in the reduce setting ) . For a single linear layer, a Gaussian likelihood with p.d. covariance matrix Σ ∈ R C × C , and a scaled sum as defined in Equation (43) as the aggregation function applied to the output of the linear function, K-FAC-reduce is exact in the reduce setting.

Just as for K-FAC-expand, we can extend this result to deep linear networks.

Proposition 2 ( Exactness of K-FAC-reduce for deep linear network in the reduce setting ) . For layer glyph[lscript] of a deep linear network (Equation (8) ), a Gaussian likelihood with p.d. covariance matrix Σ ∈ R C × C , and a scaled sum aggregation function, K-FAC-reduce is exact in the reduce setting.

for all n ∈ { 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 aggregation 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 it the reduce setting. Both settings motivate an approximation each, K-FAC-expand and K-FAC-reduce. Moreover, we presented simple cases where the approximations are exact. In Figure 1 and Figure 2 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 weight-sharing layers are used within the architecture and how this affects the approximation quality of K-FACexpand and K-FAC-reduce. Linear 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 a network that only consists of one simplified variation of the self-attention mechanism used in transformers (we ignore the softmax function, but also consider a linear projection of the output with weight matrix W V ), i.e.

We can observe that it is no longer a linear function in the input X n ∈ R R × D and that we have three linear weight-sharing layers involved in this operation. First, we consider the expand setting, i.e. the output f θ ( X n ) 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 in 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. J s Q,n,m f θ ( X n ) r for the layer with weights W Q , J s K,n,m f θ ( X n ) r for the layer with weights W K , and J s V,n,m f θ ( X n ) r for the layer with weights W V .

We can simply write the r -th row of the output of the layer with the weight matrix W Q as a function of s Q,n,r as

Therefore, we have

glyph[negationslash]

with C = P V, out and b Q,n,r,m = 0 for all m = r , which is the first assumption necessary for Lemma 3 to hold. While b Q,n,r is not the same for all n ∈ { 1 , . . . , N } , it is the same for all r ∈ { 1 , . . . , R } and hence, under the same assumptions as in Proposition 4, K-FAC-expand is exact for the layer with weights W Q in the special case of a single data point, N = 1 .

For the other two involved linear layers, we cannot express the r -th row of f θ ( X n ) as a function of the r -th row of S K/V,n , i.e. elements from all rows of S K/V,n contribute to the r -th row of the output matrix. We can also see this by directly deriving J s K/V,n,m f θ ( X n ) r which will be generally non-zero and dependent on r and m . We omit the explicit derivation by taking the partial derivatives and directly state the results. For the second layer with weight matrix W K , we have

Moreover, for the third layer with weight matrix W V , we have

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 a 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 z with factor c as the aggregation function, we can further rewrite the Jacobians occurring in Equation (39) as

where J s glyph[lscript],n,r f θ ( X n ) m are the same Jacobians we have derived for the expand case. Since according to Lemma 5 we need all b glyph[lscript],n,r to be the same for all r ∈ { 1 , . . . , R } for the first approximation to be exact under the assumptions of Proposition 6, K-FAC-reduce is only exact when N = 1 and only for the layer with weights W Q - just as K-FAC-expand in the expand setting.

When we extend this scenario to a network consisting of L blocks as defined in Equation (45), the above statements regarding the special case where K-FAC-expand and K-FAC-reduce are exact for the layer with weights W Q only hold for the last block. While we omit an 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 output S glyph[lscript],Q,n of all layers with weights W Q glyph[lscript] , besides for the layer in the last block, i.e. the layer in the L -th block with weights W Q L .

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 weightsharing layers. There are many types of GNNs and we will only explicitly cover two of them here.

Related work: node classification with GCN. We first consider a GCN layer, described in Appendix A. This specification of 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 a GNN. The only difference to the normal derivation of K-FAC is that the inputs ˜ a glyph[lscript],n to the glyph[lscript] -th layer for the node with index n now depend on its neighbourhood N ( n ) , since

Using this notation, the definition of the K-FAC GGN for node classification is simply

where X ∈ R N × D , ˜ b glyph[lscript],n := J ˜ s glyph[lscript],n f θ glyph[lscript] ( X ) ᵀ n , and ˜ s glyph[lscript],n := W glyph[lscript] ˜ a glyph[lscript],n . Again, it is important to note that we need to have access to the whole neighbourhood of x n to be able to write the n -th term of the GGN, which is why the input to the model is the matrix X containing all nodes for each loss term. Also, depending on the sparsity of ˆ C , 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 ˆ C 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 a 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

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-sharing 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 it 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 inputs take a 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 R n of graph X G n depends on the input graph itself (indicated by the index n ) and which update function within a GraphNetwork block we want to derive K-FAC-reduce for. For φ E this dimension is going to be R n = N E n , whereas it will be R n = N V n for φ V . In the case of φ u we do not have a weight-sharing dimension, as it has been reduced before this layer is applied, and we can simply apply the regular K-FAC approximation. We can define the inputs to layer glyph[lscript] of type φ E as

and as

for φ V . Correspondingly, we have b glyph[lscript],n = J S glyph[lscript],n f θ ( X G n ) ᵀ ∈ R N E n × D E × C for φ E and b glyph[lscript],n = J S glyph[lscript],n f θ ( X G n ) ᵀ ∈ R N V n × D V × C , with S glyph[lscript],n = A glyph[lscript],n W E ᵀ glyph[lscript] ∈ R N E n × D E and S glyph[lscript],n = A glyph[lscript],n W V ᵀ glyph[lscript] ∈ R N V n × D V , respectively.

Using this notation, we can approximate the GGN for layer glyph[lscript] , assuming its type is either φ E or φ V , as

analogously to Equation (39). However, there is one important difference: since R n now depends on the n -th data point, it makes a difference where we insert the scaling by 1 /R n , since ˆ A glyph[lscript] and/or ˆ B glyph[lscript] are 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 /R n in ˆ a glyph[lscript],n as in Equation (39), or in ˆ b glyph[lscript],n . Instead, to try to keep the overall scale and the weighting of the N terms balanced, we simply include a scaling by 1 / √ R n in both , ˆ a glyph[lscript],n and ˆ b glyph[lscript],n .

B.4 Practical Considerations

While we have discussed theoretically how to apply K-FAC to linear weight-sharing layers, we 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), e.g. KFAC-JAX

(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 on the implementation of KFAC-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 us to get the inputs to a 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.nn.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 et al., 2018) and Tensor Normal Training (Ren & Goldfarb, 2021) are agnostic to the architecture. Assuming the implementation of the model is appropriate, K-FAC-expand and K-FAC-reduce in their simplest form only require a minor adjustment in the code base for regular K-FAC, which is presented in Listing 1.

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 two settings applies to all linear layers with weight-sharing. Hence, using a 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, a single flag set by the user can determine if K-FAC-expand or K-FAC-reduce is 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 GraphNetwork 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 a mini-batch M of size M might have a different weight-sharing dimension R m , i.e. 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 as stated in Equation (53) and Equation (54) are simply concatenated in the first dimension, which results in a dimension of size R M := ∑ M m =1 R m . To apply K-FAC-expand here, we do not have to modify anything, besides scaling the approximation for

each mini-batch by 1 /R M instead of 1 /M . 8 To apply K-FAC-reduce, we can use a 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 is 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 a 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-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. a 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 a moving average and only compute K-FAC for a single mini-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 × C loss Hessian H f θ glyph[lscript] ( y n , f θ ( x n )) for each of the N data points through the computation graph. Instead, we use the fact that we have

and take S Monte Carlo (MC) samples from the model's predictive distribution y s ∼ p ( y | f θ ( x n ) . Taking a single sample results in a rank-1 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 MLCommons' AlgoPerf benchmark for training algorithms

The goal of the AlgoPerf benchmark for training algorithms (Dahl et al., 2023) is to measure 'training 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, a loss function, and a target performance defined by some evaluation metric. For example, training a ViT on ImageNet using the cross-entropy loss until a target validation accuracy of 77 . 309 %has been reached constitutes a 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 benchmark 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 to 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 a budget of 200 trials for every workload. The combination of algorithm and hyperparameters that reached the highest validation performance within a 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

not be confused with the baseline algorithms in the AlgoPerf benchmark, since the target-setting runs are tuned for the best possible validation performance on 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 a superior training algorithm. To provide actual evidence for this claim we would have to run a valid and optimised submission on the entire AlgoPerf benchmark or a benchmark of comparable rigour. Nevertheless, using the well-tuned reference run allows us to put K-FAC's training performance into perspective.

The two AlgoPerf workloads used in this paper are:

Graph neural network on ogbg-molpcba. The workload consists of a (binary) cross-entropy loss, the GraphNetwork instance described in Appendix A, and the ogbg-molpcba 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 to predict whether or not a 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 on ImageNet. This workload also uses the cross-entropy loss, a vision transformer 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 is 0 . 77309 and the best performing reference run algorithm is NAdamW.

C.2 Update step speed with K-FAC-expand and K-FAC-reduce

We provide the full results of the update step timing experiment 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.

C.3 Graph neural network on ogbg-molpcba

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 . 4917728606918423 , β 1 equal to 0 . 9449369031171744 , weight decay set to 1 . 285 964 054 102 592 8e -7 , and linear warmup for 3 , 000 steps and a polynomial schedule with a decay steps factor of 0 . 861509027839639 and an end factor of 1e -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 [0 . 1 , 10] and for the damping in [1e -3 , 1] . We choose the hyperparameters of the run that first reaches the validation target. This setting is a learning rate of 9 . 96871902194967 and damping of 0 . 7881965339190345 for K-FAC-expand and a learning rate of 0 . 5885756514016359 and damping

of 0 . 0579230193904011 for K-FAC-reduce. We then repeated each run five times. The Kronecker factors and the preconditioner are computed every 10 iterations and we use a single sample MC approximation of the Fisher, as explained in Appendix B.4.

C.4 Vision transformer on ImageNet

We use a training batch size of 1 , 024 , glyph[epsilon1] = 1e -8 for NAdamW, and 4 × NVIDIA V100 32 GPUs for all runs on this workload. The reference run algorithm (NAdamW) is using a learning rate of 0 . 0008445074561975979 , β 1 equal to 0 . 8895758153482813 , β 2 to 0 . 9978504782314613 , weight decay set to 0 . 08135402759553023 , 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 [1e -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 a learning rate of 0 . 0012662938340704357 and damping of 0 . 00016524019235426572 . The Kronecker factors and the preconditioner are computed every 50 iterations and we use a single sample MC approximation of the Fisher, as explained in Appendix B.4.

We also conduct a run where we update the Kronecker factors every step, use an exponential moving average with a factor ( ema_decay in ASDL ) equal to β 2 , 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 , β 1 = 0 . 7132 , β 2 = 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 of 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-FAC-expand takes about 50 hours, K-FAC-reduce reaches the target after about 37 hours. Note, that both variations are still significantly slower than the NAdamW 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-FAC makes a 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 a single marginal likelihood update step by about 50 %, as 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., a Wide ResNet 16-4 with Fixup initialisation and parameters instead of batch normalisation (Zagoruyko & Komodakis, 2016; Zhang et al., 2019b). We use a batch size of 128 , an initial learning rate of 0 . 1 cosine decayed to 1e -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 5.

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

Loss

↓

K-FAC-expand

K-FAC-reduce

Training

Validation

Validation target

20

40

60

80

100

120

Step

10

3

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

↑

10

20

30

40

50

Wall-clock time [hours]

×

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nIn 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\nLIMITATIONS:\nThis paper has no negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

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

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

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

STRENGTHS:
- The proposed method has well-grounded theoretical support.
- The proposed method has shown better results than the NAdamW/SGD optimizer.

WEAKNESSES:
My biggest concern comes from the evaluation sections.
- The proposed method is only verified on limited number of tasks and networks (4 configurations). The author should show more results with various configurations.
- 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08

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

STRENGTHS:
Very interesting perspective by looking at modern NN building blocks as weight shared linear layers. 
Solid 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.

QUESTIONS:
1. what does Exactness of K-FAC-reduce and K-FAC-expand? why is this property important?
2. 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-13

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

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.

QUESTIONS:
See above, in the weaknesses section.  Basically, contextualization of the proposed method with other similar approaches. 


-------------------------------


I 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nOne 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.\n\nQUESTIONS:\nTo 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}