[論文レビュー] Landmark Attention: Random-Access Infinite Context Length for Transformers

Amirkeivan Mohtashami, Martin Jaggi|arXiv (Cornell University)|May 25, 2023

Topic Modeling被引用数 4

ひとこと要約

本論文では、固定長の入力ブロックの取得をランドマークトークンでゲートすることで、ランダムアクセス可能で無限長のコンテキスト長を達成できるようにするLandmark Attentionという手法を提案する。ランドマーク類似度を介して注意機構が関連するブロックを取得するように訓練することで、計算量とメモリ使用量をブロックサイズに等しい割合（例：50倍）で削減でき、LLaMA 7Bが32kトークンを超える推論を実行できるようになる。これはGPT-4のコンテキスト長に相当するが、Transformer-XLと同等の性能を維持する。

ABSTRACT

While Transformers have shown remarkable success in natural language processing, their attention mechanism's large memory requirements have limited their ability to handle longer contexts. Prior approaches, such as recurrent memory or retrieval-based augmentation, have either compromised the random-access flexibility of attention (i.e., the capability to select any token in the entire context) or relied on separate mechanisms for relevant context retrieval, which may not be compatible with the model's attention. In this paper, we present a novel approach that allows access to the complete context while retaining random-access flexibility, closely resembling running attention on the entire context. Our method uses a landmark token to represent each block of the input and trains the attention to use it for selecting relevant blocks, enabling retrieval of blocks directly through the attention mechanism instead of by relying on a separate mechanism. Our approach seamlessly integrates with specialized data structures and the system's memory hierarchy, enabling processing of arbitrarily long context lengths. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number of retrieved tokens in each step. Finally, we show that fine-tuning LLaMA 7B with our method successfully extends its context length capacity to over 32k tokens, allowing for inference at the context lengths of GPT-4. We release the implementation of landmark attention and the code to reproduce our experiments at https://github.com/epfml/landmark-attention/.

研究の動機と目的

• 自己注意機構の長文処理における二次関数的メモリおよび計算コストを克服すること。
• 再帰的処理や外部リtrieverに依存せずに、注意機構のランダムアクセスの柔軟性を維持しながら、長文コンテキストの取得を可能にすること。
• 再トレーニングなしに、学習時に見られなかった長さのコンテキスト長でも任意の長さで推論を可能にすること。
• メモリ階層やFAISSのようなデータ構造とシームレスに統合し、大規模コンテキスト処理を効率的に行えるようにすること。
• 事前学習済みモデル（例：LLaMA 7B）を微調整することで、超長コンテキスト（例：32kトークン）に対応可能にすること。

提案手法

• 入力を固定長のブロックに分割し、各ブロックはそのブロックへの注意を制御するランドマークトークンで表現される。
• 推論時、ランドマークトークンへの注意スコアが、どのブロックが取得・注目されるかを決定し、任意の過去のブロックへのランダムアクセスを可能にする。
• 標準的な注意メカニズムを用いるが、クエリとランドマークキーのベクトル間の類似度に基づいてブロックの取得を条件づける。
• 以前のトークンのキャッシュ（KVキャッシュ）を維持するが、各ブロック内のランドマークでないトークンは、対応するランドマークがアクティブになるまでスワップアウトされ、必要に応じてロードされる。
• 関連するブロックの取得に効率的なデータ構造（例：FAISS）を用いて、k近傍探索を実行する。
• このメカニズムを用いて、事前学習済みモデル（例：LLaMA 7B）を微調整し、再訓練なしにコンテキスト長を拡張する。

実験結果

リサーチクエスチョン

• RQ1 再帰的処理や外部リtrieverに依存せずに、Transformerがランダムアクセス可能で無限長のコンテキスト長を達成できるか。
• RQ2 ランドマークベースの取得が、計算量とメモリ使用量を大幅に削減しつつも、注意機構の柔軟性を維持できるか。
• RQ3 この手法を用いて、既存の大規模モデルを微調整し、超長コンテキスト（例：32kトークン）を処理可能にできるか（性能劣化なしに）。
• RQ4 ランドマーク注意機構を備えたモデルの性能が、長文コンテキストにおいて再帰的モデル（例：Transformer-XL）と比較してどうなるか。
• RQ5 高度なデータ構造と統合することで、長文コンテキスト推論におけるリource使用量をさらに削減できるか。

主な発見

• Landmark Attentionにより、モデルが学習時に短いコンテキストしか見ていなくても、追加のトレーニングコストなしに任意の長さのコンテキスト長での推論が可能になる。
• 計算量とメモリ使用量がブロックサイズに等しい割合で削減される（例：50トークンブロックでは50倍の削減）ため、長文コンテキスト推論が効率的になる。
• LLaMA 7BをLandmark Attentionで微調整することで、最大32,000トークンのコンテキストから隠されたパスフレーズを高精度で抽出可能となり、GPT-4のコンテキスト長に相当する。
• 長文コンテキストタスクにおいても高い性能を維持し、32kトークンを超える長さでもフレーズ回復の精度がほぼ完璧である一方、ベースモデルは失敗する。
• FAISSのような外部データ構造と互換性があり、取得の効率性とメモリ使用量のさらなる最適化が可能になる。
• 解釈可能性が保たれており、ランドマークへの注意スコアがどのブロックがアクセスされているかを明らかにし、モデルの推論プロセスの洞察が得られる。