IO-Adam: Rethinking Memory-Efficient Adaptive Optimizers from Gradient Computation

Adaptive Moment Estimation (Adam) is one of the most popular and often the default stochastic optimizers for deep neural network training. Using first- and second-moment estimation, Adam provides adaptive learning rates for each parameter, significantly outperforming Stochastic Gradient Descent (SGD). However, as deep neural networks become larger, estimating the first and second moments consumes substantial memory. It motivates various methods to reduce memory usage for adaptive optimizers. In this paper, we propose to rethink the first and second moment estimation from a gradient computation perspective. The gradient of the weight matrix is the multiplication of the input and the gradient of the output. Instead of finding low-rank approximations of the first and second moments, as in previous work, we propose tracking the input and output gradients to efficiently estimate moments. We provide analyses of the similarities and differences between our proposed method, the widely used Adam optimizer, and previous memory-efficient optimizers designed to reduce memory usage. We conduct experiments to verify the effectiveness of our method, which reduces memory usage by up to $30$% while preserving similar performance or even improving the performance of Adam.

Adaptive optimizers, e.g., Adaptive Moment Estimation (Adam), adjust the learning rate for each parameter for fast convergence. However, these adaptive optimizers often require recording additional information during deep neural network training, e.g., first- and second-moment estimates for Adam, which can significantly increase memory usage. Various methods have been proposed to reduce Adam's memory usage. In this paper, we propose a new method that separately records second-moment estimates for the input and output gradients, thereby reducing memory usage while maintaining similar or even better performance. The gradient—the signal that tells the network how to adjust its weights—can naturally be broken down into two parts: the input entering a layer and the error signal leaving it. Instead of storing the full gradient statistics as Adam does, we store separate, much smaller summaries of the input and the output error.