Torque-Aware Momentum

We thank the reviewer for their evaluation and for recognizing the novelty and elegance of our proposed TAM method, as well as the breadth of experiments across vision and language tasks. We also appreciate the constructive suggestions, which we address in detail below.

Reviewer’s Comment: “This paper lacks convergence analysis.”

Response: We understand the reviewer's concern about convergence proof. To address this, we propose to add an argument that leverages TAM's relationship with SGDM at late training times discussed in Section 3 (“Learning Rate Transfer”). The argument goes as follows:

• As training progresses, TAM's cosine similarity term $ŝ_t$ stabilizes to a constant value $s^*$ (which we empirically observed to be approximately 0, see Appendix A.2.1).
• Under this late-time behavior, TAM's update rule effectively reduces to SGDM with a modified learning rate: $\eta^{\text {eff}}_\text{TAM} \approx \frac{1+{s}^*}{2(1-\beta)}\eta$.
• This equivalence means that in the neighborhood of optima, where $\hat{s}_t$ has stabilized, TAM inherits the well-established convergence guarantees of SGDM [1,2].
• The damping factor $(1 + \hat{s}_t)/2$ remains bounded, ensuring the effective learning rate stays within a controlled range throughout training.

This theoretical connection to SGDM, combined with our empirical evidence of $s^*$ stabilizing to $0$, ensures TAM's convergence while maintaining its enhanced exploration capabilities during early training. We've added this explanation in Section 3 in the revision.

Reviewer’s Comment: “While large oscillations are bad, oscillations are sometimes good for searching minima intuitively.”

Response: We agree that oscillations can sometimes be beneficial for exploring the loss landscape. As shown in Figure 6, TAM does not eliminate oscillations entirely. Instead, it defers abrupt sharpening events caused by large oscillations, promoting smoother exploration. Additionally, our LMC analysis highlights how TAM enhances training stability, which we argue translates to better generalization performance. To strengthen this point, we have also added more gradient norm plots to the Appendix A.2.6, as well as an analysis of $\hat{s}_t$ (cosine similarity) in Appendix A.2.1, which provides a more direct measure of oscillatory behavior. We hope these additions address your concerns.

Reviewer’s Comment: “The empirical improvements are marginal, especially for vision tasks.”

Response: We would like to emphasize that the primary goal of TAM is to enhance momentum-based optimizers by stabilizing gradient directions and improving exploration of the loss landscape. While we may expect TAM to match the performance of standard momentum methods in conventional supervised learning settings, its strengths become evident in more complex scenarios involving misaligned gradients or task adaptation:

• Fine-tuning: Fine-tuning on new datasets while retaining previous knowledge is challenging. TAM outperforms other optimizers in avoiding sharp minima and improving exploration, as shown in Fig. 3 and Fig. 4 across different models and compute budgets. Notably, AdaTAM achieves better performance on the Wikitext dataset and generalizes well on most MTEB datasets without overfitting.
• Online learning: In non-stationary environments, where maintaining plasticity to adapt to new tasks is crucial, TAM demonstrates superior adaptability across task boundaries (Fig. 5).
• Continual learning: Following Reviewer fh59’s suggestion, we also added new continual learning experiments in Appendix A.2.3, comparing SGDM and TAM on the CLEAR dataset (with 10 sequential tasks). The results confirm that TAM outperforms SGDM in settings that require both stability and plasticity.

Therefore, all these experiments highlight the advantages of TAM in stabilizing momentum updates and enhancing exploration in such settings.

Reviewer’s Comment: “The work missed a lot recent relevant works on analyzing or upgrading momentum. For example, [1,2] theoretically studied various momentum methods; [3, 4] designed momentum variants.”

Response: We appreciate the reviewer bringing these works to our attention. We have added the suggested references to the manuscript and updated the related work section to acknowledge their contributions.

[1] Yan, Y., Yang, T., Li, Z., Lin, Q., & Yang, Y. A unified analysis of stochastic momentum methods for deep learning. IJCAI 2017.

[2] Liu, Y., Gao, Y., & Yin, W. An improved analysis of stochastic gradient descent with momentum. NeurIPS 2020.

We thank the reviewer for their constructive feedback, which has helped improve our work. We believe our responses address the concerns raised, and we would appreciate it if the updated score reflects this. Please let us know if further clarifications are needed.

Reviewer’s Comment: “ While it is impressive that TAM exhibits less oscillation compared to SGDM in Figure 6, the relationship between AdaTAM and Adam is not shown ... it would be beneficial to include graphs obtained from AdaTAM as well.”

