arXiv:2501.04987cs.CL2025-01IJCAI被引 5

提出树结构压缩缓存,让大模型在长文本中保持高质量生成。

TreeKV: Smooth Key-Value Cache Compression with Tree Structures

  • 用树结构实现平滑的键值缓存压缩,无需训练。
  • 在长文本任务中,缓存减少16倍仍保持高精度。
  • 适用于生成和预填充阶段,适合资源受限场景。

高效的关键值(KV)缓存压缩对在长序列和资源受限环境下扩展基于Transformer的大语言模型至关重要。现有方法基于位置或重要性分数驱逐标记,但位置策略可能遗漏预定义区域外的关键信息,而依赖全局重要性分数的方法则导致强区域偏差,限制了缓存的整体上下文保留能力,可能影响模型在复杂任务上的表现。我们的小波分析显示,随着标记接近序列末端,其生成贡献逐渐增加,并趋于与邻近标记差异更大,表明远端到近端上下文存在平滑递增的复杂性和变异性。受此启发,我们提出TreeKV,一种直观且无需训练的树结构缓存压缩方法。TreeKV维持固定缓存大小,使模型在长文本场景下仍能生成高质量输出。与大多数压缩方法不同,TreeKV同时适用于生成和预填充阶段。在PG19和OpenWebText2的语言建模任务中,TreeKV持续优于所有基线模型,使仅训练短上下文窗口的模型也能泛化至更长窗口,实现16倍缓存压缩。在Longbench基准测试中,仅使用6%的缓存预算即达到最佳效率。

原文摘要 · Abstract (English)

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.

缓存压缩大模型树结构长文本

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。