MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866917803267719168 |
|---|---|
| author | Zuhri, Zayd Muhammad Kawakibi Adilazuarda, Muhammad Farid Purwarianti, Ayu Aji, Alham Fikri |
| author_facet | Zuhri, Zayd Muhammad Kawakibi Adilazuarda, Muhammad Farid Purwarianti, Ayu Aji, Alham Fikri |
| contents | 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 |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2406_09297 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding Zuhri, Zayd Muhammad Kawakibi Adilazuarda, Muhammad Farid Purwarianti, Ayu Aji, Alham Fikri Machine Learning 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 |
| title | MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2406.09297 |