LoLA: Low-Rank Linear Attention With Sparse Caching

要約

変圧器ベースの大手言語モデルは、長いシーケンスの推論で二次複雑さに悩まされています。
線形注意方法は効率的な代替品ですが、ソフトマックスの注意の正確な近似を提供できません。
さらに、各線形注意ヘッドにスライディングウィンドウの注意を組み込むことにより、このギャップは短いコンテキスト長タスクのために閉じることができます。
残念ながら、これらのアプローチは、「メモリの衝突」のために長い文脈から重要な情報を思い出せません。
このホワイトペーパーでは、Lolaを提案します。まばらなキャッシングを使用した低ランクの線形注意。
Lolaは、過去の連想記憶を妨げる追加のキー価値ペアを個別に保存します。
さらに、Lolaは、過去のキー価値ペアを3つの形式のメモリに分配することにより、線形注意モデルとトランスの間のギャップをさらに閉じます。（i）ローカルスライドウィンドウの最近のペア。
（ii）まばらでグローバルなキャッシュの中で、モモの困難なペア。
（iii）線形注意の再発性隠された状態における一般的なペア。
推論のみの戦略として、Lolaは、RulerのHeadle-in-a-Haystackタスクで最大8kのコンテキストの長さでパスキー検索を有効にします。
4Kコンテキストの長さでベースサブクォドラティックモデルの精度を0.6％から97.4％に高め、llama-3.18bのキャッシュよりも4.6倍小さいキャッシュを備えています。
Lolaは、1Bおよび8Bパラメーターのサブクアドラティスモデルの間で、ゼロショットの常識推論タスクで強力なパフォーマンスを示しています。
最後に、Lolaは非常に軽量のアプローチです。私たちの結果のほぼすべては、単一の消費者GPUで再現できます。

要約(オリジナル)

Transformer-based large language models suffer from quadratic complexity at inference on long sequences. Linear attention methods are efficient alternatives, however, they fail to provide an accurate approximation of softmax attention. By additionally incorporating sliding window attention into each linear attention head, this gap can be closed for short context-length tasks. Unfortunately, these approaches cannot recall important information from long contexts due to ‘memory collisions’. In this paper , we propose LoLA: Low-rank Linear Attention with sparse caching. LoLA separately stores additional key-value pairs that would otherwise interfere with past associative memories. Moreover, LoLA further closes the gap between linear attention models and transformers by distributing past key-value pairs into three forms of memory: (i) recent pairs in a local sliding window; (ii) difficult-to-memorize pairs in a sparse, global cache; and (iii) generic pairs in the recurrent hidden state of linear attention. As an inference-only strategy, LoLA enables pass-key retrieval on up to 8K context lengths on needle-in-a-haystack tasks from RULER. It boosts the accuracy of the base subquadratic model from 0.6% to 97.4% at 4K context lengths, with a 4.6x smaller cache than that of Llama-3.1 8B. LoLA demonstrates strong performance on zero-shot commonsense reasoning tasks among 1B and 8B parameter subquadratic models. Finally, LoLA is an extremely lightweight approach: Nearly all of our results can be reproduced on a single consumer GPU.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google