Adam Exploits $\ell_\infty$-geometry of Loss Landscape via Coordinate-wise Adaptivity

For SGD, the convergence speed when $\sigma = 0$ should be in the order of $O(|\nabla ^2 L(x)| _2^ {1/2} T^ {-1/4})$, which seems to be faster than either AdaSGD which is in the order of $O(d^{1/2} |\nabla^2 L(x)| _2 ^{1/2} T^{-1/4} )$ or Adam which is in the order of $O(d ^{1/2} |\nabla ^2 L(x)| _{1,1} ^{1/2} T ^{-1/4})$.

I do think the authors give an important direction that we have to find some rotation variant assumption to show the benefit of Adam, but can the proposed assumption really explain the advance of Adam compared with SGD instead of AdaSGD?

We sincerely thank the reviewer for their thoughtful feedback on paper writing and for acknowledging the significance of our work in providing a unified framework and insights into Adam-type algorithms. Below, we address the specific concerns raised:

1. Relationship between Section 3.2 and Section 3.3.

   Section 3.2 focuses on the assumptions, convergence rate, and hyperparameter choices for Adam, following a standard structure commonly seen in related work that centers specifically on Adam. We decided to present it as a standalone subsection rather than merging it into Section 3.3 to ensure that researchers interested in Adam can easily compare our results and assumptions with prior analyses. If Adam were only treated as a special case of blockwise Adam in Section 3.3, researchers would need to spend additional effort translating our general results into the specific Adam framework.

   Section 3.3, on the other hand, provides a more general framework that encompasses various optimization algorithms including AdaSGD. This broader approach highlights the distinctions among these algorithms and how these differences impact their convergence rates. By including AdaSGD as a notable special case in Section 3.3, we demonstrate how our general framework subsumes multiple algorithms and contributes to a more unified understanding of their behavior.

   This separation ensures clarity and accessibility for both researchers focused on Adam and those interested in broader generalizations.

2. Justification of (1,1)-norm and $\ell_2$ norm. Please check the general response in which we clarify the typo.

3. Explanation for table 1. You are right that the convergence of AdaSGD is related to 2-norm so we report that 2-norm (spectral norm) is 1 in the caption. Leaving the (1,1)-norm cell for AdaSGD blank might have disrupted the presentation of the table, so we chose to include it for completeness even though it is not necessary for interpreting its convergence behavior. We appreciate this observation and will clarify this in the revised manuscript to avoid any misunderstanding.

We sincerely thank the reviewer for recognizing the novelty and significance of our results and for highlighting several insightful directions to further enhance the paper. We greatly appreciate the reviewer’s constructive suggestions for expanding our work. However, due to limited time and computational resources, we were unable to conduct all the extensive experiments, such as testing more model architectures or exploring different kinds of rotations. We hope that our theoretical findings and preliminary empirical results can guide future work and inspire further exploration by the community. Below, we address the specific issues:

1. Understanding convergence result. Our convergence rate for Adam improves the dependence on the number of parameters, $d$, by replacing the constant $d$ times $\ell_2$-norm of the Hessian with the (1,1)-norm of the Hessian. For large models, where $d$ can be much larger than the total number of training steps $T$, this refinement results in a more practical and meaningful bound. Additionally, we argue that the previous $\ell_2$-norm-based bounds for Adam can be vacuous, as they do not distinguish between standard Adam and rotated Adam. Since rotated Adam has significantly worse performance in practice, the $\ell_2$-norm-based analysis does not fully capture Adam’s behavior. Our results address this limitation by using a more appropriate norm for Adam’s coordinate-wise adaptivity.

2. Smoothness assumption. As mentioned above and in our introduction (line 45-51, 366-373), we believe the standard $\ell_2$ norm is inadequate for tight analysis of Adam. The benefit of coordinate-wise adaptivity can only be captured under a non-standard $\ell_\infty$ norm geometry, which justifies the need for a new smoothness assumption. While this assumption may limit the direct comparability with other optimization methods, it provides a more accurate framework for understanding Adam’s behavior.

3. Larger model with SGD vs smaller model with Adam. Our analysis focuses on comparing Adam and SGD in terms of the number of iterations for the same problem. Comparing models of different sizes involves more complex issues related to expressiveness and scaling laws, which are outside the scope of our work. Nevertheless, we agree that exploring such trade-offs could be an interesting direction for future work.

