TreeKV:采用树结构实现平滑的键值缓存压缩
TreeKV: Smooth Key-Value Cache Compression with Tree Structures
AI总结:
针对大语言模型KV缓存现有压缩方法的信息遗漏与区域偏差问题,提出无需训练的树结构压缩方法TreeKV,可同时适用于生成与预填充阶段,在多项基准测试中以更低缓存开销实现更优性能。
摘要:
高效的键值(KV)缓存压缩对于基于Transformer的大语言模型(LLM)在长序列和资源受限场景下的扩展至关重要。现有方法根据token的位置或重要性分数进行驱逐,但基于位置的策略可能会遗漏预定义区域之外的关键信息,而依赖全局重要性分数的方法则会产生强烈的区域偏差,限制了KV缓存的整体上下文保留能力,还可能损害LLM在复杂任务上的性能。我们的小波分析表明,随着token接近序列末尾,它们对生成的贡献逐渐增加,且与相邻token的差异趋于更大,这表明从远上下文到近上下文存在着复杂度和可变性不断提升的平滑过渡。受这一观察启发,我们提出了TreeKV,这是一种直观、无需训练的方法,采用树结构实现平滑的缓存压缩。TreeKV维持固定的缓存大小,使LLM即使在长文本场景下也能输出高质量结果。与大多数压缩方法不同,TreeKV同时适用于生成阶段和预填充阶段。在PG19和OpenWebText2的语言建模任务中,TreeKV的表现始终优于所有基线模型,使经过短上下文窗口训练的LLM能够泛化到更长的窗口,同时缓存量减少16倍。在Longbench基准测试中,TreeKV以仅6%的预算实现了最佳性能,达到最优效率。
英文摘要:
Efficient key-value (KV) cache compression is critical for scaling transformer-based Large Language Models (LLMs) in long sequences and resource-limited settings. Existing methods evict tokens based on their positions or importance scores, but position-based strategies can miss crucial information outside predefined regions, while those relying on global importance scores resulting in strong regional biases, limiting the KV cache's overall context retention and potentially impairing the performance of LLMs on complex tasks. Our wavelet analysis reveals that as tokens approach the end of sequence, their contributions to generation gradually increase and tends to diverge more from neighboring tokens, indicating a smooth transition with increasing complexity and variability from distant to nearby context. Motivated by this observation, we propose TreeKV, an intuitive, training-free method that employs a tree structure for smooth cache compression. TreeKV maintains a fixed cache size, allowing LLMs to deliver high-quality output even in long text scenarios. Unlike most compression methods, TreeKV is applicable to both the generation and prefilling stages. TreeKV consistently surpasses all baseline models in language modeling tasks on PG19 and OpenWebText2, allowing LLMs trained with short context window to generalize to longer window with a 16x cache reduction. On the Longbench benchmark, TreeKV achieves the best performance with only 6\% of the budget at optimal efficiency.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。