Initializing the Layer-wise Learning Rate

We thank the reviewer for the valuable review and feedback

Relation to adaptive optimisation methods

We believe that explicit layer-wise learning rates has an effect that is different and mostly independent from adaptive optimization methods. Figure 8 in the Appendix of the updated draft shows the adaptive step size of AdamW $1/(\sqrt{\hat{v}_{t}}+\epsilon)$ when training ResNet-50 for 90 epochs with single and the proposed layer-wise learning rate. The adaptive step sizes show a low to high trend according to layer depth on both single and the proposed method. If the proposed method was an cost-effective alternative to adaptive optimisation algorithms we believe the adaptive step sizes should have flattened out.

We would like to note that the observation layer-wise learning rates can improve training was on experimental settings where adaptive methods are known to achieve lower validation accuracy. We have updated the draft to better reflect the relation with adaptive methods.

Ablation study on $T$ of Algorithm 1

We performed an ablation study on $T$ when training ResNet-50 for 90 epochs with SGD by measuring the deviation of per-layer assigned learning rates compared to the $T$ used in the paper. We report the results below and have included it in the appendix.

Additional results with multiple runs

We provide additional results that reinforce the claim that the proposed method has a different effect compared to adaptive optimization methods. We report the result of training ResNet-50 with AdamW for 90 epochs below. While we used a T of 5004, the ablation study of T above showed the choice of T can be as low as 100 and is not significant, and for single learning rate we additionally iterate with learning rate of 0 for 100 iterations. We report the average and standard deviation of three runs below, and it can be seen that layer-wise learning rates have favorable performance even when using adaptive optimization methods.

Relation with warmup

We consider learning rate warmup to be largely independent and not a setup to be replaced or removed. Our experience is that it is beneficial to both single and layer-wise learning rate, and all experiments are performed with learning rate warmup.

Regarding the complex experiment setup and confounding factors

While we agree that the experimental setup can be affected by confounding factors such as learning rate schedule and warmup, ImageNet-1k classification has been extensively studied with established baselines and reported hyperparameters. We believe evaluating on such setups is effective in demonstrating methods that work on large datasets and architectures.

We thank the reviewer for the valuable review and feedback

Details on Algorithm 1

For $T$ we find a value of 100 or higher is sufficient to obtain resonable values. $G_{i}$ is incremented $T$ times in line 8, and is used to calculate the relative learning rates in line 10. After obtaining the relative values, the normalization is performed on lines 11-13 such that the average per-parameter learning rate is one.

Choice of $\mathcal{L}^1$ over $\mathbf{g}$ and inverse of the square root of $G_i$

Algorithm 1 is developed with the interpretation that the gradient magnitude itself is a objective to be minimized, so Algorithm 1 uses the gradient magnitude directly instead of performing $\mathcal{L}^2$-norm. That said, we find using $\mathcal{L}^2$-norm over the per-parameter gradient instead of $\mathcal{L}^1$-norm assigns very similar layer-wise learning rates, with an average deviation of 0.97%, and achieves similar performance of 77.24% on ResNet-50 when trained with SGD for 90 epochs.

The square root over $G_i$ in line 10 is performed because without it we found the difference in scale of assigned learning rates varies widely and results in drastically reduced performance of 75.99%.

Comparison to $1/f_{\text{in}}$ learning rate scaling

On ResNet-50, we find $1/f_{\text{in}}$ tends to assign assigns higher learning rates to initial layers and lower learning rates to later layers, which is directly opposite to the proposed method. Under the same unit per-parameter normalization of line 11-13 in Algorithm 1, $1/f_{\text{in}}$ scaling assigns 6.295 to the first 7*7 convolution layer, and 14.459 to the directly succeding convolution layer, which is in stark contrast to the 0.038 and 0.109 assigned by Algorithm 1 in Figure 4.

