Improving Generalization and Convergence by Enhancing Implicit Regularization

We thank the reviewer for the appreciation of our work and insightful comments. Below, we address the reviewer’s questions in detail.

• W1. One weakness is that the motivation comes from the stylized example in section 2, which may be unrealistic for general deep learning optimization. For example, the idea of "sharp coordinates" only makes sense in this stylized example.

  Response: We thank the reviewer for this insightful question.

  • Clarification. We clarify that the idea behind IRE is to distinguish between "sharp/flat directions" rather than "sharp/flat coordinates". The use of a diagonal Hessian in Section 2 is only for simplification purpose. In fact, the theoretical analysis in Section 5 considers a more general setting without assuming a diagonal Hessian (see Assumption 5.1).

  • Practical Considerations. In the practical implementation described in Section 3, we indeed use the diagonal Hessian to approximate the full Hessian, aiming to reduce computational costs. The diagonal Hessian is commonly acknowledged as a reasonable approximation of the full Hessian in deep neural networks and is frequently employed in the design of deep learning optimizers, as noted in references [1][2]. For a more detailed discussion, please refer to lines 156-162 in our manuscript.

• Q1. Shouldn't this algorithm also speed up optimization? The algorithm uses a larger learning rate on some coordinates, which should also have an optimization effect. (This would not be visible in the analysis setting of manifold-of-global-minima because no optimization is occurring in that setting.)

  Response: We thank the reviewer for this insightful comment and will include a discussion on this issue in our revised version. The motivation of IRE is to speed up the sharpness reduction, which only requires to increase learning rate along completely flat (zero-curvature) directions. However, practical implementation may also increase learning rate along directions with small but non-zero curvatures, which can further speed up loss convergence. However, explaining why this approach can provide a significant acceleration is unclear.

• Q2. How would you define IRE at the most abstract level? If we define an abstract version of preconditioned gradient descent as $\theta_{t+1}=\theta_t-Pg_t$ (where vanilla GD with learning rate $\eta$ corresponds to $P=\eta I$), is the idea that it's better to use $P_2$ rather than $P_1$ provided that $P_2\succeq P_1$ and both are locally stable i.e. $\lambda_{\max}(PH)<2$ for $P\in\{P_1,P_2\}$?

  Response: We are grateful to the reviewer for sharing this high-level perspective, which can help explain the acceleration of loss convergence observed with IRE. Given a base optimizer with update $\theta_{t+1}=\theta_t-P_1g_t$, we aim to improve it by selecting a better $P_2=QP_1$, while ensuring $\lambda_{\max}(P_2 H)<2$ to maintain training stability. Let $P_{flat}$ be the projection matrix into the flat (zero-curvature) directions. In our IRE approach, we choose $P_2=(I+\mu P_{flat}) P_1$, which meets the condition $P_2\succeq P_1$ mentioned by the reviewer. However, we remark that our original motivation of this specific $P_2$ is to accelerate effective dynamics on the manifold of global minima, thereby accelerating the sharpness reduction.

[1] George et al. Fast approximate natural gradient descent in a Kronecker-factored eigenbasis. (NeurIPS 2018)

[2] Liu et al. Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. (ICLR 2024)

We appreciate the reviewer's recognition of our work and helpful comments. Below, we offer detailed responses to the reviewer’s questions:

• W&Q1. Comparison with related works.

  Response: We thank the reviewer for this question and will provide a more detailed comparison with related algorithms in the revised version. First, we clarify that the mechanism behind IRE in accelerating sharpness reduction is fundamentally different from that of SAM. Moreover, IRE can be integrated with generic base optimizers, including SGD, Adam, and SAM. In contrast, the SAM variants, including Kwon et al. (2021) and Liu et al. (2022), are specifically designed to improve SAM.

• W&Q2. Limitations and failure cases.

  Response: Thank the reviewer for the insightful question. We acknowledge that diagonal Hessian may not always provide accurate second-order information for identifying the flat (zero-curvature) directions. Recall that IRE iterates as $\theta_{t+1} = \theta_t - \eta (1+\kappa P) g_t$, where $P$ represents the projection to flat directions. When $\kappa=0$, it recovers the base optimizer. For problems where diagonal Hessian provides relatively accurate second-order information, a large $\kappa$ can be used. Conversely, for problems where diagonal Hessian is less informative, we can set $\kappa$ to be small or even to zero. This approach allows to control over the influence of whether the diagonal Hessian approximation is effective. This idea is analogous to the damping technique adopted in Newton's methods with an approximate Hessian. Specifically, we tune $\kappa\in${1,2} for CNNs, $\kappa\in${2,3,4} for Llama pre-training, and $\kappa\in${20,50} for ViTs.

