Initializing the Layer-wise Learning Rate

We thank the reviewer for detailed and constructive feedback. We hope to sufficiently address your concerns and questions below.

W1 weak motivation

Motivating the paper as simply mitigating ADAM's computation time is not enough

• We would like to clarify that our motivation is not necessarily pertained to replacing or mitigating a particular aspect of Adam, but to answer a more broader question of whether assigning non-adaptive static layer-wise learning rates can be a practical or effective stategy.

The gradient information at initialization has no information of the curvature at any point in the loss landscape. So, the fundamental motivation for the proposed algorithm seems to be missing.

• We appreciate the reviewers concern in the proposed method does not have a strong theoretical connection, which may mean the presented results do not generalize.
• However, in a situation where the connection between existing theory and empirical performance remains unclear, we believe approaching from different directions constitutes a valid motivation.
• For example, [1] as mentioned by reviewer NtQ7 investigates the sign descent property of Adam and present evidence that counter the idea ADAM outperforms SGD on transformers due to a more robust estimate of the descent direction. Given evidence of the effectiveness of sign-descent optimizers [2], we consider exploring methods of that adjust the learning rate that at a layer-wise level in a static manner a reasonable motivation.

W2 No substantial improvement

The improvement on top of SGD or ADAM-W seems to be marginal and it is uncertain that whether it is solely because of the layerwise learning rate or the interaction of this layerwise lr with other tricks such as cosine scheudling or small batch size. It's hard to justify this algorithm unless studied in isolation without any other regularization tricks.

• We agree the presented improvement on top of SGD or ADAM-W is not enough to ascertain that the layer-wise learning rate is beneficial given the many potential factors that could have affected the results.
• The overfitting experiments on Table 2, 3 and in Table 8, 9 are settings designed to explicitly factor out the many potential tricks. It is a setting where we do not perform cosine scheduling and where we measure the train loss of the small dataset, such that factors that pertain to generalization or regularization would not affect the result. Arguably the small batch size trick remains, but in practical training a small batch size is likely to be relavent than performing a full gradient descent. The experiments performed on 4 optimizers (AdamW, Adam-mini, Lion, Sophia) demonstrate that the layer-wise method consistently improve train loss.
• In particular, the additional experimenents with Sophia [3], which were performed during rebuttal phase, show that the layer-wise learning is also beneficial to a Hessian-based optimzier. It further suggests the existence of additional factors that can be leveraged to improve performance even when using computation-efficient second-order methods. Given the highly specialized architectures which modern neural network training is concerned with, how architecture interacts with training makes it a high priority for investigation.

various sharpness regularizers like SAM or even noise injection that outperforms their base optimizers (SGD or ADAM-W) by a significant margin. In this context the use of such layerwise learning rate is obsolete and unnecessary in the long run.

• SAM and noise injection methods are methods that focus on improve generalization, at the cost of increased difficulty in training the original objective. Our experience with SAM is that while it improves validation accuracy there is a noticeable degradation, i.e. increase in train loss. Our understanding is that both methods that improve either convergability or generalization are significant in the long run.

[1] Kunstner, Frederik, et al. "Noise Is Not the Main Factor Behind the Gap Between Sgd and Adam on Transformers, But Sign Descent Might Be." The Eleventh International Conference on Learning Representations.

[2] Chen, Xiangning, et al. "Symbolic discovery of optimization algorithms." Advances in neural information processing systems 37 (2023).

[3] Liu, Hong, et al. "Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training." The Twelfth International Conference on Learning Representations.

We thank the reviewer for detailed and constructive feedback. We hope to sufficiently address your concerns and questions below.

Q1, Q2 and Q3: Regarding the combinations with different weight initializations with the proposed method

• Many weight initialization methods previously proposed would have been evaluated on settings where a single learning rate is used. This means that the performance of such methods could have tuned to be favorable at such settings. As such, when performing ablation studies in Table 2, we focused on the most basic initialization schemes that were initially proposed, that is, the activation preserving and gradient preserving fan-in and fan-out initialization that samples weights from $N(0, 1/f_{in})$ and $N(0, 1/f_{out})$
• We further provide results from additional initialization experiments below, where we experiment with downweighting the residual branches $h(x) = x + \beta f(x)$. We would like to note the $\beta$ scaling is incorporated as a multiplicative scalar rather than just adjusting the variance from which the residual branch weight was sampled from.
• As it is a initialization scheme concerned with scalability with depth, we provide results on the 774M GPT-2 model that has triple the depth (36) compared to 124M GPT-2 model (12). Due to computational constraints, we provide the train loss (↓) results from small scale overfitting experiments when trained with Lion.

• Results show the method further improves convergability on top of the additional $\beta$ scaling, demonstrating that the method can further benefit from advancements in initialization schemes. We updated the manuscript to include the result in Table 5 and the accompanying description.

Q4: number of prior training steps needed to approximating gradient trajectory