For ViT-S/16, $1/f_{\text{in}}$ scaling assigns identical learning rates to all the query, key and value layers in the self attention block, while the proposed method assigns higher learning rates to the query and key layers. We believe the practice of normalizing query and key layers in self attention [1][2] further validates the empirical significance of our method.

[1] Henry, Alex, et al. "Query-Key Normalization for Transformers." Findings of the Association for Computational Linguistics: EMNLP 2020. 2020.

[2] Dehghani, Mostafa, et al. "Scaling vision transformers to 22 billion parameters." International Conference on Machine Learning. PMLR, 2023.

Results under different learning rates

We provide additional results when training ResNet-50 with half and double the base learning rate for 90 epochs with SGD below. It shows that the proposed method performs well over various learning rates.

We thank the reviewer for the valuable review and feedback

Motivation of proposed Algorithm 1

Algorithm 1 adjusts the layer-wise learning rate to regularize architecture-induced convergence bias, measured as the gradient magnitude at initialization. It is motivated from an observation that layer-wise learning rate can further improve training of SGD in settings where adaptive methods are considered to have lower validation accuracy due to higher generalization gap.

Effect and performance of different architectures

Inspecting the assigned learning rates show that the convergence bias of Swin-T and ConvNeXt-T resembles that of vision transformer more than convolutional networks. While there could be various factors, we think lack of hyperparameter tuning could be a major reason for the lack of performance gain on Swin-T and ConvNeXt-T when using AdamW. Given how sensitive performance is to learning rates, we believe achieving competitive performance when learning rates of layers can differ by an order of magnitude is still a very interesting phenomenon.

Swin-T and ConvNeXt-T are complex architectures that incorporate design choices of both convolutional and transformer architectures. Swin-T incorporates hierarchical feature maps which is similar to feature map hierarchy in convolutional networks, and ConvNeXt-T stems from introducing vision transformer designs on convolutional networks. We mainly focus on ResNet and vision transformer as they are representative of the two main architecture family in vision tasks, and we believe performance on SGD is a good indication of training stablity and generalizability.

In Figure 6, single seems better than proposed method

In terms of final accuracy we agree Figure 6 shows using a single learning rate is better in such settings. However CIFAR-100 differs widely from ImageNet-1k in terms of dataset size, resolution and difficulty, and we believe methods that are beneficial on larger dataset size and difficulty is of interest even if it isn't necessarily beneficial on smaller tasks.

Ablation studies to validate the influence of hyper-parameters in Algorithm 1

We performed an ablation study on $T$ when training ResNet-50 for 90 epochs with SGD by measuring the deviation of per-layer assigned learning rates compared to the $T$ used in the paper. We report the results below and have included it in the appendix.

We thank the reviewer for the valuable review and feedback

Regarding the simplicity of the method and comparison with LARS, LAMB

We would like to emphasize that while LARS, LAMB are also layer-wise methods, the underlyring mechanisms and interpretation are very different in that LARS and LAMB relies on per-layer gradient normalization and additional weight norm scaling, while the proposed method is concered with architecture-induced convergence bias. In practice it means the learning rates of LARS and LAMB is insensitive to gradient of other layers, and due to the additional weight norm scaling layers initialized 0 are in principle not updated.

We have provided results of layer-wise gradient normalization techniques on ResNet-50 with SGD on Table 2. We found that the batchnorm scale and bias layers had to be excluded from LAMB scaling for LAMB to achieve competitive performace, perhaps due to its additional weight norm scaling. In that sense, we believe the proposed method can be considered simpler as the step size remains dependent only the gradient and does not modifiy the optimizer.

Training cost of Algorithm 1

We performed an ablation study on $T$ when training ResNet-50 for 90 epochs with SGD by measuring the deviation of per-layer assigned learning rates compared to the $T$ used in the paper. We report the results below and have included it in the appendix. We also ran a baseline experiement of 91 epochs, but found the final validation accuracy to be lower that the reported 76.80%, suggesting that other factors are more dominant compared to a single epoch difference.