Compressing Gradient Optimizers via CountSketches
[edit]
Proceedings of the 36th International Conference on Machine Learning, PMLR 97:59465955, 2019.
Abstract
Many popular firstorder optimization methods accelerate the convergence rate of deep learning models. However, these algorithms require auxiliary variables, which cost additional memory proportional to the number of parameters in the model. The problem is becoming more severe as models grow larger to learn from complex, largescale datasets. Our proposed solution is to maintain a linear sketch to compress the auxiliary variables. Our approach has the same performance as the fullsized baseline, while using less space for the auxiliary variables. Theoretically, we prove that countsketch optimization maintains the SGD convergence rate, while gracefully reducing memory usage for largemodels. We show a rigorous evaluation on popular architectures such as ResNet18 and TransformerXL. On the 1Billion Word dataset, we save 25% of the memory used during training (7.7 GB instead of 10.8 GB) with minimal 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 minibatch size 3.5x using our countsketch optimizer.
Related Material


