arXiv · 2410.16103
LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics
Abstract
We introduce LDAdam, a memory-efficient optimizer for training large models, that performs adaptive optimization steps within lower dimensional subspaces, while consistently exploring the full parameter space during training. This strategy keeps the optimizer's memory footprint to a fraction of the model size. LDAdam relies on a new projection-aware update rule for the optimizer states that allows for transitioning between subspaces, i.e., estimation of the statistics of the projected gradients. To mitigate the errors due to low-rank projection, LDAdam integrates a new generalized error feedback mechanism, which explicitly accounts for both gradient and optimizer state compression. We prove the convergence of LDAdam under standard assumptions, and show that LDAdam allows for accurate and efficient fine-tuning and pre-training of language models. Code is available at https://github.com/IST-DASLab/LDAdam
Explore related subjects
Keep this discovery
Thomas Robert, Mher Safaryan, Ionut-Vlad Modoranu, Dan Alistarh. 2024-10-21. LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics. https://arxiv.org/abs/2410.16103
Cite the original work for its findings. Save a collection to share your selection of sources.