Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Lin, Wu, Dangel, Felix, Eschenhagen, Runa, Bae, Juhan, Turner, Richard E., Makhzani, Alireza
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913532749021184
author Lin, Wu
Dangel, Felix
Eschenhagen, Runa
Bae, Juhan
Turner, Richard E.
Makhzani, Alireza
author_facet Lin, Wu
Dangel, Felix
Eschenhagen, Runa
Bae, Juhan
Turner, Richard E.
Makhzani, Alireza
contents 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)
format Preprint
id arxiv_https___arxiv_org_abs_2402_03496
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective
Lin, Wu
Dangel, Felix
Eschenhagen, Runa
Bae, Juhan
Turner, Richard E.
Makhzani, Alireza
Machine Learning
Optimization and Control
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)
title Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective
topic Machine Learning
Optimization and Control
url https://arxiv.org/abs/2402.03496