• W&Q3. Experimental setup and implementation details.

  Response: Thank the reviewer for this question. For image classification tasks, the models including ResNets and ViT-T/S follows the ones in Muller et al. (2023). For LLM pre-training, we use the Llama models from HuggingFace. For image data pre-processing, we apply the default normalization technique: transforms.Normalize(mean, std). The search ranges of hyperparameters, as well as the data augmentation methods are specified in Appendix B. In the revised version, we will carefully incorporate all these experimental details in the main text.

• W&Q4. Comparisons with state-of-the-art optimizers.

  Response: Thank the reviewer for this question. We clarify that IRE is a generic framework for boosting optimizer's implicit bias and can be integrated with general optimizers. Due to space and resource limit, we have only compared IRE with the popular algorithm including SGD, SAM and AdamW. Specifically, we show that IRE can improve the generalization in image classification task and convergence of AdamW in LLM pre-training. As for other optimizers mentioned by the reviewer, Adagrad has been replaced by Adam(W) in most fields; second-order algorithms like K-FAC and Shampoo is computational expensive and it is less popular in training large models.

  About regularization, IRE is proposed to enhance implicit regularization and thus, orthogonal to those explicit regularization, including weight decay, dropout and label smoothing. It is worth mentioning that our experiments incorporate various regularization techniques such as weight decay and label smoothing. Please refer to the Appendix for more details.

• W&Q5. Downstream performance evaluation.

  Response: Thank the reviewer for this constructive comment. It is valuable to explore the performance of the pretrained models in downstream tasks. However, due to the limited time for response, we could not fully evaluate the downstream tasks for pretrained models, but we conducted preliminary experiments to explore this issue.

  • Notably, "Liu et al. (2023) Same pretraining loss..." points out a strong correlation between the flatness and downstream task performance among models. Specifically, it implies that for models with the same pre-training loss, flatter solutions yield better performance on downstream tasks.
  • Inspired by this observation, we evaluated the flatness, measured by the trace of Hessian, of models pre-trained by AdamW and AdmIRE. Due to time constraints, we only focus on the experiments in Fig.3 (left), i.e., training Llama (60M) on wiki-103 dataset using either AdamW or AdmIRE. The experimental results, shown in the attached PDF in our Response to All Reviewers, demonstrate that AdmIRE not only achieves the same loss in only half the iterations required by AdamW, but also the solutions found by AdmIRE are significantly flatter than that found by AdamW.

  In the revised version, we will include a comprehensive evaluation of downstream tasks.

• W&Q6. Theoretical analysis.

  Response: Thanks the reviewer for the question.

  • We clarify that Assumption 5.1 essentially only assume: 1) the loss $\mathcal{L}(\cdot)$ is $C^4$ smooth and 2) the global minima are connected. For over-parametrized models, this connectivity assumption on the minima manifold has been empirically verified in works such as Draxler et al. (2018) and Garipov et al. (2018), and theoretically supported in Cooper (2018). We refer to lines 267-270 of our manuscript for more details.
  • Our main theoretical results are explained in lines 286-369, where we provide an intuitive interpretation: IRE accelerates the "effective dynamics" of base optimizers in flat directions, achieving faster reduction of sharpness; moreover, IRE maintains stability by not affecting movement along sharp directions. This dual effect ensures IRE can improve generalization performance without compromising training stability.
  • As for other optimization techniques, our experiments have incorporated momentum, weight decay, and the adaptive method (AdamW). However, these techniques poses new challenges to theoretical analysis, which we leave as future work.

We appreciate the reviewer’s recognition of our work and helpful comments. Below, we offer detailed responses to the reviewers questions.

