Paper: Initializing the Layer-wise Learning Rate (ICLR-2024)

Reviewer comment (Writing):
(clarity) Layer-wise learning rates, which are probably also addressed in the first edition of tricks of the trade (1998), were probably the only way to make deep networks trainable. The main advantage of adaptive optimisers has always been that the painfull process of finding learning rates for each layer is no longer necessary.

Author response:
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.

Reviewer comment (Experiments):
(clarity) It is unclear how important the choice for $T$ in algorithm 1 is. In the experiments, $T$ is the number of batches in one epoch, but there are no ablations for different choices of $T$.

Author response:
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.

$T$ | Max deviation % | Average deviation % | Validation accuracy %|
-----------   | :-------------:| :-------------:| :-------------:|
1    | 11.02            | 2.66                 | 77.07               | 
10   | 3.51             | 0.69                 |  77.01              | 
100  | 1.01             | 0.18                 |   77.04             | 
1000 | 0.32             | 0.06                 |   77.07             |
5004 | 0               | 0                     |    77.15             |

Reviewer comment (Reproducibility):
In Algorithm 1, the $G_i$ are apparently incremented $T$ times, but they are not normalized by $T$ or any other quantity depending on $T$. This raises two questions: how should $T$ be chosen, or should the $G_i$ be normalized somewhere?

Author response:
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.

Reviewer comment (Theory):
In Algorithm 1, line 8, there is no explanation for why the $\mathcal{L}^1$-norm over $\mathbf{g}$ is used, instead of the $\mathcal{L}^2$-norm or any other norm.

Author response:
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.

Reviewer comment (Theory):
In Algorithm 1, line 10, there is no justification for choosing the inverse of the square root of $G_i$, instead of the inverse of $G_i$ or any other quantity.

Author response:
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%.

Reviewer comment (Theory):
In Algorithm 1, lines 11-13, there is no justification provided for the computation performed there.

Author response:
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.

Reviewer comment (Reproducibility):
The reviewer asks for the authors’ perspective on the paper *Neural tangent kernel: Convergence and generalization in neural networks*, Jacot et al., 2018. In that paper, each weight tensor is scaled by $1/\sqrt{f_{\text{in}}}$, which, combined with a unique learning rate for all layers, is equivalent to the “normal” setting (without scaling) using a per-layer learning rate proportional to $1/f_{\text{in}}$. The reviewer asks how the learning rates computed by the authors compare to these.

Author response:
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.

Reviewer comment (Experiments):
Regarding experimental results, the reviewer asks whether the results are consistent when the learning rate is changed, and whether the proposed method performs better than the “single lr method” in any circumstance.

Author response:
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.

Initial LR | Single | Layer-wise | 
-----------   |  :-------------:| :-------------:| 
0.05 | 76.61    | **76.67** | 
0.1 | 76.80    | **77.12** | 
0.2 | 76.33    | **76.90** |

Reviewer comment (Experiments):
In Algorithm 1, the reason of choice of T and corresponding ablation study is missing which may be vital to the performance of the proposed algorithm.

Author response:
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.  
$T$ | Max deviation % | Average deviation % | Validation accuracy %|  
-----------   | :-------------:| :-------------:| :-------------:|  
1    | 11.02            | 2.66                 | 77.07               |   
10   | 3.51             | 0.69                 |  77.01              |   
100  | 1.01             | 0.18                 |   77.04             |   
1000 | 0.32             | 0.06                 |   77.07             |  
5004 | 0               | 0                     |    77.15             |

Reviewer comment (Writing):
What is the motivation of proposed Algorithm 1?

Author response:
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.

Reviewer comment (Experiments):
Why layer-wise learning rate scheme performs not so good on Swin-T and ConvNeXt-T when using AdamW? ResNet-50 and SGD are no longer mainstream models or algorithms in 2023. What are the impacts of proposed algorithms on different model structures and modules?

Author response:
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.

Reviewer comment (Presentation):
Figure 6 is not intuitive, even seems that single way is better than proposed methods.

Author response:
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.

Reviewer comment (Experiments):
More ablation studies are needed to validate the influence of different hyper-parameters in Algorithm 1.

Author response:
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.  
$T$ | Max deviation % | Average deviation % | Validation accuracy %|  
-----------   | :-------------:| :-------------:| :-------------:|  
1    | 11.02            | 2.66                 | 77.07               |   
10   | 3.51             | 0.69                 |  77.01              |   
100  | 1.01             | 0.18                 |   77.04             |   
1000 | 0.32             | 0.06                 |   77.07             |  
5004 | 0               | 0                     |    77.15             |

Reviewer comment (Experiments):
You need to compare the proposed method with more layer-wise optimization method, such as LARS and LAMB. I noticed that the main baseline is SGD and Adam. and these methods are not layer-wose methods. Although their performance is very strong, I think LRAS / LAMB can also further improve the performance of SGD / Adam. To better illustrate the performance gain of your method, maybe you should provide these results on layer-wise method.

Author response:
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.

Reviewer comment (Experiments):
I would like to ask the training cost of Algorithm 1. Since the method need to estimate the gradient magnitude and other methods don't need it. Therefore, I would like to ask the cost, such as time. In addition, the proposed need to use T steps in Algorithm 1 and that means we need more steps to finish the training with the proposed method. If we add these T steps to the baselines, such as SGD and Adam, whether we can further improve their performance?

Author response:
We thank the reviewer for the valuable review and feedback

> 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. 
$T$ | Max deviation % | Average deviation % | Validation accuracy %|
-----------   | :-------------:| :-------------:| :-------------:|
1    | 11.02            | 2.66                 | 77.07               | 
10   | 3.51             | 0.69                 |  77.01              | 
100  | 1.01             | 0.18                 |   77.04             | 
1000 | 0.32             | 0.06                 |   77.07             |
5004 | 0               | 0                     |    77.15             |