Ring Attention with Blockwise Transformers for Near-Infinite Context

要約

トランスフォーマーは、多くの最先端の AI モデルに最適なアーキテクチャとして登場し、幅広い AI アプリケーションにわたって優れたパフォーマンスを発揮します。
ただし、トランスフォーマーによって課されるメモリ要求により、長いシーケンスを処理する能力が制限されるため、複雑な環境でビデオ、アクション、その他の長い形式のシーケンスやモダリティを利用する際に課題が生じます。
我々は、ブロック単位のトランスフォーマーによるリング アテンション (リング アテンション) という新しいアプローチを提案します。これは、セルフ アテンションとフィードフォワードのブロック単位の計算を利用して、キーと値のブロックの通信をブロック単位のアテンションの計算で完全にオーバーラップさせながら、複数のデバイスに長いシーケンスを分散します。
私たちのアプローチにより、近似に頼ったり、追加の通信や計算のオーバーヘッドを発生させたりすることなく、従来のメモリ効率の高い Transformer によって達成可能であったものよりも、最大デバイス数倍長いシーケンスのトレーニングと推論が可能になります。
言語モデリングと強化学習タスクに関する広範な実験により、数百万のトークンコンテキストサイズを許容し、パフォーマンスを向上させるというアプローチの有効性が実証されました。

要約(オリジナル)

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 posing challenges in utilizing videos, actions, and other long-form sequences and modalities in complex environments. We present a novel approach, Ring Attention with Blockwise Transformers (Ring Attention), which leverages blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices while fully overlapping the communication of key-value blocks with the computation of blockwise attention. Our approach enables training and inference of sequences that are up to device count times longer than those achievable by prior memory-efficient Transformers, without resorting to approximations or incurring additional communication and computation overheads. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of our approach in allowing millions of tokens context size and improving performance.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google