Ring Attention with Blockwise Transformers for Near-Infinite Context

要約

トランスフォーマーは、多くの最先端の AI モデルに最適なアーキテクチャとして登場し、幅広い AI アプリケーションにわたって優れたパフォーマンスを発揮します。
ただし、トランスフォーマーによって課されるメモリ要求により、トランスフォーマーが長いシーケンスを処理する能力が制限されるため、拡張されたシーケンスや長期的な依存関係を含むタスクには課題が生じます。
我々は、セルフ アテンションのブロック単位の計算を利用して、キーと値のブロックの通信をブロック単位のアテンションの計算で同時にオーバーラップさせながら、複数のデバイスに長いシーケンスを分散する独特のアプローチであるリング アテンションを提案します。
メモリ効率を維持しながらより長い入力シーケンスを処理することで、Ring Attendance は、以前のメモリ効率の高い Transformer よりもデバイス カウント倍長いシーケンスのトレーニングと推論を可能にし、個々のデバイスによって課されるメモリ制約を効果的に排除します。
言語モデリング タスクに関する広範な実験により、大きなシーケンス入力サイズを許可し、パフォーマンスを向上させる上での Ring Attendance の有効性が実証されています。

要約(オリジナル)

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, Google