Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

要約

リニア アテンションは、従来のソフトマックス アテンションに代わる有望な代替手段として最近登場した効率的なアテンション メカニズムです。
リニア アテンションは、線形の計算複雑さでトークンを処理できるため、理論上、速度を犠牲にすることなく無制限の長さのシーケンスを処理できます。つまり、一定のメモリ消費量でさまざまなシーケンス長に対して一定のトレーニング速度を維持できます。
ただし、累積加算 (cumsum) の問題により、現在の線形注意アルゴリズムは因果関係のある設定では理論上の利点を実証できません。
この論文では、線形アテンションの理論的な計算上の利点を実現する最初の線形アテンションの実装である Lightning Attendance-2 を紹介します。
これを実現するために、タイリングの考え方を活用し、線形アテンション計算でブロック内コンポーネントとブロック間コンポーネントを個別に処理します。
具体的には、ブロック内には従来のアテンション計算メカニズムを利用し、ブロック間には線形アテンション カーネル トリックを適用します。
GPU ハードウェアを最大限に活用するために、前方手順と後方手順の両方を通じてタイリング技術が採用されています。
私たちはアルゴリズムを Triton に実装して、IO を認識し、ハードウェアに優しいものにしています。
さまざまなモデルサイズとシーケンス長に対してさまざまな実験が行われます。
Lightning Attendee-2 は、入力シーケンスの長さに関係なく、一貫したトレーニングと推論の速度を維持し、他のアテンション メカニズムよりも大幅に高速です。
ソース コードは https://github.com/OpenNLPLab/lightning-attention で入手できます。

要約(オリジナル)

Linear attention is an efficient attention mechanism that has recently emerged as a promising alternative to conventional softmax attention. With its ability to process tokens in linear computational complexities, linear attention, in theory, can handle sequences of unlimited length without sacrificing speed, i.e., maintaining a constant training speed for various sequence lengths with a fixed memory consumption. However, due to the issue with cumulative summation (cumsum), current linear attention algorithms cannot demonstrate their theoretical advantage in a causal setting. In this paper, we present Lightning Attention-2, the first linear attention implementation that enables linear attention to realize its theoretical computational benefits. To achieve this, we leverage the thought of tiling, separately handling the intra-block and inter-block components in linear attention calculation. Specifically, we utilize the conventional attention computation mechanism for the intra-blocks and apply linear attention kernel tricks for the inter-blocks. A tiling technique is adopted through both forward and backward procedures to take full advantage of the GPU hardware. We implement our algorithm in Triton to make it IO-aware and hardware-friendly. Various experiments are conducted on different model sizes and sequence lengths. Lightning Attention-2 retains consistent training and inference speed regardless of input sequence length and is significantly faster than other attention mechanisms. The source code is available at https://github.com/OpenNLPLab/lightning-attention.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google