Compressing Gradient Optimizers via Count-Sketches

Abstract

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.

Initially submitted to WWW 2019 (November 2018)

References in corpus (10)

• Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Exploring the Limits of Language Modeling
• Training Deep Nets with Sublinear Memory Cost
• Pointer Sentinel Mixture Models
• Learning Intrinsic Sparse Structures within Long Short-Term Memory
• Multiplicative LSTM for sequence modelling
• Deep Networks With Large Output Spaces
• Large Scale Language Modeling: Converging on 40GB of Text in Four Hours
• Extreme Classification in Log Memory

Cited by in corpus (4)

• From Distributed Machine Learning to Federated Learning: A Survey
• FetchSGD: Communication-Efficient Federated Learning with Sketching
• 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed
• 1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed