REVISITING LARS FOR LARGE BATCH TRAINING GENERALIZATION OF NEURAL NETWORKS

1. Comments on whether TVLARS would be useful in datasets that are large in modern concept: We thank the reviewer for a very constructive comment.

• We agree that challenging datasets are required to verify a good algorithm. However, due to the short rebuttal phase, it is impossible for us to generate the results according to the challenging dataset such as ImageNet for the revised paper. We will implement and add the results according to the ImageNet in our public repository in the future. Besides, we want to discuss that the TinyImageNet can be considered as a challenging dataset for top AI conferences. Despite its volume compared to ImageNet, TinyImageNet is still considered as a challenging dataset and widely used in various research which has been accepted in top AI conferences (e.g., CVPR, ICCV, ICLR, ICML, NeurIPS). For instance:

  • Stochastic Marginal Likelihood Gradients using Neural Tangent Kernels (ICML 2023)
  • Diffusion Model as Representation Learner (ICCV 2023)
  • Online Continual Learning through Mutual Information Maximization (ICML 2022)
  • Delving into Out-of-Distribution Detection with Vision-Language Representations (Neural IPS 2022)
  • The Heterogeneity Hypothesis: Finding Layer-Wise Differentiated Network Architectures (CVPR 2021)

• Moreover, in various LBL research, the authors use less challenging dataset to evaluate the LBL performance (e.g., Cifar100, Cifar10, TIMIT, MNIST). For instance:

  • On Large Batch Training for Deep Learning: Generalization Gap and Sharp Minima (ICLR 2017)

• In addition, TinyImageNet is a challenging classification task that it contains 200 categories with 500 samples for each class. However, each sample's size is 64 $\times$ 64, which makes the model harder to extract the useful features from the image. Moreover, inheriting from ImageNet, the images are sampled in various scenarios, but with a much lower number of training samples, the model trained by using this data set will face heavily the out of distribution and long-tailed learning problem as discussed in the above papers.

2. Why we should drop the warm-up: We the reviewer for the helpful comment. We acknowledge that the main problem is the mutual information between Figures 1 and 2. Furthermore, we want to discuss that we use Figure 2 to show the relationship between $\Vert w \Vert$, $\Vert \nabla w\Vert$, and the test loss in different scenarios. By doing so, we can go to the conclusion that the stable convergence behavior in LBL come along with some characteristics (e.g., the stable reduction in $\Vert w \Vert$, and $\Vert \nabla w\Vert$) (mentioned in Section 3.2).

• Per your comment, we have removed Figure 1 (as it is already shown in Appendix H).

We thank the reviewer for their time and efforts, as well as their valuable comments.

As per the reviewer's suggestions, we have made the following revisions to our paper:

We have revised the abstract along with the introduction in order to:

• We provided references to sharp minimizers along with a brief definition for the reader's comprehension.
• We have explained briefly the potential issues of warm-up techniques when incorporating them into LARS and LAMB to tell why we have to propose our new algorithm TVLARS.

We have revised the empirical analysis (section 2) and methodology (section 3) in order to:

• We highlighted the difference between our TVLARS and WALARS to address your comments about the clarification of the paper’s motivation and contributions.
• We conducted further empirical experiments according to layer-wise learning rate annealed by warmup strategy to give a more thorough understanding about the shortcomings of using warm-up in LARS and LAMB.
• We demonstrated the algorithm more systematically and clearly defined the algorithm’s ingredients.
• We revised and improved the paper consistency in using notations and abbreviations.

1. Why we should drop the warm-up: We totally agree that warmup is a very simple and important method while our proposed method requires parameter tuning to get great performance. However, we want to discuss that:

