[論文レビュー] ABC: Attention with Bounded-memory Control

Hao Peng, Jungo Kasai+6 more|arXiv (Cornell University)|Oct 6, 2021

Topic Modeling被引用数 7

ひとこと要約

本稿では、記憶サイズを制約することで線形計算量を達成する、ABC（Bounded-Memory Controlを備えたアテンション）という統一的フレームワークを紹介する。ABC$_{\text{MLP}}$ は、学習可能で文脈に依存するメモリ制御機構であり、特に長文系列において、従来手法に比べて精度と効率の両面で優れている。標準のトランスフォーマーと比較して、精度の損失は最小限に抑えられる。

ABSTRACT

Transformer architectures have achieved state-of-the-art results on a variety of sequence modeling tasks. However, their attention mechanism comes with a quadratic complexity in sequence lengths, making the computational overhead prohibitive, especially for long sequences. Attention context can be seen as a random-access memory with each token taking a slot. Under this perspective, the memory size grows linearly with the sequence length, and so does the overhead of reading from it. One way to improve the efficiency is to bound the memory size. We show that disparate approaches can be subsumed into one abstraction, attention with bounded-memory control (ABC), and they vary in their organization of the memory. ABC reveals new, unexplored possibilities. First, it connects several efficient attention variants that would otherwise seem apart. Second, this abstraction gives new insights--an established approach (Wang et al., 2020b) previously thought to be not applicable in causal attention, actually is. Last, we present a new instance of ABC, which draws inspiration from existing ABC approaches, but replaces their heuristic memory-organizing functions with a learned, contextualized one. Our experiments on language modeling, machine translation, and masked language model finetuning show that our approach outperforms previous efficient attention models; compared to the strong transformer baselines, it significantly improves the inference time and space efficiency with no or negligible accuracy loss.

研究の動機と目的

• 多様な効率的アテンション手法を、同じ抽象化の下に統合し、見た目には異なる手法同士の背後にある共通の構造的原則を明らかにすること。
• 長文系列における標準的トランスフォーマーの二次的計算コストを、メモリサイズを制限することで解消しつつ、性能を維持すること。
• ヒューリスティックに基づく手法に代わる、学習可能で文脈に即応したメモリ制御戦略を開発し、効率性と精度のトレードオフを改善すること。
• 従来、因果的アテンションでは適用不可能とされてきた手法（例：Linformer）が、ABCフレームワークを通じて適合可能であることを示すこと。
• 標準の自己アテンションの実用的で即座に置き換え可能な代替手段を提供し、精度に損失を生じさせることなく、著しい速度向上とメモリ使用量の削減を実現すること。

提案手法

• ABCは、固定サイズのメモリにキー・バリュー対を格納し、クエリへのアクセスをメモリ管理戦略で制御する、制限付きメモリシステムとしてアテンションを定式化する。
• このフレームワークは、スパースアテンション、ローカルアテンション、Linformerといった既存手法を、それらのメモリ管理戦略を異なる制御メカニズムとして抽象化することで統一的に扱う。
• ABC$_{\text{MLP}}$ は、文脈的入力特徴に基づいて、どのトークンをメモリに格納し、どこに格納するかを決定する、学習可能なニューラルネットワークベースのコントローラーを導入する。
• コントローラーは、メモリへのライト優先度を計算するために多層パーセプトロン（MLP）を用い、動的で文脈に依存するメモリ管理を可能にする。
• ABCは因果的および非因果的アテンションの両方をサポートし、線形計算量を維持しながら、効率的な自己回帰的デコードとシーケンス符号化を実現する。
• アテンション計算では、制限付きメモリ上でのソフトマックス正規化されたリード操作を用い、定数時間・定数空間のアクセスオーバーヘッドを保証する。

実験結果

リサーチクエスチョン

• RQ1 多様な効率的アテンション手法は、それらの共通する構造的原則を明らかにする、同一の抽象化の下に統合可能か？
• RQ2 ヒューリスティックに基づくメモリ管理に代わる、学習可能で文脈に即応したメモリ制御戦略は、精度と効率の両面で優れているか？
• RQ3 従来、因果的アテンションでは適用不可能とされてきた手法（例：Linformer）は、ABCフレームワークを通じて適合可能か？
• RQ4 メモリサイズの変更が、さまざまなアテンションタイプやタスクにおけるモデル性能にどのように影響を与えるか？
• RQ5 ABC$_{\text{MLP}}$ は、標準のトランスフォーマー性能に匹敵または上回りつつ、推論時間とメモリ使用量を削減できるか、その程度はどの程度か？

主な発見

• ABC$_{\text{MLP}}$ は、RFA よりも最大1.25倍速く、T2R よりも20％速く、著しいメモリオーバーヘッドの削減を達成している。
• 8-8のメモリサイズ（クロスおよび因果的）を用いたABC$_{\text{MLP}}$ は、最高性能を示した設定からわずか1.1 BLEUポイントの差に抑えられ、高い効率-精度トレードオフを示している。
• テキスト符号化の文脈では、ABC$_{\text{MLP}}$ はメモリ使用量を50％削減し、推論速度を標準のトランスフォーマー基準モデルと比較して1.5倍速くしている。
• 従来、因果的アテンションでは適用不可能とされていたLinformerが、ABCの抽象化を通じて適合可能であることが示され、機械翻訳タスクで競争力のある性能を示した。
• ABC$_{\text{MLP}}$ は、他のABCバージョンに比べて、より小さなメモリサイズで高い精度を達成しており、文脈に即応したメモリ制御の利点を裏付けている。
• ABC$_{\text{MLP}}$ の効率的利点は、長文系列において最も顕著であり、長文脈処理アプリケーションにおける強い漸近的優位性を示している。