On the Parameterization of Second-Order Optimization Effective towards the Infinite Width

We would like to thank the reviewer for the helpful comments and suggestions.

In the experiments, for example in Table 3, there isn’t any standard deviation next to the accuracies.

Thank you for the valuable advice. We have added data with standard deviations for Table 3 in Appendix G.2 (Table S.2). We initially avoided showing the standard deviation to prevent the tables and figures from becoming overly dense and impairing readability. However, since the purpose of Table 3 is the comparison of accuracy, it will be indeed beneficial to include standard deviations. Thus, we have now incorporated them in Table S.2 and it clarifies that the superiority of $\mu$P increases for wider networks. Please note that not all data have standard deviations due to limited time for rebuttal.

The paper mentions some conditions for the stable feature learning that comes from the previous works. However, I think it needs more context for a reader who is not aware of previous work.

Thank you for your kind suggestion. We have added an additional explanation of the conditions to Appendix A.1. We hope you will also refer to Appendix A.2, which describes conditions A.1 and A.2 in more detail.

I was wondering if we change the layer width for different layers what part of your analysis would change?

Thank you for your insightful question. Actually, if all widths approach infinity in the same order, our results remain the same! In more detail, if we define the model width as $M_l = \alpha_l M$ and consider the limiting case where $M$ becomes infinitely large while the coefficients $\alpha_l$(>0) remain constant, the obtained $\mu$P is the same. This is because the $\mu$P is based on order evaluation and the constant coefficients play no essential role. We have added an explanation of this case in the footnote of Appendix A.1.

If our response satisfies the reviewer, we appreciate it if you would consider adjusting their score accordingly.

We would like to thank the reviewer for the helpful suggestions and for correctly acknowledging the strengths of our work.

What do the authors think about apply to other second order methods, such as Newton, Gauss-Newton and generalized Gauss-Newton?

Thank you for the insightful question. In deep learning, for standard training objectives such as MSE loss and cross-entropy loss, the Gauss-Newton matrix is equivalent to the Fisher Information Matrix [1]. Thus, K-FAC, which is an approximation method for Natural Gradient Descent, can also be regarded as an approximation method for the Gauss-Newton method.

Furthermore, we can derive the $\mu$P for the Gauss-Newton method in a similar way to K-FAC and Shampoo. Through some calculations, we found that the $\mu$P for the (layerwise) Gauss-Newton method is the same as in Proposition 4.1. In more detail, $b_l$ and $c_l$ are the same as in K-FAC and damping terms are the same as in Shampoo. Thus, the current work gives a unified perspective covering various individually developed second-order methods. We have included an explanation for this in Appendix A.4.3.

[1] New insights and perspectives on the natural gradient method, James Martens, 2020

Could the authors clarify the meaning of "transferring the learning rate" in Section 5.2 and in Figure 6, and clarify the claims of the corresponding paragraphs?

Thank you for your very helpful remarks. We should emphasize that in the previous work of $\mu$P, learning rate transfer is defined as satisfying both "we can set the optimal learning rate without depending on the order of width" and "wider is better". Since our explanation was insufficient at this point, we have added this definition in Section 5.2.

But, when looking at Figure 6, this is barely true: actually, the transfer is better with K-FAC SP than with muP in the MLP case.

Thank you for the question. In Figure 6, the learning rate does not transfer in SP. More specifically, when training MLP with K-FAC SP, the optimal learning rate for width = 16384 is lower than that for width = 128 (which is marked by star symbols in the figure). Consequently, if we use the optimal learning rate determined for a width of 128 to train a model with a width of 16384, there is a significant drop in accuracy compared to using the optimal learning rate for the latter width. In contrast, in the K-FAC $\mu$P, we can transfer a learning rate, allowing the model with width = 16384 to be effectively trained using the optimal learning rate for width = 128. Furthermore, "wider is better" holds as well.

We would like to thank the reviewer for the helpful comments and suggestions.

Answer to Question

I take the main benefit of the paper to be the actual derived scaling laws ... Does this work introduce any new theoretical results/mathematical techniques of wider interest/applicability...?

As you appropriately understand, the main benefit of this paper is the actual derived scaling laws. However, we think that our theoretical approach also includes some novelty in the theory. First, we provided a unified derivation that is easier to apply to 1-st order graident compared to the previous work and also applicable to our 2nd order gradients. In more detail, as is summarized in the beginning of Section A.2, we showed the proof focusing on the order evaluation of $h_{l,1}$ rather than the kernel value $||h_{l,1}||^2$. Indeed, this does not involve developing new mathematical tools, so from that perspective, the novelty might be low. However, we believe that formulating derivations in a unified manner that is easily applicable to various problems is also a significant contribution in theory.

Second, a technical novelty in applying muP to second-order optimization is the usage of push-through identity(Lemma A.7). This gives a unified derivation over K-FAC and Shampoo. Note that it also covers (layer-wise) Gauss-Newton gradient, which was requested by Reviewer Xrk9 and has been added in Section A.4.3. As far as we know, these 2nd-order methods have been analyzed only in an individual manner. Therefore, it is interesting that we can show new theoretical results and techniques that are widely applicable to these seemingly different methods.

Reply to Weakness

Theory is only given for least-squares losses, although the numerics suggest the results hold for other losses.

Thank you for your valuable suggestion. As you say, the results of Proposition 4.1 in this paper can easily be extended to the cross-entropy loss. Technically speaking, we need to change only the definition of $\chi$ in (S.21) and the coefficient matrix of order 1 appearing in $B_l$ matrices, which does not affect the order of $\delta_l$ and $h_l$. We have added an explanation in Appendix A.5 that the muP is the same among all MSE and cross-entropy losses.

I found the mathematical presentation fairly informal... This is a presentational comment, not a critique of the correctness of the work.

Our focus was the derivation of the actual scaling, so we prioritized a structure that is somewhat informal but more accessible to readers outside of mathematics, rather than adhering strictly to a mathematically formal composition (writing everything in the format of definitions, lemmas, propositions, and theorems). However, we agree that at least we need to clearly mention any assumption or new definition unique to our work. To avoid potential misunderstanding in the main text, we have added references to the assumptions used in the Appendix. We have also added a detailed explanation of the 'valid', as defined by us, in the derivations within the Appendix (Definition A.6).

While the empirical results show merit, their real benefit would only be clear if applied for larger-scale networks than those tested.

We also agree that doing more experiments on large-scale networks will further verify the benefit of this work. We hope that the appearance of the current work will inspire the community and subsequent work will challenge individual large-scale networks including Transformers as is mentioned in Section 6.

If our response satisfies the reviewer, we appreciate it if you would consider adjusting their score accordingly.