• The existing warm-up strategy comprises two phases: 1) a gradual rise in the scaling ratio of the learning rate from 0 to $\frac{B}{B_\textrm{base}}$, and 2) a gradual decay in the scaling ratio of the learning rate from $\frac{B}{B_\textrm{base}}$ back to the lower threshold $\gamma_\textrm{min}$. From these work flow, we have two observations:

  • First of all, we recognize that phase 1) is redundant. When the scaling ratio $\gamma^k_t$ is excessively small, especially in the early stages, the learning process fails to prevent the memorization of noisy data [1], [2]. Additionally, if the model becomes trapped in the sharp minimizers during the warm-up phase, the steepness of these minimizers prevents the model from escaping and further loss landscape exploration. As a consequence, this process consumes a significant amount of computation resources and learning time.
  • Secondly, as the decay ratio of the warmup technique in LARS is fixed by $\frac{1}{2}\Big[1 + \cos\Big(\frac{t-de}{T-de}\pi\Big) \Big]$, where t is the current time step and the T is the total timestep. Therefore, the decay characteristics are the same among all datasets and model architectures. As a result, when applying to different datasets (with different sharp minimizers distributions [3]) and model architectures, the learning is not adaptable and tunable to achieve the best performance.

• As per your comment, we have:

  • Revised the introduction, LARS analysis, and proposed method to indicate the drawbacks of warmup in LARS and emphasize our TVLARS motivations and contributions.
  • We conducted further empirical experiments according to the layer-wise learning rate value to reinforce our observations about the redundant learning phase when incorporating warm-up into LARS and LAMB. The experiments are demonstrated in Appendix J.

2. Hyper-parameter tuning needs to be tuned more which is less convenient to use: In our proposed method, there are several parameters that do not need tuning. However, to conduct a fair comparison between TVLARS, LARS, and LAMB, they are fixed at specific values.

• In TVLARS, the layer-wise learning rate $\gamma_t^k$ is annealed by the formula: $\phi_t = (\alpha + e^{\psi_t})^{-1} + \gamma_{\min}$ where $\psi_t = \lambda(t - d_{\rm e})$. A fair comparison is conducted when the target learning rate of the three optimizers and the final lower threshold $\gamma_{\rm min}$ are the same, which means $\alpha$ is always set to $1$ and $\gamma_{\rm min}$ is set to $\frac{\mathcal{B}}{\mathcal{B}_{\rm base}} \times 0.001$.
• As a result, only one parameter needs tuning: $\lambda$. As we have stated above, the use of $\lambda$ is crucial as the decay characteristics may not be the same among datasets and model architectures. Therefore, using $\lambda$ can not only make the optimizer adaptive to different scenarios but also to different sizes of batch, as we made an ablation test on this in Section 5.2.1.
• Finally, using those parameters, upper and lower boundaries are constructed to make the training more stable.

3. The accuracy is still too low according to LARS/LAMB:

• As we can see from Table 1, our TVLARS has higher accuracy in almost every case of classification tasks. Moreover, in the cases that our proposed method did not dominate the WALARS, our performance gap between WALARS is trivial (i.e., $<0.3\%$). Especially, in SSL tasks, our TVLARS outperforms LARS and LAMB in every case. Notably, in many cases, we achieve from 5 to 10% higher than LARS and LAMB in terms of accuracy.
• Moreover, as we mentioned in answer to questions Q1, and Q2, our TVLARS method does not only outperform warmup LARS and LAMB in terms of accuracy but also achieves faster convergence. This significant improvement comes from the cancellation of the redundant ratio scaling phase (as mentioned in Figure 1), which makes the statistical model stuck at the sharp minimizer.

4. Try to use different initialization methods and analyze their results. We thank the reviewer for the constructive comment.

• We agree that the layer-wise optimization method is sensitive to the initialization method of weights. In response to your comments, we have carried out further comprehensive experiments under different model initialization methods and the results have been added in the revised manuscript. As we can see from the extensive ablation test, our proposed TVLARS significantly outperforms the two baselines. Besides, we want to discuss more about our implementation and first manuscript. We want to reveal why we can claim that our proposed method is carefully considered in terms of fine-tuning.
• First of all, we claim that we do use the same method of weight initialization as the authors of the LARS and LAMB method, which is the Xavier Uniform sampling method LARS, LAMB. However, taking this challenge as a chance, we conducted more experiments with different methods of weight initialization including Xavier Uniform, Xavier Normal, Kaiming He Uniform and Kaiming He Normal, which are the most popular methods nowadays. After conducting those experiments, we found out that the methods of weight initialization do not affect the performance of the training model in LBL scenarios as the results are close to the first ones.
• Among all weight initialization methods, the method that we used is known to be the most popular and widely used ere are some official implementations that use LARS optimizer with the same weight initialization method as our implementation.
  • BarlowTwins
  • VICReg
  • SimCLR