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

要約

トランスフォーマーの自己回帰推論は、Key-Value (KV) キャッシュから大きな恩恵を受けますが、モデル サイズ、バッチ サイズ、シーケンスの長さが大規模に増大するにつれて、大きなメモリ ボトルネックが発生する可能性があります。
KV 共有をトランス層全体に拡張する新しいアプローチである、Multi-Layer Key-Value (MLK​​V) 共有を導入して、Multi-Query Attendant (MQA) および Grouped-Query Attendant (GQA) で可能であった以上にメモリ使用量を削減します。
アップトレーニングされた Pythia-160M バリアントを使用したさまざまな NLP ベンチマークと推論メトリクスの評価では、MLKV がパフォーマンスの損失を最小限に抑えながらメモリ使用量を大幅に削減し、KV キャッシュ サイズを MQA と比較して 6 分の 1 に削減することが実証されました。
これらの結果は、変圧器モデルを大規模に効率的に導入できる MLKV の可能性を浮き彫りにしています。
https://github.com/zaydzuhri/pythia-mlkv でコードを提供します。

要約(オリジナル)

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

arxiv情報

提供元, 利用サービス

arxiv.jp, Google