Despite the significant success of large language models (LLMs), their extensive memory requirements pose challenges for deploying them in long-context token generation. The substantial memory footprint of LLM decoders arises from the necessity to store all previous tokens in the attention module, a requirement imposed by key-value (KV) caching. In this work, our focus is on developing an efficient compression technique for the KV cache. Empirical evidence indicates a significant clustering tendency within key embeddings in the attention module. Building on this key insight, we have devised a novel caching method with sublinear complexity, employing online clustering on key tokens and online $\ell_2$ sampling on values. The result is a provably accurate and efficient attention decoding algorithm, termed SubGen. Not only does this algorithm ensure a sublinear memory footprint and sublinear time complexity, but we also establish a tight error bound for our approach. Empirical evaluations on long-context question-answering tasks demonstrate that SubGen significantly outperforms existing and state-of-the-art KV cache compression methods in terms of performance and efficiency.
翻译:尽管大语言模型取得了显著成功,但其庞大的内存需求对部署长期上下文中的令牌生成提出了挑战。大语言模型解码器的大量内存占用源于注意力模块中需存储所有先前令牌的键值缓存。本文致力于开发一种高效的键值缓存压缩技术。经验证据表明,注意力模块中的键嵌入具有显著的聚类趋势。基于这一关键发现,我们提出了一种具有亚线性复杂度的新型缓存方法,该方法通过对键令牌进行在线聚类并对值进行在线$\ell_2$采样实现。由此产生的注意力解码算法SubGen不仅保证了亚线性内存占用和亚线性时间复杂度,而且我们为该方法建立了严格的误差界。长期上下文问答任务的实证评估表明,SubGen在性能和效率上显著优于现有及最先进的键值缓存压缩方法。