On the Parameterization of Second-Order Optimization Effective towards the Infinite Width
Authors: Satoki Ishikawa
Ryo Karakida
Venue: ICLR 2024 poster (2024)
Keywords: Deep learning
Second-order optimization
K-FAC
Feature learning
Infinite width
Maximum update parameterization
Primary area: optimization
TL;DR: We uncover the parameterization of second-order optimization methods that enable feature learning in infinite-width neural networks and demonstrate its benefits.

Abstract: Second-order optimization has been developed to accelerate the training of deep neural networks and it is being applied to increasingly larger-scale models. In this study, towards training on further larger scales, we identify a specific parameterization for second-order optimization that promotes feature learning in a stable manner even if the network width increases significantly. Inspired by a maximal update parametrization, we consider a one-step update of the gradient and reveal the appropriate scales of hyperparameters including random initialization, learning rates, and damping terms. Our approach covers two major second-order optimization algorithms, K-FAC and Shampoo, and we demonstrate that our parametrization achieves higher generalization performance in feature learning.
In particular, it enables us to transfer the hyperparameters across models with different widths.

Published as a conference paper at ICLR 2024
ON THE PARAMETERIZATION OF SECOND-ORDER OPTI- MIZATION EFFECTIVE TOWARDS THE INFINITE WIDTH
Satoki Ishikawa
Ryo Karakida
Department of Computer Science
Artificial Intelligence Research Center
Tokyo Institute of Technology, Japan r riverstcne@ria gsic wwwechhc........................ ac. jp
AIST, Japan
karakida ryoBaist.go......
ABSTRACT
Second-order optimization has been developed to accelerate the training of deep neural networks and it is being applied increasingly larger-scale models. In this study, towards training on further larger scales, we identify a specific parameteriza- tion for second-order optimization that promotes feature learning in a stable manner even if the network width increases significantly. Inspired by a maximal update parameterization, we consider one-step update of the gradient and reveal the ap- propriate scales of hyperparameters including random initialization, learning rates, and damping terms. Our approach covers two major second-order optimization algorithms, K-FAC and Shampoo, and we demonstrate that our parameterization achieves higher generalization performance in feature learning. In particular, it enables us transfer the hyperparameters across models with different widths.
INTRODUCTION
Second-order optimization has been attracting considerable interest in improving the training effi- ciency of neural networks (Amari, 1998; Pascanu & Bengio, 2014). It accelerates the convergence of gradient dynamics (Martens & Grosse, 2015; Gupta et al., 2018; Li, 2017) and can optimize neural networks that are especially hard train due highly distorted parameter landscape (Martens et al., 2018; Zhang et al., 2022a; Beyer et al., 2022). Because of these successes in improving training efficiency, increasingly being applied to even larger-scale models year by year (Osawa al., 2023b; Pauloski et al., 2021; Shi et al., 2021; Zhang et al., 2023; 2022b; Anil et al., 2021).
In general, gradient methods depend on some we need careful HP tuning to achieve better performance. The most trriigtforwwad approach is to train the model multiple times and search for better HP, but for large-scale models, the computational cost is high even for single training, making very challenging to find it. In particular, second-order optimization methods possess not only a learning rate but also a damping term which requires careful tuning (Pascanu & Bengio, 2014; Martens, 2020).
(HPs), including a learning rate, and
How can we obtain quantitative insight into appropriate settings of HP that hold on large scales? Usually, appropriate scale of HPs depends on the model size and we need scale them up down depending on the network width or depth (Park et al., 2019; Iyer et al., 2023; Dinan et al., 2023). One approach to obtaining robust scale of HPs that universally works in large models is to consider a large limit of the model size (Schoenholz et al., 2017; Lee et al., 2018; Xiao et al., 2018; Luo et al., 2021). In particular, to obtain preferable HP that works in first-order gradient descent, Yang & Hu (2021) proposed Maximum Update Parameterization (MUP; μP). They analyzed an infinite width limit of neural networks and successfully identified HPs that prevent the parameter update from vanishing or exploding and ensured stable progress of training, which is independent of the width. The μPP also enables us to avoid the lazy regime, where the model parameters remain sufficiently close to the initialization and allow feature learning at the maximum scale of the update in all layers. Moreover, Yang et al. (2021) have demonstrated that since the juP works effectively towards the infinite width, we re-use the HP including learning rate obtained in a smaller model to the large-scale one. To date, however, these studies have been limited to first-order and entry-wise optimization (Yang & Littwin, 2023).

---

Published as a conference paper at ICLR 2024
In this work, we consider infinite width limit and propose jPP for second-order optimization including two major methods; K-FAC (Martens & Grosse, 2015) and Shampoo (Gupta et al., 2018). Our main contributions can be summarized as follows:
We consider a one-step update of the second-order optimization methods in the infinite- limit and reveal the HPs (i.c., scales of random initialization, learning rates, damping terms) that allow the μP (in Section 4.1). We especially clarify that the stable feature learning requires specific scales of learning rates; K-FAC works for constant learning rates whereas Shampoo requires scaling depending on the width. Regarding the damping terms, we find that classical heuristic scaling of Shampoo satisfies the μP condition while that of K-FAC requires re-scaling (in Section 4.2).
In practice, the last layer's weight is sometimes initialized not by random initialization but by zero. By carefully considering the one-step update, we find that while the zero initialization allows feature learning in the usual first-order gradient, it can cause an approach to the network parameter corresponding to the neural network Gaussian process (NNGP) in the case of K-FAC (in Section 4.3). This can be regarded as a novel implicit bias specific to K-FAC that appears the infinite-width limit.
We empirically verify the effectiveness of our proposed parameterization in the training of various neural networks. In particular, it enables us to transfer optimal learning rates and damping terms from narrow models to wider ones (in Section 5.2, 5.3).
Thus, this work provides quantitative insights that will serve as foundation for scaling second-order optimization towards the learning of even larger models in the future.
2
RELATED WORK
Feature learning: In general, it is highly non-trivial that the large limit of the model allows stable learning under the widely-used standard parameterization (SP). In the infinite width limit, we may have unstable dynamics (i.e., vanishing/cxploding gradient depending on the width) or, more non- trivially, the lazy regime (a.k.a. neural tangent kernel regime) (Chizat et al., 2019; Jacot et al., 2018). While the learning dynamics in the lazy regime progress in stable manner, the parameters remain sufficiently close to the initialization, and the network is essentially approximated by linear model. This means that no feature learning appears, thus there is growing interest in under what conditions feature learning progresses outside of the lazy regime (Woodworth et al., 2020; Geiger et al., 2021; Luo et al., 2021; Bordelon & Pehlevan, 2022). Based on an order evaluation of parameter updates in the middle of training, Yang & Hu (2021) proposed μP for stochastic gradient descent (SGD) that realizes feature learning in the infinite-widthe limit. Some experimental studies demonstrated the utility of yP Yang al. (2021); Vyas al. (2023). The theory is also generalized to entry-wise adaptive gradient methods including Adam (Littwin & Yang, 2023). The second-order optimization does not belong to the class analyzed in these existing studies, and thus the current work is the first to challenge this problem. Note that the second-order optimization in the lazy regime has been investigated by some work (Zhang et al., 2019a; Cai al., 2019; Karakida & Osawa, 2020).
Second-order optimization: The preconditioned gradient can speed up the convergence of training dynamics. Natural gradient descent (NGD) (Amari, 1998) is a classical example of second-order optimization which preconditions the gradient by the inverse of the Fisher information matrix. However, since the computation of the inverse is computationally demanding, some approximation is required. K-FAC (Martens & Grosse, 2015; Grosse & Martens, 2016) is such an approximation for deep neural networks and has been frequently employed for training large-scale models where acceleration is particularly important (Osawa et al., 2023b; Pauloski et al., 2020; 2021; Shi et al., 2021; Zhang et al., 2023; 2022b). Shampoo is another commonly used second-order optimization method (Gupta al., 2018) and achieves fast convergence (Anil et al., 2021). Note that aecond-order optimization methods contain a damping term. Careful selection of such HPs is known to be important for the success of training (Martens, 2020; Zhang et al., 2019h Gao et al., 2020).
2

---

