[论文解读] Memory-Efficient Adaptive Optimization

Rohan Anil, Vineet Gupta+2 more|arXiv (Cornell University)|Jan 30, 2019

Advanced Neural Network Applications被引用 19

一句话总结

本文提出 SM3，一种内存高效的自适应优化方法，通过仅维护二阶统计量的压缩低秩近似，显著降低了 Adam 和 Adagrad 等自适应优化器的内存开销。该方法在保持相近或更优收敛性能的同时，使模型规模和批量大小显著增大，在大规模语言模型和图像分类任务中，训练时间最高可实现两倍加速，同时保持收敛性保证。

ABSTRACT

Adaptive gradient-based optimizers such as Adagrad and Adam are crucial for achieving state-of-the-art performance in machine translation and language modeling. However, these methods maintain second-order statistics for each parameter, thus introducing significant memory overheads that restrict the size of the model being used as well as the number of examples in a mini-batch. We describe an effective and flexible adaptive optimization method with greatly reduced memory overhead. Our method retains the benefits of per-parameter adaptivity while allowing significantly larger models and batch sizes. We give convergence guarantees for our method, and demonstrate its effectiveness in training very large translation and language models with up to 2-fold speedups compared to the state-of-the-art.

研究动机与目标

• 为解决 Adam 和 Adagrad 等自适应优化器带来的高内存开销问题，这些开销限制了大规模训练中的模型大小和批量大小。
• 开发一种方法，在大幅降低内存消耗的同时保留逐参数自适应的优势。
• 通过释放内存以支持更大的批量大小，实现对超大模型（尤其是自然语言处理和视觉任务）的训练。
• 在凸在线优化设置下提供理论收敛性保证。
• 通过实证结果证明，该方法在减少内存使用和加快实际运行时间的同时，实现相近或更优的收敛性能。

提出的方法

• SM3 不存储每个参数的完整二阶矩，而是维护二阶梯度统计量的低秩近似。
• 该方法采用一种随机且内存高效的更新规则，通过秩一近似压缩自适应统计量。
• 通过应用对角缩放因子，在最小化存储空间的同时保持逐参数自适应能力。
• 该算法设计为与标准深度学习框架兼容，仅需极少的代码修改。
• 通过 Fisher 信息矩阵的结构化近似，实现类似自然梯度的更新。
• 该方法能动态适应梯度模式，而无需预先了解激活结构。

实验结果

研究问题

• RQ1 一种内存高效的自适应优化器能否在减少内存使用的同时，保持与 Adam 和 Adagrad 相当的收敛性能？
• RQ2 SM3 所节省的内存能在多大程度上用于在分布式训练中提升批量大小？
• RQ3 对二阶统计量进行低秩近似是否会影响大规模模型中的收敛速度？
• RQ4 与现有的内存高效方法（如 Adafactor 和 Shampoo）相比，SM3 在内存使用、速度和精度方面表现如何？
• RQ5 在固定计算预算下，SM3 是否能实现比标准自适应优化器更快的实际运行时间收敛？

主要发现

• 在 WMT’14 en→fr 翻译任务中，SM3 以比 Adam 少 40% 的内存使用量，实现了 40.50 的 BLEU 分数，同时支持更大的批量大小。
• 在 BERT-Large 语言建模任务中，当批量大小加倍时，SM3 在比 Adam 少 35% 的实际运行时间内达到了 70% 的掩码语言模型准确率。
• 在批量大小为 2048 时，SM3 每核心仅使用 6.02 GiB 内存，与 Adam 在 1024 批量大小下的内存使用量相当，但可扩展至 2^16 而不触及内存限制。
• 在 ImageNet 上使用 AmoebaNet-D 模型，SM3 达到了 78.71% 的 top-1 准确率和 94.31% 的 top-5 准确率，与当前最先进性能持平。
• 与 Adam 相比，SM3 在大规模模型上将内存使用量减少了 40%，同时保持了相似的收敛速度和准确率。
• SM3 的单步时间比 Adam 快 3%，当批量大小加倍时，收敛的实际运行时间最多减少 50%。