SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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, C., 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., Vinyls, 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.
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  $\mathcal{G} \coloneqq (\mathcal{V}, \mathcal{E})$ , where  $\mathcal{V}$  is the set of  $N$  nodes and  $\mathcal{E}$  the set of edges. The edges can be encoded relative to the nodes in an adjacency matrix  $\mathbf{C} \in \mathbb{R}^{N \times 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  $\ell$ -th GCN layer is defined as
$$
f _ {\boldsymbol {\theta} _ {\ell}} (\boldsymbol {X}) = \phi \left(\hat {\boldsymbol {C}} \boldsymbol {X} \boldsymbol {W} _ {\ell} ^ {\intercal}\right) \tag {11}
$$
which is identical to a regular dense linear layer from Section 2.2, but the input matrix  $\mathbf{X} \in \mathbb{R}^{N \times P_{\ell, \mathrm{in}}}$ , which has the  $N$  node features  $\mathbf{x}_n$  of size  $P_{\ell, \mathrm{in}}$  stacked in the rows, is first transformed by the normalised adjacency matrix  $\hat{\mathbf{C}} := (\mathbf{D} + \mathbf{I}_N)^{-\frac{1}{2}}(\mathbf{C} + \mathbf{I}_N)(\mathbf{D} + \mathbf{I}_N)^{-\frac{1}{2}}$ , where  $\mathbf{D}$  is the diagonal node degree matrix of the graph and  $\mathbf{I}_N$  is the  $N \times N$  identity matrix.
$$
\tilde {\boldsymbol {x}} _ {n} := \sum_ {j = 1} ^ {N} \hat {\boldsymbol {C}} _ {n j} \boldsymbol {x} _ {j} = \sum_ {j \in \mathcal {N} (n)} \hat {\boldsymbol {C}} _ {n j} \boldsymbol {x} _ {j}, \tag {12}
$$
we can express the forward pass for a single node and layer as
$$
f _ {\boldsymbol {\theta} _ {\ell}} \left(\tilde {\boldsymbol {x}} _ {n}\right) = \phi \left(\boldsymbol {W} _ {\ell} \tilde {\boldsymbol {x}} _ {n}\right), \tag {13}
$$
where  $\mathcal{N}(n) \coloneqq \{j \in \{1, \dots, N\} | \hat{C}_{nj} \neq 0\}$  is the neighbourhood of the node with index  $n$ . Notably, the forward pass for a single node  $\boldsymbol{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  $\mathcal{N}(n)$ . Moreover, we can now see that the forward pass through the linear layer, i.e. the matrix multiplication of the weight matrix  $\boldsymbol{W}_{\ell}$  with the transformed input  $\tilde{\boldsymbol{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  $\mathcal{G} = (\pmb{u},\mathcal{V},\mathcal{E})$  where  $\pmb{u}\in \mathbb{R}^{D_u}$  are the global features of the graph, and  $\mathcal{V}$  and  $\mathcal{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  $\mathbb{X}_n^G\coloneqq (\pmb{x}_n^u,\pmb {X}_n^V,\pmb {X}_n^E,\pmb {r}_n,\pmb {s}_n)$ , with global features  $\pmb{x}_n^u\in \mathbb{R}^{D_u}$ , node features  $\pmb {X}_n^V\in \mathbb{R}^{N_n^V\times D_V}$ , and edge features  $\pmb {X}_n^E\in \mathbb{R}^{N_n^E\times D_E}$  for all  $n = 1,\dots ,N$ . The two vectors  $\pmb {r}_n\in \mathbb{R}^{N_n^E}$  and  $\pmb{s}_n\in \mathbb{R}^{N_n^E}$  contain the indices of the receiving and sending nodes of each edge, respectively. Using these indices, we define  $\pmb {X}_{n,r_n}^V\in \mathbb{R}^{N_n^E\times D_V}$  and  $\pmb {X}_{n,\pmb {s}_n}^V\in \mathbb{R}^{N_n^E\times D_V}$  which contain the node features  $\pmb {X}_n^V$  at indices  $\pmb{s}_n$  and  $\pmb{r}_n$ , respectively. Note, that these graph inputs unfortunately cannot trivially be batched by stacking them, since the number of nodes  $N_{n}^{V}$  or edges  $N_{n}^{E}$  are not necessarily the same for all  $n\in \{1,\ldots ,N\}$ .
A GraphNetwork block updates the 3-tuple  $(\pmb{x}_n^u, \pmb{X}_n^V, \pmb{X}_n^E)$  by using three update functions  $\phi$ ,
$$
\boldsymbol {X} _ {n} ^ {E} \gets \phi^ {E} (\boldsymbol {X} _ {n} ^ {E}, \boldsymbol {X} _ {n, \boldsymbol {r} _ {n}} ^ {V}, \boldsymbol {X} _ {n, \boldsymbol {s} _ {n}} ^ {V}, \boldsymbol {x} _ {n} ^ {u})
$$
$$
\boldsymbol {X} _ {n} ^ {V} \leftarrow \phi^ {V} \left(\boldsymbol {X} _ {n} ^ {V}, \tilde {\boldsymbol {X}} _ {n} ^ {E}, \boldsymbol {x} _ {n} ^ {u}\right) \tag {14}
$$
$$
\boldsymbol {x} _ {n} ^ {u} \leftarrow \phi^ {u} \left(\boldsymbol {x} _ {n} ^ {u}, \bar {\boldsymbol {X}} _ {n} ^ {V}, \bar {\boldsymbol {X}} _ {n} ^ {E}\right),
$$
and three permutation-invariant aggregation functions  $\rho$
$$
\tilde {\boldsymbol {X}} _ {n} ^ {E} \leftarrow \rho^ {E \rightarrow V} (\boldsymbol {X} _ {n} ^ {E})
$$
$$
\bar {\boldsymbol {X}} _ {n} ^ {E} \leftarrow \rho^ {E \rightarrow u} \left(\boldsymbol {X} _ {n} ^ {E}\right) \tag {15}
$$
$$
\bar {\boldsymbol {X}} _ {n} ^ {V} \gets \rho^ {V \rightarrow u} (\boldsymbol {X} _ {n} ^ {V}).
$$
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 \in \{1, \dots, N\}$ :
1. Update edges  $\mathbf{X}_n^E$  with  $\phi^E(\mathbf{X}_n^E, \mathbf{X}_{n,\mathbf{r}_n}^V, \mathbf{X}_{n,\mathbf{s}_n}^V, \mathbf{x}_n^u)$ .
2. Aggregate updated edges over all nodes in  $\tilde{\pmb{X}}_n^E\in \mathbb{R}^{N_n^V\times D_E}$  using  $\rho^{E\to V}(X_n^E)$ .
3. Update nodes  $\mathbf{X}_n^V$  using  $\phi^V(\mathbf{X}_n^V, \tilde{\mathbf{X}}_n^E, \mathbf{x}_n^u)$ .
4. Aggregate updated edges over all graphs in  $\bar{\pmb{X}}_n^E\in \mathbb{R}^{D_E}$  using  $\rho^{E\to u}(X_n^E)$ .
5. Aggregate updated nodes over all graphs in  $\bar{\pmb{X}}_n^V\in \mathbb{R}^{D_V}$  using  $\rho^{V\to u}(\pmb {X}_n^V)$
6. Update global features  $\pmb{x}_n^u$  with  $\phi^u(\pmb{x}_n^u, \bar{\pmb{X}}_n^V, \bar{\pmb{X}}_n^E)$ .
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  $\phi$  and an element-wise sum for the aggregation functions  $\rho$ . 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_{n}^{u}$  after they are updated by the last GraphNetwork block.
To be more precise, the update functions are in this case specified as
$$
\begin{array}{l} \phi^ {E} \left(\boldsymbol {X} _ {n} ^ {E}, \boldsymbol {X} _ {n, \boldsymbol {r} _ {n}} ^ {V}, \boldsymbol {X} _ {n, \boldsymbol {s} _ {n}} ^ {V}, \boldsymbol {x} _ {n} ^ {u}\right) := \operatorname {c o n c a t} \left(\boldsymbol {X} _ {n} ^ {E}, \boldsymbol {X} _ {n, \boldsymbol {r} _ {n}} ^ {V}, \boldsymbol {X} _ {n, \boldsymbol {s} _ {n}} ^ {V}, \operatorname {r e p e a t} _ {N _ {n} ^ {E}} \left(\boldsymbol {x} _ {n} ^ {u}\right)\right) \boldsymbol {W} ^ {E ^ {\intercal}} \\ \phi^ {V} \left(\boldsymbol {X} _ {n} ^ {V}, \tilde {\boldsymbol {X}} _ {n} ^ {E}, \boldsymbol {x} _ {n} ^ {u}\right) := \operatorname {c o n c a t} \left(\boldsymbol {X} _ {n} ^ {V}, \tilde {\boldsymbol {X}} _ {n} ^ {E}, \operatorname {r e p e a t} _ {N _ {n} ^ {V}} \left(\boldsymbol {x} _ {n} ^ {u}\right)\right) \boldsymbol {W} ^ {V ^ {\intercal}} \tag {16} \\ \phi^ {u} (\pmb {x} _ {n} ^ {u}, \bar {\pmb {X}} _ {n} ^ {V}, \bar {\pmb {X}} _ {n} ^ {E}) := \pmb {W} ^ {u} \mathrm {c o n c a t} (\pmb {x} _ {n} ^ {u}, \bar {\pmb {X}} _ {n} ^ {V}, \bar {\pmb {X}} _ {n} ^ {E}) \\ \end{array}
$$
with  $\pmb{W}^{E}\in \mathbb{R}^{D_{E}\times (D_{E} + 2D_{V} + D_{u})},\pmb{W}^{V}\in \mathbb{R}^{D_{V}\times (D_{V} + D_{E} + D_{u})}$ , and  $\pmb{W}^u\in \mathbb{R}^{D_u\times 3D_u}$ .
Note, that this is a simplification since in reality, the update functions  $\phi$  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  $\phi^E$  and  $\phi^V$ , i.e. for updating the edge and node features. However, since the number of edges  $N_{n}^{E}$  and the number of nodes  $N_{n}^{V}$  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  $\phi^u$  used to update the global features  $\boldsymbol{x}_n^u$  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<sup>6</sup>
$$
\begin{array}{l} \boldsymbol {F} (\boldsymbol {\theta}) = - \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} \mid f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right))} [ \nabla_ {\boldsymbol {\theta}} ^ {2} \log p (\boldsymbol {y} \mid f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)) ] \tag {17} \\ = \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} [ \nabla_ {\boldsymbol {\theta}} \log p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) (\nabla_ {\boldsymbol {\theta}} \log p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))) ^ {\intercal} ]. \\ \end{array}
$$
Notably, the labels  $\mathbf{y}$  are samples from the predictive distribution of the model and are not the empirical labels from the data. Replacing  $\mathbf{y}$  with  $\mathbf{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
$$
\boldsymbol {F} (\boldsymbol {\theta}) \approx \operatorname {d i a g} \left(\boldsymbol {F} \left(\boldsymbol {\theta} _ {1}\right), \dots , \boldsymbol {F} \left(\boldsymbol {\theta} _ {\ell}\right), \dots , \boldsymbol {F} \left(\boldsymbol {\theta} _ {L}\right)\right) \in \mathbb {R} ^ {P \times P}, \tag {18}
$$
where  $\pmb{F}(\pmb{\theta}_{\ell}) \in \mathbb{R}^{P_{\ell} \times P_{\ell}}$  and  $\mathrm{diag}(\cdot, \dots, \cdot)$  build a block-diagonal matrix with the input matrices as blocks.
To derive K-FAC, we first note that the pre-activation for layer  $\ell$  and the  $n$ -th data point  $\pmb{x}_n$  can be expressed as  $\pmb{s}_{\ell,n} = \pmb{W}_{\ell}\pmb{a}_{\ell,n}$ , with  $\pmb{W}_{\ell} \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times P_{\ell,\mathrm{in}}}$  and  $\pmb{a}_{\ell,n} \in \mathbb{R}^{P_{\ell,\mathrm{in}}}$ , the input to the  $\ell$ -th layer (or equivalently, the activation of the  $\ell-1$ -th layer). We have omitted an explicit bias parameter  $\pmb{b}_{\ell}$  since it can always be subsumed in  $\pmb{W}_{\ell}$ . Hence, by applying the chain rule, the gradient of the loss w.r.t. the weights of the  $\ell$ -th layer can be written as  $\nabla_{\pmb{W}_{\ell}}\mathcal{L}(\pmb{y}, f_{\pmb{\theta}}(\pmb{x}_n)) = \nabla_{\pmb{s}_{\ell,n}}\mathcal{L}(\pmb{y}, f_{\pmb{\theta}}(\pmb{x}_n))\pmb{a}_{\ell,n}^{\top} =: \pmb{g}_{\ell,n}\pmb{a}_{\ell,n}^{\top} \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times P_{\ell,\mathrm{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.
$$
\begin{array}{l} \boldsymbol {F} (\boldsymbol {\theta} _ {\ell}) = \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} [ \operatorname {v e c} \left(\nabla_ {\boldsymbol {W} _ {\ell}} \mathcal {L} (\boldsymbol {y}, f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))\right) \operatorname {v e c} \left(\nabla_ {\boldsymbol {W} _ {\ell}} \mathcal {L} (\boldsymbol {y}, f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))\right) ^ {\top} ] (19a) \\ = \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} \left[ \operatorname {v e c} \left(\boldsymbol {g} _ {\ell , n} \boldsymbol {a} _ {\ell , n} ^ {\top}\right) \operatorname {v e c} \left(\boldsymbol {g} _ {\ell , n} \boldsymbol {a} _ {\ell , n} ^ {\top}\right) ^ {\top} \right] (19b) \\ = \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} \mid f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} [ (\boldsymbol {a} _ {\ell , n} \otimes \boldsymbol {g} _ {\ell , n}) (\boldsymbol {a} _ {\ell , n} ^ {\intercal} \otimes \boldsymbol {g} _ {\ell , n} ^ {\intercal}) ] (19c) \\ = \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} \left[ \boldsymbol {a} _ {\ell , n} \boldsymbol {a} _ {\ell , n} ^ {\intercal} \otimes \boldsymbol {g} _ {\ell , n} \boldsymbol {g} _ {\ell , n} ^ {\intercal} \right] (19d) \\ \approx \underbrace {\left[ \frac {1}{N} \sum_ {n = 1} ^ {N} \boldsymbol {a} _ {\ell , n} \boldsymbol {a} _ {\ell , n} ^ {\intercal} \right]} _ {=: \boldsymbol {A} _ {\ell}} \otimes \underbrace {\left[ \sum_ {n = 1} ^ {N} \mathbb {E} _ {\boldsymbol {y} \sim p (\boldsymbol {y} | f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}))} [ \boldsymbol {g} _ {\ell , n} \boldsymbol {g} _ {\ell , n} ^ {\intercal} ] \right]} _ {=: \boldsymbol {G} _ {\ell}}, (19e) \\ \end{array}
$$
where  $A_{\ell} \in \mathbb{R}^{P_{\ell,\mathrm{in}} \times P_{\ell,\mathrm{in}}}$  and  $G_{\ell} \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times P_{\ell,\mathrm{out}}}$ . For this derivation, we have used three convenient properties of the Kronecker product (using matrices  $A, B, C, D$  with appropriate dimensions):  $\operatorname{vec}(ABC) = (C^{\intercal} \otimes A)\operatorname{vec}(B)$  and  $(A \otimes B)^{\intercal} = A^{\intercal} \otimes B^{\intercal}$  for Equation (19c), and  $(A \otimes B)(C \otimes D) = AC \otimes BD$  for Equation (19d).
6More generally, it is defined with an expectation over  $x\sim p(x)$  as well.
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  $\mathcal{O}(P_{\ell,\mathrm{in}}^2 P_{\ell,\mathrm{out}}^2)$  to  $\mathcal{O}(P_{\ell,\mathrm{in}}^2 + P_{\ell,\mathrm{out}}^2)$  and the computational complexity of the preconditioning of the gradient with the approximate Fisher from  $\mathcal{O}(P_{\ell,\mathrm{in}}^3 P_{\ell,\mathrm{out}}^3)$  to  $\mathcal{O}(P_{\ell,\mathrm{in}}^3 + P_{\ell,\mathrm{out}}^3)$ , since
$$
\begin{array}{l} \boldsymbol {F} (\boldsymbol {\theta} _ {\ell}) ^ {- 1} \boldsymbol {g} (\boldsymbol {\theta} _ {\ell}) \approx (\boldsymbol {A} _ {\ell} \otimes \boldsymbol {G} _ {\ell}) ^ {- 1} \boldsymbol {g} (\boldsymbol {\theta} _ {\ell}) \\ = \operatorname {v e c} \left(\boldsymbol {G} _ {\ell} ^ {- 1} \nabla_ {\boldsymbol {W} _ {\ell}} \mathcal {L} (\boldsymbol {y}, f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \boldsymbol {A} _ {\ell} ^ {- 1}\right) \tag {20} \\ \end{array}
$$
with  $\pmb{g}(\pmb{\theta}_{\ell}) = \mathrm{vec}\left(\nabla_{\pmb{W}_{\ell}}\mathcal{L}(\pmb{y},f_{\pmb{\theta}}(\pmb{x}_{n}))\right)$  and the property  $(A\otimes B)^{-1} = A^{-1}\otimes B^{-1}$ .
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_{\theta_{\ell}}(\boldsymbol{x}_n) \coloneqq J_{\theta_{\ell}}f_{\theta}(\boldsymbol{x}_n) = J_{\boldsymbol{s}_{\ell,n}}f_{\theta}(\boldsymbol{x}_n)J_{\theta_{\ell}}\boldsymbol{s}_{\ell,n} \in \mathbb{R}^{C \times P_{\ell}}$  as the Jacobian of the model outputs w.r.t. the parameters of the  $\ell$ -th layer and  $\Lambda(f_{\theta}(\boldsymbol{x}_n)) \coloneqq H_{f_{\theta}}\mathcal{L}(\boldsymbol{y}_n, f_{\theta}(\boldsymbol{x}_n)) \in \mathbb{R}^{C \times C}$  as the Hessian of the loss w.r.t. the model outputs. Now we can write  $s_{\ell,n} = W_{\ell}\boldsymbol{a}_{\ell,n} = (\boldsymbol{a}_{\ell,n}^{\intercal} \otimes \mathbf{I}_{P_{\ell,\text{out}}})\mathrm{vec}(\boldsymbol{W}_{\ell})$  and with this we have  $J_{\theta_{\ell}}\boldsymbol{s}_{\ell,n} = \boldsymbol{a}_{\ell,n}^{\intercal} \otimes \mathbf{I}_{P_{\ell,\text{out}}}$ . Additionally, by defining  $\boldsymbol{b}_{\ell,n} \coloneqq J_{\boldsymbol{s}_{\ell,n}}f_{\theta}(\boldsymbol{x}_n)^{\intercal} \in \mathbb{R}^{P_{\ell,\text{out}} \times C}$  as the transposed Jacobian of the model outputs w.r.t. to the pre-activations of the  $\ell$ -th layer, we note that  $J_{\theta_{\ell}}(\boldsymbol{x}_n)^{\intercal} = (\boldsymbol{a}_{\ell,n}^{\intercal} \otimes \mathbf{I}_{P_{\ell,\text{out}}})^{\intercal}\boldsymbol{b}_{\ell,n} = \boldsymbol{a}_{\ell,n} \otimes \boldsymbol{b}_{\ell,n}$ .
Replacing the (transposed) Jacobians in the definition of the GGN by this expression, we have