• W1. The improvements of IRE in Table 1 and 7 are not as significant as expected. It makes one wonder about its usefulness for CNN networks and its oversensitivity to hyperparameters.

  Response: We thank the reviewer for this insightful comment.

  • CNNs vs ViTs. It is important to emphasize that CNNs are designed with a strong image prior, which limits the potential for further improvement through regularization. Consequently, it is not unexpected that the improvement achieved by IRE for CNNs is less significant. This aligns with observations in previous regularization techniques, such as ASAM [1] and SAM-ON [2], which similarly shows limited improvements for CNNs. For instance, ASAM improves SAM by only 0.26% points when training WRN-28-10 on CIFAR-100, and SAM-ON improves SAM by just 0.17% points when training ResNet-50 on ImageNet. In contrast, ViTs, which have a much weaker image prior, benefit more from regularization, as shown in Table 3.
  • Sensitivity to hyperparameters. 1) For image classification tasks, the optimal hyperparameters for IRE do vary depending the dataset and model. Nevertheless, Table 7 demonstrates that IRE consistently improves performance even without tuning the hyper-parameters, although the improvement is modest. 2) In contrast, when training Llama models, IRE with a fixed hyperparameter $\gamma=0.6$ performs effectively across varying model and data sizes, as shown in Fig.3.

• W2. A big concern for me is that the experiment in Figure 3 does not appear to have converged yet, the rate of convergence in the early part of an experiment does not equate to the rate of convergence throughout the training process, as well as a more detailed analysis of the performance of the model after convergence should have been added to demonstrate that the final position of convergence is good.

  Response: We thank the reviewer for raising this concern.

  • Loss convergence in LLM pre-training. It is important to note that in LLM pre-training, the training loss never decrease significantly due to the huge size of dataset. For example, as reported in [3], the official Llama pre-training achieves only a final loss around 1.55, despite the model size being 65B.
  • Final loss v.s. performance. In LLM pre-training, it is widely observed that the performance on downstream tasks is predominantly determined by the final pre-training loss and has little correlation with other pre-training components such as model types. We refer to [4] for a detailed analysis of this issue. Therefore, the primary focus in LLM pre-training is to reduce the final loss as much as possible, given specific computational and data resources.
  • Final loss. Our final loss aligns with expectations based on our model scale. For instance, for our baseline, our LLama (229M) was trained on openwebtext by AdamW for 100k steps, resulting in a final loss of 2.835. Fig.1(d) in [5] show that GPT models of varying sizes were trained on openwebtext by AdamW for 100k steps, achieving losses of 2.915 for GPT-small (125M) and 2.695 for GPT-middle (355M). Our loss of 2.835 for Llama (229M) (the baseline) is consistent with these results. Notably, our proposed AdmIRE achieves the same loss of 2.835 with only 50k iterations, which is half the number of iterations required by AdamW.

• W3. Judging from the code, the cost of each training step is at least twice that of SGD.

  Response: We clarify that the cost of our IRE takes twice that of base optimizers only during the steps of estimating the diagonal Hessian and updating the mask. However, these steps are triggered only once every $K$ steps (see Alg. 1). In all our experiments, we set $K=10$ (see l.168-l.172), so the average cost of IRE is only 1.1 times that of the base optimizer. Table 4 further verify this estimate empirically.

• Q1. Could IRE help in the finetuning phase of LLMs, Including both convergence speed and convergence position properties?

  Response: We thank the reviewer for this constructive question. To address this, we have conducted a new Supervised Fine-Tuning (SFT) experiment. Specifically, we finetune pretrained Llama2-7B with LoRA on Stanford's Alpaca dataset. The results, shown in the attached PDF in our Global Response to All Reviewers, demonstrate that IRE can accelerate the convergence in SFT. This is consistent with the advantage of IRE in pretraining. Unfortunately, we do not have enough time to evaluate the performance of the convergent models on downstream tasks. In the revised version, we will include complete results for SFT.

• L1. The author did not provide information on the computing resources they used.

  Response: Thank you for this reminder. For Section 4.1 (image classification), the experiments on Cifar-10/100 were conducted using a single A800 GPU, while the experiments on ImageNet were conducted using 4 A800 GPUs. Details regarding the computing resources for the experiments in Section 4.2 are provided in Appendix B (l.664 and l.678).

[1] Kwon et al. ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. (ICML 2021)

[2] Muller et al. Normalization Layers Are All That Sharpness-Aware Minimization Needs. (NeurIPS 2023)

[3] Touvron et al., LLaMA: Open and Efficient Foundation Language Models. (2023)

[4] Du et al. Understanding Emergent Abilities of Language Models from the Loss Perspective. (2024)

[5] Liu et al. Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. (ICLR 2024)