Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

要約

キーバリュー (KV) キャッシュは、トランスフォーマーベースの自己回帰大規模言語モデル (LLM) のデコードを高速化する上で重要な役割を果たします。
ただし、シーケンスの長さが長く、バッチ サイズが大きい場合、KV キャッシュの保存に必要なメモリの量が法外に大きくなる可能性があります。
トランスフォーマーの発明以来、KV キャッシュのサイズを削減するために発見された最も効果的な介入の 2 つは、マルチクエリ アテンション (MQA) とその一般化であるグループ化クエリ アテンション (GQA) です。
MQA と GQA は両方とも、複数のクエリ ヘッドが 1 つのキー/値ヘッドを共有できるようにアテンション ブロックの設計を変更し、精度の低下を最小限に抑えながら個別のキー/値ヘッドの数を大幅に削減します。
この論文では、隣接するレイヤー間でキーとバリューのヘッドを共有することによって、マルチクエリ アテンションをさらに一歩進めることが可能であり、クロスレイヤー アテンション (CLA) と呼ばれる新しいアテンション デザインが得られることを示します。
CLA を使用すると、未修正の MQA とほぼ同じ精度を維持しながら、KV キャッシュのサイズをさらに 2 倍削減できることがわかりました。
1B および 3B パラメータ モデルをゼロからトレーニングする実験では、CLA が従来の MQA で可能であったメモリと精度のトレードオフをパレート的に改善し、他の方法で可能であるよりも長いシーケンス長と大きなバッチ サイズでの推論を可能にすることを実証しました。

要約(オリジナル)

Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to store the KV cache can become prohibitive at long sequence lengths and large batch sizes. Since the invention of the transformer, two of the most effective interventions discovered for reducing the size of the KV cache have been Multi-Query Attention (MQA) and its generalization, Grouped-Query Attention (GQA). MQA and GQA both modify the design of the attention block so that multiple query heads can share a single key/value head, reducing the number of distinct key/value heads by a large factor while only minimally degrading accuracy. In this paper, we show that it is possible to take Multi-Query Attention a step further by also sharing key and value heads between adjacent layers, yielding a new attention design we call Cross-Layer Attention (CLA). With CLA, we find that it is possible to reduce the size of the KV cache by another 2x while maintaining nearly the same accuracy as unmodified MQA. In experiments training 1B- and 3B-parameter models from scratch, we demonstrate that CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA, enabling inference with longer sequence lengths and larger batch sizes than would otherwise be possible

arxiv情報

提供元, 利用サービス

arxiv.jp, Google