SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We can express the Jacobian of the  $r$ -th row of the model output  $f_{\pmb{\theta}}(\pmb{x}_n) \in \mathbb{R}^{R \times C}$  w.r.t.  $\pmb{\theta}_{\ell}$  as  $\pmb{J}_{\pmb{\theta}_{\ell}}(\pmb{x}_n)_r = \sum_{m=1}^{R} \pmb{J}_{\pmb{s}_{\ell,n,m}} f_{\pmb{\theta}}(\pmb{x}_n)_r \pmb{J}_{\pmb{\theta}_{\ell}} \pmb{s}_{\ell,n,m}$ . Since the weights  $\pmb{\theta}_{\ell}$  are shared across the weight-sharing dimension of size  $R$ , we can write the  $r$ -th row of  $\pmb{S}_{\ell,n}$  as  $\pmb{s}_{\ell,n,r} = \pmb{W}_{\ell} \pmb{a}_{\ell,n,r}$  and we have  $\pmb{J}_{\pmb{\theta}_{\ell}} \pmb{s}_{\ell,n,r} = \pmb{a}_{\ell,n,r}^{\top} \otimes \mathbf{I}_{P_{\ell,\mathrm{out}}}$ , as for a regular linear layer. We denote  $\pmb{b}_{\ell,n,r,m} := \pmb{J}_{\pmb{s}_{\ell,n,m}} f_{\pmb{\theta}}(\pmb{x}_n)_r^{\top}$ .
Hence, we have  $J_{\theta_{\ell}}(\pmb{x}_n)_{r}^{\mathsf{T}} = \sum_{m=1}^{R} \pmb{a}_{\ell,n,m} \otimes \pmb{b}_{\ell,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 \cdot R$  loss terms as a Kronecker product without any approximation. The Jacobians  $J_{\theta_{\ell}}(\pmb{x}_n)_{r}$  could be approximated with a Kronecker product of sums, but this requires access to  $\sum_{m=1}^{R} \pmb{b}_{\ell,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
$$
\begin{array}{l} \mathbf {G G N} \left(\boldsymbol {\theta} _ {\ell}\right) = \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \binom {R} {m = 1} \boldsymbol {a} _ {\ell , n, m} \otimes \boldsymbol {b} _ {\ell , n, r, m} \Bigg) \Lambda \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \left(\sum_ {m ^ {\prime} = 1} ^ {R} \boldsymbol {a} _ {\ell , n, m ^ {\prime}} ^ {\intercal} \otimes \boldsymbol {b} _ {\ell , n, r, m ^ {\prime}} ^ {\intercal}\right) \tag {6} \\ \approx \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \sum_ {m = 1} ^ {R} \left(\boldsymbol {a} _ {\ell , n, m} \otimes \boldsymbol {b} _ {\ell , n, r, m}\right) \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \left(\boldsymbol {a} _ {\ell , n, m} ^ {\intercal} \otimes \boldsymbol {b} _ {\ell , n, r, m} ^ {\intercal}\right), \\ \end{array}
$$
where all terms with  $m \neq m'$  are ignored. Consequently, we can apply the regular K-FAC approximation over  $N \cdot R$  terms instead of the usual  $N$  terms, resulting in what we call the  $K$ -FAC-expand approximation:
K-FAC-expand
$$
\mathbf {G} \hat {\mathbf {G}} \mathbf {N} _ {\boldsymbol {\theta} _ {\ell}} ^ {\text {e x p a n d}} := \underbrace {\left[ \frac {1}{N R} \sum_ {n = 1} ^ {N} \sum_ {m = 1} ^ {R} \boldsymbol {a} _ {\ell , n , m} \boldsymbol {a} _ {\ell , n , m} ^ {\intercal} \right]} _ {= A _ {\ell}} \otimes \underbrace {\left[ \sum_ {n = 1} ^ {N} \sum_ {m = 1} ^ {R} \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n , r , m} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \boldsymbol {b} _ {\ell , n , r , m} ^ {\intercal} \right]} _ {= B _ {\ell}} \tag {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
$$
f _ {\boldsymbol {\theta}} (\boldsymbol {x}) = \boldsymbol {W} _ {L} \dots \boldsymbol {W} _ {\ell} \dots \boldsymbol {W} _ {1} \boldsymbol {x} = \boldsymbol {W} \boldsymbol {x}, \tag {8}
$$
where  $\pmb{x} \in \mathbb{R}^{D}$  and  $\pmb{W}_{L} \in \mathbb{R}^{C \times P_{L,\mathrm{in}}}$ ,  $\pmb{W}_{\ell} \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times P_{\ell,\mathrm{in}}}$  (with  $P_{\ell,\mathrm{in}} = P_{\ell-1,\mathrm{out}}$ ), and  $\pmb{W}_{1} \in \mathbb{R}^{P_{1,\mathrm{out}} \times D}$ . Decomposing a single weight matrix  $\pmb{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  $\pmb{X} \in \mathbb{R}^{R \times 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  $\ell$  of a deep linear network defined as in Equation (8) and a Gaussian likelihood with p.d. covariance matrix  $\pmb{\Sigma} \in \mathbb{R}^{C \times 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.  $\mathcal{L}_{\mathrm{reduce}}(f_{\theta},\mathcal{D})\coloneqq -\sum_{n = 1}^{N}\log p(\pmb {y}_n|f_\theta (\pmb {x}_n))$ , and thus the weight-sharing dimension must have been reduced somewhere in the forward pass of the neural network  $f_{\theta}$ . 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.
Although the work on transformers does not scale by  $1 / R$  in Equation (7), see Appendix B.2.1 for details.
Figure 2: Visualisation of K-FAC-expand and K-FAC-reduce in the reduce setting. This is similar to Figure 1, but for the reduce setting, where K-FAC-reduce is exact and K-FAC-expand is not (Section 3.3).
Since  $\mathbf{A}_{\ell,n} \in \mathbb{R}^{R \times P_{\ell,\mathrm{in}}}$  is now a matrix, we have  $\mathbf{S}_{\ell,n} = \mathbf{A}_{\ell,n} \mathbf{W}_{\ell}^{\mathsf{T}} \in \mathbb{R}^{R \times P_{\ell,\mathrm{out}}}$ . Hence,  $\mathbf{J}_{\theta_{\ell}} \mathbf{S}_{\ell,n}$  and  $\mathbf{J}_{\mathbf{S}_{\ell,n}} f_{\theta}(\mathbf{x}_n)$  are now both multi-dimensional arrays. Luckily, we can avoid dealing with this directly by writing  $\mathbf{J}_{\theta_{\ell}} f_{\theta}(\mathbf{x}_n) = \sum_{r=1}^{R} \mathbf{J}_{\mathbf{s}_{\ell,n,r}} f_{\theta}(\mathbf{x}_n) \mathbf{J}_{\theta_{\ell}} \mathbf{s}_{\ell,n,r}$ , where  $\mathbf{s}_{\ell,n,r} \in \mathbb{R}^{P_{\ell,\mathrm{out}}}$  is the  $r$ -th row of  $\mathbf{S}_{\ell,n}$  and  $\mathbf{s}_{\ell,n,r} = \mathbf{W}_{\ell} \mathbf{a}_{\ell,n,r}$ . Using this equivalence, we can approximate the (transposed) Jacobians in the GGN for layer  $\ell$  as
$$
\boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) ^ {\top} = \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \boldsymbol {b} _ {\ell , n, r} \approx \frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n, r}. \tag {9}
$$
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
K-FAC-reduce
$$
\mathbf {G} \hat {\mathbf {G}} \mathbf {N} _ {\theta_ {\ell}} ^ {\text {r e d u c e}} :=
$$
$$
\underbrace {\left[ \frac {1}{N R ^ {2}} \sum_ {n = 1} ^ {N} \left(\sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n , r}\right) \left(\sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n , r} ^ {\mathsf {T}}\right) \right]} _ {= \hat {\boldsymbol {A}} _ {\ell}} \otimes \underbrace {\left[ \sum_ {n = 1} ^ {N} \left(\sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n , r}\right) \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {X} _ {n}\right)\right) \left(\sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n , r} ^ {\mathsf {T}}\right) \right]} _ {= \hat {\boldsymbol {B}} _ {\ell}}. \tag {10}
$$
To ensure that K-FAC-reduce is exact in the simple setting of Proposition 1, we have to choose an appropriate aggregation function  $z: \mathbb{R}^{R \times P_{\ell, \mathrm{out}}} \to \mathbb{R}^{P_{\ell, \mathrm{out}}}$ . A simple and relevant case for which this holds is a scaled sum, i.e.  $z(\pmb{S}_{\ell,n}) = c \sum_{r=1}^{R} \pmb{s}_{\ell,n,r}$  with  $c \in \mathbb{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  $\ell$  of a deep linear network (Equation (8)), a Gaussian likelihood with p.d. covariance matrix  $\Sigma \in \mathbb{R}^{C\times 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_{\ell}$  costs  $\mathcal{O}(NRP_{\ell,\mathrm{in}}^2)$  for K-FAC-expand and only  $\mathcal{O}(NP_{\ell,\mathrm{in}}(P_{\ell,\mathrm{in}} + R))$  for  $\hat{A}_{\ell}$  and K-FAC-reduce. The same holds for  $B_{\ell}$  and  $\hat{B}_{\ell}$ , with complexities of  $\mathcal{O}(NRCP_{\ell,\mathrm{out}}(C + P_{\ell,\mathrm{out}}))$  and  $\mathcal{O}(NCP_{\ell,\mathrm{out}}(C + P_{\ell,\mathrm{out}} + R))$ . The space complexity of K-FAC's overhead also differs by a factor of  $R$ , with  $\mathcal{O}(NR(P_{\ell,\mathrm{in}} + P_{\ell,\mathrm{out}}))$  for K-FAC-expand and  $\mathcal{O}(N(P_{\ell,\mathrm{in}} + P_{\ell,\mathrm{out}}))$  for K-FAC-reduce.
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 MLCCommons 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 target-setting runs but the warmup and expected number of steps are multiplied by 0.75 to account for the fact that second-order methods tend to require fewer steps. Additionally, we tune the learning rate and damping for K-FAC via random search and set all other hyperparameters 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
Table 1: Timing [s] of an update step with the two K-FAC variants for a Wide ResNet on CIFAR-10.
K-FAC	Batch size
128	256	512	1024	2048
expand	0.24	0.38	0.75	1.36	OOM
reduce	0.17	0.24	0.43	0.63	1.17
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 \times$  (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 \pm 2,070$  steps, whereas K-FAC-expand takes  $16,106 \pm 863$  and K-FAC-
5https://bit.ly/algoperf
Figure 3: Training results for a graph neural network on ogbg-molpcba. Both K-FAC variations require  $\approx 50\%$  of the steps which almost directly translates into reduced wall-clock time. The K-FAC statistics are updated every 10 steps. We show runs with five different random seeds for each method.
reduce  $15,566 \pm 429$  steps, about  $51\%$  and  $49\%$  of the reference run, respectively. In wall-clock time, the reference run takes about  $2.33 \pm 0.22$  hours, whereas K-FAC-expand takes  $1.32 \pm 0.11$  and K-FAC-reduce  $1.23 \pm 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 target-setting 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 \pm 1119$  steps or about  $26.10 \pm 0.24$  hours. K-FAC-expand only takes  $87,464 \pm 245$  steps or  $19.57 \pm 0.13$  hours, and K-FAC-reduce  $92,550 \pm 731$  steps or  $20.28 \pm 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 \times$  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.
Figure 4: Training results for a vision transformer on ImageNet. The K-FAC statistics are updated every 50 steps. Due to amortising the costs, the reduced number of steps to the target translates into reduced wall-clock time. We show runs with three different random seeds for each method.
Therefore, we compare K-FAC-reduce and K-FAC-expand in this setting following the setup of Daxberger 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.
Table 2: Performance of K-FAC variants for marginal likelihood maximisation on CIFAR-10 with a Wide ResNet with and without data augmentation (DA).
K-FAC	DA	NLL ↓	Acc [%] ↑	Time [%] ↓
expand	X	0.42	88.9	100
✓	0.24	92.5
reduce	X	0.70	86.7	50.5
✓	0.35	93.5
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).
The authors thank Kazuki Osawa for fruitful discussions and support with using and extending the ASDL library. The authors also thank Wu Lin for pointing out a mistake in the original equation of the K-FAC-expand approximation. Finally, the authors thank Felix Dangel for helpful comments on this manuscript. Runa Eschenhagen is supported by ARM and the Cambridge Trust. Alexander Immer is funded by the Max Planck ETH Center for Learning Systems (CLS). Richard E. Turner is supported by Google, Amazon, ARM, Improbable and EPSRC grant EP/T005386/1. Frank Schneider and Philipp Hennig are funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) in the frame of the priority programme SPP 2298 "Theoretical Foundations of Deep Learning" - HE 7114/5-1, as well as the DFG Cluster of Excellence "Machine Learning - New Perspectives for Science", EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg.