> Training large transformer models still poses a challenge for second-order methods. (...)
We fully agree that training large models still poses challenges for second-order methods, probably mostly due to the memory overhead. We believe this could be addressed by considering sparse structures to approximate the Kronecker factors, see e.g. [1]. However, this is beyond the scope of our paper.
While we could not provide full training results for GPT-2 due to the time constraints of the rebuttal, we agree that this will be an interesting experiment.
> I understand that the authors' focus is on generalizing K-FAC to modern neural network architectures. (...)
The first-order baseline we show per experiment is the best one among AdamW, NAdamW, SGD with Nesterov momentum, and SGD with heavy ball momentum, as described in appendix C.1.
Comparing the K-FAC variations to other second-order methods is definitely interesting and could increase the appeal to some readers, but we do not think that this is necessary for the central claims and contributions of the paper. However, we are happy to add at least one second-order baseline for the camera-ready version, e.g. Shampoo.
> The method K-FAC is specifically developed to approximate the Fisher Information Matrix (FIM). In general cases, the Generalized Gauss-Newton (GGN) is not equivalent to the FIM. (...)
We agree that the two approximations should be distinguished in the case where the FIM and the GGN do not coincide. For the cases where they do, which include common loss functions like the cross-entropy and mean squared error loss we consider in the paper, the two approximations are exactly identical. Hence why we decided to also talk about "K-FAC to approximate the GGN". We will add a clarifying sentence to the paper.
> There appears to be a mistake in equation 26. (...) In general cases, $(A \\otimes B) D (A^T \\otimes B^T) \\neq (AA^T) \\otimes (BDB^T)$, where $A \\in \\mathbb{R}^{p \\times 1}$, $B \\in \\mathbb{R}^{b \\times C}$, $D \\in \\mathbb{R}^{C \\times C}$, and $C > 1$. Since Lemma 3 relies on equation 26, please provide clarification on this matter.
Sure. The crucial detail is that $A \\in \\mathbb{R}^{p \\times 1}$ is a vector, as you correctly define. In this case the equality indeed holds. To see this, we can rephrase your claim to $(A \\otimes B) D \\neq (A \\otimes BD),$ since if this was an equality we could write $(A \\otimes BD) (A^T \\otimes B^T) = (AA^T) \\otimes (BDB^T)$ due to the mixed-product property of the Kronecker product. We will now show that this inequality is in fact an equality. We have
$$(A \\otimes B) D = \\begin{bmatrix} A_1B \\ ... \\ A_pB \\end{bmatrix}^T D = \\begin{bmatrix} A_1BD \\ ... \\ A_pBD \\end{bmatrix}^T = (A \\otimes BD),$$
where $A_1, ..., A_p$ are the scalar components of the vector $A$.
> The case in KFAC-expand and the case in [13] are different. In [13], the authors do not consider the loss with N·R terms. (...)
In [13] the authors derive K-FAC-expand for CNNs by simply explicitly stating the independence assumptions it implies and not by arguing about the exactness in the deep linear case.
We have to distinguish between the setting of K-FAC-expand, i.e. a loss with NR terms, and the approximation itself. As explained in the paper, we can also use K-FAC-expand in the reduce setting and K-FAC-reduce in the expand setting -- which is also what we do in the experiments, since we apply both approximations in the same setting. Crucially, K-FAC-expand was the only approximation known for both settings. However in the reduce setting, K-FAC-expand cannot be motivated by making statements about the exactness in the deep linear case with a Gaussian likelihood. In contrast, we can motivate K-FAC-reduce in this way in the case of a CNN for image classification (assuming average pooling as the aggregation function). We state this in lines 240-242 of the paper, below Proposition 2.
To provide some intuition on why K-FAC-expand and K-FAC-reduce can both always be used for linear weight-sharing layers, please refer to the Python code snippet in Listing 1 in the appendix. Here you can see that the quantities needed to calculate the K-FAC approximation will have the same shape in both settings, allowing us to arbitrarily choose one of the two K-FAC variations to deal with the weight-sharing dimension.
> I noticed the idea in "Eva(...)" is very similar to the KFAC-reduce. (...)
Thanks for the reference. The difference is that the approximation in Eva sums over the mini-batch dimension, whereas K-FAC-reduce takes a sum/mean over the weight-sharing dimension. This is quite different, since Eva does not have a clear motivation in terms of the connection to the exact Fisher/GGN, leads to rank-1 Kronecker factors, and is completely unrelated to weight-sharing.
[1] R. Grosse, J. Bae, C. Anil et al. Studying Large Language Model Generalization with Influence Functions. Preprint (arXiv: 2308.03296), 2023.