Reproduce “Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression” on OpenResearch

About this paper

Ruijie Miao, Zhiming Wang, Wang Li, Shiwei Wu, Shufan Liu, Yanbing Jiang, Tong Yang

The MixedDimKV method introduces a fine-grained, adaptive key-value (KV) cache compression scheme for large language models, significantly reducing memory footprint and improving inference efficiency for long-context tasks. This approach enables performance comparable to full attention on Llama-3-8B-Instruct while reducing peak memory usage to 66.44% and total decoding latency to 61.31% for 128K contexts.

Key insights

• Fine-grained, adaptive allocation of latent dimensions to individual tokens, moving beyond binary token eviction, can significantly optimize information retention under memory constraints.
• A robust accuracy loss score that accounts for both attention significance and value magnitudes provides a more accurate measure for prioritizing which tokens and dimensions to compress.
• The optimal dimension allocation problem can be efficiently solved for long-context scenarios using a bisection search due to the negligible duality gap, ensuring scalability without prohibitive computational overhead.

Summary by alphaXiv.