Published as a conference paper at ICLR 2024
3 PRELIMINARIES
This section explains the second-ordes optimization in L-layered fully-connected neural networks:
Wih-- + bt. h1 = ) (1 1....... L),
(1)
where we define weight matrices Wi IR Mi-1, bias terms b, and activations hy E RMM We
set the width of the hidden layer to Mi = M (1 = I,...,L L - 1) for clarity, but this does not lose the generality of the following analysis. of-) is differentiable and polynomially-bounde activation function and theoretical works in uP usually assume either Tanh function or o-GELU function if necessary (Yang & Hu, 2021). Let (æ,, y/) be a pair of input and target training sample. For simplicity, we consider the mean squared loss function for a mie-dimensional target: [()) = 1 -E y - fo [=1 E IR, where 0 denotes a vector of all parameters and fe 7y is the output of the deep neural network. is straightforward generalize the following results to the cases of multi-classes and the cross-entropy loss as mentioned in Section A.5.
3.1 OVERVIEW OF SECOND-ORDER OPTIMIZATION
Second-order optimization is an algorithm that updates parameters by a preconditioned gradient: 0u+1 = 0, - (C(e,)
(10,)) + +I) Ve, C(+,), where is a learning rate, C(0) is the curvature matrix and p is the damping term. Usually, this inverse is computationally demanding and hard to use. Therefore, the following seminal works have introduced smaller preconditioners for each layer and updated rules in matrix form.
K-FAC: Natural gradient descent (NGD) the case where C' is given by the Fisher information matrix. K-FAC approximates this C by the Kronecker product of two matrices (Martens & Grosse, 2015; Grosse & Martens, 2016). Its update rule in the matrix form is given by
Wi.++ W1G ry B: PR, EM w. (A + PA 1) EA
where C eB 1. The preconditioning matrices are given by B E[ō,5; 1 and A, = E (h 1 where 8; = n. and E[|] is an average over training samples. General exponents е and eB are introduced here because some work only considers to use part of preconditioners like (E éB) (1,0) (Benzing. 2022: Amid et al., 2022). The size of damping terms is usually determined in a heuristic manner as mentioned in Section 6.3.
Shampoo: Gupta et al. (2018) proposed the following update rule as second-order optimization;
W1,++ Wt, r1L !) 2/2 T
V ((,)) + PR. -//2
and R. E 1: 8 with 6, V., C. In Shampoo, = 1/2 is applied. If we neglect the non-diagonal entries of the preconditioners, this is similar to Adam and AdaGrad.
where L, 16
3.2 ABC-PARAMETERRIATTOON
ABC- ABC-paraameterization scales parameters by the width as follows (Yang & Hu, 2021):
Wi = w//M", wi ~ NO, o²2 /M22h, T = n//
The yP an abc- abbeparraettrizzaiinn that induces feature learning in every layer in a model with infinite widths. In short, the previous work characterizes, feature learning by
Ah := hie huo =(1)',
(5)
where (--) denotes the Big Theta notation for the order evaluation with respect to the width. Note that for the lazy regime, have Ah, o(1) in hidden layers. The previous work found that the feature learning (5) appears for some specific conditions. In particular, the following condition of Wi updated maximally plays fundamental role:
AWth (().
(6)
Precisely speaking. this feature learning (Definition H.9) especially satisfying a stability condition (Thcorem H.6) Yang & Hu (2021). Note that as the previous work, the last equality represents the coordinate size where ((M") means / /l @(M)) for E R
3

---

Published as a conference paper at ICLR 2024
Table 1: μP for K-FAC and Shampoo.
Input weights & all biases
Output weights
Hidden weights
SP
b = 0, 0
b = 1/2. 0
b=1/2, c= 0
SGD (e =
b = 0, = -1
b= 1=
b-1/2. 0
Shampoo (e
b 0, c = -1/2
b= 1, 1/2
b= 1/2, c 0
K-FAC (CA,B =
b b= 0,c=0
b= 1,
b= 1/2 0
Yang & Hu (2021) analyzed feature learning in the one-step gradient under this condition. Note that they also obtain mathematical expression of forward and backward propagated signals for the first-order gradient at general time steps by the Tensor Program. For the derivation of the μP,, we focus only on the first one-step gradient as is explained in Section A.2 of Appendices.
In the following sections, we will set = 0 for all layers in the same way as in the first-order case (Yang et al., 2021). The previous work has demonstrated that the μPP of the first-order gradient is scale-invariant to the constant shift (a, b, c) t (a, + (k. -k, -kk). As we will show later, the μPP of coddrr optimization also has this indeterminacy and we can eliminate it by aj
PREFERABLE SCALING OF HPS IN SECOND-ORDER OPTIMIZATION
4.1 μP FOR SECOND-ORDER OPTIMIZATION
In this section, we derive the μP by considering the first one-step update of second-order optimization. We suppose that the damping term px (X {A, B, L, R)) satisfies
px /xx /
with a positive constant p'x For simplicity, suppose a one-step update from the initialization where the gradient and all preconditioners are computed on the same samples (Assumption A.5). We also suppose common assumptions used in yP for the first-order gradient (Assumptions A.),A.4). When the eigenvalues of the preconditioning matrices have equal width-dependent scale the damping term (e.g., PA e([[A]))) we say that the second-order optimization is valid (Definition A.6). Then, obtain the following.
Proposition 4.1 (uP of second-order parameterization) Consider the first one-step update of K-FAC and Shampoo in the infinite width limit. The second-order optimization becomes valid for
1=1
d^ = [-1 1<lS 1-1
00 l=
dp,
dia dR,
1 <<l< (8)
[1 1 1<l<L'
1=
It admits the μP for feature learning
0
1=1
EB - 1 1=1
b
1/2
1<1<
q
en es l<l<,,
1=L
I e EA 1=b
where we set a₁ and setting C = cB = corresponds Shampoo.
Rough sketch of derivation. The detailed and comprehensive derivation is presented in Section A. Here, us briefly explain the derivation of yP for K-FAC. The μP has two conditions to be satisfied (A.1, For an infinitesimal one-step update, these conditions have explicit and tractable expressions as described in Section A.2. They are applicable to both first-order and second-order optimization methods. To check the conditions,
use the push- yush-through identity:
AW1h--1= = 1/M 2a+0 B + PBl) - 8diag(x) (A -
1/M22++0888 5 5 PD diag(x) Ththhhhhh -
(10)
where we omitted the layer index and X means an error vector. In the second line, the size of the inverse matrix is independent of the width and this expression enables us to carry out the order evaluation of AW1h-11 Then, juP's conditions become
(2eg l) bL) 0, 20 + q + CA + cg - (2eg - l)(ar + by) = O,
201 + C1 ee
(11)
(12)
20 C7 + CA = 0, ar + by 1 0.
(13)
4

---

Published as a conference paper at ICLR 2024
KFAC SP Heuristies
KFAC muP Heuristics
KFAC muP Restaled
Shampoo SP
Shsmpoo mup
2a
4h (MLP))
29
2-2
2
2
29
2/2
29
2
2
2
2/2
25
4h (Myytiee
29
2-5
2
2
2 width
24
2 widch
2
2
2
2 width
2
width
width
Input layer
output layer
hidden loyers
Figure 1: /P achieves feature learning across the width. In SP (Pytorch's Default), Ah; in each layer exhibits dependence on the width. For K-FAC, the default setting of the damping (heuristics) does not satisfy the condition of μP and we need to utilize the rescaled one is explained in Section 4.2. We train 3-layer MIP with CIFAR in the first line and Myrtle-5 with CIFAR1O in the second line. This result does not depend on exponential moving averages or activation (Appendix.D).
Figure 2: The obtained parameterization consistent throughout the training. The order the curvature matrix in K-FAC does not change with time. The input layer is pro- portional M whereas the output layer is proportional to 1 1/M, which is natural order terms of the μP of SGD. We trained a 3-layer MLP on FashionMNIST dataset.
imput layer
output layer
width
Grad Norm Ratio
21
218
Epurch
Epooch
Fixing a constant shift by setting a; = we obtain the result.
Here, let explain some rather technical details of the result. First, note that we derived the pP from the one-step update. This is the same as in the original work on jPP where the one-step update determines the μP for the whole (inductive) t-step updates. In fact, effectiveness of yP for / > 1. Second, we focus on the case where the preconditioning matrices become valid. Even if p takes much larger value than these matrices, the gradient may realize the feature learning under an appropriate scaling because it reduces to the first-order gradient. However, such unusual switching between the first and second-order optimization outside the scope of the current work. Thus, we refer to the setting of this proposition as jP of second-order parameterization.
empirically confirm the
Table summarizes the μP for K-FAC and Shampoo. by is common characteristic of μPP for all methods whereas and depend on e. One interesting point is that just by setting b, = 1, K-FAC with μP works for q 0. That is, K-FAC does not require sealing of the learning rate in contrast to the first-order gradient (SGD), which requires depending on the width. Moreover, (9) indicates that this is unique to the K-FAC (e 1). other words, K-FAC's preconditioning effectively achieves the JP's scaling of learning rates of the first-order gradient. Note that we can also extend μPP to the Gauss-Newton method without using the K-FAC approximation (Appendix A.4.3). As side note, we also summarize the parameterization for the lazy regime in Section I.
Figure empirically confirms that μP realizes the feature learning (5) in the training of MLP and Myrtle-5. The order of the features was kept during the training, as is shown in Figure 2. Figure 1 also justifies the μP in CNN. In the CNN model, width represents the number of channels (Xiao et al., 2018: Yang et al., 2021). Although /P is derived the MIP and K-FAC for CNN includes
5

---

Published as a conference paper at ICLR 2024
SGD (Train Acc)
KFAC (Train Aec)
Shampoo (Train Acc)
84
logg 6 (ontpoot layer)
82
80
80
78
76
76
136
123 136
100
Iteration
Iteration
Iteation
Figure 3: K-FAC converges to the NNGP solution when the variance of the last layer is close to zero. When by in the last layer is increased (other parameters are fixed to yP), K-FAC can converge to the NNGP solution in one step. Therefore, when by, is increased, it converges to kernel solution, which limits b which feature learning can occur.
an additional approximation addition K-FAC for MIP, we empirically observe that the same parameterization as MLP is also true for CNN.
4.2 JUSTIFICATION OF DAMPING HEURISTICS
Implementations in practice often adjust the damping term of K-FAC using the following heuristics:
PA: 0=-,(=1//B))
M
O(V M. M2(az
(14)
M
tr (r)
that is ((M) for the {a, b, c} given in Proposition 4.1. This heuristics is consistent with the valid damping scales (8) in hidden layers but not in the input and output layers. This causes Ah; to decay when using damping heuristies even a, b, e are set to μP settings. See Section A.6 for more details. To overcome this problem, we propose to use the following rescaled damping satisfying the valid damping scale in μP:
Rc == p' hr
PB
Re = ir(d, 5z).
(15)
This rescaling is useful because it enables explicitly expressing the dampings of all layers in unified manner. Furthermore, it provides the heuristic scales the preconditioners (i.e., trace) as proportional coefficients which expected ensure stable learning dynamics. If damping is set consistent with μP, Ah; neither decays nor grows with respect to width as shown in Figure (3rd column).
In contrast to the K-FAC, we found that standard heuristics of the damping, i.e., a constant multiple of the largest eigenvalue of L,, RG is consistent with the μP in Proposition 4.1 and requires no modification towards the infinite width. See Section A.6 for more detail.
4.3 IMPLICIT BIAS OF K-FAC TOWARDS NNGP ZERO INITIALIZATION
Up to here, supposed the most common setting where the output weight is given by random (Gaussian) initialization. However, some recent implementations utilized zero initialization on the head (Yang et al., 2021; Wightman, 2019; Botev & Martens, 2022). This corresponds to by. cc, which also induces feature learning after the second step in SGD. The reason why feature learning be explained as follows. When the last layer is initialized with zero, the
occurs even at weights after 1-step SGD update are
W1.== = mhh,0y (l= WLo (1 L).
(16)
Since WT,t=1 = (1/MM), the weights after 1-step update can be regarded as weights initialized by IPP. Thus feature learning begins from the second step. This also true for by > However, in K-FAC, when the last layer is initialized with zero, single update results in the weight:
WrA- = nh 1,0h-1101011100 + pa hooy
L), Wi < L).
(17)
Interestingly, this represents an NNGP solution (Lee et al., 2018). When the model can deviate from the NNGP solution, feature learning begins in the training process, similar to SGD. However, when

---

Published as a conference paper at ICLR 2024
Table 2: Test accuracy for differ- ent batch sizes shows the consis- tent results (3-layer CNN on Fash- ionMNIST. The whole dataset size
Batch Size
Optimizer
16
64
256
1024
0.5
82.60
80.91
78,86
74,50
66.83
SGD
1.0
83.62
83.61
82.10
77.99
73.40
set to 1024). The case of br. consistently performs better than
64.0
83.98
83.82
82.60
79.53
74.63
0.5
83.18
84.17
83.75
84.07
80.25
by = 0.5 and 1y = 64 in K- F-FAC
K-FAC
1.0
83.84
84.16
84.29
84.33
83.21
but not in SGD, Learning rates are
64.0
81.56
82.72
82.63
79.51
75,37
tuned for each hatch size.
Analogz Ace (Word2Vec)
40
width
width
Thi Loo
20
10
5000
5000
10a
10
10a
10a
10
10²
Beeation
Beration
Iteracionn
Iterattion
Figure 4: Wider models learn well under jP throughout training. Using μPP training proceeds equally across widths. In μP, the loss is lower for wider widths throughout training. (Left) We trained CBOW on WikiText2 by Shampoo with various widths. (Right) We trained ResNet18 on CIFAR100 by K-FAC while increasing the number of channels from to 16.
the batch size is close to the full batch size or the learning rate is set considerably low, moving away from the NNGP solution becomes challenging and does incur feature learning. Since K-FAC is often used for large batch training, we must be careful about its behavior on by >> 1.
Figure confirms this phenomenon. In this experiment, we trained a 3-layer CNN on FashionMNIST with MSE loss. For simplicity, the dataset size is reduced to 1024 and we are training in full batch. After sufficient learning. the training accuracy is highest only for by fixed at the initial kernel solution when by >> 1, we can observe that accuracy at by. >> is lower than by = However, SGD and Shampoo achieve almost the same accuracy for b, and by > 1. Table 2 shows that when by >> (br 64), its accuracy while there no decrease in accuracy by setting b, = 64 SGD. In addition, b, (nP) achieves higher accuracy than by 0.5 (SP) across different batch sizes. We empirically observed that the difference in accuracy between by = 1 and by = 64 decreases as the batch size decreases. The same behavior also be observed when using cross-entropy loss (Appendix.F.) F.2..
Since the parameters are
K-FAC is consistently lower than by, = 1
Note that our purpose is not to fully deny the current usage of zero initialization relatively finite neural networks. Our finding claims that the current default settings do not necessarily work well with large models, and careful attention is required.
5 EXPERIMENTS
5.1 μP IN WIDE NEURAL NETWORKS
μP can invoke feature learning equally across width. This enables a wider model performs better when trained in μP and given the same HPs (Yang et al., 2021). This section demonstrates that the above statement is also true for second-order optimization.
This is toy model used in the previous work of pP (Yang & Hu, 2021) to check the advantage of feature learning. We train CBOW on Wikitext2 dataset with Shampoo and evaluate its embedding vectors by a Word Analogy task2. Figure 4 shows that in SP the accuracy decreases as the width increased. In the infinite width limit of SP, the embedding layer is fixed at initialization, and the accuracy does not increase from the churn rate. However, in pP, increasing the width does not
Word2Vee:
2K-FAC is very sensitive to by, and found it difficult optimize word2vec in K-FAC. This could be explained by Section 4.3. Its details descrihed in Appendix E.5.
7

---

Published as a conference paper at ICLR 2024
VGGT9 (CTOO)
ResNet18 (CT00)
ResNer50 (INet)
K-FAC
Shampoo
K-FAC
Shampoo
K-FAC
width
SP
μP
SP
μP
SP
μP
SP
uP
SP
μP
1
64.96
+0.31
63.86
-0.28
66.81
+0.23
66.42
+0.03
61.94
+0.17
2
72.08
+0.65
70.55
-0.66
72.16
+0.28
69.99
+0.22
62.00
+0.12
4
76.38
+0.22
74.61
+0.80
74.38
+0.88
73.85
+0.42
75.64
+0.26
8
78.27
+0.31
76.83
+0.65
74.96
+1.98
76.11
+1.04
78.63
+0.13
16
78.11
+0.56
74.26
+4.00
77.91
+0.61
Table 3: Test accuracies with different widths. JoP has a higher accuracy compared with SP in models with large widths. The learning rate set slightly small to enlarge the effect of infinite width. Missing points are owing to the too-long computational time. The results for ResNet18 (CIFAR1O) and ResNet50 (CIFARION 100) are in the Appendix.G.I.
decrease the accuracy. These results highlighted that the pP for second-orde optimization enables feature learning even in infinite- wiith models.
ResNet: We evaluate the test accuracy for SP and uP with VGG19 on CIFARI00, ResNet18 on CIFAR and ResNer50 on ImageNet. The number of channels in the middle layers is scaled from the original mod- els. Original models correspond to width 4. HPs are fixed regardless of width. Table indicates that the accuracy can consistently im- proved by μ.P. Figure shows the learning curves for SP and μP with
Test Acc
78
SP
different widths. In SP, the test loss for width = 16 (wider model) is not necessarily higher than that for width 1 (narrower model). However, in μP,, wider models consistently match or outperform narrower models. This is more evident in the early stages of training.
muP
74
2-4
LR
Figure 5: μPP consistently achieves higher accuracy for various learning rates. (ResNet50 on ImageNet)
The advantage of μP generally holds regardless of a fixed learning rate. Figure 5 shows that regardless of the fixed learning rate, jP achieves higher test accuracy compared with SP. As a side note, the difference between SP and μP particularly large when the learning rate small.
5.2 LEARNING RATE TRANSFER
jPP can transfer learning rate for models with different widths. Learning rate transfer, defined in previous study, means that we can set the optimal learning rate to remain constant relative to the order of width, and that "wider is better" holds at this optimal learning rate. We confirm this across three different architectures, MLP, CNN, and ResNet in Figure 6. In the SP for SGD, the optimal learning rate decreases as the width increases. Similarly, when the damping heuristics of Equation (1) are used for SP in K-FAC, there is a shift in the optimal learning rate. These shifts in the optimal learning rate jPP, which is a stable parameterization.
are caused by lack of stability these parameterizations. the optimal learning rate is fixed with respect to width. One can also confirm that the wider is better for fixed HP. For instance, if MI.P trained with SP in K-FAC, test accuracy at width 16384 lower than width = 512. In contrast, if one uses jP for training, accuracy consistently increases with the width; thus, ,P works as better parameterization to promote feature learning.
5.3 DAMPING TRANSFER
The damping term another HP that requires careful tuning in second-order optimization. μPP, we can re-use the damping obtained in small model to the large-scale one. As shown in Figure 7, if one optimizes CNN with K-FAC using the damping heuristics (Eq. 14) SP (i.e., the default setting), the optimal damping value increases as the width increases. In contrast, when damping is consistent with JPP can be transferred across the widths.
CONCLUSION
We proposed a desirable parameterization inspired by μP for second-order optimization to achieve stable feature learning in infinite-wice neural networks. As the computational resources en-
8

---

Published as a conference paper at ICLR 2024
SGD SP
SGD mup K-FAC SP(leurristie) K-FAC SP
K-FAC muP
width
256
60
1024
2048
Test
40
4096
8192
-10
10
-10
-10
(CNN)
80
60
16
32
64
256
(ReaNett
85
width
80
75
-5
5
0
5
-
logg
LR
loga
LR
LR
Figure 6: pP allows the learning rate (1)) to transfer across widths. Using pPP, one transfer the learning rate concerning width. KFAC with SP, the heuristic damping obstructs the transfer of learning rates. With pP, the transfer succeeds and the wider models perform better They trained by MSE loss with 1024 samples. This learning rate transfer also holds for the full dataset (Appendix.F.) F.3). In addition, pP enables this transfer Shampoo (Appendix.F)) and FOOF (Appendix.F.2) as well.
Heuristies (SP) Rescaled (SP) Rescaled (muP)
Figure 7: Under the rescaled damping of
90
16 K-FAC, the optimal damping fixed re- gardless of width. 3-layer CNN is trained on FashionMNIST by MSE loss (Full dataset size). The same transfer can also be observed for MLP (Appendix.H.2),
Test Ace
80
64
60
-20
-20
0
20
logg damping log, damping log, damping
riched more and the model becomes larger, there is no guarantee that the current default settings of initialization, learning rate, and damping terms work well in a similar way to smaller models. In pursuit of such training of wider models, we empirically confirmed the effectiveness of the μP in commonly used second-order methods and the necessary modification of the HPs. is also one of the advantages of the proposed parameterization that when using μP, HPs tuned in smaller model can be transferred to a larger one.
Limitation and future direction. The current work and the framework of μP focus on the wide neural networks as large-scale models. It would be interesting to investigate the parameterization effective for the depth. Some studies have scaled weight initialization with respect to depth (Shoeybi et al., 2019; Radford et al., 2019). Transformer is a widely used architecture in modern deep learning; however second-order optimization has not yet been established well, and we have still limited empirical observation in both pertaining (Pauloski et al., 2022; Osawa et al., 2023b) and fine-tuning (Ding et al., 2023). Moreover, the finite width effect may non-negligible the Transformer; thus, will need more careful study focused on it (Dinan et al., 2023; Wortsman et al., 2023). From theoretical perspective, mathematically rigorous evaluation of feature learning in general steps is curious well. Tensor Program (Yang & Littwin, 2023) is strong candidate for such evaluation, although is currently not applicable to the second-order parameterization and analytical evaluation still limited even in the first-order gradient. We expect that our proposed parameterization and empirical observation will serve as a foundation for providing deeper insight into second-order optimization both in theory and practice.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The authors thank Rio Yokota for his helpful comments and acknowledge the funding support from JST FOREST (Grant Number: JPMJFR226Q) and JSPS KAEKENHI (Grant Number: 22H051 16, 23K16965).
REFERENCES
Amari. Natural Gradient Works Efficiently in Learning. Neural Computation, 10:251-276, 1998.
Shunichi Amari, Jimmy Ba, Roger Baker Grosse, Xuechen Li, Atsushi Nitanda, Taiji Suzuki, Denny Wu, and Ji Xu. When does preconditioning help or hurt generalization? In International Conference on Learning Representations, 2021.
Ehsan Amid, Rohan Anil, and Manfred Warmuth, Locoprop: Enhancing backprop via local loss optimization. In International Conference on Artificial Intelligence and Statistics, pp. 9626-9642.
PMLR, 2022.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable Second Order Optimization for Deep Learning. arX1v:2002.09016 2021.
Frederik Benzing. Gradient descent neurons and its link approximate second-order optimization. In International Conference on Machine Learning, pp. 1817-1853 1853. PMLR, 2022.
Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov. Knowledge distillation: good teacher patient and consistent. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10925-10934, 2022.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks, Advances in Neural Information Processing Systems, 35:32240-32256. 2022.
Aleksandar Botev and James Martens. google-des eeoogee nd/ -facjjax
KFAC-JAX, 2022. URL httpt:///gtthbu....... com/
Tianle Cai, Ruiqi Gao, Jikai Hou, Siyu Chen, Dong Wang, Di He, Zhihua Zhang, and Liwei Wang. Gram-gauss-newton method: Learning overparameterized neural networks for regression problems. 1rX.v: 2019.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization. arXiv:2304.02334, 2023.
Ning Ding, Qiaosen Wang, Yulin Chen, Pengjun Xie, Zhiyuan Liu, Hai-Tao Zheng, and Maosong Sun. Applying second order optimization deep transformers with parameter-efficien tuning. In OpenReview, 2023.
Kai-Xin Gao, Xiaolei Liu. Zheng-Hat Hai Huang, Min Wang, Zidong Wang, Dachuan Xu, and F. Yu. A trace-restricted frottcker-factored approximation to natural gradient. In AAAI Conference on Artificial Intelligence, 2020.
Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Perspective: A phase diagram for deep learning unifying jamming, feature learning and lazy training. arXiv:2012.15110, 2020.
Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physics Reports, 924:1-188 2021.
Roger Grosse and James Martens. A Kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582, PMLR, 2016.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
10

---

Published as a conference paper at ICLR 2024
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dongsung Huh. Curvature-corrected learning dynamics in deep neural networks. In Proceedings of the 37th International Conference Machine Learning. PMLR, 2020.
Gaurav Iyer, Boris Hanin, and David Rolnick. Maximal initial learning rates deep relu networks. In International Conference on Machine Learning, pp. 14500-14530. PMLR, 2023.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Ryo Karakida and Kazuki Osawa. Understanding approximate fisher information for fast convergence of natural gradient descent wide neural networks. Advances in neural information processing systems, 111100991-90011 2020.
Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Universal statistics of Fisher information in deep neural networks: Mean field approach. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 1032-1041. PMLR, 2019.
Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Pathological spectra of the Fisher information metric and variants in deep neural networks. Neural Computation, 33(8):2274-2307 2021.
Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural nehworks: Tricks of the trade, pp. 9-50. 1998.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. ICLR 2018, arXiv:1711.0016 2018.
Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE transactions on neural networks and learning systems, 1954): 1466, 2017.
Etai Littwin and Greg Yang. Adaptive optimization in the SünftyS-width limit. In The Eleventh International Conference on Learning Representations, 2023.
Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer ReLU neural networks -nfinite-width limit. The Journal of Machine Learning Research, 22(1):3327-3375 2021.
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-ffcctored approximate curvature. In International conference on machine learning, pp. 2408-2417 PMLR, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Cronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representa- tions in vector space. XXXiv:/301.377 2013.
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler. ASDL: A unified interface for gradient preconditioning in PyTorch. a3X5v: 04684, 2023a.
Kazuki Osawa, Shigang Li, and Torsten Hoefler. Pipefisher: Efficient training of large language models using pipelining and fisher information matrices. Proceedings of Machine Learning and Systems, 5. 2023b.
Daniel Park, Jascha Sohl-Dickstein Quoe Le, and Samuel Smith. The effect of network width stochastic gradient descent and generalization: an empirical study. In International Conference on Machine Learning, pp. 5042-5051. PMLR, 2019.

---

Published as a conference paper at ICLR 2024
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. ICLR 2014, arXiv:1301.358,, 2014.
J Gregory Pauloski, Zhao Zhang, Lei Huang. Weijia Xu, and Ian T Foster. Convolutional neural network training with distributed K-FAC. SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-12. IEEE, 2020.
J Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, lan Foster, and Zhao Zhang. Kaisa: an adaptive second-order optimizer framework for deep neural networks. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14, 2021.
J Gregory Pauloski, Lei Huang, Weijia Xu, Kyle Chard, Ian T Foster, and Zhao Zhang. Deep neural network training with distributed K-FAC. JEEE Transactions on Parallel and Distributed Systems, 33(12):3616-366 3627, 2022.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Demmark, 7(15):510, 2008
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ blog, 1(8):9, 2019.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dicksteii Deep information propagation. ICLR2017 arXiv:1611.0122. 2017.
Vaishaal Shankar, Alex Fang. Wenshuo Guo, Sara Fridovich-Keil Jonathan Ragan-Kelley, Ludwig Schmidt, and Benjamin Recht. Neural kernels without tangents. In International conference on machine learming, pp. 8614-8623, PMLR, 2020.
Shaohuai Shi, Lin Zhang, and Bo Li. Accelerating distributed K-FAC with smart parallelism of computing and communication tasks. In 2021 IEEE 41st International Conference on Distributed Computing Systems (ICDCS), pp. 550-560. IEEE, 2021.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. 10X.::9009000553 2019.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learinng networks are consistent across widths at realistic scales. arXiv:2305.18411, 2023.
Ross Wightman. Pytorch image models.
com/
pytorch- ietrhh--imme-mmddeee 2019.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pp. 3635-3673. PMLR, 2020.
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities. urXiv:2309.14322, 2023.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel S Schoenholz, and Jeffrey Pennington. Dynamical isometry and mean field theory CNNs: How to train 10,000-layer vanilla convolu- tional neural networks. In Proceedings of International Conference Machine Learning (ICML), pp. 5393-5402, 2018.
Greg Yang. Tensor programs II: Neural tangent kernel for any architecture. arXiv:2006.14554 2020.
Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-widtt neural net- works. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 11727-11737. PMLR, 2021.
12

---

Published as a conference paper at ICLR 2024
Greg Yang and Etai Littwin. Tensor programs IVb: Adaptive optimization in the Solnite-widthe limit. arXiv:2308.01814, 2023.
Greg Yang. Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Guodong Zhang, James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for over-parameterrize neural networks. In Advances in Neural Information Processing Systems (NeurIPS), pp. 8080-8091, 2019a.
Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. In International Conference Learning Representations, 2019b.
Guodong Zhang, Alcksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. In International Conference on Learning Representations, 2022a.
Lin Zhang, Shaohuai Shi, Wei Wang, and Bo Li. Scalable K-FAC training for deep neural networks with distributed preconditioning. IEEE Transactions on Cloud Computing, 2022b.
Lin Zhang, Shaohuai Shi, and Bo Li. Accelerating distributed K-FAC with efficient collective communication and scheduling. In IEEE INFOCÓM 2023-IEEE Conference Computer Com- munications, pp. 1-10. IEEF, 2023.
13

---

Published as a conference paper at ICLR 2024
Appendices
A DEVIATION OF MUP
In this section, we first explain basic conditions and assumptions for μP in Section A.I. Next, we provide a derivation of μP for the first-order gradient in Section A.2. This deviation is based one-step update and its perturbation, which is easily generalized to the cases of K-FAC (in Section 1.4.1) and Shampoo (in Section A.4.2).
A.1 SETTING OF pP
Here, let overview the conditions of μP proposed in Yang & Hu (2021), As we described in Section 3.2, the previous work defined the feature learning by
Ahi := ht, - h4,0 = (
(S.1)
Note that in the appendices, we avoid using bold fonts of vectors and matrices order to prevent messy appearance. Intuitively speaking, the activation varies with the constant order to generate features in each layer. This is crucial for distinguishing the learning from the lazy regime (Jacot et al., 2018; Chizat et al., 2019). In the lazy regime, the model can converge to global minimum but the trained parameters remain sufficiently close to their initial values and the model is equivalent to a linearized model. In other words, changes in the hidden layer are infinitesimal, i.e., Ah Af = ((1) ) in the infinite width limit. Thus, the condition S.1 is required to avoid the lazy regime. For clarity, we set the width of the hidden layers the same and take the infinite width limit 2:
while
Mi M, M
11=1,...
(S.2)
To realize the feature learning, they proposed to the abe- parameterization satisfying the following two conditions:
Condition A.1 (W₂ updated maximally).
AW 0(1),
(S.3)
where AW == Wi,- W2o-
This first condition plays fundamental role in feature learning. As described in Section H.7 of Yang & Hu (2021) and we will explain more detail in the next subsection, this condition naturally appears when one considers the infinitesimal change of the parameter (i.e., 0,)). The previous work also requires the following condition for pP:
Condition A.2 (W, initialized maximally).
W W1oAAm--1==00)
(S.4)
This condition plays fundamental role for determining the variance of the weights in the last layer. Combining this condition with Condition A. makes the output change of order 1, that is, Jt - Jo = AWr.thr.-1. + Wr.0Dhr.-1.e = O(1). The yP is derived from these two conditions. Subsequent studies have empirically verified that the successfully facilitates feature learning in various settings of actual training (Yang et al., 2021; Vyas et al., 2023).
Yang & Hu (2021) also assume Assumption A.3. 01,0 14,0 = @(1) (K< D), So =o u1,0 = 0(1).
The input and output of the activation function are assumed to be of the order of 1. This assumption has been commonly employed in the training of neural networks (LeCun et al., 1998). Additionally,
OOur analysis remains applicable even if the widths of different layers vary, long all widths approach infinity in the same order. More precisely, even if the network width is given by Mj = M,, with each > 0 heing constant, the μP remains the same in the infinite width limit of M
14

---

Published as a conference paper at ICLR 2024
the μP setting allows that the network output f can be close to zero at initialization. This assumption immediately leads to
a1 bj 0,
(S.5)
+ b, 1/2 (1<|<L)
(S.6)
a7 + by 1/2.
(S.7)
as is described in Theorem H.6 of Yang & Hu (2021). We also have Assumption A.4. For feature learning, we suppose or. + by. 1/2.
This is the Assumption H.23 of Yang & Hu (2021). This assumption used to avoid technical subtleties potentially caused by az + bb. 1/2: for random initialization, fo goes to for az. + by, > 1/2 while becomes Gaussian process with a non-zero variance for ag + by = 1/2. Roughly speaking, this assumption ensures that backpropagated error converges deterministic constant (y - fo) - y in the infinite width limit and makes the derivation more clear by avoiding potential correlation between the error and trained features.
A.2 PP's s CONDITIONS OF OTE-STEP UPDATE AND PERTURBATION
As preparation for the second-order optimization, let explain the pP's conditions written in an infinitesimal one-step update. The essentially same perturbation argued in Section H.7 of Yang & Hu (2021), but their explicit evaluation focuses the development of feature kernels, i.e., ||//,1 and the precise evaluation this value requires much- sauchhhomplicatee argument. Since the current work interested only in simpler problem on the order evaluation of hu,1 it will be more informative and clearer to show the perturbation argument specific to h4,1 directly. As you can see below, this clarifies an explicit connection between feature learning and Conditions A.1 A.2. Express the first one-step update of the weight by
AWL1 22 Gi,
(S.8)
G::= == PAV w/ CoPB,
(S.9)
where PA and Pe are preconditioners. Note that Vw, Co depends only on the initialization (t 0). We have
W1th1-111 = (W1,0 + M2a+ta G1) G))(u-111()))
(S.10)
where we remarked the dependence of (Wi,1...,, W111,) on v by u1-1, (7).
If the derivative of the one-step update with respect to the learning rate is (1), this implies that feature learning appears. While Yang & Hu (2021) explicitly evaluated 0, ||/ II y= let us
evaluate here the following quantities which have an explicit expansion by the chain rule with respect to ro:
1
0z B. =((w) Maatta G
Σ
v(lkk)
(<< L).
(S.11)
Duik
fi= f 10== M2a G1 W
(S.12)
where
ull = 2 kt. GI-shi---200
(S.13)
We can find connections between 15.11,S.122 and Conditions (A.1, A.2) as follows:
1
(WW, M3ar+a M G
L),
(S.14)
and
tW Wr.
(S.15)
That is, under the perturbation with small v/,, the conditions reduce
15

---

Published as a conference paper at ICLR 2024
Condition A.1'
M2a+0 11h ((1)
(S.16)
Condition A.2'
W WL,0O 1/=00 (())
(S.17)
After all, (S.16,S.13 17) under the perturbation. If Conditions A.1' and A.2' hold, it implies (S.11, S.12) of ((1) and then the feature learning (S.1) appears. One say that the uP is the "maximal" update because induces all layers trained as much as possible.
that the Conditions of μP (A.1.A.2) reduces to explicit representations
A.3 CASE OF FIRST-ORDER OPTIMIZATION
As an exercise for preparing second-order optimization, we first evaluate Conditions A.1' and A.2' for the first-order gradient with the preconditioners PA Pn I. The second-order cases shown later.
On Condition A.1'. To avoid cumbersome notation, we omit the initialization index as long as it doesn't t lead to confusion. We can express the gradient G by
G dydiag())4
(S.18)
where forward and backward signals are given by matrix forms h4 81 E RMxx Note that the backward signals are given by the chain rule:
8 = d' (mi o (W (L<I-1),
(S.19)
5y = Cn
(S,20)
where we defined en == (1, 1) IR". We also introduced an error vector
x y f R
(S.21)
As is pointed out in the previous works, the important point of Gjh1-1 (in other words, AWth-11 is that G2 and h not independent. We have
G = (4)T
(S.22)
Here, let us introduce
KAA : NTTTT//MM KP = 8, 6r.
(S.23)
For the random initialization, these two matrices are well known as intermediate components of the neural tangent kernel, as discussed in Jacot et al. (2018); Yang (2020); Karakida & Osawa (2020). They are also known as essential components of the Fisher information matrix, which characterizes the geometric structure of the parameter space (Karakida et al., 2019: 2021). Assume that the Then, in the infinite width limit,
activation function and its derivatives are colynomially-boundee these two matrices become deterministic constants independent of the width. In more detail, the tensor program ensures almost sure convergence of the moments composed of the forward and backpropagated signals (h; and 6,) (Yang, 2020). The l-layer component of the neural tangent kernel is given by K", K%2 where denotes the Hadamard product. The kernel (K() ) is @(1). Since we suppose the abe-parameterization, - @(W,) @(1 (b.)). This leads to
(KB B mm = O (11/M2((tttbu--))
(S.24)
for < 1 < L, that is, the coordinate size is given by
5 θ(1 0(1/A/laztbb))
(S.25)
After all, we have
Og (AW,, 1/=0 M2ar+a--1
-
(S.26)
16

---

Published as a conference paper at ICLR 2024
with
201 + (al bL) (1 1),
2a + q 1 + (az + bb) (1 (1<l< 1), 77
(S.27)
20,, Cr. - ()=
where we used the fact that X is constant vector of (1) in the infinite width limit under Assumption A.4. Compared to r. r1 does not include -1 because M, ( and Ao (1/MM) The
case of rr. does not include (z, + by. because 6,. (1).
On Condition A.2'. We can represent Eq. (S.17) by
0, (Wr.o
WL,0 (0' (uL--) (B, -o (AW
(S.28)
Note that
Wr,0lø' o 0. (AWL-1,hh-22 =
1
- Meal-1teL--1 G GL-1h
(S.29)
- eM ( ( (L-
(S.30)
Because 5 has the coordinate size of Eq. (S.25) and the product with eM means the summation over M, we obtain
8y
6(1/M)
(S.31)
Finally, from Eqs. (S.27) and (S.31), the pP is given by
20 cı + (ar. + br.) = 0,
2a + q + (at b1) O,
(S.32)
20z + cl 1 = 0,
0r. br. O.
A.4 CASE OF SECOND-ORDER OPTIMIZATION
What we need to do here is to find the abc- parameterization satisfying Conditions A.1' and A.2' for the preconditioners PA and PB (S.9) of the second-order optimization methods. We assume following:
Assumption A.S. The gradient of the loss V C, preconditioners Pa and PB are calculated on the same batch of training samples the first one step.
Furthermore, we consider the situation where the second-order optimization is valid, defined as follows.
Definition A.6 (Valid second-order optimization). Suppose that each preconditioning matrix X has damping term Px as X + Px We define the second-order optimization as valid if the eigenvalues of X have an equal width-dependent scale to the damping term, that is, px = (1()
This valid situation is rational because it prevents the effect of preconditioner from vanishing in the infinite width limit. also avoids the damping term's 's faster convergence to zero, a scenario that can potentially lead to numerical instability when computing inverses for rank-deficient preconditioners.
A.4.1 μP FOR K-FAC
On Condition A.1': For K-FAC, we have
1
G
(B, + Pn, 1) -n 6ydia((x) + EA
(S.33)
Maarta
22
where
A = hibe
B::= 5,0
(S.34)
17

---

Published as a conference paper at ICLR 2024
As technical remark, define (or implement) the metric matrix Bi using ō; vu fo for derivations. If the implementation is based on the derivative on w7, we have to multiply 1/M to Bt. Since the results essentially the same, are using more readable notations here without ay.
Under Assumption A.5, for ea, E {0, 1), we can apply the push-through identity, i.c., (1 XY)-IX ( (XX+ as follows.
(i) For < 1 < L, We have
M G
(S.35)
M2a 51 (KP PB. PB. - diag(x)) PA,,, .T) - A
K
(S.36)
noteworthy that a similar way to Eq. (S.26), K2 converges to a deterministic value in the infinite width limit. Since K/* is deterministic constant of ((1), its eigenvalues also ((1). This means that the damping where the second-order optimization becomes valid is dA, = -1. As side note, for d,, 1, the damping term in (K₁ PA MT) becomes dominant and the contribution of the preconditioner vanishes in the infinite width limit. In contrast, we may take dA. 1 if KJ* is positive-definite Similarly, we have dæ, 2(ar + 11) 1.
have B, (AW,,
= -(1/M²)) with
ri = 2ay + - 122B - l) + by) 0.
(S.37)
(ii) For 7. = L, 82 has the coordinate size ( and this means 82 (1 0g + pB|)) = O(1). Therefore, we easily obtain da, = - dB, = and
TL = 2az + eλ 1 0.
(S.38)
(iii) For / note that the input ho has the coordinate size ((1) and (he ho + pAl) -ea ca O(1).
Therefore, we have da, 1, dBg 2(az + bl) 1 and
r1 2a1 + +
+ ee (2eg l) + b1) 0.
(S.39)
On Condition 2', We have
W
1
=em(ōL-19 Maa
GL-1hL-2)
(S.40)
Moo(2(aztba)-1)--.
( (81-2
(S.41)
Maattr-1
with K₁ given by Eq. (S.36). To satisfy Condition 2', we need
ar. + br. - 0.
(S.42)
Finally, by combining Conditions A.I' and A.2', we obtain the paramcterizaaioon of pP
201 + ci + en (2en (a + br.) O, + ep 1 (2e l) + bb) 0,
20 + + +
(S.43)
20z + +ea
0, ar. + br. 0.
By setting a; = 0, we can climinate the indeterminacy of the parameterization and obtain Proposition 4.1. By setting e en 0, recover the case of the first-order gradient method.
18

---

Published as a conference paper at ICLR 2024
A.4.2 μP FOR SHAMPOO Here, let consider general e 0.
On Condition A.I": For Shampoo, we have
M2arta G M2 (Li + PL, (daa)) + 1) h--1, (S.44)
where we use &t defined Eq. (S.19).
(i) For 1 <1 < L, In a similar way to the case of K-FAC, we use the push-through identity. The update is given by
Maatta G 22 (
PL, 1), 1)-022 +/diii--//d
hi-1(hi- 6ydiag(x)h ( + PR: 1) h-11-
e/2
(S.45)
Note that we may multiply 1/M2a: the metric matrices L and R4 when their implementation is based on the derivative on wt.
Under Assumption A.5, we use the following push-through identity here: Lemma A.7 (e.g., Petersen al. (2008)). For any analytic function and real mutrices X and g(XY)x=3X X).
(S.46)
First, let us consider
Q Q= (ia 6 + PR. e/2 h1-1-
(S.47)
Let us express the largest eigenvalue of X by || 22 we have [|h=-diag(x)6] (ddiag())/]] = [[dia((X)6] 6,diag( h1-1/12
(S.48)
(S.49)
By setting PRi-: = (az + bL) and taking certain constant PR. diag(x))) 5,diag(X)) (ooagg()))(_ll///p)
satisfying
1, the inverse matrix in has a matrix series expan- sion that converges. Thus, we can use Lemma A.7 and obtain
Q h4 h (d (xdia(())/_sh-1 + PR:_1
(S.50)
Applying the same argument the La side of (S.45), We obtain
2artta G ta bi (diag x) by di- 61 + I) ()e diig))(())))
diag(x)
h h (diag 5,diag( + P)
(S.51)
In similar way Eq. (S.36), obtain
r1 20; a + 2e - (2c 1)(az bz).
(S.52)
(ii) For = 1, L, Similar to the case with K-FAC, by carefully handling R, for 1 = L and Lo for 1 0, we immediately find
ri 2a C1 + (2c 1)(a2 + bb),
(S.53)
P7 = 2ag + cL -
(S.54)
On Condition 2'. In the same way as in the case of K-FAC, we have
at + bn 11-11 0.
(S.55)
In summary, we obtain the thchearramttrizzationn of μuP as
2a 1 + t
-2e () by)
2a + q 2e (2e - 1) + br.) = 0,
(S.56)
2ar, + Ct, -e-1
ar + by - 1 0.
19

---

Published as a conference paper at ICLR 2024
A.4.3 GAUSS-NEWTON METHOD
From the standpoint of computational cost, practical applications use approximate second-order optimization methods. For K-FAC and Shampoo two smaller pre-conditioners are applied to the first-order gradient in matrix form. Although other forms of second-order optimization are so often utilized in deep learning, mentioning methods beyond K-FAC and Shampoo could be valuable, particularly from the perspective of demonstrating the theoretical applicability of these techniques.
Let us consider Gauss-Newtor method in a layer-wise manner. We choose this setting because the K-FAC and Shampoo are layer-wise and it makes the comparison easier. It also straightforward to apply similar order evaluation to the case without layer-wise approximation. The vector form of the Gauss-Newton gradient is given by
(V
Va C + Ml) - VV W. C (J7 diag (rg(x2)) x2)x + pt -JT
(S.57)
= Ji (diag(x)) JiJJi + +p(T)
(S.58)
where J, = V / as x M Jacobian matrix. Put XI == (diag(x²)) + px Then, the matrix form of the gradient (S.58) is given by
V C 6idiag((r))
(S.59)
under a 0. This corresponds G: in (S.18) and can easily apply the same order evaluation shown in Section The point is that J is the NTK matrix for the 1-th layer and we have
JJJ = (6:67) (h1-1h7-1).
(S.60)
From Eq. (S.25) and ar + bb, = 1, we can see that JIJí is ((1) for 1 <1 < L, OLLMM) for and (M) for l = L and determines the order of X. Thus, the pP for the layer-wise Gauss-Newton method is as follows:
1=1
1=1
by =
1/2 <<
-=
d =
0
<<l<
(S.61)
1=
-1
1= L
where PI= = Pf/Ma, Thus, bt and a are the same as in K-FAC and iI is the same as in Shampoo.
A.S EXTENSION TO OTHER LOSSES
Until now, we have considered the MSE loss with a one-dimensional target. Extending these results to multi-dimensional target cases is straightforward. Suppose that the target sample y is a C- dimensional vector, with C O(1). For a mileti-dimensionall target, its gradient (S.18) is given by
G d1ag(x) @ hT)).
(S.62)
We defined 8, 18(1) ...., S1 where 8(4) denotes the backward signal from the k-th output unit and
= RCC
(S.63)
Since (1, h }h (1, & h ht). Consequently, we can easily obtain the same jP for multi-dimensional targets.
apply the same arguments those shown in Section A.2.
Similarly, we obtain the same jP for the case of C--clas cross-entropy loss. For SGD and Shampoo, we just need to replace the error vector in the gradient (S.18) with
x = y - (f) E TR"C
(S.64)
where y is the one-hot encoded true label vector and off) denotes a softmax function with the input logits f. For each input sample H;, the softmax is defined by off((z)) = exp(ff(xx))) / exp(fx (x;)) for i = 1,..., n and k k= 1, ..., C. Thus, from the same argument as for MSE loss with a dimen-ddimesionnll target, we can obtain the same pp.
For K-FAC, note that the Fisher information matrix is given with the following backward part:
B = SAAT
(S.65)
20

---

Published as a conference paper at ICLR 2024
where A(o) is nC nC' block diagonal matrix which is composed of C' x C' block matrices; of off(2)))((((x)) (2 = I,...; n). Because A depends only on f and its size is independent of the width, its contribution amounts merely constant multiplication. In addition, the only difference from the MSE case lies in the backward part. We can easily employ the push-through identity
+ pp.) - - (A + pp,) Therefore, we can obtain the same jP as those presented Proposition 4.1.
(6
(S.66)
A.6 DAMPING HEURISTICS
When using damping heuristics in K-FAC, second-order optimization does not become valid. In addition, Aly decay with width if a, b, e are set to jP settings. This mechanism can be explained as follows. Consider the damping of the input layer determined by Eq. (14).
M tr(B,) P ((M(o2th))1))))
PA ( (= 1/pn))
(S.67)
V.
According to Eq. (8), the appropriate damping scales are PΛ = (11) and pB = 0(1AA/(extbb))) Thus, if we use damping heuristics, second-order optimization is not valid because PA and pg both larger than the appropriate damping scale. Since the order of the damping term is larger than the appropriate damping scaling and Pa PB 1, AW, bg is of the same order as SGD. Thus the scale of AW,ho will be as follows:
(AW ho, , = /ataMMaattttttttt
(S.68)
Thus, Ah1 will decay when = 0 and by 2 O.
On the other hand, in the case of Shampoo, even we use the damping of heuristics (determined by constant multiple of the maximum eigenvalue), second-order optimization becomes valid. This mechanism can be explained as follows. The following three matrices all share the largest eigenvalue.
L + pT = -5th1-1 1 + +pT,
(S.69)
Ri + pl = 14-18 8 + pl,
(S.70)
Ki pl = -hT-,/y-11 1 pl.
(S.71)
Since the matrix size of K₁ is independent of width, the maximum and mean eigenvalues of Ki are of the same order with respect to width. Thus, the maximum eigenvalues of R; and L, are of the same order the mean eigenvalue of K,
B EXPERIMENTAL DETAILS
B.
MODELS
We considered the following models for vision tasks.
MLP:
We considered a 3-layer multilayer perceptron (MLP) with ReLU activation. The MLP models do not include bias.
We considered 3-layer CNN with ReLU activation. The models consist of a two-layer convolution and a linear layer. We trained with different hidden widths where the width represents the output dimension of the first layer. Max pooling is applied after the activation function.
CNN:
Myrtle-5:
We considered Myrtle family (Shankar al., 2020) example for CNN. We trained with different hidden widths where the width represents the output dimension of the first layer.
We considered ResNet18 and ResNet50 (He et al., 2016) with different hidden widths where widths represent the number of channels. We used the existing implemen- tation for training CIFARIO. In addition, we used models from Pytorch Image models (Wightman, 2019) for ImageNet training.
ResNet:
hhtpps:// thub. www.mom//ooooo uoguel ph-m1 (hhmmrr//uutee rg/Cutout
21

---

Published as a conference paper at ICLR 2024
In addition these four models, we also considered CBOW as a language model.
We used CBOW model for the word embedding task (Mikolov et al., 2013). The CBOW model is network consisting of an embedding layer that embeds words and linear layer that predicts context words. The models do not include an activation function.
CBOW:
B.2 DETAILS OF FIGURES
This section explains the experimental setting for each figure. In all experiments, we implemented second-order optimization based the ASDL library (Osawa et al., 2023a). Note that p in the Re P below represents a proportionality constant for damping as PA Mi tr tr(h,,hh h1-1) and PB
P. Mi 1r(6, 8t).
Figure the upper graph, we trained a 3-layer MLP on the MNIST dataset with cross-entropy loss. In the second graph, we trained Myrtle-5 the CIFARIO with cross-entropy loss. Both graphs represent Ah; at the 10th iteration. The training sets have been reduced to 256 samples, and trained models by full-batch training. We apply data augmentation.
Figure 2
We trained a 3-layer MLP on FNST with 77 = 0.001, = 1. Its parameterization follows yP settings. We apply data augmentation.
Figure 3 We trained a 3-layer CNN on FashionMNIST with different by for 50 epochs. This graph shows the ratio between the gradient and the preconditioned gradient, Layer-wise learning rate follows the scaling of in muP in Eq. 9. The learning rate is tuned by grid search. Its settings are as follows.
Learning rate 2 where E
-20
Damping term p for K-FAC:
Damping term for Shampoo: 1e-3
In Table.2, we experimented with exactly the same settings as above, but with different batch sizes.
Figure 4 (Lefft) We trained Context as Bag-of-Wond (CBOW) on WikiText2 by Shampoo with 0.1, p = 10-6 and r = 0.1. We test the embeddings by the word analogy task on WordSim353. We used words that appeared more than 128 times in the dataset.
Figure (Right) We trained ResNet18 CIFARI00 by K-FAC with v 0.003, = 10 and T= 0.5. We use cross-entropy loss with label smoothing. We apply RandomCrop. RandomHor izontalFlip, AutoAugment, and Cutout data augmentation. In Table 3, experimented with exactly the same settings when training ResNet18 on CIFARI00.
Figure We trained ResNet50 on ImageNet by K-FAC with p = 0.001 and 0.05. We use cross-entropy loss with label smoothing. We apply RandomCrop, HhandomHorizontalFIin and Cutout as data augmentation. In addition, to prevent instability the training, we used gradient clipping in Grosse & Martens (2016). In Table 3, we experimented with exactly the same settings when training ResNet50 on ImageNet.
In the first row, trained 3-layer MLP on MNIST for 20 epochs. The number of samples reduced to 1024 and trained by MSE loss. In training with K-FAC, we used p = 0.001 for heuristics damping and p 1 for rescaled damping.
Figure 6
In the second row, we trained a 3-layer CNN on FashionMNIST for 50 epochs. The number of samples is reduced 1024 and trained by MSE loss. In training with K-FAC, we used p 1 for heuristics damping and p = 100 for rescaled damping.
In the last row, trained Re FashionMNIST for 20 epochs. The number of samples is reduced to 1024 and trained by MSE loss. In training with K-FAC, we used p = 10 for heuristics damping and p 10 for rescaled damping.
22

---

Published as a conference paper at ICLR 2024
We trained 3-layer CNN FashionMNIST for 50 epochs with 1) 0.003. The number of samples is reduced to 1024 and trained by MSE loss.
Figure 7
Table We trained VGG19 CIFAR 100 for 300 epochs, ResNet18 CIFAR 100 for 300 epochs and ResNer50 on ImageNet for 55 epochs. For VGG19 training with K-FAC, we set the learning rate " 0.01, for ResNet18 training with shampoo, set 77 0.001, and for ResNet50 training with K-FAC, we set 77 0.2. In all other settings, 0.003.
C ADDITIONAL DESCRIPTION OF SECOND-ORDER OPTIMIZATION
C.I K-FAC
Natural gradient descent (Amari, 1998) is an optimization algorithm that preconditions the gradient by the inverse of the Fisher information matrix. Its update rule is given by
1u+1 = -T (F((.) pl) VL.).).
(S.72)
where
F(0)= P
V log pe (t')
V log pe (t' "
(S.73)
is the Fisher information matrix. Here, the general exponent While 1 is commonly used, NGD with e < 1 has also been explored in some studies (Huh, 2020; Amari et al., 2021).
K-FAC (Martens & Grosse, 2015; Grosse & Martens, 2016) an approximation algorithm for NGD.
approximates F((O,) by Kronecker product of two matrices
(O;) pl = (Br + pBl) (Ai-1 + PAl),
(S.74)
where 816; and A, E hehy the estimation of the K-FAC curvature matrix.
Exponential moving average is often utilized to stabilize
1 {A()) (1) + (1 - E) hihy
(S.75)
B((+1) =B(() + (1 - { 8,6;
(S.76)
C.2 SHAMPOO
Adaptive gradient descent utilizes the square root of batched empirical Fisher
Fbutch (0) BBaopan VLs(O)VC,(0)T]]
(S.77)
instead F(0,) where VLs( = [8] E( V log
Pe (t x:) represents the mini-batch gradient. Shampoo approximates this batched empirical Fisher by Kronecker product of two matrices (Gupta al., 2018).
Fhutch (0) + pl (Ri pnT) (I1-1 PTT),
(S.78)
where L, E h 1 and Ri 1 81h1-1 J. Summation is often utilized to stabilize the estimation of the Shampoo curvature matrix while exponential moving average is also well utilized.
((+1) (4) E(8 T 0-11 J
(S.79)
R (++)) = RQ(") + 1 8,11-1
(S.80)
This summation prevents - and R(" from becoming a zero matrix, which ensures stable learning an infinite width limit.
23

---

Published as a conference paper at ICLR 2024
cmA=0.00 0.01
ema-=
cma 0.1
cmA= 0.3
CmA
AA SP (K-FAG)
29
AA XOUL (K-FAC)
2-2
24
212
2
2
ema=0.01
ema=0.00 0.00
ema=0011
ema=0.3
ema=1
At SP
Ah mup (Sbsspooo)
2-2
212
2
width
width
widch
width
width
Figure S.1 Exponential moving average does not affect order evaluation. Exponential moving average has almost no effect on the behavior of Ah. More precisely, regardless of the exponential moving average, A decays in the input and hidden layers with increasing width when trained with SP, but it does not decay when trained with μPP We trained a 3-layer MLP CIFARIO. In K-FAC we used rescaled damping.
D ADDITIONAL EXPERIMENTS ON EMPIRICAL VERIFICATION OF μP
D.1 ADDITIONAL EXPERIMENTS Ah
Exponential moving average If we take an exponential moving average over the curvature matrix, we cannot use the push-through identity Eq. 10. However, even this case, the behavior of Ah is consistent with the order evaluation. Figure S.1 shows that the exponential moving average has almost no effect on the order of Ah. This shows that even if use an exponential moving average, uP is still reasonable.
KFAC SP Heuristics
KFAC muP Heuristics
KFAC muP Rescaled
Shampoo SP
Shampoo muP
MN Manay
2-2
29
212
29
232
29
212
29
232
29
212
Figure S.2 juP achieves feature learning MLP with tanh activation trained 3-layer MLP on MNIST using cross-entropy loss. Even if the activation function for the MLP is tanh, we observed the same behavior as in Figure I.
24

---

Published as a conference paper at ICLR 2024
Activation Function The behavior of Ah is not affected by the activation function. We trained 3-layer MLP with tanh activation in Figure S.2 It shows the same behavior as when we trained 3 3-layer MLP with relu activation (Figure 1).
E ADDITIONAL EXPERIMENTS ON THE VARIANCE OF THE LAST LAYER
E.1 TRAIN CURVE
Figure S.3 shows the learning curves for the experiment in Figure 3. K-FAC can reach the NNGP solution in step. Thus, can observe that K-FAC reaches high test accuracy from the first step for b = 64. However, the accuracy at = 1 is higher than that at b = 64 since the training process for b 64 stays at the initial value.
Note that the optimal learning rate of L = 64 is smaller than that of b 0.5 or b = 1. This because the NNGP solution is sharp local solution, and large learning rate makes the subsequent behavior unstable.
(A) SGD
(B) K-FAC
(C) Shampoo
80
80
(outpuut layer)
Test Acc
60
60
0.5
60
40
40
20
40
20
40
20
40
epochs
epochs
epochs
Figure S.3 Train Curve Comparison for different by. K-FAC achieves high accuracy in one step when ,y 64 since K-FAC can achieve NNGP solution in one step. However, K-FAC can not escape from the NNGP solution. In this figure, we trained 3-layer CNNs Fashion-MNIST with batch size 1024.
E.2 CHOICE OF LOSS FUNCTION
The experiment in Figure is performed with MSE loss to clarify the connection with NNGP solution. This result is easily extended from MSE loss cross-entropy loss. Figure S.4 shows that even with a cross-entropy loss, K-FAC shows higher accuracy only when by is at or near in the last layer.
SGD (Train Acc)
KFAC (Train Acc)
Shampoo (Train Ace)
-00
-2.0
84 -2.0
Losg 6 (output laver)
00.5
0055
82
0.5
10
78
76
16
37
40
76
106
136
16
A0
61
91
106
121
1316
21
60
61
10
006
1906
Iteration
Figure S.4 K-FAC converges NNGP solution even the cross-entropy loss. experimented with a 3-lyerr CNN, moving only the last layer by from μP. The loss function is cross-entropy. K-FAC shows that the regions of high accuracy are concentrated only around b 1.
25

---

Published as a conference paper at ICLR 2024
E.3 CHOICE OF MODEL ARCHITECTURE
The experiment in Figure 3 was performed with 3-layer CNN model. This result can be generalized to other architectures. Figure S.5 indicates that accuracy at ,,, is higher than 6,. = 0.5. this example, train accuracy for by. > is consistently higher than the train aceuracy for b, 0.5. In ResNet-18 18, the maximum value appears to be obtained at by = 2, but this may be due to the constant factor rather than order, which occurs when experiments are conducted with finite widths.
Train Acc(K-FAC)
Test (co(K-FAC)
95
95
model
MLP CNN
80
85
80
80
ReaNet18
70
70
2-1
2
2
2
2
bg. (ontput layar)
bg (output layor)
Figure S.5 μP is effeetive in various architecture. According to the K-FAC μPP, by. = has a higher accuracy than by 0.5. This is consistent across various models. MLPs trained with MNIST, while other models are trained with FashionMNIST.. The number of samples is reduced 1024. MLP and CNN trained in full batches, while Myrtle and ResNet18 trained in mini-batches with a batch size of 128.
E.4
ON THE EFFECT OF MOMENTUM
SGD (Train Acc)
KFAC (Train Acc)
Shampoo (Train Acc)
2.0
lozg 5 (output layez)
80
0.0
0.0
0.0
10
66
78
76
16
31
46
70
106
123
130
16
31
40
11
01
106
121
136
F6
31
46
11
78
106
136
Iteration
Iteration
Iteration
Figure S.6 K-FAC converges to NNGP solution even without momentum. We experimented with 3-layer CNN, moving only the last layer by from μP. The loss function is cross-entropy. K-FAC shows that the regions of high accuracy are concentrated only around b 1.
Figure S.7 : juP is effective across batch-size without momentum We trained FNSH with 3-layer CNN with different batch sizes. We did not use momentum in this figure. The values represent the accuracy of the training dataset.
Batch Size
Optimizer
16
64
256
1024
0.5
80.52
78,69
75.05
67.94
49.70
SGD
1.0
82.85
80.41
77.58
73.12
64.30
64.0
83.59
81.23
77.25
73.63
70.53
0.5
77.60
79.66
83.94
82.14
78.63
K-FAC
1.0
79.10
81.69
83.92
83.16
80.27
64.0
76.06
76.95
77.28
76.04
75.94
26

---

Published as a conference paper at ICLR 2024
Figure S.6 and Figure S.7 represents the results for training 3-layer CNN without momentum. Since momentum is known to have the effect of escaping from the saddle point, is expected that is more difficult to escape from the NNGP solution without momentum. However, there was almost no effect of removing momentum on the relationship between accuracy and by. while the overall accuracy is lowered by removing momentum.
E.5
WORD2VEC
KFAC (Analogy Acc)
Shampoo (Analogy Acc)
0.75 7.250.25
40
0.75
40
b (output layer)
35
35
1.5
1.5
30
16.0 4.0
4.0
30
25
16.0
25
1228
256.0
512.0
10210
2048.0 0
4096. .
8 8192.00
128. 0
256.0
512.0
1022A
20045 D
40080 0
$192.0
220
width
width
Figure S.8 Effect of by on Accuracies of Word Analogy In K-FAC, only limited range of by has
high accuracy, whereas Shampoo has nearly 40% accuracy over a wide range by. et is correctly set to match the μP setting.
Feature learning is especially important in word2vec, which requires careful tuning of by in K-FAC. Figure S.8 examines the effect by. on accuracy K-FAC and Shampoo. In this figure, HPs are tuned with the following settings. Note that we tuned the variance of weights at width=128 since K-FAC sensitive to the variance the last layer.
learning rate = (1e-1, le-2, le-3
embedding weight multiplier = (1, le-1, le-2)
output weight multiplier= = 2, 4, 8, 16)
As shown in Figure S.8, K-FAC achieves high accuracy only around b, = 1. On the other hand, Shampoo achieves high accuracy wide range of by. This example implies that while K-FAC learns features in Vord2Vee, K-FAC requires more careful tuning of b,, than Shampoo
F ADDITIONAL EXPERIMENTS ON LEARNING RATE TRANSFER
F.I SHAMPOO
Figure S.9 shows that μP for Shampoo allows the learning rate tuned in a small model to be re-used in larger model. In Shampoo, the learning rate is stable even in SP while in jPP, the wider model achieves higher accuracy over wide range of learning rates.
F.2 FOOF
Since we introduced general exponents and EB, we can consider jP for FOOF The pP for FOOF immediately derived from Proposition 9.
0
1 -
1
==1
(-1 1<15L
61
1/2
L, 1
1 < 1
-1 1 <l<
==
df
1 1=1
0 ==L
(S.81)
27

---

Published as a conference paper at ICLR 2024
SP MLP
muP MLP
SP CNN
muP CNN
width
80
widch
80
128
16
Test
256
Test
60
512
70
1024
2015
$152
-10
log2 learning rate
logg learning
learning rate
logg learuing rate
Myrtle
Myrtlc
ResNet18
mul! ResNet18
80
width
80
sitth
Test Act.
60
126
256
40
5:22
1024
20
10
0
-0
-10
-10 logg learning rate
-10 loga leazning rate
logg learning rate
logg learning rate
Figure S.9; ,P for Shampoo allows the learning rate (17") to transfer across widths. We trained various models with SP and μP by Shampoo. We trained 3-layer MLP on MNIST, 3-layer CNN on MashionMNIST, Myrtle-5 on FashionMNIST and ResNet 8 on FashionMNIST. The number of samples reduced to 1024.
Figure S.10 shows that pP for FOOF enables the learning rate to transfer across the width as well as in SGD K-FAC. The learning rate may be transferred across widths since the SP for FOOF is stable parameterization. In addition, we can observe that accuracy obtained by training with jP is higher than SP, especially in wide models.
SIP MLP
mup MLP
SP CNN
mup CNN
90
Test Acc
64
128
256
512
64
1024
125
4006
8192
lo82 learning
logo learning
logo learning
loga learning
Figure S.10 μP for FOOF allows the learning rate (17) to transfer across widths. We trained 3- 3-layer MLP on MNIST and a 3-layer CNN on MNSS. The number of samples is reduced to 1024.
F.3 LEARNING RATE TRANSFER WHEN THE SAMPLE SIZE IS FULL
In Figure 6, we reduced the training samples. By reducing the dataset size, known to behave more closely to infinite- width models, has often been in papers examining the thcoretical aspects of second-ordes optimization and feature lcarning(Geige) et al., 2020; Karakida when training the full dataset, its learning rate
models are
& Osawa, 2020; Amari al., 2021). However, can be transferred using juP as shown in Figure S.11 In addition, we can observe that the learning rate is correctly transferred when optimize ResNet50 on ImageNet. As shown in Figure S.1. the optimal learning rate is fixed at 0.2 regardless of the width.
F.4 ACTIVATION FUNCTION LOSS FUNCTION
In Figure 6, trained 3-yer MLP with ReLU activation by MSE Loss. Similar observations be made when using Tanh activation or cross-entrope loss as shown in Figure 1.1 Note that squashing activation functions such tanh and softmax are not recommended in some experiments in Yang
28

---

Published as a conference paper at ICLR 2024
SGD SP
SGD mup K-FAC SP(leuriitiie) K-FAC SP
K-FAC muP
width
256
80
1024
Test
70
2048
4096
-10
-10
-10
-10
-10
16384
Ingg
loga
Irgg
loga 1.R:
loga I.R
SGD SP
SGD muP K-FAC (Heursssss))))
K-FAC
K-FAC muP
wifth
80
32
128
256
10
10
logg LR
logg LR
logg
Jong
logg LR
Figure S.11 j&P allows the learning rate (7/) to transfer across widths(Full Dataset) Using JP, one can transfer the learning rate with respect to width. We trained a 3-layer MLP on FNSSI and 3-layer CNN FashionMNIST.. contrast to Figure 6, the sample size been reduced. In KFAC with SP, we use damping heuristics which prevent the transfer of learning rate.
Learning Rate
Table S.1 optimal lcarning rate does not
width
0.025
0.050
0.100
0.200
0.400
shift in training ResNet50 on ImageNet.
0.5
69.26
70.07
70.19
70.56
70.15
The optimal learning rate is fixed at 0.2 and does not shift when scaling the width.
1.0
74.35
74.67
75.54
75.92
75.32
2.0
77.45
78.07
78.47
78.79
47.64
al. (2021) D.3. However, the properties the hyperparameter landscape for each parameterization in Figure S.12 are almost the same as in Figure 6.
G ADDITIONAL EXPERIMENTS ON TRAINING WIDE MODELS
G.1 ANOTHER DATASET AND MODELS
In Figure 3, we compare the test accuracy with SP and jP and find jP takes a higher accuracy than SP in wide models. This tendency is also true for other settings. As shown in Figure S.13 and Figure S.14, ,PP achieves higher valuation accuracy than SP in many cases. Specifically, the difference between SP and μP is larger for wider models, and the difference is often larger in the early stages of learning.
G.2 ERROR BAR
Table.S.2 shows the results of the comparison of the accuracy of SP and jPP with standard deviation. There is significant difference in accuracy between SP and uP regardless of the effect of seed.
G.3 DISTANCE FROM INITIAL WEIGHT
The relative distance from their initial weights value AW/IW| is often used in the analysis of infinite- width models. A WW/WW/ is always zero throughout training, the model will fall into a lazy regime. As shown in Figure S.15 in the early stages of training with SP, AW/|W| of wide-width models is almost zero. As the training proceeds, the weights begin to deviate from their initial values, which be considered finite-width effects. The wider the width, the longer the time that AW/[W] almost zero. This suggests that the weights do not move away from their initial values in the width limit. In pP, however, the point which weights start
away from their initial
29

---

Published as a conference paper at ICLR 2024
SGD SP
SGD muP K-FAC SP(Heuristies) K-FAC SP
K-FAC muP
width
64
Test Acc
80
256
60
512
1024
40
2048
4096
100
10
5
-10
5
8192
logg LR
loga LR
logg LR
LR
log LR
SGD SP
SGD muP K-FAC SP(Heuristies) K-FAC SP
K-FAC muP
width
64
Tert Acc
256
512
70
20:68
60
4096
10
10
10
10
10
Ingg 1.R
logg LR
L.R
lodg LR
loga LRR.
16384
SGD SP
SGD mup K-FAC SP(Heuristies)
K-FAC SP
K-FAC muP
width
Teat Acc (ReiLU xent)
90
256
1024
2048
4096
-10
55
10 -
-10 5
- 10 5
10 5
logg LR
loga LR
loga LR
logg LR
logg LR
16338
Figure S.12 : Learning rate transfer across different activation functions and loss functions. We trained a 3-layer MLP on MNIST. We used tanh activation with MSE loss in the first row, tanh activation with cross-entropy loss in the second row, and ReLU activation with cross-entropy loss in the last row. We trained MNIST dataset and the sample size was reduced 1024.
Epochs=100 100
Epochs=200
Epochs=300 300
K-FAC
Shampoo SP μP
K-FAC
Shampoo SP/ JPP
K-FAC
Shampoo
wid
SP / μP
SP/ μP
SP/ μP
SP uP
89.54/40.00
88.27/40.43
91.35/+0.00
91.01/-0.25
91.97/+0.00
91.57/-0. 14
2
92.10/+0.26
91.31/+0.40
93.92/+0.30
93.62/+0.14
94.32/+0.03
93.95/+0.05
93.64/+0.57
93.32/40.22
95.11/40.45
95.19/-0.15
95.57/+0.14
95.56/-0.10
93.56/40.74
94.02/40.33
95.44/40.20
95.81/+0.19
95.73/+0.44
96.09/-0.04
16
93.00/41.36
94.81/40.96
95.04/40.66
96.28/+0.14
95.31/40.72
96.49/+0.24
Figure S.13 Test Accuracies of ResNet18 on CIFAR1O. We trained ResNet18 on CIFAR100 with "
0.003, p' 10. jP has higher accuracy than SP in models with wide widths.
values remains the same, even as the width increases, which implies that training can proceed even at an infinite width limit.
H ADDITIONAL EXPERIMENTS ON DAMPING TRANSFER
H.1 ON FIXED DAMPING SCALING
In this paper, we determine damping mainly by Eq. 15 as this is consistent with jPP. However, even damping determined by Eq. 7, damping is still consistent with μP. In Figure S.16. we trained 3-layer MLP with damping determined by Eq. 7. Its learning rate is fixed at 0.01. d' and Pg are tuned using the following grid:
PA 2", where E 10,1,2,... 10)
30

---

Published as a conference paper at ICLR 2024
Epochs=100)
Epochs=200
Epochs=300
K-FAC
Shampoo
K-FAC
Shampoo
K-FAC
Shampoo
wid
SP/ / μP
uP
SP / μP
SP/ uP
SP / μP
SP μP
68.53/+1.88
65.36/-0.31
71.55/+1.88
68.49/+0.30
72.28/+1.93
69.20/-0.02
2
71.28/43.66
69.28/+0.54
73.85/+2.99
72.24/+0.11
74.30/43.23
73.05/+0.20
69,45/47.87
72.81/40.78
73,39/45.83
75,59/40.38
73.79/45.91
76.57/40.31
8
61.15/417.01
74.87/40.28
68.13/411.63
77.10/41.64
68.56/412.09
78.08/41.74
Figure S.14 Test Accuracies of ResNet50 CIFAR100. We trained ResNet50 CIFARI00
with = 0.003, p' = 10. juP has higher accuracy than SP in models with wide widths. Note that in the training of K-FAC, the standard deviation of the last layer width=1 is tuned carefully and it is 1/16 of the default SP value.
VGG19 (C100)
ResNet18 (C100)
Shampoo
K-FAC
Shampoo
width
SP
jrP
SP
μP
SP
μP
63.30 (6.26)
63.01 (0.44)
67. 67.02
67.00
67.01 (0.33)
66.99
70.10(0.23)
70.21 (0.29)).29)
71.88
(2.31(0.39)
71.68 70.83
4
(2.45)
75.04(0.18)
(4.09
75.00 75.06(0.25)
73.80
75. (5.28((oo (0.34)
8
76.51 (0.38)
77. 77.30
74.90
76.92
76.23
78. 78.39 (0.17)
16
77.89(0.18)
78.35
(4.29(0.41)
78.12 (0.25)
78.05
80. 80.24
Table S.2 : Test accuracies with different widths (with error bar). The results are averaged over 5 random seeds, with standard deviation shown in the brackets. The settings are the same as in Table.3.
P'B = = where E {0, -0, -2.... -00}
Figure S.16 shows that the optimal damping value for wide width models given at (da, dp) = (-1, 1) in Eq. 7. This is consistent with the scaling in Eq. 9.
There are two reasons why we determine damping using Eq. 15 instead of Eq. 7. The first reason is that the optimal values for A and d are generally different, resulting in one extra hyperparameter. The second reason is that the eigenvalues of the curvature matrix may change over training and the optimal values for P'A and I'g may change during training(Zhangg et al., 2019b). Figure S.16 shows that the mean eigenvalue of A, and Bi change during the training and the damping determined by Eq. 15 and Eq. 14 is adjusted according the change of mean eigenvalue.
VGG19 CIFAR100 (IR=0.03)
ReaNer50 GIFAR100 (IRR=0.01)
SP
muf
SP
mnu
widch
300
widch
400
200
100
0
2000 4000 6000 0
2000 4000 6000
10000
10000
Iteration
Iteration
Iteration
Figure S.15 For SP, AW/(W] in the early stages of training depends on the model width In training with SP, AW/IW| depends on the width, especially in the early stages of training. Specifically, as the width increases, the period during which AWt/WW z O gets longer. This suggests that the infinite width limit with SP, AW/WW remains 0 throughout training. other hand, in jPP, the point at which weights start to move away from their initial values remains the same, even as the width increases (Left) We trained VGG19 on CIFAR 100 using K-FAC, (Right) We trained ResNet50 on CIFAR 100 using K-FAC.
the
31

---

Published as a conference paper at ICLR 2024
Damping Scaling (MLP)
Figure S.16 Damping term needs to be scaled as
(dx:da)
90
(42.dg)-(1. 13
dA = -1 and dg = 1 (Eq. 7). We trained a 3-layer MLP on MNIST with K-FAC. The number of samples is reduced to 1024 and trained by MSF loss. Validation
(fx.dg)=(1--1 -1)
Test
80
(dx.dn)-(11.)) 1)
(dx.dg)m(11,-1)
80
(4xid))=0000) 0)
accuracy is highest at d = and dg =1
2
2
2
logg width
MeanEig A
MeanLig B
damping A
damping B
109
10
CIFATLOO
10a
10
10-2
10a
10
10-2
10-4
109
10
20
40
20 Epoch
40
20 Epoch
40
20 Epoch
40
Epoch
MeanEig A
MeanEig B
damping A
damping B
CIFARLOD
10"
10
10a
10
10
10-3
10a
10
10-2
10a
10-2
20
40
20
40
20
20
Epoch
Epoch
Epoch
Epoch
Figure S.17 Mean eigenvalues and the heuristics damping terms vary during training. We trained ResNet18 on CIFARIO and CIFARIOC by K-FAC. (Left) Transition of the mean eigenvalues of A and B; in each layer during the training process. We can observe that the mean eigenvalues change significantly during the training. This implies that PA and Pg need to be modified during the training to retain their second-order properties for numerical stability. Note that the transition of damping determined by Eq. 15 is consistent with the transition of mean eigenvalues. (Right) Transition of the damping determined by Eq. 14. The trend of damping determined by Eq. 14 is generally consistent with that of mean eigenvalues.
H.2 DAMPING TRANSFER ON MLP
In Section 5.3, have observed that when using heuristics damping, the optimal damping value shifts as the width increases. We have also confirmed that rescaling damping prevents this shift. This phenomenon is observed not only in CNN but also in MLP as shown in Figure S.18 In the second column of Figure S.18, the dataset size is not reduced, but again the damping is transferred by using rescaling damping.
I LAZY PARAMETERIZATTOON
In the lazy regime, the network output changes by an order of ( but feature learning does not occur, i.e., FILL > O. Consider the uniform parameterization (UP), which is easy to characterize and known to include two important parameterization methods; jPP and NTK parameterization (Yang & Hu, 1).1). It is defined by rıel. = r and Wr. satisfying Conditions A. 1 and A.2. For K-FAC, we from these conditions and thus
obtain 20 + er. + 1 = 0, as, + by. =
2a + c (2r 1) 2r 1,
2a + q + eΛ en (2r 1) 2r;
(S.82)
2ar, + C1 +
0,
01 + by = -
32

---

layer MLP on MNIST and a 3-layer Published as a conference paper at ICLR 2024
Heuristics (SP) Rescaled (SP) Rescaled (muP)
width
90
Test Acc.
128
266
Figure S.18 We can transfer
damping in MLP We trained a 3-
2048
60
4006
MLP on FashionMNNIST with K-
8192
FAC. When training MNIST, the number of samples is reduced to
log2 damping log2 damping log2 damping
1024 while in FashionMNIST we trained MLP by full-dataset. When using heuristics damping, the maxi- 128
Heuristics (SP) Rescaled (SP) Rescaled (muP)
width
90
Test Acc
mum damping value that does not
256
diverge increases as the width in-
512
creases.
70
60
4096
8192
20
100 log2 damping loga damping
loga damping
We usually suppose r = 1/2 for NTK parameterization. By fixing the shift invariance by aj 0, we have
= 0, b =, b 1/2.
ci>1 = -
(S.83)
This means that K-FAC SP achieves the lazy regime for the constant learning rates of @(1). In contrast, the first-order gradient (e = 0) in SP requires the re-scaled learning rates of (11MM) for the lazy regime. Given that the naive setting of default hyperparameters in implementation often assumes SP and constant learning rates, it can be said that K-FAC is more likely to suffer from the lazy regime compared to (S)GD. is also noteworthy that the obtained pate-mmaatttiitiiionn is consistent with the parameterization used in parameterization used in the previous work K-FAC in the NTK regime (Karakida & Osawa, 2020).
Similarly, for Shampoo, we obtain
= 0, Ct> = 1 - e, bi = 0,
(S.84)
bi> = 1/2.
SGD SP
SGD UP
SGD muP
K-FAC SP
K-FAC muP
width
(ReiLU mox)
90
128
Tost Acc
80
512
1024
70
4096
6192
10
0
10
1U
10
10
16384
loga LR
logg LR
LR
LR
LR
Figure S.19 UP allows the learning rate to transfer while the accuracy of UP is lower than j.P. We trained a 3-yayer MLP on MNIST and the sample size was reduced to 1024. Note that in K-FAC SP is equal to the UP.
33

---

Published as a conference paper at ICLR 2024
Table S.3: Lazy parameterization for SGD, K-FAC, FOOF and Shampoo
Input weights & all biases
Output weights
Hidden weights
SGD
b b= 0, c
= 1/2, 1
b= ==1/
Shampoo K-FAC FOOF
b 0, e b b= 0, c
1/2, 1/2
b b=1/2,c 1/2
1/2, 0
b 1/122 =0 0
b= 0, c
1/2,c = 0
b 1/2,c =
34


Reviews:
[{"review_id": "408mNLqdTJ", "reviewer": "Reviewer_ZUFF", "summary": "The authors derive scaling laws for two second-order optimization methods (K-FAC (Martens & Grosse, 2015) and Shampoo (Gupta et al, 2018)), based on the muP approach from (Yang & Hu, 2021). This approach has previously been used for first-order methods, including with momentum. It is used to determine how the weight initialization and per-layer learning rate for deep neural networks should scale based on layer width. This allows 'parameter transfer', where good hyperparameters can be trained on small networks and suitably scaled up for wider networks. The derived scalings are successfully tested on some simple example problems.", "questions": "To me, it appears that the extent of the theoretical contribution is not huge (largely based on combining existing methods), and I take the main benefit of the paper to be the actual derived scaling laws (i.e. choices of $a$, $b$, $c$, etc) rather than the overall approach/techniques.  Does this work introduce any new theoretical results/mathematical techniques of wider interest/applicability, or is the theoretical approach quite restricted to the specific methods being considered?\n\n=========== AFTER REBUTTAL ===========\nThe new version of the work includes a mild strengthening of the mathematical formatlity in Appendix A, plus new results deriving scaling laws for layer-wise Gauss-Newton and other loss functions such as cross-entropy. This I think improves the paper by showing that their adaptation of muP to second-order methods yields a systematic approach rather than needing ad-hoc approaches, giving it wider interest.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is a new application of the muP framework to a useful class of methods. I think the new scaling laws are useful and beneficial to the community. The authors apply a successful approach for deriving scaling laws to the case of specific second-order networks. The results are clearly stated and yield new insights into the merit of existing choices of damping parameter for K-FAC, and the use of zero initial weights for the last layer. The paper is reasonably clearly written and the presented experiments thorough.", "weaknesses": "I found the mathematical presentation fairly informal, both in terms of the exposition of the approach (Prop 4.1 has \"the optimization becomes valid for...\", for example) and the actual derivation in Appendix A. This is a presentational comment, not a critique of the correctness of the work. \n\nTheory is only given for least-squares losses, although the numerics suggest the results hold for other losses. If the authors could derive these laws - and Appendix A seems to suggest it might not be a huge piece of work - this would help the paper. \n\nWhile the empirical results show merit, their real benefit would only be clear if applied for larger-scale networks than those tested. I understand the practical issues around doing this and recognize that it isn't feasible for all researchers, but those results would strengthen the paper.\n\n=========== AFTER REBUTTAL ===========\nThe new sections on layer-wise Gauss-Newton and extension to other loss functions are helpful here. I would suggest the authors mention the Gauss-Newton work somewhere in the main text (just like they did for the alternative loss functions), as I couldn't see clear reference to this extra contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QmBTAYoUF3", "reviewer": "Reviewer_Xrk9", "summary": "This paper is a continuation of [1] in the context of second-order optimization. The initial idea consists in tuning simultaneously several hyperparameters (HPs) related to neural network (NN) training based on the following heuristic: after one training step, the pre-activations of each layer should have made a move of order $1$. More specifically, this heuristic should hold when the number of neurons of each layer tends to infinity. In order-$1$ optimization, the HPs include: variance of the initialization distribution of the weights, scaling of the weights, scaling of the learning rate. In the present work, the authors propose to tune the HPs specific to order-$2$ optimization, such as the damping parameters.\n\nIn practice, the authors focus on two order-$2$ optimization techniques (K-FAC and Shampoo), and test experimentally their heuristic with multilayer perceptrons (MLPs), convolutional NNs (CNNs) and ResNets. The authors find out that their method leads to better test losses, especially for NNs with very large layers. Moreover, one choice of learning rate seems to lead to results stable w.r.t. the size of the layers.\n\n[1] *Tensor programs IV: Feature learning in infinite-width neural networks*, Yang and Hu, 2021.\n\n=== Edit ===\nI acknowledge the other reviews and the authors' answers. I keep my rating (8: Accept).", "questions": "What do the authors think about apply $\\mu \\mathrm{P}$ to other second order methods, such as Newton, Gauss-Newton and generalized Gauss-Newton? \n\nCould 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?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "## Originality\n\nGeneralizing the heuristic of [1] to second order methods is new.\n\n## Clarity\n\nOverall, the paper is clear for a reader who is already aware of the works of Greg Yang (see [1]). \n\n## Quality\n\n* This work is technically correct.\n* The initial heuristic is tested (Fig. 1): $\\Delta h$ behaves as expected when using $\\mu \\mathrm{P}$.\n* In order to compare fairly $\\mu \\mathrm{P}$ with SP, the authors have performed a sanity check (Fig. 5): the results are stable for a wide range of learning rates.\n* The experiments show that $\\mu \\mathrm{P}$ is better than SP, especially when the number of neurons per layer tends to infinity. This result was expected and is explained by the \"infinite width limit assumption\".", "weaknesses": "## Section 5.2 and Figure 6\n\nThere is a mix of lack of clarity and unconvincing results in this section.\n\nFirst, the authors claim in Section 5.2 that $\\mu \\mathrm{P}$ allows better \"transfer\" of learning rate than SP. The term \"transfer\" is not so clear to me, but apparently, it means that the optimal learning is stable by width change. But, when looking at Figure 6, this is barely true: actually, the transfer is better with K-FAC SP than with $\\mu \\mathrm{P}$ in the MLP case. \n\nSecond, in the legend of Figure 6, the authors claim something else: according to the context, \"transferring the learning rate\" means that, with a fixed l.r., the results get better and better as the width of the layers increases. This claim is more convincing, but differs substantially from the claim provided in Section 5.2. \n\n## Significance\n\nThe narrowness of the focus on K-FAC and Shampoo is a limitation of this work. What do the authors think about Newton, Gauss-Newton and generalized Gauss-Newton? \n\nBesides, the results presented in Table 1 are somewhat disappointing: basically, the usual NTK scaling corresponds to $\\mu \\mathrm{P}$ for all the hidden weights, in all the settings (SGD, Shampoo, K-FAC). The only difference lies in the input and output weights. Naturally, the authors cannot do anything about that.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "djuIB2tkAb", "reviewer": "Reviewer_SW4H", "summary": "This paper considers finding the optimal hyperparameters for two second-order methods KFAC and Shampoo when training a neural network with a large width or when the width is infinite. These hyperparameters include the learning rate and the initialization hyperparameters. This optimal setting could help the network to learn features in a stable manner. This work is built upon the previous works that find the optimal setting for the first-order optimization methods.", "questions": "I was wondering if we change the layer width for different layers what part of your analysis would change?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Analyzing the optimal hyper-parameters for the second-order methods is provided in this work for the first time. Empirically they show that their suggested values for the hyperparameters help to find useful features when they increase the width of the layers i.e. their experimental results confirm their theoretical analysis. As a side contribution, they show that for a specific parameter setting KFAC has an implicit bias toward NNGP.", "weaknesses": "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. \nIn the experiments, for example in Table 3, there isn’t any standard deviation next to the accuracies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)