SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbf {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}) \\ = \sum_ {n = 1} ^ {N} \left(\boldsymbol {a} _ {\ell , n} \otimes \boldsymbol {b} _ {\ell , n}\right) \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \left(\boldsymbol {a} _ {\ell , n} \otimes \boldsymbol {b} _ {\ell , n}\right) ^ {\intercal} \\ = \sum_ {n = 1} ^ {N} \left(\boldsymbol {a} _ {\ell , n} \boldsymbol {a} _ {\ell , n} ^ {\mathsf {T}}\right) \otimes \left(\boldsymbol {b} _ {\ell , n} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \boldsymbol {b} _ {\ell , n} ^ {\mathsf {T}}\right) \tag {21} \\ \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} \boldsymbol {b} _ {\ell , n} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \boldsymbol {b} _ {\ell , n} ^ {\intercal} \right]} _ {=: \boldsymbol {B} _ {\ell}}. \\ \end{array}
$$
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  $\mathbf{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_{\theta_{\ell}}(\mathbf{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 \cdot R$  terms, which corresponds to assuming  $N \cdot R$  i.i.d. examples,
The Expand Setting
$$
\mathcal {L} _ {\text {e x p a n d}} \left(f _ {\boldsymbol {\theta}}, \mathcal {D}\right) := - \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \log p \left(\boldsymbol {y} _ {n, r} \mid f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right), \tag {22}
$$
where  $f_{\pmb{\theta}}(\pmb{x}_n)_r$  is the  $r$ -th row of the model output  $f_{\pmb{\theta}}(\pmb{x}_n) \in \mathbb{R}^{R \times C}$  and  $\pmb{y}_{n,r}$  is the  $r$ -th row of the label  $\pmb{Y}_n \in \mathbb{R}^{R \times 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  $\pmb{x}_n$  to have an additional weight-sharing dimension since we only require that the input to the  $\ell$ -th layer has this additional dimension, i.e.  $\pmb{A}_{\ell,n} \in \mathbb{R}^{R \times D}$ . This obviously does not exclude the case where  $\pmb{x}_n$  already has this weight-sharing dimension, e.g. a sequence of tokens in translation tasks.
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. the parameters  $\pmb{\theta}_{\ell}$  as
$$
\left(\boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) _ {i j} = \sum_ {m = 1} ^ {R} \sum_ {p = 1} ^ {P _ {\ell , \text {o u t}}} \frac {\partial f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r i}}{\partial \boldsymbol {S} _ {\ell , n , m p}} \frac {\partial \boldsymbol {S} _ {\ell , n , m p}}{\partial \boldsymbol {\theta} _ {\ell , j}} \tag {23}
$$
or in matrix form
$$
\boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \left(\boldsymbol {x} _ {n}\right) _ {r} = \sum_ {m = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {s} _ {\ell , n, m}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \boldsymbol {s} _ {\ell , n, m}. \tag {24}
$$
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}^{\intercal} \otimes \mathbf{I}_{P_{\ell,\mathrm{out}}}$ , as for regular K-FAC (c.f. Appendix B.1). We denote  $\pmb{b}_{\ell,n,r,k} := \pmb{J}_{\pmb{s}_{\ell,n,k}}f_{\pmb{\theta}}(\pmb{x}_n)_r^{\intercal}$ . Hence, we have
$$
\begin{array}{l} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \left(\boldsymbol {x} _ {n}\right) _ {r} ^ {\intercal} = \left(\sum_ {m = 1} ^ {R} \boldsymbol {b} _ {\ell , n, r, m} ^ {\intercal} \left(\boldsymbol {a} _ {\ell , n, m} ^ {\intercal} \otimes \mathbf {I} _ {P _ {\ell , \text {o u t}}}\right)\right) ^ {\intercal} \tag {25} \\ = \sum_ {m = 1} ^ {R} \boldsymbol {a} _ {\ell , n, m} \otimes \boldsymbol {b} _ {\ell , n, r, m}. \\ \end{array}
$$
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 \cdot R$  independent examples and deriving the approximation in the same way as we would with  $N$  examples (c.f. Appendix B.1),
$$
\begin{array}{l} \mathbf {G G N} \left(\boldsymbol {\theta} _ {\ell}\right) = \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \left(\boldsymbol {x} _ {n}\right) _ {r} ^ {\intercal} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \left(\boldsymbol {x} _ {n}\right) _ {r} \tag {26} \\ = \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \left(\sum_ {m = 1} ^ {R} \boldsymbol {a} _ {\ell , n, m} \otimes \boldsymbol {b} _ {\ell , n, r, m}\right) \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). \\ \end{array}
$$
However, we cannot directly write each of the  $N \cdot R$  loss terms as a Kronecker product without any approximation. One approach could be to approximate each  $J_{\theta_\ell}(\boldsymbol{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  $\sum_{m=1}^{R} b_{\ell,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
$$
\begin{array}{l} \mathbf {G G N} \left(\boldsymbol {\theta} _ {\ell}\right) \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) \tag {27} \\ = \sum_ {n = 1} ^ {N} \sum_ {m = 1} ^ {R} \left(\boldsymbol {a} _ {\ell , n, m} \boldsymbol {a} _ {\ell , n, m} ^ {\intercal}\right) \otimes \left(\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), \\ \end{array}
$$
where we ignore all the terms with  $m \neq m'$ , which allows us to express each of the  $N \cdot R$  terms as a Kronecker product. Consequently, we can apply the regular K-FAC approximation over  $N \cdot R$  terms instead of just  $N$  terms as usual. We call the resulting approximation  $K$ -FAC-expand:
However, the authors of the work on transformers do not discuss this extension of K-FAC at all; although Zhang et al. (2018) do explicitly discuss the tying of the embedding and linear output layer weights.
$$
\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]} _ {= \boldsymbol {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 {28}
$$
There is one simple case where the exact expression in Equation (26) is identical to the approximation in Equation (27). When  $\pmb{b}_{\ell,n,r,m} = \mathbf{0}$  for all  $r \neq m$ , both expressions are equivalent to
$$
\sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \left(\boldsymbol {a} _ {\ell , n, r} \otimes \boldsymbol {b} _ {\ell , n, r, r}\right) \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \left(\boldsymbol {a} _ {\ell , n, r} ^ {\intercal} \otimes \boldsymbol {b} _ {\ell , n, r, r} ^ {\intercal}\right). \tag {29}
$$
With other words, when  $f_{\theta}(\pmb{x}_n)_r$  is independent of all pre-activations  $s_{\ell,n,m}$  with  $m \neq 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 weight-sharing 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_{\ell,n,r,m} \neq \mathbf{0}$  for  $r \neq m$  and is also what has been used by previous work on transformers (Grosse et al., 2023; Ottawa 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.
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 \cdot 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  $\boldsymbol{a}_{\ell,n,r}$  are the same for all  $n \in \{1,\dots,N\}$  and  $r \in \{1,\dots,R\}$ . We do not state these types of conditions explicitly from now on.
Lemma 3 (Sufficient condition for exactness of K-FAC-expand in the expand setting). Let  $C_{\ell} \in \mathbb{R}^{P_{\ell, \mathrm{out}} \times P_{\ell, \mathrm{out}}}$  be a constant matrix for layer  $\ell$ . If  $\pmb{b}_{\ell, n, r, m} = \mathbf{0}$  for all  $r \neq m$  and  $\pmb{b}_{\ell, n, m, m} \Lambda(f_{\theta}(\pmb{x}_n)_m) \pmb{b}_{\ell, n, m, m}^{\intercal} = C_{\ell}$  for all  $n \in \{1, \dots, N\}$  and  $m \in \{1, \dots, R\}$ , then the K-FAC approximation in Equation (7) is equal to the exact GGN/Fisher of the  $\ell$ -th layer in the expand setting.
Proof. As mentioned before, when  $\pmb{b}_{\ell,n,r,m} = \mathbf{0}$  for all  $r \neq 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  $\ell$  from there. We have
$$
\begin{array}{l} \left(\frac {1}{N R} \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \boldsymbol {a} _ {\ell , n, r} ^ {\intercal}\right) \otimes \left(\sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n, r, r} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \boldsymbol {b} _ {\ell , n, r, r} ^ {\intercal}\right) \\ = \left(\frac {1}{N R} \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \boldsymbol {a} _ {\ell , n, r} ^ {\intercal}\right) \otimes \left(N R C _ {\ell}\right) \tag {30} \\ = \left(\sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} \boldsymbol {a} _ {\ell , n, r} ^ {\intercal}\right) \otimes \boldsymbol {C} _ {\ell} \\ = \sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {R} \left(\boldsymbol {a} _ {\ell , n, r} \boldsymbol {a} _ {\ell , n, r} ^ {\intercal}\right) \otimes \left(\boldsymbol {b} _ {\ell , n, r, r} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {r}\right) \boldsymbol {b} _ {\ell , n, r, r} ^ {\intercal}\right), \\ \end{array}
$$
where we have used the assumption that  $\pmb{b}_{\ell,n,m,m}\pmb{\Lambda}(f_{\pmb{\theta}}(\pmb{x}_n)_m)\pmb{b}_{\ell,n,m,m}^{\top} = \pmb{C}_{\ell}$  is the same for all  $n \in \{1,\dots,N\}$  and  $m \in \{1,\dots,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  $\Sigma \in \mathbb{R}^{C\times C}$ ,  $K$ -FAC-expand is exact in the expand setting.
Proof. We can write  $f_{\pmb{\theta}}(\pmb{x}_n)_r = \pmb{W}_{\ell}\pmb{x}_{n,r}$  and hence  $\pmb{b}_{\ell,n,r,m} = \mathbf{0}$  for  $r \neq m$ . Moreover, we have  $\Lambda(f_{\pmb{\theta}}(\pmb{x}_n)_r) = \pmb{\Sigma}^{-1}$  and  $\pmb{b}_{\ell,n,m,m} = \mathbf{I}_C$  ( $P_{\ell,\mathrm{out}} = C$  for a single layer). Hence,
$$
\boldsymbol {b} _ {\ell , n, m, m} \boldsymbol {\Lambda} (\boldsymbol {f} _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}) _ {r}) \boldsymbol {b} _ {\ell , n, m, m} ^ {\intercal} = \mathbf {I} _ {C} \boldsymbol {\Sigma} ^ {- 1} \mathbf {I} _ {C} = \boldsymbol {\Sigma} ^ {- 1}
$$
for all  $n \in \{1, \dots, N\}$  and  $m \in \{1, \dots, 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 _ {\boldsymbol {\theta}} (\boldsymbol {x}) = \boldsymbol {W} _ {L} \dots \boldsymbol {W} _ {\ell} \dots \boldsymbol {W} _ {1} \boldsymbol {x} = \boldsymbol {W} \boldsymbol {x}, \tag {31}
$$
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). We adopt the notation of Bernacchia et al. (2018) and define
$$
\boldsymbol {W} _ {\ell} ^ {a} := \boldsymbol {W} _ {L} \dots \boldsymbol {W} _ {\ell + 1} \tag {32}
$$
as the product of the weight matrices ahead of  $W_{\ell}$  and
$$
\boldsymbol {W} _ {\ell} ^ {b} := \boldsymbol {W} _ {\ell - 1} \dots \boldsymbol {W} _ {1} \tag {33}
$$
as the product of the weight matrices behind of  $\mathbf{W}_{\ell}$ . Hence, we can write  $f_{\theta}(\pmb{x}) = \pmb{W}_{\ell}^{a}\pmb{W}_{\ell}\pmb{W}_{\ell}^{b}\pmb{x}$ . Note, that now
$$
\boldsymbol {a} _ {\ell , n, r} = \boldsymbol {W} _ {\ell} ^ {b} \boldsymbol {x} _ {n, r} \in \mathbb {R} ^ {P _ {\ell , \text {i n}}} \tag {34}
$$
$$
\boldsymbol {b} _ {\ell , n, r, r} = \boldsymbol {W} _ {\ell} ^ {a ^ {\intercal}} \in \mathbb {R} ^ {P _ {\ell , \text {o u t}} \times C}. \tag {35}
$$
Using these insights, we can now easily state the result for deep linear networks.
Proof. We can write  $f_{\pmb{\theta}}(\pmb{x}_n)_r = \pmb{W}\pmb{x}_{n,r}$  and hence  $\pmb{b}_{\ell,n,r,m} = \mathbf{0}$  for  $r \neq m$ . We have  $\pmb{\Lambda}(f_{\pmb{\theta}}(\pmb{x}_n)_r) = \pmb{\Sigma}^{-1}$  and  $\pmb{b}_{\ell,n,r,r} = \pmb{W}_{\ell}^{a^{\intercal}}$ . Hence,  $\pmb{b}_{\ell,n,m,m}\Lambda(f_{\pmb{\theta}}(\pmb{x}_n)_m)\pmb{b}_{\ell,n,m,m}^{\intercal} = \pmb{W}_{\ell}^{a^{\intercal}}\pmb{\Sigma}^{-1}\pmb{W}_{\ell}^{a}$  for all  $n \in \{1,\dots,N\}$  and  $m \in \{1,\dots,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
$$
\mathcal {L} _ {\text {r e d u c e}} \left(f _ {\boldsymbol {\theta}}, \mathcal {D}\right) := - \sum_ {n = 1} ^ {N} \log p \left(\boldsymbol {y} _ {n} \mid f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right), \tag {36}
$$
where the crucial observation is that 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 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  $\pmb{x}_n$  and labels  $\pmb{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  $\mathbf{A}_{\ell,n} \in \mathbb{R}^{R \times P_{\ell,\mathrm{in}}}$  is now a matrix, we have  $S_{\ell,n} = \mathbf{A}_{\ell,n} \mathbf{W}_{\ell}^{\mathsf{T}} \in \mathbb{R}^{R \times P_{\ell,\mathrm{out}}}$ . Hence,  $J_{\theta_{\ell}} S_{\ell,n}$  and  $J_{S_{\ell,n}} f_{\theta}(x_n)$  are now both multi-dimensional arrays. Luckily, we can simplify this by writing
$$
\left(\boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) _ {i j} = \sum_ {r = 1} ^ {R} \sum_ {p = 1} ^ {P _ {\ell , \text {o u t}}} \frac {\partial f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) _ {i}}{\partial \boldsymbol {S} _ {\ell , n , r p}} \frac {\partial \boldsymbol {S} _ {\ell , n , r p}}{\partial \boldsymbol {\theta} _ {\ell , j}}, \tag {37}
$$
$$
\boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) = \sum_ {r = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {s} _ {\ell , n, r}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \boldsymbol {s} _ {\ell , n, r}, \tag {38}
$$
where  $s_{\ell ,n,r}\in \mathbb{R}^{P_{\ell ,\mathrm{out}}}$  is the  $r$  -th row of  $S_{\ell ,n}$  and  $s_{\ell ,n,r} = W_{\ell}a_{\ell ,n,r}$
Using this equivalence we can approximate the GGN for layer  $\ell$  as
$$
\begin{array}{l} \mathbf {G G N} (\boldsymbol {\theta} _ {\ell}) = \sum_ {n = 1} ^ {N} \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} (\boldsymbol {x} _ {n}) ^ {\top} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} (\boldsymbol {x} _ {n}) \\ = \sum_ {n = 1} ^ {N} \left(\sum_ {r = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {s} _ {\ell , n, r}} f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \boldsymbol {s} _ {\ell , n, r}\right) ^ {\intercal} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \left(\sum_ {r = 1} ^ {R} \boldsymbol {J} _ {\boldsymbol {s} _ {\ell , n, r}} f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}) \boldsymbol {J} _ {\boldsymbol {\theta} _ {\ell}} \boldsymbol {s} _ {\ell , n, r}\right) \\ = \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} \\ \approx \sum_ {n = 1} ^ {N} \underbrace {\left[ \frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n , r} \right]} _ {=: \hat {\boldsymbol {a}} _ {\ell , n}} \otimes \underbrace {\left[ \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n , r} \right]} _ {=: \hat {\boldsymbol {b}} _ {\ell , n}} \boldsymbol {\Lambda} (f _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n})) \left[ \frac {1}{R} \sum_ {r = 1} ^ {R} \boldsymbol {a} _ {\ell , n, r} ^ {\intercal} \right] \otimes \left[ \sum_ {r = 1} ^ {R} \boldsymbol {b} _ {\ell , n, r} ^ {\intercal} \right] \\ = \sum_ {n = 1} ^ {N} \left(\hat {\boldsymbol {a}} _ {\ell , n} \hat {\boldsymbol {a}} _ {\ell , n} ^ {\mathsf {T}}\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} ^ {\mathsf {T}}\right) \\ \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(\boldsymbol {x} _ {n}\right)\right) \hat {\boldsymbol {b}} _ {\ell , n} ^ {\mathsf {T}} \right]} _ {=: \hat {\boldsymbol {B}} _ {\ell}}, \tag {39} \\ \end{array}
$$
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
$$
\mathrm {G} \hat {\mathrm {G}} \mathrm {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} ^ {\intercal}\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} ^ {\intercal}\right) \right]} _ {= \hat {\boldsymbol {B}} _ {\ell}} \tag {40}
$$
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.