[论文解读] Sub-Linear Memory: How to Make Performers SLiM

Valerii Likhosherstov, Krzysztof Choromański+3 more|arXiv (Cornell University)|Dec 21, 2020

Topic Modeling参考文献 31被引用 5

一句话总结

本文提出了一种针对Performer模型的内存高效反向传播算法，通过分片处理序列实现次线性内存扩展，当C=1时实现O(1)内存使用。该方法无需近似即可计算精确梯度，从而在内存受限设备（如智能手机或旧款GPU）上实现训练或微调的完全向后兼容性。

ABSTRACT

The Transformer architecture has revolutionized deep learning on sequential data, becoming ubiquitous in state-of-the-art solutions for a wide variety of applications. Yet vanilla Transformers are notoriously resource-expensive, requiring $O(L^2)$ in serial time and memory as functions of input length $L$. Recent works proposed various linear self-attention mechanisms, scaling only as $O(L)$ for serial computation. We perform a thorough analysis of recent Transformer mechanisms with linear self-attention, Performers, in terms of overall computational complexity. We observe a remarkable computational flexibility: forward and backward propagation can be performed with no approximations using sublinear memory as a function of $L$ (in addition to negligible storage for the input sequence), at a cost of greater time complexity in the parallel setting. In the extreme case, a Performer consumes only $O(1)$ memory during training, and still requires $O(L)$ time. This discovered time-memory tradeoff can be used for training or, due to complete backward-compatibility, for fine-tuning on a low-memory device, e.g. a smartphone or an earlier-generation GPU, thus contributing towards decentralized and democratized deep learning.

研究动机与目标

• 解决长序列Transformer模型训练时的高内存开销，尤其是在内存受限设备上的问题。
• 克服标准Transformer模型因O(L²)内存复杂度导致的批量大小减小的限制。
• 实现在内存受限设备（如智能手机或旧款GPU）上对Performer模型进行训练或微调。
• 在降低内存消耗的同时保持向后兼容性和精确梯度计算。
• 探索时间-内存权衡机制，使用户能够通过增加并行计算时间来控制内存使用。

提出的方法

• 将输入序列按大小为C的分片处理，其中C控制内存与时间的权衡。
• 使用算法1以顺序方式逐片执行前向和反向传播，仅存储损失值和中间缓冲区。
• 在反向传播过程中构建符号表示Φ^(n)，以实现无需近似的精确梯度计算。
• 在每个分片内重用Dropout掩码的随机种子，以支持随机计算。
• 与现有优化技术（如梯度累积和可逆层）集成，以进一步降低内存使用。
• 通过计算精确梯度确保向后兼容性，使该方法适用于训练和微调。

实验结果

研究问题

• RQ1 我们能否在不引入近似误差的情况下，将Performer模型的内存消耗降低至O(L)以下？
• RQ2 在长序列训练中，内存使用与并行计算时间之间的权衡关系是什么？
• RQ3 所提出的方法是否可用于在内存受限设备上进行微调，同时保持模型精度？
• RQ4 该方法是否与现有深度学习优化技术（如梯度累积和可逆层）兼容？
• RQ5 基于分片的反向传播是否与全批量训练相比保持了梯度精度？

主要发现

• 当C=1时，所提出的算法在训练期间实现O(1)内存使用，仅具有O(L)时间复杂度，使在内存极小的设备上训练成为可能。
• 在复制任务中（L=512），内存使用从全批量的0.0449 GB降至C=64时的0.0374 GB，每次迭代时间从0.0474秒增加到0.1228秒。
• 在Penn Treebank数据集上（L=1024），内存使用从全批量的0.300 GB降至C=256时的0.231 GB，每次迭代时间从0.1377秒增加到0.3060秒。
• 在Enwik8数据集上（L=4096），内存使用从全批量的1.513 GB降至C=1366时的0.909 GB，每次迭代时间从0.4598秒增加到0.8654秒。
• 实验结果表明，所有方法的性能在统计上无法区分，证实该算法计算了精确梯度并保持了模型精度。
• 该方法完全向后兼容，支持在内存受限设备上进行微调，从而推动去中心化和普惠化的深度学习。