• The layer-wise gradient is measured as the average per parameter gradient magnitude of a layer at initialization. As the network is yet to have trained on a single minibatch, there is little deviation in the calculated values with respect to the actual iterations taken to measure to gradient. We recommend using a T of 100 and is the value used in the nlp experiments as it would still make the estimation more accurate and would take neglibible time compared to the entire training duration under most settings.

Q5: connection between using T=1 and the empirical study for ResNet50 in Figure 1

• The empirical study in Figure 1 is tracking and visualizing the step size modifier of various optimizers (AdamW, Lars and Lamb) when compared to SGD and AdamW. We show the per parameter values that is averaged across a layer and a single epoch. The value of T used refers to the number of minibatches used to measure the gradient at initalization for the proposed method, and is not performed in the empirical study in Figure 1. We are unsure of the question and would be glad to answer upon further clarification.

Q6: experiments with optimizers for large batch training (Lars, Lamb)

• Lars and Lamb are optimizers designed for large batch training, and while are also called "layer-wise learning rate" their motivation and methodology are different making them rather independent from our method.
• We have added figures for Lamb in the empirical study of Figure 1, which shows that for ResNet50, both Lars and Lamb show little layer-wise patterns. In particular, the common implementation of Lamb assigns very low values to the conv/linear layers while batchnorm layers are not modified, creating a large discrepancy in the assigned learning rates. On the other hand, our method gracefully handles the batchnorm and layernorm layers which are naturally assigned similar values to adjacent layers. We consider investigating the effect of assigning large learning rates to batchnorm or layernorm layers beyond the scope of our work.

We thank the reviewer for detailed and constructive feedback. We hope to sufficiently address your concerns and questions below.

W1 and Q4: Novelty compared to [1] and [2]

• The proposed method's layer-wise learning rate adjust the "macro" level step size such that the per-layer step sizes become significantly different from sign descent. As the assigned learning rates are non-adaptive values, it affects the entire training duration and is an additional multiplicative factor on top of the optimizer specific step size adjustments, which we will refer to as "micro" level step size adjustment.

• In contrast, the layer-wise learning rate of [2] is effectively using layer-wise average $1/\sqrt{v}$ instead of per-parameter $1/\sqrt{v}$, and mentions its similarity to Adam as a major characteristic. It is effectively a replacement/approximation to the "micro" level step size adjustment already performed by Adam.

W1: evidence provided in this paper to support the effectiveness of the proposed methods is primarily intuitive

• We provide extensive experiments that the "macro" level adjustment is effective over varying types of architectures and optimizers, that show strong empirical evidence the effectiveness is not limited to a particular optimizer or architecture.

W2: lacks a theoretical justification for the presented method, Q1: rationale behind the proposed method's effectiveness in conjunction with Lion optimizer, Q4: novelty and distinction from [4]

• Due to analytical convenience, theory behind initialization has been mostly focus at initialization to investigate architectural factors as it is a special case where it is not overshadowed by effects that arise during training. We demonstrate that initialization characteristics can be leveraged to improve training stability through the assignment of layer-wise learning rates at initialization. While theoretical justification would be beneficial, we believe empirical evidence of the effectiveness of the techniques is also valuable, and that initialization characteristics can carry on to prolonged training has also been demonstrated in [5].

• The effectiveness in conjunction with Lion further demonstrates that the "macro" level adjustment according to initialization characterisctics accounts for the architectural factors that is difficult to measure from in-training time gradients, and by extension makes it different the "micro" level step size adjustments. It shows the benefit of the additional multiplicative factor performed in conjunction with optimizers that have sign-descent properties (Lion, Adam), and aligns with the claims of [4] that having uniform changes despite large changes in (in-training) time gradients is a contribution to the success of Adam.

W2: The authors assert that the proposed methods can enhance model generalization

• The main argument is that the proposed method can improve convergability, which we demonstrate by showing that layer-wise methods significantly improve train loss when overfitting on small dataset when trained with AdamW and Lion in Table 2 and 3. We performed additional experiments that further demonstrate improved convergability with different optimizers (Adam-mini, Sohpia) in Table 8, 9 in the appendix and when scaling model depth in Table 5.

W3: The authors could enhance their evaluation by comparing their proposed method with more cutting-edge optimizers, such as Adam-Mini [2] and Sophia [3]

• Per reviewer's request, we provide additional experiments on Adam-mini and Sophia. The table below shows the lowest train loss when sweeping learning rates in small scale overfitting experiments and the validation loss when trained for 9.89B tokens.

• The layer-wise scheme improves train loss for all optimizers in small scale experiments, demonstrating the different effect of assigning static learning rates compared to Adam-mini. Interestingly, in 9.89B training the validation loss is slightly behind for both Adam-mini and AdamW, demonstrating that Adam-mini's layer-wise learning scheme is effectively approximating Adam. In fact, our results suggest that there is also a slight performance degradation when approximating AdamW through Adam-mini. The method also results in a noticable improvement for Sophia in 9.89B training.

[5] He, Bobby, et al. "Understanding and Minimising Outlier Features in Neural Network Training." arXiv preprint arXiv:2405.19279 (2024)