[论文解读] Compressing Gradient Optimizers via Count-Sketches

Ryan Spring, Anastasios Kyrillidis+2 more|arXiv (Cornell University)|Feb 1, 2019

Advanced Neural Network Applications参考文献 32被引用 5

一句话总结

本文提出使用 Count-Sketch 数据结构对 Adam、Momentum 和 AdaGrad 等一阶优化器中的辅助变量进行压缩，以减少大规模深度学习模型中的内存占用。通过利用嵌入层和 Softmax 层中的稀疏性，该方法在保持收敛速度和模型性能的同时，将内存使用量最多减少 25%，并在包含 4950 万个类别的极端分类任务中将训练速度提升 38%。

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.

研究动机与目标

• 为解决 Adam 和 Momentum 等一阶优化器中随模型规模增长的辅助参数带来的内存开销问题。
• 通过减少深度学习中优化器状态的内存消耗，实现更大的批量大小和更具表现力的模型。
• 利用嵌入层和 Softmax 层中的结构稀疏性，高效压缩优化器状态，且不造成性能下降。
• 提供一种理论基础坚实的、内存高效的优化器状态替代方案，保持与 SGD 相同的收敛速率。
• 在大规模语言建模和极端分类任务中展示实际效益，且精度损失可忽略。

提出的方法

• 使用 Count-Sketch 数据结构在嵌入层和 Softmax 层中维护二阶优化器状态（如 Adam 的速度和动量）的压缩表示。
• 仅存储一个小型固定大小的 Sketch（例如 [3, 266, 1024]），而非完整的优化器状态，从而将二阶矩估计的内存使用量最多减少 99%。
• 通过随机哈希函数将梯度更新映射到桶中，增量式地更新 Sketch，实现常数时间操作。
• 在参数更新过程中通过查询 Sketch 来近似完整优化器状态，误差通过理论界得到控制。
• 通过调整哈希函数数量和桶的数量，实现 Sketch 大小与精度之间的平稳权衡。
• 将基于 Sketch 的优化器直接集成到训练流水线中，用压缩版本替代标准优化器状态。

实验结果

研究问题

• RQ1 Count-Sketch 是否可用于压缩深度学习中的一阶优化器辅助变量，而不会降低收敛性或性能？
• RQ2 基于 Sketch 的优化器在大规模模型上的内存占用和训练时间与标准 Adam 相比如何？
• RQ3 在使用压缩优化器状态时，批量大小最多可增加多少，对训练速度有何影响？
• RQ4 在理论分析中，基于 Sketch 的优化器是否能保持与标准一阶方法相同的收敛速率？
• RQ5 该方法是否能有效应用于包含超过 4900 万个类别的超大规模分类任务？

主要发现

• 在 1-Billion Word 语言建模任务中，Count-Sketch 优化器将内存使用量减少了 25%（从 11.7 GB 降至 8.6 GB），且未造成精度或性能损失。
• 在包含 4950 万个类别的极端分类任务中，使用 Count-Sketch 优化器将小批量大小提高 3.5 倍，使训练时间减少了 38%。
• Count-Sketch 优化器的 Recall@100 达到 0.6889，而 Adam 基线为 0.6881，表明性能下降可忽略不计。
• 通过使用 1% 大小的 Sketch 压缩二阶矩估计，模型每参数的内存成本从 4 GB 降低至 2.6 GB（减少 35%）。
• 理论分析证实，Count-Sketch 优化器保持了与全精度基线相同的 SGD 收敛速率。
• 该方法实现了内存使用量与近似误差之间的实用权衡，且计算开销极低。