When Can You Get Away with Low Memory Adam?

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Kalra, Dayal Singh, Kirchenbauer, John, Barkeshli, Maissam, Goldstein, Tom
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866915202807627776
author Kalra, Dayal Singh
Kirchenbauer, John
Barkeshli, Maissam
Goldstein, Tom
author_facet Kalra, Dayal Singh
Kirchenbauer, John
Barkeshli, Maissam
Goldstein, Tom
contents Adam is the go-to optimizer for training modern machine learning models, but it requires additional memory to maintain the moving averages of the gradients and their squares. While various low-memory optimizers have been proposed that sometimes match the performance of Adam, their lack of reliability has left Adam as the default choice. In this work, we apply a simple layer-wise Signal-to-Noise Ratio (SNR) analysis to quantify when second-moment tensors can be effectively replaced by their means across different dimensions. Our SNR analysis reveals how architecture, training hyperparameters, and dataset properties impact compressibility along Adam's trajectory, naturally leading to $\textit{SlimAdam}$, a memory-efficient Adam variant. $\textit{SlimAdam}$ compresses the second moments along dimensions with high SNR when feasible, and leaves when compression would be detrimental. Through experiments across a diverse set of architectures and training scenarios, we show that $\textit{SlimAdam}$ matches Adam's performance and stability while saving up to $98\%$ of total second moments. Code for $\textit{SlimAdam}$ is available at https://github.com/dayal-kalra/low-memory-adam.
format Preprint
id arxiv_https___arxiv_org_abs_2503_01843
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle When Can You Get Away with Low Memory Adam?
Kalra, Dayal Singh
Kirchenbauer, John
Barkeshli, Maissam
Goldstein, Tom
Machine Learning
Disordered Systems and Neural Networks
Adam is the go-to optimizer for training modern machine learning models, but it requires additional memory to maintain the moving averages of the gradients and their squares. While various low-memory optimizers have been proposed that sometimes match the performance of Adam, their lack of reliability has left Adam as the default choice. In this work, we apply a simple layer-wise Signal-to-Noise Ratio (SNR) analysis to quantify when second-moment tensors can be effectively replaced by their means across different dimensions. Our SNR analysis reveals how architecture, training hyperparameters, and dataset properties impact compressibility along Adam's trajectory, naturally leading to $\textit{SlimAdam}$, a memory-efficient Adam variant. $\textit{SlimAdam}$ compresses the second moments along dimensions with high SNR when feasible, and leaves when compression would be detrimental. Through experiments across a diverse set of architectures and training scenarios, we show that $\textit{SlimAdam}$ matches Adam's performance and stability while saving up to $98\%$ of total second moments. Code for $\textit{SlimAdam}$ is available at https://github.com/dayal-kalra/low-memory-adam.
title When Can You Get Away with Low Memory Adam?
topic Machine Learning
Disordered Systems and Neural Networks
url https://arxiv.org/abs/2503.01843