MLKV: メモリ効率の良いTransformerデコーディングのための多層キー・バリューヘッド

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

June 13, 2024

著者: Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji

cs.AI

要旨

トランスフォーマーの自己回帰型推論は、Key-Value（KV）キャッシュの活用によって大きな恩恵を受けますが、モデルサイズ、バッチサイズ、シーケンス長がスケールアップするにつれて、メモリのボトルネックが深刻化する可能性があります。本論文では、Multi-Layer Key-Value（MLKV）共有という新たなアプローチを提案します。これは、トランスフォーマーの層を跨いでKV共有を拡張し、Multi-Query Attention（MQA）やGrouped-Query Attention（GQA）を超えるメモリ使用量の削減を実現します。アップトレーニングされたPythia-160Mのバリエーションを用いた様々なNLPベンチマークおよび推論メトリクスでの評価により、MLKVが性能の低下を最小限に抑えつつメモリ使用量を大幅に削減し、MQAと比較してKVキャッシュサイズを最大6分の1にまで縮小できることが示されました。これらの結果は、MLKVがスケールアップしたトランスフォーマーモデルの効率的なデプロイメントにおいて大きな可能性を秘めていることを示しています。コードはhttps://github.com/zaydzuhri/pythia-mlkvで公開しています。

English

Auto-regressive inference of transformers benefit greatly from Key-Value (KV) caching, but can lead to major memory bottlenecks as model size, batch size, and sequence length grow at scale. We introduce Multi-Layer Key-Value (MLKV) sharing, a novel approach extending KV sharing across transformer layers to reduce memory usage beyond what was possible with Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Evaluations on various NLP benchmarks and inference metrics using uptrained Pythia-160M variants demonstrate that MLKV significantly reduces memory usage with minimal performance loss, reducing KV cache size down to a factor of 6x compared to MQA. These results highlight MLKV's potential for efficient deployment of transformer models at scale. We provide code at https://github.com/zaydzuhri/pythia-mlkv