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

Alham Fikri Aji; Ayu Purwarianti; Muhammad Farid Adilazuarda; Zayd Muhammad Kawakibi Zuhri

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

Not yet reviewed by Pith; the record is open.

Re-run · record.json Download PDF Read on arXiv ↗

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2406.09297 v3 pith:QQQUQSKZ submitted 2024-06-13 cs.LG

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

Zayd Muhammad Kawakibi Zuhri , Muhammad Farid Adilazuarda , Ayu Purwarianti , Alham Fikri Aji This is my paper

classification cs.LG

keywords memorymlkvkey-valuesizetransformerattentionefficientinference

verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

0 comments

read the original 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

discussion (0)

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

SOS-LoRA: Static Orthogonal-Subspace Low-Rank Adaptation with Fixed Multi-Scale Scaling
cs.LG 2026-06 conditional novelty 5.0

A LoRA update split into several fixed, differently-scaled low-rank experts with orthogonal input directions improves fine-tuning accuracy at the same parameter count.