Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2605.14037 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914564903272448 |
|---|---|
| author | Szilvasy, Gergely Faysse, Manuel Lomeli, Maria Douze, Matthijs Mazaré, Pierre-Emmanuel Cabannes, Loïc Yih, Wen-tau Jégou, Hervé |
| author_facet | Szilvasy, Gergely Faysse, Manuel Lomeli, Maria Douze, Matthijs Mazaré, Pierre-Emmanuel Cabannes, Loïc Yih, Wen-tau Jégou, Hervé |
| contents | Under modern test-time compute and agentic paradigms, language models process ever-longer sequences. Efficient text generation with transformer architectures is increasingly constrained by the Key-Value cache memory footprint and bandwidth. To address this limitation, we introduce Self-Pruned Key-Value Attention (SP-KV), a mechanism designed to predict future KV utility in order to reduce the size of the long-term KV cache. This strategy operates at a fine granularity: a lightweight utility predictor scores each key-value pair, and while recent KVs are always available via a local window, older pairs are written in the cache and used in global attention only if their predicted utility surpasses a given threshold. The LLM and the utility predictor are trained jointly end-to-end exclusively through next-token prediction loss, and are adapted from pretrained LLM checkpoints.
Rather than enforcing a fixed compression ratio, SP-KV performs dynamic sparsification: the mechanism adapts to the input and typically reduces the KV cache size by a factor of $3$ to $10\times$, longer sequences often being more compressible. This leads to vast improvements in memory usage and decoding speed, with little to no degradation of validation loss nor performance on a broad set of downstream tasks. Beyond serving as an effective KV-cache reduction mechanism, our method reveals structured layer- and head-specific sparsity patterns that we can use to guide the design of hybrid local-global attention architectures. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_14037 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Self-Pruned Key-Value Attention: Learning When to Write by Predicting Future Utility Szilvasy, Gergely Faysse, Manuel Lomeli, Maria Douze, Matthijs Mazaré, Pierre-Emmanuel Cabannes, Loïc Yih, Wen-tau Jégou, Hervé Machine Learning Computation and Language I.2.6; I.2.7 Under modern test-time compute and agentic paradigms, language models process ever-longer sequences. Efficient text generation with transformer architectures is increasingly constrained by the Key-Value cache memory footprint and bandwidth. To address this limitation, we introduce Self-Pruned Key-Value Attention (SP-KV), a mechanism designed to predict future KV utility in order to reduce the size of the long-term KV cache. This strategy operates at a fine granularity: a lightweight utility predictor scores each key-value pair, and while recent KVs are always available via a local window, older pairs are written in the cache and used in global attention only if their predicted utility surpasses a given threshold. The LLM and the utility predictor are trained jointly end-to-end exclusively through next-token prediction loss, and are adapted from pretrained LLM checkpoints. Rather than enforcing a fixed compression ratio, SP-KV performs dynamic sparsification: the mechanism adapts to the input and typically reduces the KV cache size by a factor of $3$ to $10\times$, longer sequences often being more compressible. This leads to vast improvements in memory usage and decoding speed, with little to no degradation of validation loss nor performance on a broad set of downstream tasks. Beyond serving as an effective KV-cache reduction mechanism, our method reveals structured layer- and head-specific sparsity patterns that we can use to guide the design of hybrid local-global attention architectures. |
| title | Self-Pruned Key-Value Attention: Learning When to Write by Predicting Future Utility |
| topic | Machine Learning Computation and Language I.2.6; I.2.7 |
| url | https://arxiv.org/abs/2605.14037 |