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.

通过引入线性草图的压缩辅助变量方法，可以在减小成本的情况下，实现与原算法相同的性能，这在大规模深度学习模型的训练中具有潜在的应用价值。