[論文レビュー] Simple linear attention language models balance the recall-throughput tradeoff

Simran Arora, Sabri Eyuboglu+7 more|arXiv (Cornell University)|Feb 28, 2024

Topic Modeling被引用数 4

ひとこと要約

この論文は、線形アテンションとスライディングウィンドウアテンションを組み合わせることで、リcallと推論効率のバランスをとったハイブリッド言語モデルアーキテクチャであるBasedを提案する。ウィンドウサイズと特徴次元を調整することで、Basedはメモリ使用量を著しく削減しながら、完全なアテンションレベルのリcallを達成し、FlashAttention-2に比べ24倍高いスループットを実現する。MambaなどのSOTAモデルに比べ、リcall集中的なタスクで6.22ポイントの精度向上を達成した。

ABSTRACT

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.

研究の動機と目的

• 自己回帰的言語モデルにおけるリcall性能と推論効率のトレードオフを解消すること。
• 異なるアーキテクチャにおいて、モデル状態サイズとリcall能力のパレートフロンティアを特定・走査すること。
• 標準アテンションと同等の高リcallを維持しながら、効率的代替手法（例：Mamba、RWKV）と同様にメモリ消費を低減するモデルを設計すること。
• 実世界のスループットベンチマークで競争力を持つよう、I/Oに配慮した最適化を実現すること。
• 線形アテンションとスライディングウィンドウアテンションを組み合わせることで、個々のコンponentよりも優れたトレードオフ性能を達成できることを示すこと。

提案手法

• 同じモデルレイヤー内で線形アテンションとスライディングウィンドウアテンションを組み合わせたハイブリッドアーキテクチャ、Basedを導入すること。
• 可変なスライディングウィンドウサイズ（例：64, 128）と線形アテンションの特徴次元を用いて、状態サイズとトレードオフ特性を制御すること。
• 生成処理中のメモリアクセスオーバーヘッドを低減し、スループットを向上させるためにI/Oに配慮したアルゴリズム的最適化を適用すること。
• Adam、コサインスケーリング、BFloat16精度を用いた標準的なトレーニングレシピに従い、最大1.3Bパラメータのモデルを訓練すること。
• スライディングウィンドウコンponentにはロータリー位置エンコーディングを適用し、線形コンponentには位置エンコーディングを適用しないこと。
• ブロック排他的多項式解析を実装し、関連的リcallタスクを解くために十分なモデル深さが必要であることを理論的に裏付けること。

Figure 1: Based overview. Combining linear attention with tiny sliding window softmax attention (e.g., 64 or 128 tokens in width) enables improved recall accuracy with limited efficiency overhead vs. smaller tile sizes. ( Left ) Time to execute Cutlass GEMMs ( $y$ ) vs. sliding window attention size

実験結果

リサーチクエスチョン

• RQ1 Mambaのような効率的代替手法の小さな状態サイズを維持しながら、標準アテンションと同等のリcall性能を達成できるか？
• RQ2 異なるアーキテクチャクラス間で、モデル状態サイズとリcall能力の根本的トレードオフは何か？
• RQ3 線形アテンションとスライディングウィンドウアテンションを組み合わせることで、両コンponentを個別に使用するよりも優れたパレート最適解が得られるか？
• RQ4 I/O効率の良いアルゴリズムを設計することで、ハイブリッドアテンションモデルが実世界のスループットベンチマークでFlashAttentionに匹敵または劣らない性能を発揮できるか？
• RQ5 線形アテンションと局所アテンションメカニズムを組み合わせることで、個別に使用する場合よりも優れた長距離かつ正確なトークン相互作用が可能になるか？

主な発見

• Basedは、Mambaに比べて実世界のリcall集中的タスクで6.22ポイント高い精度を達成し、小さな状態サイズであるにもかかわらず優れたリcall能力を示した。
• 1.3Bパラメータのモデルで、1024トークンの自己回帰的生成において、FlashAttention-2に比べ24倍高いスループットを達成した。
• ハイブリッドアーキテクチャは、リcallとメモリのトレードオフのパレートフロンティアを効果的に走査し、わずかな効率的オーバーヘッドで完全なアテンションレベルのリcallを回復した。
• 線形アテンション単体では、局所的トークン比較における精度不足のため、関連的リcallタスクに失敗するが、スライディングウィンドウアテンションはウィンドウサイズの制限に縛られる。
• 理論的分析により、データに依存しないBaseConvモデルは、関連的リcallを解くために少なくとも⌊log(2p)⌋層を必要とすることが示され、このようなモデルにおける十分な深さの必要性が裏付けられた。
• I/Oに配慮した最適化により、ハイブリッドアテンションコンponentを備えたBasedは、最適化済みアテンション実装のスループットに匹敵または上回る性能を発揮した。

Figure 2: Throughput (memory) - recall tradeoff. $x$ -axis shows state size (bytes) during generation; $y$ -axis shows accuracy on the MQAR recall task [ 1 ] . For each architecture, we train several models varying hyperparameters that affect the recurrent state size ( e.g. model dimension). The plo