Ring Attention with Blockwise Transformers for Near-Infinite Context

要約

トランスフォーマーは、多くの最先端AIモデルで選択されるアーキテクチャとして登場し、幅広いAIアプリケーションで卓越した性能を発揮している。しかし、Transformerが要求するメモリは、長いシーケンスを処理する能力を制限するため、長時間のシーケンスや長期的な依存関係を含むタスクに課題をもたらす。我々は、自己注意のブロック単位の計算を活用し、キー・バリュー・ブロックの通信とブロック単位の注意の計算を同時にオーバーラップさせながら、長いシーケンスを複数のデバイスに分散させる、独自のアプローチであるリング・アテンションを提示する。リング・アテンションは、メモリ効率を維持しながら長い入力シーケンスを処理することで、メモリ効率の高い従来のトランスフォーマーよりもデバイス数倍長いシーケンスの学習と推論を可能にし、個々のデバイスによって課されるメモリ制約を効果的に排除する。言語モデリングタスクに関する広範な実験により、大きなシーケンス入力サイズを可能にし、パフォーマンスを向上させるRing Attentionの有効性が実証されている。

要約(オリジナル)

Transformers have emerged as the architecture of choice for many state-of-the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transformers limit their ability to handle long sequences, thereby creating challenges for tasks involving extended sequences or long-term dependencies. We present a distinct approach, Ring Attention, which leverages blockwise computation of self-attention to distribute long sequences across multiple devices while concurrently overlapping the communication of key-value blocks with the computation of blockwise attention. By processing longer input sequences while maintaining memory efficiency, Ring Attention enables training and inference of sequences that are device count times longer than those of prior memory-efficient Transformers, effectively eliminating the memory constraints imposed by individual devices. Extensive experiments on language modeling tasks demonstrate the effectiveness of Ring Attention in allowing large sequence input size and improving performance.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL