arXiv:2411.05787cs.CL2024-11ACL被引 9

让大模型生成长文本时更高效准确,动态更新小缓存避免信息丢失。

RefreshKV: Updating Small KV Cache During Long-form Generation

  • 生成时动态更新小键值缓存,结合全上下文与子集注意力
  • 在长文本生成任务上性能优于传统压缩方法,速度相当
  • 适合需要长时间生成的场景,如故事创作、报告撰写

长序列输入下生成长文本是大语言模型(LLMs)极为耗算力的推理任务。现有加速方法通过构建较小的键值(KV)缓存,避免模型对长序列重复计算注意力,但仅适用于短序列生成,长文本生成时性能急剧下降。多数方法一次性压缩缓存,过早丢弃后续可能有用的上下文。我们提出新方法 RefreshKV,生成过程中灵活交替使用全上下文注意力和部分输入注意力。每次全注意力后,根据整个输入的注意力模式更新小缓存。该方法在现成的 LLM 上实现与淘汰式方法相当的速度提升,同时在多种长文本生成任务中显著改善性能。此外,采用该推理设置继续预训练可进一步提升模型表现。

原文摘要 · Abstract (English)

Generating long sequences of tokens given a long-context input is a very compute-intensive inference scenario for large language models (LLMs). One prominent inference speed-up approach is to construct a smaller key-value (KV) cache, relieving LLMs from computing attention over a long sequence of tokens. While such methods work well to generate short sequences, their performance degrades rapidly for long-form generation. Most KV compression happens once, prematurely removing tokens that can be useful later in the generation. We propose a new inference method, RefreshKV, that flexibly alternates between full context attention and attention over a subset of input tokens during generation. After each full attention step, we update the smaller KV cache based on the attention pattern over the entire input. Applying our method to off-the-shelf LLMs achieves comparable speedup to eviction-based methods while improving performance for various long-form generation tasks. Lastly, we show that continued pretraining with our inference setting brings further gains in performance.

长文本生成推理优化注意力机制

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