SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma 5 (Sufficient condition for exactness of K-FAC-reduce in the reduce setting). Let  $D_{\ell,n} \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times C}$  be a constant matrix for layer  $\ell$  and data point  $\pmb{x}_n$ . Further, let  $C_\ell \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times P_{\ell,\mathrm{out}}}$  be a constant matrix for layer  $\ell$ . If it holds for each  $n$  that  $\pmb{b}_{\ell,n,r} = D_{\ell,n}$  for all  $r \in \{1, \dots, R\}$  and  $\hat{\pmb{b}}_{\ell,n}\pmb{\Lambda}(f_\theta(\pmb{x}_n))\hat{\pmb{b}}_{\ell,n}^\top = C_\ell$  for all  $n \in \{1, \dots, N\}$ , then the K-FAC-reduce approximation in Equation (10) is equal to the exact GGN of the  $\ell$ -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
$$
\begin{array}{l} \sum_ {n = 1} ^ {N} \left(\frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n, r}\right) \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \left(\frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n, r}\right) ^ {\intercal} \\ = \sum_ {n = 1} ^ {N} \left(\frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes R \boldsymbol {D} _ {\ell , n}\right) \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \left(\frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes R \boldsymbol {D} _ {\ell , n}\right) ^ {\intercal} \tag {41} \\ = \sum_ {n = 1} ^ {N} \left(\sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \boldsymbol {b} _ {\ell , n, r}\right) \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \left(\sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \boldsymbol {b} _ {\ell , n, r}\right) ^ {\intercal}, \\ \end{array}
$$
where we have used the assumption that for each  $n$ , we have  $\pmb{b}_{\ell,n,r} = \pmb{D}_{\ell,n}$  for all  $r \in \{1, \dots, R\}$ . Now we consider the second approximation in Equation (39). Analogously, we have
$$
\begin{array}{l} \left(\frac {1}{N} \sum_ {n = 1} ^ {N} \hat {\boldsymbol {a}} _ {\ell , n} \hat {\boldsymbol {a}} _ {\ell , n} ^ {\intercal}\right) \otimes \left(\sum_ {n = 1} ^ {N} \hat {\boldsymbol {b}} _ {\ell , n} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \hat {\boldsymbol {b}} _ {\ell , n} ^ {\intercal}\right) \\ = \left(\frac {1}{N} \sum_ {n = 1} ^ {N} \hat {\boldsymbol {a}} _ {\ell , n} \hat {\boldsymbol {a}} _ {\ell , n} ^ {\intercal}\right) \otimes N \boldsymbol {C} _ {\ell} \tag {42} \\ = \sum_ {n = 1} ^ {N} \left(\hat {\boldsymbol {a}} _ {\ell , n} \hat {\boldsymbol {a}} _ {\ell , n} ^ {\intercal}\right) \otimes \left(\hat {\boldsymbol {b}} _ {\ell , n} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \hat {\boldsymbol {b}} _ {\ell , n} ^ {\intercal}\right), \\ \end{array}
$$
where we have used that  $\hat{\pmb{b}}_{\ell,n}\pmb{\Lambda}(f_{\pmb{\theta}}(\pmb{x}_n))\hat{\pmb{b}}_{\ell,n}^{\top} = \pmb{C}_{\ell}$  for all  $n \in \{1, \dots, N\}$ .
Until now, we did not have to explicitly take the aggregation function  $z: \mathbb{R}^{R \times P_{\ell, \mathrm{out}}} \to \mathbb{R}^{P_{\ell, \mathrm{out}}}$  into account, since its Jacobian is simply subsumed in  $\pmb{b}_{\ell, 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  $\pmb{J}_{\pmb{s}_{\ell, n, r}} \pmb{z}_{\ell, n}$  with  $\pmb{z}_{\ell, n} := z(\pmb{S}_{\ell, n}) \in \mathbb{R}^{P_{\ell, \mathrm{out}}}$  is the same for all  $r \in \{1, \dots, 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.
$$
\begin{array}{l} z \left(\boldsymbol {S} _ {\ell , n}\right) = c \sum_ {r = 1} ^ {R} \boldsymbol {s} _ {\ell , n, r} \\ = c \boldsymbol {S} _ {\ell , n} ^ {\mathsf {T}} \mathbf {1} _ {R} \tag {43} \\ = \left(\mathbf {1} _ {R} ^ {\intercal} \otimes c \mathbf {I} _ {P _ {\ell , \text {o u t}}}\right) \operatorname {v e c} \left(\boldsymbol {S} _ {\ell , n} ^ {\intercal}\right) \\ = \left(\mathbf {1} _ {R} ^ {\intercal} \otimes c \mathbf {I} _ {P _ {\ell , \text {o u t}}}\right) K ^ {(R, P _ {\ell , \text {o u t}})} \operatorname {v e c} \left(\boldsymbol {S} _ {\ell , n}\right) \\ \end{array}
$$
with  $c\in \mathbb{R}$  and the commutation matrix
$$
\boldsymbol {K} ^ {\left(R, P _ {\ell , \text {o u t}}\right)} := \sum_ {r = 1} ^ {R} \sum_ {p = 1} ^ {P _ {\ell , \text {o u t}}} \left(\boldsymbol {e} _ {R, r} \boldsymbol {e} _ {P _ {\ell , \text {o u t}}, p} ^ {\intercal}\right) \otimes \left(\boldsymbol {e} _ {P _ {\ell , \text {o u t}}, p} \boldsymbol {e} _ {R, r} ^ {\intercal}\right), \tag {44}
$$
where  $e_{i,j}$  is the  $j$ -th canonical vector of dimension  $i$ . This is a linear function in  $\operatorname{vec}(S_{\ell,n})$  and we have  $J_{s_{\ell,n},r}z_{\ell,n} = c\mathbf{I}_{P_{\ell,\mathrm{out}}}$  for all  $r \in \{1, \ldots, 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  $\operatorname{vec}(S_{\ell,n})$ , because as soon as we have a weighted sum with weights  $c_r \in \mathbb{R}$  and they are not the same for all  $r \in \{1, \ldots, R\}$ , the Jacobians  $J_{s_{\ell,n},r}z_{\ell,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  $\pmb{\Sigma} \in \mathbb{R}^{C \times 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.
Proof. We have  $\boldsymbol{\Lambda}(f_{\boldsymbol{\theta}}(\boldsymbol{x}_n)) = \boldsymbol{\Sigma}^{-1}$  and  $\boldsymbol{b}_{\ell,n,r} = (J_{\boldsymbol{z}_{\ell,n}}f_{\boldsymbol{\theta}}(\boldsymbol{x}_n)\boldsymbol{J}_{\boldsymbol{s}_{\ell,n,r}}\boldsymbol{z}_{\ell,n})^\top = c\mathbf{I}_C$  for all  $r \in \{1, \ldots, R\}$  and  $n \in \{1, \ldots, N\}$  ( $P_{\ell,\mathrm{out}} = C$  for a single layer). Hence,  $\hat{\boldsymbol{b}}_{\ell,n}\boldsymbol{\Lambda}(f_{\boldsymbol{\theta}}(\boldsymbol{x}_n))\hat{\boldsymbol{b}}_{\ell,n}^\top = c^2 R^2\mathbf{I}_C\Sigma^{-1}\mathbf{I}_C = c^2 R^2\Sigma^{-1}$  for all  $n \in \{1, \ldots, N\}$ . Therefore, the desired result follows from Lemma 5.
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  $\ell$  of a deep linear network (Equation (8)), a Gaussian likelihood with p.d. covariance matrix  $\pmb{\Sigma} \in \mathbb{R}^{C \times C}$ , and a scaled sum aggregation function,  $K$ -FAC-reduce is exact in the reduce setting.
Proof. We have  $\Lambda(f_{\theta}(\pmb{x}_n)_r) = \pmb{\Sigma}^{-1}$  and  $\pmb{b}_{\ell,n,r} = (J_{\pmb{z}_{\ell,n}}f_{\theta}(\pmb{x}_n)\pmb{J}_{\pmb{s}_{\ell,n,r}}\pmb{z}_{\ell,n})^{\intercal} = c\pmb{W}_{\ell}^{a^{\intercal}}$  for all  $r \in \{1, \ldots, R\}$  and  $n \in \{1, \ldots, N\}$ . Hence,
$$
\hat {\boldsymbol {b}} _ {\ell , n} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \hat {\boldsymbol {b}} _ {\ell , n} ^ {\intercal} = c ^ {2} R ^ {2} \boldsymbol {W} _ {\ell} ^ {a ^ {\intercal}} \boldsymbol {\Sigma} ^ {- 1} \boldsymbol {W} _ {\ell} ^ {a}
$$
for all  $n\in \{1,\ldots ,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 \cdot 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-FAC-expand 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^{\bar{V}}$ ), i.e.
$$
f _ {\boldsymbol {\theta}} \left(\boldsymbol {X} _ {n}\right) = \underbrace {\boldsymbol {X} _ {n} \boldsymbol {W} ^ {Q} {} ^ {\intercal}} _ {=: \boldsymbol {S} _ {Q, n}} \underbrace {\boldsymbol {W} ^ {K} \boldsymbol {X} _ {n} ^ {\intercal}} _ {=: \boldsymbol {S} _ {K, n} ^ {\intercal}} \underbrace {\boldsymbol {X} _ {n} \boldsymbol {W} ^ {V} {} ^ {\intercal}} _ {=: \boldsymbol {S} _ {V, n}}. \tag {45}
$$
We can observe that it is no longer a linear function in the input  $\mathbf{X}_n \in \mathbb{R}^{R \times 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_{\theta}(\mathbf{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_{\theta}(X_n)_r$  for the layer with weights  $W^Q$ ,  $J_{s_{K,n,m}}f_{\theta}(X_n)_r$  for the layer with weights  $W^K$ , and  $J_{s_{V,n,m}}f_{\theta}(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  $\mathbf{W}^Q$  as a function of  $s_{Q,n,r}$  as
$$
f _ {\boldsymbol {\theta}} \left(\boldsymbol {X} _ {n}\right) _ {r} = \boldsymbol {s} _ {Q, n, r} ^ {\intercal} \boldsymbol {S} _ {K, n} ^ {\intercal} \boldsymbol {S} _ {V, n}. \tag {46}
$$
$$
\boldsymbol {J} _ {\boldsymbol {s} _ {Q, n, r}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {X} _ {n}\right) _ {r} = \boldsymbol {S} _ {V, n} ^ {\intercal} \boldsymbol {S} _ {K, n} = \boldsymbol {b} _ {Q, n, r} ^ {\intercal} \in \mathbb {R} ^ {C \times P _ {K, \text {o u t}}}, \tag {47}
$$
with  $C = P_{V,\mathrm{out}}$  and  $\pmb{b}_{Q,n,r,m} = \mathbf{0}$  for all  $m\neq r$ , which is the first assumption necessary for Lemma 3 to hold. While  $\pmb{b}_{Q,n,r}$  is not the same for all  $n\in \{1,\dots ,N\}$ , it is the same for all  $r\in \{1,\ldots ,R\}$  and hence, under the same assumptions as in Proposition 4, K-FAC-expand is exact for the layer with weights  $\pmb{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_{\theta}(\mathbf{X}_n)$  as a function of the  $r$ -th row of  $\mathbf{S}_{K / V,n}$ , i.e. elements from all rows of  $\mathbf{S}_{K / V,n}$  contribute to the  $r$ -th row of the output matrix. We can also see this by directly deriving  $\mathbf{J}_{\mathbf{s}_{K / V,n,m}}f_{\theta}(\mathbf{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  $\mathbf{W}^K$ , we have
$$
\boldsymbol {J} _ {\boldsymbol {s} _ {K, n, m}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {X} _ {n}\right) _ {r} = \boldsymbol {s} _ {V, n, m} \boldsymbol {s} _ {Q, n, r} ^ {\intercal} \in \mathbb {R} ^ {C \times P _ {Q, \text {o u t}}} \tag {48}
$$
Moreover, for the third layer with weight matrix  $W^V$ , we have
$$
\boldsymbol {J} _ {\boldsymbol {s} _ {V, n, m}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {X} _ {n}\right) _ {r} = \boldsymbol {s} _ {Q, n, r} ^ {\intercal} \boldsymbol {s} _ {K, n, m} \mathbf {I} _ {C} \in \mathbb {R} ^ {C \times C}. \tag {49}
$$
This means that the assumption of Lemma 3 that the  $R$  elements along the weight-sharing dimension are independent does not hold, since the Jacobians depend on  $r$  and  $m$ . The approximation leads to an inexact computation, even though only linear layers are involved and 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
$$
\begin{array}{l} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} z (f _ {\boldsymbol {\theta}} (\boldsymbol {X} _ {n})) = \sum_ {r = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {s} _ {\ell , n, r}} z (f _ {\boldsymbol {\theta}} (\boldsymbol {X} _ {n})) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \boldsymbol {s} _ {\ell , n, r} \\ = \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \boldsymbol {b} _ {\ell , n, r} \tag {50} \\ = \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \otimes \left(c \sum_ {m = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {s} _ {\ell , n, r}} f _ {\boldsymbol {\theta}} (\boldsymbol {X} _ {n}) _ {m}\right), \\ \end{array}
$$
where  $J_{s_{\ell,n,r}} f_{\theta}(X_n)_m$  are the same Jacobians we have derived for the expand case. Since according to Lemma 5 we need all  $b_{\ell,n,r}$  to be the same for all  $r \in \{1, \dots, 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  $\mathbf{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  $\mathbf{S}_{\ell ,Q,n}$  of all layers with weights  $\mathbf{W}_{\ell}^{Q}$ , besides for the layer in the last block, i.e. the layer in the  $L$ -th block with weights  $\mathbf{W}_L^Q$ .
This shows that even without explicit nonlinear activation functions, the self-attention mechanism in transformer models breaks the two approximations. Hence, it is not inherently clear how useful it is to consider the corresponding approximation in the expand and reduce setting. This is especially relevant given that we know that the computational complexity of K-FAC-reduce is generally smaller than of K-FAC-expand and given the similar downstream optimisation performance we report in Section 4.
B.3.2 K-FAC for GNNs
Beyond transformers, we have introduced GNNs as a class of models that also use linear weight-sharing layers. There are many types of GNNs and we will only explicitly cover two of them here.
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  $\tilde{\pmb{a}}_{\ell,n}$  to the  $\ell$ -th layer for the node with index  $n$  now depend on its neighbourhood  $\mathcal{N}(n)$ , since
$$
\tilde {\boldsymbol {a}} _ {\ell , n} := \sum_ {j \in \mathcal {N} (n)} \hat {\boldsymbol {C}} _ {n j} \boldsymbol {a} _ {\ell , j} \in \mathbb {R} ^ {P _ {\ell , \text {i n}}}. \tag {51}
$$
Using this notation, the definition of the K-FAC GGN for node classification is simply
$$
\begin{array}{l} G G N (\boldsymbol {\theta} _ {\ell}) = \sum_ {n = 1} ^ {N} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} (\boldsymbol {X}) _ {n} ^ {\intercal} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {X}) _ {n}) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} (\boldsymbol {X}) _ {n} \\ \approx \underbrace {\left[ \frac {1}{N} \sum_ {n = 1} ^ {N} \tilde {\boldsymbol {a}} _ {\ell , n} \tilde {\boldsymbol {a}} _ {\ell , n} ^ {\mathsf {T}} \right]} _ {=: \tilde {\boldsymbol {A}} _ {\ell}} \otimes \underbrace {\left[ \sum_ {n = 1} ^ {N} \tilde {\boldsymbol {b}} _ {\ell , n} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} (\boldsymbol {X}) _ {n}\right) \tilde {\boldsymbol {b}} _ {\ell , n} ^ {\mathsf {T}} \right]} _ {=: \tilde {\boldsymbol {B}} _ {\ell}}, \tag {52} \\ \end{array}
$$
where  $\mathbf{X} \in \mathbb{R}^{N \times D}$ ,  $\tilde{\pmb{b}}_{\ell,n} \coloneqq J_{\tilde{\pmb{s}}_{\ell,n}} f_{\pmb{\theta}_{\ell}}(\mathbf{X})_n^{\mathrm{T}}$ , and  $\tilde{\pmb{s}}_{\ell,n} \coloneqq W_{\ell} \tilde{\pmb{a}}_{\ell,n}$ . Again, it is important to note that we need to have access to the whole neighbourhood of  $\pmb{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  $\mathbf{X}$  containing all nodes for each loss term. Also, depending on the sparsity of  $\hat{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  $\hat{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 A1goPerf 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  $\mathbb{X}_n^G$  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  $\phi^E$  this dimension is going to be  $R_{n} = N_{n}^{E}$ , whereas it will be  $R_{n} = N_{n}^{V}$  for  $\phi^V$ . In the case of  $\phi^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  $\ell$  of type  $\phi^E$  as
$$
\boldsymbol {A} _ {\ell , n} = \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) \in \mathbb {R} ^ {N _ {n} ^ {E} \times \left(D _ {E} + 2 D _ {V} + D _ {u}\right)} \tag {53}
$$
and as
$$
\boldsymbol {A} _ {\ell , n} = \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) \in \mathbb {R} ^ {N _ {n} ^ {V} \times \left(D _ {V} + D _ {E} + D _ {u}\right)} \tag {54}
$$
for  $\phi^V$ . Correspondingly, we have  $\pmb{b}_{\ell,n} = \pmb{J}_{\pmb{S}_{\ell,n}}f_{\pmb{\theta}}(\mathbb{X}_n^G)^\top \in \mathbb{R}^{N_n^E\times D_E\times C}$  for  $\phi^E$  and  $\pmb{b}_{\ell,n} = \pmb{J}_{\pmb{S}_{\ell,n}}f_{\pmb{\theta}}(\mathbb{X}_n^G)^\top \in \mathbb{R}^{N_n^V\times D_V\times C}$ , with  $\pmb{S}_{\ell,n} = \pmb{A}_{\ell,n}\pmb{W}_\ell^{E^\intercal} \in \mathbb{R}^{N_n^E\times D_E}$  and  $\pmb{S}_{\ell,n} = \pmb{A}_{\ell,n}\pmb{W}_\ell^{V^\intercal} \in \mathbb{R}^{N_n^V\times D_V}$ , respectively.
Using this notation, we can approximate the GGN for layer  $\ell$ , assuming its type is either  $\phi^E$  or  $\phi^V$ , as
$$
\begin{array}{l} \mathbf {G G N} (\boldsymbol {\theta} _ {\ell}) = \sum_ {n = 1} ^ {N} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} (\mathbb {X} _ {n} ^ {G}) ^ {\intercal} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\mathbb {X} _ {n} ^ {G})) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} (\mathbb {X} _ {n} ^ {G}) \\ \approx \sum_ {n = 1} ^ {N} \underbrace {\left[ \sum_ {r = 1} ^ {R _ {n}} \frac {1}{\sqrt {R _ {n}}} \boldsymbol {a} _ {\ell , n , r} \right]} _ {=: \hat {\boldsymbol {a}} _ {\ell , n}} \otimes \underbrace {\left[ \sum_ {r = 1} ^ {R _ {n}} \frac {1}{\sqrt {R _ {n}}} \boldsymbol {b} _ {\ell , n , r} \right]} _ {=: \hat {\boldsymbol {b}} _ {\ell , n}} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\mathbb {X} _ {n} ^ {G}\right)\right) \left[ \sum_ {r = 1} ^ {R _ {n}} \frac {1}{\sqrt {R _ {n}}} \boldsymbol {a} _ {\ell , n, r} \right] \otimes \left[ \sum_ {r = 1} ^ {R _ {n}} \frac {1}{\sqrt {R _ {n}}} \boldsymbol {b} _ {\ell , n, r} \right] ^ {\intercal} \\ \approx \underbrace {\left[ \frac {1}{N} \sum_ {n = 1} ^ {N} \hat {\boldsymbol {a}} _ {\ell , n} \hat {\boldsymbol {a}} _ {\ell , n} ^ {\mathsf {T}} \right]} _ {=: \hat {\boldsymbol {A}} _ {\ell}} \otimes \underbrace {\left[ \sum_ {n = 1} ^ {N} \hat {\boldsymbol {b}} _ {\ell , n} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\mathbb {X} _ {n} ^ {G}\right)\right) \hat {\boldsymbol {b}} _ {\ell , n} ^ {\mathsf {T}} \right]} _ {=: \hat {\boldsymbol {B}} _ {\ell}}, \tag {55} \\ \end{array}
$$