CAME: Confidence-guided Adaptive Memory Efficient Optimization

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 4442–4453 July 9-14, 2023 ©2023 Association for Computational Linguistics CAME: Confidence-guided Adaptive Memory Efficient Optimization Yang Luo1*, Xiaozhe Ren2, Zangwei Zheng1, Zhuo Jiang1, Xin Jiang2, Yang You1 1School of Computing, National University of Singapore 2Noah’s Ark Lab, Huawei {yangluo,zangwei,jiangz,youy}@comp.nus.edu.sg {renxiaozhe,jiang.xin}@huawei.com Abstract Adaptive gradient methods, such as Adam and LAMB, have demonstrated excellent perfor- mance in the training of large language models. Nevertheless, the need for adaptivity requires maintaining second-moment estimates of the per-parameter gradients, which entails a high cost of extra memory overheads. To solve this problem, several memory-efficient optimizers (e.g., Adafactor) have been proposed to obtain a drastic reduction in auxiliary memory usage, but with a performance penalty. In this pa- per, we first study a confidence-guided strategy to reduce the instability of existing memory efficient optimizers. Based on this strategy, we propose CAME to simultaneously achieve two goals: fast convergence as in traditional adaptive methods, and low memory usage as in memory-efficient methods. Extensive exper- iments demonstrate the training stability and superior performance of CAME across various NLP tasks such as BERT and GPT-2 training. Notably, for BERT pre-training on the large batch size of 32,768, our proposed optimizer attains faster convergence and higher accuracy compared with the Adam optimizer. The im- plementation of CAME is publicly available1. 1 Introduction Robust training of large language models (LLMs) often relies on adaptive gradient-based optimiza- tion methods (Li et al., 2022; Kingma and Ba, 2015; Zhuang et al., 2020). Through the use of cumula- tive second-order statistics, these methods adapt the per-parameter learning rate and demonstrate supe- rior convergence speed during the training process of LLMs. However, the remarkable performance of adaptive methods incurs an extra cost of mem- ory usage indeed. For example, Adam requires to preserve the first moment estimate and second *Work was done when Yang Luo was an intern at Huawei Noah’s Ark Lab. 1https://github.com/huawei-noah/Pretrained-Language- Model/tree/master/CAME Figure 1: Visualization of Non-negative Matrix Factor- ization (NMF). Generally, NMF reduces the memory requirements from O(nm) to O(n + m). In this paper, we focus on the special case of rank-1 factors. raw moment estimate of each gradient in order to tune the learning rate for each parameter, which inevitably triples the memory usage concerning the optimizer states. Besides, with the growing size of the model, LLMs are becoming increasingly ex- pensive in terms of memory, and the limitation of memory is gradually emerging as a main bottleneck for training LLMs. Many existing memory-efficient optimizers at- tempt to store second-order statistics with sublin- ear memory requirement while retaining the excep- tional convergence property of adaptivity (Shazeer and Stern, 2018; Anil et al., 2019). Adafactor opti- mizer achieves remarkable memory cost reduction by applying the non-negative matrix factorization algorithm (Lee and Seung, 2000) to factorize the accumulator matrix for squared gradients into two rank-1 factors as shown in Figure 1, where the memory requirement for the original matrix V de- creases from O(nm) …

Content truncated to 500 words