Understanding Adam Requires Better Rotation Dependent Assumptions
Tianyue H. Zhang, Lucas Maes, Alan Milligan, Alexia Jolicoeur-Martineau, Ioannis Mitliagkas, Damien Scieur, Simon Lacoste-Julien, Charles Guille-Escuret
Abstract
Despite its widespread adoption, Adam's advantage over Stochastic Gradient Descent (SGD) lacks a comprehensive theoretical explanation. This paper investigates Adam's sensitivity to rotations of the parameter space. We observe that Adam's performance in training transformers degrades under random rotations of the parameter space, indicating a crucial sensitivity to the choice of basis in practice. This reveals that conventional rotation-invariant assumptions are insufficient to capture Adam's advantages theoretically. To better understand the rotation-dependent properties that benefit Adam, we also identify structured rotations that preserve or even enhance its empirical performance. We then examine the rotation-dependent assumptions in the literature and find that they fall short in explaining Adam's behaviour across various rotation types. In contrast, we verify the orthogonality of the update as a promising indicator of Adam's basis sensitivity, suggesting it may be the key quantity for developing rotation-dependent theoretical frameworks that better explain its empirical success.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 87b10d23-1b9e-43a4-b0f2-00a3c50e0cb8Cited by top-tier papers5
- Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its PreconditionerRuna Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E. Turner et al.NeurIPS 2025 · 24 citations
- Muon in Associative Memory Learning: Training Dynamics and Scaling LawsKaifei Wang, Binghui Li, Han Zhong, Pinyan Lu et al.ICML 2026 · 7 citations
- Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis RotationHyunji Jung, Sungbin Shin, Namhoon LeeICML 2026 · 2 citations
- Implicit Bias of Per-sample Adam on Separable Data: Departure from the Full-batch RegimeBeomhan Baek, Minhak Song, Chulhee YunICLR 2026 · 2 citations
- Adam Exploits ℓ∞-geometry of Loss Landscape via Coordinate-wise AdaptivityShuo Xie, Mohamad Amin Mohamadi, Zhiyuan LiICLR 2025
Builds on17
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn et al.ICLR 2021 · 21,477 citations
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra et al.NeurIPS 2022 · 5,493 citations
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank ProjectionJiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang et al.ICML 2024 · 433 citations
- Why are Adaptive Methods Good for Attention Models?Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim et al.NeurIPS 2020 · 397 citations
Related papers
- Why Transformers Need Adam: A Hessian PerspectiveYushun Zhang, Congliang Chen, Tian Ding, Ziniu Li et al.NeurIPS 2024 · 149 citations
- On the Optimization and Generalization of Two-layer Transformers with Sign Gradient DescentBingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou et al.ICLR 2025
- VectorAdam for Rotation Equivariant Geometry OptimizationSelena Ling, Nicholas Sharp, Alec JacobsonNeurIPS 2022 · 28 citations
- Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language ModelsFrederik Kunstner, Alan Milligan, Robin Yadav, Mark Schmidt et al.NeurIPS 2024 · 100 citations
- Noise Is Not the Main Factor Behind the Gap Between Sgd and Adam on Transformers, But Sign Descent Might BeFrederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, Mark SchmidtICLR 2023 · 5 citations