Response: Thank you for your suggestion. We have added a plot in the Appendix A.2.6 comparing AdaTAM and Adam (Figure 9-right), showing the evolution of gradient norm (Grad Norm) using the same setup as in Fig. 6, with a fixed learning rate of 0.001 for both adaptive optimizers. The results demonstrate that while both optimizers experience a high gradient norm in the first epoch, AdaTAM consistently maintains a lower gradient norm as training progresses, indicating that the damping effect is effectively controlling large gradients in AdaTAM.

Reviewer’s Comment: “the threshold for similarity is set at 1%-point ... It is recommended to lower the threshold for similarity to around 0.2%-point.”

Response: Thank you for your recommendation. We have updated the results in Figure 4 to use a 0.2%-point similarity threshold and moved the previous 1%-point results to the Appendix A.2.7. With this change, we observe that, except for RoBERTa-large and BERT-base fine-tuned for 10 epochs, AdaTAMW generally matches or outperforms AdamW in most settings.

Reviewer’s Comment: “There is no convergence proof for TAM and AdaTAM. ”

Response: We understand the reviewer's concern about convergence proof. To address this, we propose to add an argument that leverages TAM's relationship with SGDM at late training times discussed in Section 3 (“Learning Rate Transfer”). The argument goes as follows:

• As training progresses, TAM's cosine similarity term $ŝ_t$ stabilizes to a constant value $s^*$ (which we empirically observed to be approximately 0, see Appendix A.2.1).
• Under this late-time behavior, TAM's update rule effectively reduces to SGDM with a modified learning rate: $\eta^{\text {eff}}_\text{TAM} \approx \frac{1+{s}^*}{2(1-\beta)}\eta$.
• This equivalence means that in the neighborhood of optima, where $\hat{s}_t$ has stabilized, TAM inherits the well-established convergence guarantees of SGDM [1,2].
• The damping factor $(1 + \hat{s}_t)/2$ remains bounded, ensuring the effective learning rate stays within a controlled range throughout training.

This theoretical connection to SGDM, combined with our empirical evidence of $s^*$ stabilizing to $0$, ensures TAM's convergence while maintaining its enhanced exploration capabilities during early training. We've added this explanation in Section 3 in the revision.

Reviewer’s Comment: “Why does TAM diverge at epoch 0 in the middle figure of Figure 6?”

Response: We would like to clarify that TAM does not diverge at epoch 0 in the middle figure of Fig 6. The gradient norm starts at a higher value for all methods, including TAM, and decreases quickly after the first epoch before increasing again in later stages due to abrupt sharpening. To address the confusion, we have updated Figure 6 to improve clarity.

Reviewer’s Comment: “It would be useful to understand how performance changes with hyperparameters, as seen in Adam.”

Response: Thank you for this suggestion. We conducted an ablation study on ResNet18, varying gamma while keeping all other hyperparameters fixed for CIFAR10 and CIFAR100. Our experiments indicate that varying gamma has minimal impact on the overall behavior of the optimization trajectories. The results, included in the Appendix A.2.4, demonstrate that even with changes in gamma, TAM consistently outperforms other baselines. The results are as follows:

CIFAR10

CIFAR100

Thank you for your detailed and thoughtful review. We appreciate your recognition of the strengths of our work, particularly in applying the damping effect based on cosine similarity, the diverse set of experiments on language models, and the observed impact of TAM on gradient stability. We also appreciate the suggestions and have provided our responses to the concerns below:

Reviewer’s Comment: “TAM and AdaTAM(W) are evaluated using a limited set of models, ... It is recommended to use other models in the image domain ... how performance would change if the TAM approach were used from the pre-training stage of the LM.”

Response: Thank you for your suggestion. In response, we have expanded our experiments to include additional models. Specifically, we have evaluated TAM with MobileNet on CIFAR-10/100 and ViT fine-tuning (pre-trained on ImageNet-1k). We perform the learning rate grid search to obtain the best performing setup. Details of these experiments are included in the revised version of the paper in Appendix A.2.2. We also provide the results below:

CIFAR10 + MobileNet:

CIFAR100 + MobileNet:

CIFAR10 + ViT:

CIFAR100 + ViT:

Overall, we observe that TAM and AdaTAM either match the performance or outperform the baselines similar to results in Table 1.

We also would like to inform the reviewer that we are also currently running experiments on training GPT, which we will update the paper with when they are finished.

Reviewer’s Comment: “In Table 1, the performance of TAM on the ImageNet dataset is quite margina ...”

