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

Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji
6/13/2024
cs.LG

Official code: zaydzuhri/pythia-mlkv (35 stars, last commit 2024-06-14).

Abstract

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

View on arXivView PDF

Code Implementations(1)

zaydzuhri/pythia-mlkvOfficial100%

Multi-Layer Key-Value sharing experiments on Pythia models

352Python, PowerShellDec 24, 20232 years ago

Similar Papers

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

88% match1 repo172024

Layer-Condensed KV Cache for Efficient Inference of Large Language Models

87% match1 repo1562024

Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads

87% match2025

Low-Rank Key Value Attention

86% match2026

LLMCache: Layer-Wise Caching Strategies for Accelerated Reuse in Transformer Inference

85% match2025

KV-Latent: Dimensional-level KV Cache Reduction with Frequency-aware Rotary Positional Embedding

85% match1 repo32025

Cite this paper

@article{zuhri2024mlkv,
  title  = {MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding},
  author = {Zayd Muhammad Kawakibi Zuhri and Muhammad Farid Adilazuarda and Ayu Purwarianti and Alham Fikri Aji},
  year   = {2024},
  eprint = {2406.09297},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2406.09297v3}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2406.09297.svg)](https://codeofpaper.com/paper/2406.09297)

Discussion