2024/02/05 by Lin Wu, Lin, Wu, Felix Dangel +9 · 4 citations
Mathematics · #Advanced Optimization Algorithms Research #FOS: Computer and information sciences #FOS: Mathematics #Iterative Methods for Nonlinear Equations #Machine Learning (cs.LG) #Numerical methods in inverse problems #Optimization and Control (math.OC)
paper · pdf · doi:10.48550/arxiv.2402.03496
openalex publication_date 2024/02/05 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
Adaptive gradient optimizers like Adam(W) are the default training algorithms for many deep learning architectures, such as transformers. Their diagonal preconditioner is based on the gradient outer product which is incorporated into the parameter update via a square root. While these methods are often motivated as approximate second-order methods, the square root represents a fundamental difference. In this work, we investigate how the behavior of adaptive methods changes when we remove the root, i.e., strengthen their second-order motivation. Surprisingly, we find that such square-root-free adaptive methods close the generalization gap to SGD on convolutional architectures, while maintaining their root-based counterpart's performance on transformers. The second-order perspective also has practical benefits for developing non-diagonal methods that can incorporate arbitrary curvature approximations through the concept of preconditioner invariance. In contrast to root-based methods like Shampoo, root-free counterparts work well and fast with half-precision since they do not require numerically unstable matrix root decompositions and inversions. Overall, our findings provide new insights into the development of adaptive methods and raise important questions regarding the overlooked role of adaptivity in their success. (experiment code: https://github.com/yorkerlin/remove-the-square-root optimizer code: https://github.com/f-dangel/sirfshampoo)