Kv cache compression method retains long and off context query strength
Cartridges++: KV Cache Compression without Off-Context Derailment
Machine Learning
Summary
Working with long documents on large language models uses a lot of computer memory and time, so people try to compress the memory it needs. The authors study a popular compression method called Cartridges, which works well for questions about the original document but struggles when questions are unrelated to that document. They find a trade-off between remembering document details and handling general queries. They propose Cartridges++, which modify the compression to keep the benefits for document queries while also preserving the model’s ability to answer off-topic questions well.
What this means in practice
- •For llm deployment engineers: Implement efficient memory compression for long documents that also maintains model quality on unrelated queries.
- •For cloud ai infrastructure teams: Optimize serving costs by reducing KV cache memory footprint without losing generalization on off-context user queries.
Authors
Sonia Laguna, Joao Monteiro, Marco Cuturi, Pierre Ablin, Eleonora Gualdoni
Abstract
Serving long documents to a Large Language Model (LLM) repeatedly is expensive: computations grow with context length, and the memory footprint of the key-value (KV) cache balloons. Compressed KV (CKV) representations aim to mimic the cache of a document and are typically computed once and for all, ahead of inference time. Methods to obtain CKVs range from drop mechanisms that reduce their number of columns, to learned approaches. Among the latter, Cartridges have emerged as a leading compression method, learning compact KV representations through distillation on relevant Q/A pairs. While existing evaluations focus primarily on whether Cartridges and other CKVs yield approximately similar responses to document-related, on-context queries, we investigate the crucial deployment question of whether they can handle off-context queries, something the native KV representation is particularly good at, thanks to the mechanics of attention. We observe a fundamental trade-off: while Cartridges perform better for on-context queries, heuristic-variants preserve better the original LLM's ability to operate off-context. We measure this through their capability to avoid context contamination in their response, retain general knowledge, and follow instructions. We propose Cartridges++, simple modifications to cartridges that retain off-context abilities at small or negligible cost. The router variant decides at inference time whether the query should use the learned long-context memory, while the data-mixing variant allocates a small fraction of training Q/As to queries outside the reference long document. Our study shows that assessing CKVs on document utility alone can mask substantial degradation in broader model capabilities, yet those issues can be fixed with benign changes to CKV inference or training.