Landmark Attention: Random-Access Infinite Context Length for Transformers

要約

Transformers は自然言語処理において目覚ましい成功を収めてきましたが、そのアテンション メカニズムの大量のメモリ要件により、より長いコンテキストを処理する能力は制限されていました。
リカレント メモリや検索ベースの拡張などの以前のアプローチでは、アテンションのランダム アクセスの柔軟性 (つまり、コンテキスト全体で任意のトークンを選択する機能) が損なわれるか、関連するコンテキストの検索に別のメカニズムに依存していました。
モデルの注意と一致すること。
この論文では、ランダムアクセスの柔軟性を維持しながら完全なコンテキストへのアクセスを可能にし、コンテキスト全体に注意を向けるのとよく似た新しいアプローチを紹介します。
私たちのメソッドは、ランドマーク トークンを使用して入力の各ブロックを表し、関連するブロックを選択するためにそれを使用するようにアテンションをトレーニングします。これにより、別個のメカニズムに依存するのではなく、アテンション メカニズムを通じて直接ブロックを取得できるようになります。
私たちのアプローチは、特殊なデータ構造とシステムのメモリ階層をシームレスに統合し、任意の長いコンテキスト長の処理を可能にします。
私たちの方法は、各ステップで取得されるトークンの数を大幅に削減しながら、Transformer-XL と同等のパフォーマンスを得ることができることを実証します。
最後に、私たちの方法で LLaMA 7B を微調整すると、そのコンテキスト長の容量が 32,000 トークンを超えるまで拡張され、GPT-4 のコンテキスト長での推論が可能になることを示します。
ランドマーク アテンションの実装と実験を再現するコードを https://github.com/epfml/landmark-attention/ でリリースします。

要約(オリジナル)

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/.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google