Response: We would like to emphasize that the primary goal of TAM is to enhance momentum-based optimizers by stabilizing gradient directions and improving exploration of the loss landscape. While we may expect TAM to match the performance of standard momentum methods in conventional settings, its strengths become evident in more complex scenarios involving misaligned gradients or task adaptation:

• Fine-tuning: Fine-tuning on new datasets while retaining previous knowledge is challenging. TAM outperforms other optimizers in avoiding sharp minima and improving exploration, as shown in Fig. 3 and Fig. 4 across different models and compute budgets. Notably, AdaTAM achieves better performance on the Wikitext dataset and generalizes well on most MTEB datasets without overfitting.
• Online learning: In non-stationary environments, where maintaining plasticity to adapt to new tasks is crucial, TAM demonstrates superior adaptability across task boundaries (Fig. 5).
• Continual learning: Following Reviewer fh59’s suggestion, we also added new continual learning experiments in Appendix A.2.3, comparing SGDM and TAM on the CLEAR dataset (with 10 sequential tasks). The results confirm that TAM outperforms SGDM in settings that require both stability and plasticity.

Therefore, all these experiments highlight the advantages of TAM in stabilizing momentum updates and improving exploration in such settings.

We’d like to thank the reviewer for the overall positive evaluation of our paper, specifically the appreciation of the novelty, robustness and wide applicability of our method, and the clarity of our exposition. We appreciate the suggestions and have provided our responses below:

Reviewer’s Comment: “Despite the fact that TAM is computationally efficient, the requirement to compute the cosine similarity between gradients and momentum might introduce non-trivial implementation complexity in certain training frameworks.”

We would like to clarify that TAM only requires the computation of the cosine similarity between two precomputed vectors (gradient and momentum), which involves a simple normalized dot product. This operation is computationally efficient and induces negligible overhead in modern training frameworks. Furthermore, TAM does not require any additional memory beyond storing the gradient and momentum vectors, which are already maintained in standard optimizers like SGDM.

Please note that we also provide a runtime time comparison in training BERT models in Appendix A.2.8.

Reviewer’s Comment: “While comparisons with the standard optimizers and a few existing appoachs are robust, additional benchmarks against more recent optimizers concentrating on gradient stability would strengthen the claims furthermore.”

Thank you for the suggestion. We have compared TAM with AngularGrad, the closest baseline targeting gradient stability. If you have specific recommendations for additional recent optimizers to benchmark against, we would be happy to consider them in future experiments.

Reviewer’s Comment: “The paper illustrates TAM's effectiveness on various tasks, but its performance in more demanding and complex training scenarios, such as very large LLMs (e.g., 70B), could be further explored.”

Thank you for the suggestion. While running experiments on very large-scale models like 70B LLMs is well beyond our current computational resources, we have included evaluations across a range of model sizes, including BERT variants, to demonstrate TAM's scalability and effectiveness. We believe these experiments provide sufficient evidence of TAM's robustness and general applicability across diverse scenarios.

Reviewer’s Comment: “Would TAM still perform well in non-stationary environments, such as continual learning, where the gradient directions frequently change or shift?”

Yes, TAM is particularly well-suited for non-stationary settings, as exploring the loss surface and preserving previously learned features are critical for maintaining a balance between plasticity and stability in continual learning. In Fig. 5, we demonstrate TAM's superior ability to adapt to new tasks across different levels of difficulty (task boundaries).

To further validate this, in response to the reviewer’s comment, we conducted continual learning experiments using a pre-trained ResNet-50 on the CLEAR dataset, which consists of 10 sequential tasks [1]. We compare TAM with SGDM across various learning rates on top of two continual learning setups: (i) Naive and (ii) Learning without forgetting (LwF) [2] which is a well-known continual learning method. We summarize the results obtained using the best-performing setup below:

Detailed results are provided in Appendix A.2.3 of the revision. These results again show the effectiveness of TAM in a challenging setting where a model is required to maintain both stability and plasticity.

[1] Xiaohui Zhang, Jiangyan Yi, Jianhua Tao, Chenglong Wang, and Chu Yuan Zhang. Do you remember? overcoming catastrophic forgetting for fake audio detection. In International Conference on Machine Learning, pp. 41819–41831. PMLR, 2023.

[2] Zhizhong Li and Derek Hoiem. Learning without forgetting. IEEE transactions on pattern analysis and machine intelligence, 40(12):2935–2947, 2017.

We appreciate the review and its valuable feedback. We hope our responses address the concerns raised. If so, we kindly request this be reflected in the score. Please feel free to reach out with any further questions, and we will respond promptly.

Thanks for the authors' response which solves my concerns to a considerable extent. After consideration, I decide to maintain my score.