Compressing Gradient Optimizers via Count-Sketches

By Ryan Spring, Anastasios Kyrillidis, Vijai Mohan, Anshumali Shrivastava. 2019

41 citations recorded by the source corpus.

Many popular first-order optimization methods (e.g., Momentum, AdaGrad, Adam) accelerate the convergence rate of deep learning models. However, these algorithms require auxiliary parameters, which cost additional memory proportional to the number of parameters in the model. The problem is becoming more severe as deep learning models continue to grow larger in order to learn from complex, large-scale datasets. Our proposed solution is to maintain a linear sketch to compress the auxiliary variables. We demonstrate that our technique has the same performance as the full-sized baseline, while using significantly less space for the auxiliary variables. Theoretically, we prove that count-sketch optimization maintains the SGD convergence rate, while gracefully reducing memory usage for large-models. On the large-scale 1-Billion Word dataset, we save 25% of the memory used during training (8.6 GB instead of 11.7 GB) by compressing the Adam optimizer in the Embedding and Softmax layers with negligible accuracy and performance loss. For an Amazon extreme classification task with over 49.5 million classes, we also reduce the training time by 38%, by increasing the mini-batch size 3.5x using our count-sketch optimizer.

Subfields: Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Optimization & Training

Citations in the corpus

Memory-Efficient Adaptive Optimization for Large-Scale Learning (2019, methodology): “Spring et al. [25] provide an alternative way to reduce memory costs, making use of the Count-Sketch data structure [7] to maintain a compressed approximation to the auxiliary variables.”
1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed (2021, unknown): “Previous work has applied a range of techniques such as quanti-zaiton, sparsiﬁcation, and sketching (Alistarh et al., 2017; Agarwal et al., 2018; Spring et al., 2019; Ye and Abbe, 2018; Shi et al., 2021).”
FetchSGD: Communication-Efficient Federated Learning with Sketching (2020, unknown): “Spring et al. (2019) also propose using sketches for distributed optimization.”