CAME : Confidence-guided Adaptive Memory Efficient Optimization (2023)

Abstract: Adaptive gradient methods, such as Adam and LAMB, have demonstrated excellent performance 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 paper, 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 experiments 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 implementation of CAME is publicly available.

## introduction robust training of large language models (llms) often relies on adaptive gradient-based optimization methods @xcite @xcite . through the use of cumulative second-order statistics, these methods adapt the per-parameter learning rate and demonstrate superior convergence speed during the training process of llms. however, the remarkable performance of adaptive methods incurs an extra cost of memory usage indeed. for example, adam requires to preserve the first moment estimate and second 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 expensive in terms of memory, and the limitation of memory is gradually emerging as a main bottleneck for training llms. many existing memory-efficient optimizers attempt to store second-order statistics with sublinear memory requirement while retaining the exceptional convergence property of adaptivity @xcite . adafactor optimizer achieves remarkable memory cost reduction by applying the non-negative matrix factorization algorithm @xcite 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 decreases from o(nm) to o(n + m). whereas, it is observed that adafactor suffers a performance degradation in the training of large language models universally compared with conventional adaptive gradient-based optimization methods. the reason for this phenomenon is adafactor inevitably introduces some errors that cause instability in training deep networks due to the operation of nonnegative matrix factorization. in addition, in the case of large-batch training that aims to accelerate the training of deep neural networks, the memory consumption of each machine (gpu/tpu) is much higher than general batch size training, which further imposes a grave constraint on the performance of the trained model. in comparison to standard training tasks, largebatch training presents more challenges for optimizers. empirically, when the mini-batch size increases after a certain point (e.g. 1024), the test accuracy of the converged solution decreases significantly compared with the baseline @xcite . to our knowledge, there is currently no work related to memory-efficient optimizers for large-batch training. motivated by these challenges, we firstly study a confidence-guided strategy catered to alleviate the instability of adafactor by calculating the confidence of the generated update at each training step. on the basis of the adaptation strategy, we propose a novel came optimizer that saves nearly the same memory footprint as existing memoryefficient optimizers while attaining faster convergence and superior generalization performance. to further assess the scalability of our proposed algorithm, we consider an additional challenging experiment -performing large-batch training on bert using came optimizer. contributions of our paper can be summarized in the following: ## related work memory efficient adaptive optimization memory efficient optimizers maintain the benefits of standard per-parameter adaptivity while significantly reducing memory footprint. adafactor @xcite proposes to reconstruct a low-rank approximation of the exponentially smoothed accumulator at each training step that is optimal with respect to the generalized kullback-leibler divergence. sm3 @xcite divides the elements in the second-order gradient matrix into sets by the observed similarity of the elements, and each item in the generated approximation matrix is the minimum of the maximum value of each set in which it is located. the methods mentioned above behave poorly in the training of large language models and converge slowly, which raises a significant challenge for memory-efficient optimization methods. large batch training a large-batch training scheme is preferred in distributed machine learning because of its ability to increase parallelism by enhancing large-scale cluster utilization. it has seen growing interest in recent years in large-batch training @xcite @xcite . in particular, a layer-wise adaptive learning rate algorithm lars @xcite ) is proposed to scale the batch size to 32k for resnet-50. based on lars, lamb optimizer @xcite can finish the bert training in 76 minutes through tpu v3 pod. despite the success of these approaches for bert models, the much larger batch size highly boosts the gpu usage which is prohibitively expensive and inaccessible to most researchers. moreover, training with a large batch size incurs additional challenges @xcite . large-batch training is prone to converge to sharp local minima, since the number of interactions will decrease when the batch size is increased if the number of epochs is fixed, which causes a wide gap in generalization of the model @xcite . traditional methods seek to narrow the generalization gap by carefully tuning hyperparameters, such as learning rate, momentum, and label smoothing, to narrow the generalization gap @xcite @xcite ). yet there have been few attempts to reduce memory usage in large-batch training, and the underlying challenge remains unclear. ## method in this section, we firstly provide a brief description of the adafactor optimizer and discuss the errors contained in the update of adafactor (erroneous update). we further study a confidence-guided strategy and introduce the proposed came in detail in light of the strategy. ## experiments in this section, we present extensive comparisons with existing optimizers on training tasks of three important large language models: bert @xcite , gpt-2 @xcite and t5 (raffel et al., 2022). ## conclusion in this paper we propose a novel memory-efficient optimizer called came, which supports adaptive confidence-based updating guided by the residual between predicted update and generated update. came achieves a considerable improvement compared to existing memory-efficient optimizers in the training of large language models, with an ignorable extra memory footprint. moreover, came shows comparable convergence to adam and lamb with huge memory reduction. in particular, came has proven effective for large-batch training, which serves as an advantageous extension to memory-efficient optimizers. we hope our work will provide insight into memory reduction 0 10 20 30 40 50 steps(k) 10 20 30 40 50 60 70 80 90 100 validation perplexity adam adafactor came figure 11: validation perplexity of t5 language model. came demonstrates similar convergence speed to adam.