SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\hat{A}_{\ell}$  and/or  $\hat{B}_{\ell}$  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  $\hat{a}_{\ell,n}$  as in Equation (39), or in  $\hat{b}_{\ell,n}$ . Instead, to try to keep the overall scale and the weighting of the  $N$  terms balanced, we simply include a scaling by  $1 / \sqrt{R_{n}}$  in both,  $\hat{a}_{\ell,n}$  and  $\hat{b}_{\ell,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
1 # Check if there even is a weight-sharing dimension; if not, the Kronecker  
2 # factors can directly be calculated.  
3 if in_data.ndim == 3:  
4 # Mini-batch size M, weight-sharing dimension R, feature dimension  $P_{\ell, \text{in/out}}$ .  
5 M, R, P_in = in_data.shape  
6 P_out = out_grads.shape[2]  
7 if approximation == 'expand':  
8 # Flatten the weight-sharing dimension into the mini-batch dimension.  
9 in_data = in_data.view(M*R, P_in) / math.sqrt(R)  
10 out_grads = out_grads.view(M*R, P_out)  
11 elif approximation == 'reduce':  
12 # Reduce the weight-sharing dimension with mean and sum.  
13 in_data = in_data.mean(dim=1)  
14 out_grads = out_grads.sum(dim=1)  
15 # Calculate Kronecker factors  $A_{\ell} / \hat{A}_{\ell}$  and  $B_{\ell} / \hat{B}_{\ell}$ .  
16 A = torch/matmul(in_data.T, in_data) / M  
17 B = torch/matmul(out_grads.T, out_grads)
Listing 1: Illustration of K-FAC-expand and K-FAC-reduce with code. This piece of code calculates the approximations on one mini-batch and for one layer. We receive the inputs to the layer, in_data, from a forward hook and the gradients of the loss w.r.t. the outputs of the layer, out_grads, from a backward hook. Here we assume that only one additional weight-sharing dimension exists and that the first dimension is always the mini-batch dimension. Note that the dimensions of in_data and out_grads are the same, independently of the setting we are in (c.f. Section 3.1). This illustrates why we can choose to use either approximation, K-FAC-expand or K-FAC-reduce, in each setting. The actual implementation in ASDL is very similar for torch nn. Linear modules, the logic is just separated into multiple functions and allows for multiple weight-sharing dimensions. However, the adjustments for the GraphNetwork are more involved.
(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 K-FAC-expand and K-FAC-reduce within ASDL, which we also use for the experiments in Section 4. K-FAC is implemented using forward and backward hooks, which allow 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  $\mathcal{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_{\mathcal{M}} := \sum_{m=1}^{M} 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_{\mathcal{M}}$  instead of  $1 / M$ .<sup>8</sup> 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  $\mathbf{A}$  and  $\mathbf{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 \times C$  loss Hessian  $H_{f_{\theta}}\ell(\mathbf{y}_n,f_{\theta}(\mathbf{x}_n))$  for each of the  $N$  data points through the computation graph. Instead, we use the fact that we have
$$
\mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} \mid f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} [ \nabla_ {f _ {\boldsymbol {\theta}}} \log p (\boldsymbol {y} \mid f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \nabla_ {f _ {\boldsymbol {\theta}}} \log p (\boldsymbol {y} \mid f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) ^ {\top} ] = \boldsymbol {H} _ {f _ {\boldsymbol {\theta}}} \ell (\boldsymbol {y} _ {n}, f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \tag {56}
$$
and take  $S$  Monte Carlo (MC) samples from the model's predictive distribution  $\mathbf{y}_s \sim p(\mathbf{y} | f_\theta(\mathbf{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
As explained below Equation (55), we could also choose a different way of scaling here. We choose this one as it enables a simple implementation in contrast to K-FAC-reduce.
Table 3: Results from Table 1 with standard errors over 10 epochs.
K-FAC	Batch size
128	256	512	1024	2048
expand	0.24 ± 0.01	0.38 ± 0.02	0.75 ± 0.03	1.36 ± 0.05	OOM
reduce	0.17 ± 0.01	0.24 ± 0.02	0.43 ± 0.03	0.63 ± 0.04	1.17 ± 0.08
not be confused with the baseline algorithms in the AlgoPerf benchmark, since the target-setting runs are tuned for the best possible validation performance 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  $350\mathrm{k}$  examples and almost  $44\mathrm{k}$  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.
Table 4: Timing of K-FAC GGN approximation for GPT-2 on full DART dataset.
K-FAC	Absolute time [min] ↓	Relative time [%] ↓
expand	9.55 ± 0.13	100
reduce	6.68 ± 0.04	70
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,  $\beta_{1}$  equal to 0.9449369031171744, weight decay set to 1.2859640541025928e-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  $1\mathrm{e} - 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,  $\epsilon = 1\mathrm{e} - 8$  for NAdamW, and  $4\times$  NVIDIA V100 32 GPUs for all runs on this workload. The reference run algorithm (NAdamW) is using a learning rate of 0.0008445074561975979,  $\beta_{1}$  equal to 0.8895758153482813,  $\beta_{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  $\beta_{2}$ , update the preconditioner every 10 steps, set the damping to  $1\mathrm{e} - 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  $2\mathrm{e} - 3$ ,  $\beta_{1} = 0.7132$ ,  $\beta_{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  $1\mathrm{e} - 6$  and weight decay of  $5\mathrm{e} - 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.
Figure 5: Training results for a vision transformer on ImageNet. This is similar to Figure 4, but the K-FAC statistics are updated every step and different hyperparameters are used. Due to K-FAC's overhead, the wall-clock time is not reduced in this setting. Moreover, the discrepancy in speed between K-FAC-expand and K-FAC-reduce becomes apparent.
Table 5: Results from Table 2 with standard errors over three random seeds. The update time is the average time per full K-FAC Laplace approximation to the marginal likelihood, which requires a full data set pass with computation and eigendecomposition of the Kronecker factors.
K-FAC	Data Augmentation	NLL ↓	Accuracy [%] ↑	Update Time [s] ↓
expand	X	0.422 ± 0.013	88.89 ± 0.24	196 ± 22
✓	0.244 ± 0.004	92.52 ± 0.12
reduce	X	0.703 ± 0.012	86.71 ± 0.13	99 ± 23
✓	0.352 ± 0.008	93.50 ± 0.06