4. General $\beta_1$. Please refer to the beginning of Section 3 (lines 151–157).

   Similar to previous work (Défossez et al., 2022), our analysis could be extended to the most general case of Adam, where both $\beta_1$, $\beta_2$ are non-zero, but the rate becomes strictly worse than the RMSProp (the case of $\beta_1=0$), as there will be some extra polynomials of $\frac{1}{1-\beta_1}$. We decide not to include the result for the most general case, on one hand for ease of presentation, and on the other hand, because such result can not explain the optimization benefit of momentum ($\beta_1>0$) in practice and does not add any insight on the benefit of Adam. We hypothesize that we are missing some important features of loss landscape of transformers in the theoretical assumptions and we leave this for future work.

We sincerely thank the reviewer for acknowledging the importance of the rotation invariance property towards understanding gap between Adam and SGD and the value of our results in guiding the design of new optimization algorithms. We greatly appreciate the constructive feedback and will address the concerns in detail below.

1. Justification of (1,1)-norm and $\ell_2$ norm. See general response in which we clarify the typo.

2. Comparison with Zhang et al. (2024a).

   Both our work and Zhang et al. (2024a) aim to highlight Adam’s advantages over SGD under specific properties of the Hessian matrix. At a high level, our conclusions align, as $\\| \nabla^2 L(x) \\|_{1,1}$ is typically smaller than $d \\| \nabla^2 L(x) \\|_2$ when there is block heterogeneity in the Hessian matrix $\nabla^2 L(x)$. The quadratic loss example in Section 4.1 of our work serves as an extreme example of such block heterogeneity, supporting this shared insight.

   There is difference in terms of the specific results. The theoretical result in Zhang et al. (2024a) is for strongly convex quadratic loss functions and only holds for the degenerate case $\beta_2 = 1$ in Adam. Moreover, as shown in Zhang et al. (2024a), $\beta_2=1$ is crucial in their analysis because any $\beta_2<1$ will lead to non-convergence of Adam. We believe this is not a very realistic regime for Adam analysis because, in deep learning practice, $\beta_2$ is typically chosen to be strictly smaller than 1. Setting $\beta_2=1$ leads to suboptimal performance because the denominator in the update rule of Adam becomes fixed at its initialization, causing Adam to degenerate into SGD with constant diagonal preconditioner. In contrast, our analysis focuses on the convergence of Adam in the standard non-convex stochastic optimization setting and is valid for a wide range of $\beta_2$ values.

3. Non-rotation-invariance of neural networks. We appreciate this insightful question. We believe theoretically understanding the advantage of the standard coordinate system is an important and interesting future direction while we will discuss our initial thoughts on it below. The lack of rotation invariance in deep neural networks can be attributed to both architecture and data:

   Architecture: The hierarchical structure and the properties of specific components like activation functions and normalization layers can cause non-rotation-invariance. For example, coordinate-wise nonlinear activation functions, such as ReLU, introduce dependencies that are not preserved under rotation. These elements are tightly coupled to the original coordinate system, and arbitrary rotations can disrupt this alignment.

   Data: Specific data formats may also contribute to non-rotation-invariance. For example, in language modeling, each row in the embedding matrix corresponds to different tokens so the embedding matrix itself is not randomly rotated. The imbalance in training data distribution may also attribute to the asymmetry of model parameters and cause Adam to optimize faster than SGD as observed by Kunstner et al. 2024. This explanation is consistent with the quadratic example we provide in Sec 4.1, where different diagonal values corresponds to the frequency of different classes. Such imbalance only allows faster convergence of Adam over SGD in the standard coordinate system, and the optimization advantage disappears when we apply Adam in the rotated coordinate system.

We kindly ask the reviewer to consider revisiting their score in light of our responses, as we believe we have addressed the main concerns about the relationship between different norms. We also welcome the reviewer to raise any additional questions or suggestions to help us further improve our work.

Kunstner, Frederik, et al. "Heavy-tailed class imbalance and why adam outperforms gradient descent on language models." arXiv preprint arXiv:2402.19449 (2024).