[論文レビュー] 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

ひとこと要約

この論文では、O(1)のメモリ使用量を達成するため、C=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)を構築する。
• 各スライス内で、ドロップアウトマスクのランダムシードを順方向と逆方向の両方で再利用することで、確率的計算をサポートする。
• 勾配蓄積や可逆レイヤーなどの既存の最適化技術と統合し、さらにメモリ消費量を削減する。
• 正確な勾配を計算することで、後方互換性を確保し、トレーニングおよび微調整の両方に適した方法を実現する。

実験結果

リサーチクエスチョン

• RQ1 近似誤差を導入せずに、O(L)未満のメモリ消費量を実現できるか？
• RQ2 長時間シーケンスのトレーニングにおいて、メモリ使用量と並列計算時間のトレードオフはどのようなものか？
• RQ3 提案手法は、低メモリデバイスでの微調整に使用可能であり、モデルの精度を維持できるか？
• RQ4 勾配蓄積や可逆レイヤーなどの既存のディープラーニング最適化技術と互換性があるか？
• RQ5 フルバッチ学習と比較して、スライスベースの逆誤差伝搬は勾配の正確性を保持しているか？

主な発見

• C=1のとき、提案されたアルゴリズムはトレーニング中にO(1)のメモリ使用量を達成し、時間計算量はO(L)に留まるため、最小限のメモリを備えたデバイスでもトレーニングが可能になる。
• コピータスク（L=512）では、メモリ使用量がフルバッチ時（0.0449 GB）からC=64のとき（0.0374 GB）に低下し、1イテレーションあたりの時間は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秒に増加した。
• 実験的結果から、すべての手法間で区別できない性能が得られ、アルゴリズムが正確な勾配を計算し、モデルの精度を維持していることが確認された。
• この手法は完全に後方互換性があり、低メモリデバイスでの微調整を可能にし、分散型で民主的（デモクラティック）なディープラーニングを実現する。