Simple linear attention language models balance the recall-throughput tradeoff

要約

最近の研究では、注意に基づいた言語モデルがリコール時に優れていることが示されています。これは、以前にコンテキストで見られたトークンで世代を接地する能力です。
ただし、注意ベースのモデルの効率は、KV-Cacheの積極的なメモリ消費による推論中にボトルネックされています。
この作業では、リコールを妥協することなく、言語モデルの効率を改善できるかどうか（たとえば、メモリ消費を減らすことで）を探ります。
実験と理論を幅広いアーキテクチャに適用することにより、モデルの状態サイズとリコール能力の間の重要なトレードオフを特定します。
注意の効率的な代替手段（例：H3、Mamba、RWKV）は固定サイズの再発状態を維持しているが、リコールに苦労することを示しています。
線形とスライドのウィンドウの注意を組み合わせたシンプルなアーキテクチャを基にして提案します。
さまざまなベースのウィンドウサイズと線形注意特徴の寸法により、状態サイズをダイヤルして、リコールメモリトレードオフカーブのパレートフロンティアを通過し、一方の端での注意の完全な品質と、他方の注意相対者の小さな状態サイズを回復できます。
1.3Bパラメーターまでの言語モデルをトレーニングし、ベースが困惑の最も強力な亜二次モデル（MAMBAなど）と一致し、実際のリコール集約型タスクで6.22の精度ポイントでそれらを上回ることを示します。
線形注意の実装は、最適化された標準的な注意実装よりも効率が低いことがよくあります。
1.3Bパラメーターモデルを使用して1024トークンを生成するときに、Flashattention-2よりも言語生成で24倍高いスループットを有効にするIO認識アルゴリズムを開発するために、1.3Bパラメーターモデルを使用して1024トークンを開発します。
この作業のコードは、https：//github.com/hazyresearch/badedに提供されています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google