论文
IndexMem：用于长上下文 LLM 推理的学习 KV 缓存驱逐和潜在内存
IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inference
摘要
人们越来越期望 大语言模型 (LLM) 在长上下文上运行，但标准的 softmax 注意力会产生随序列长度线性增长的 KV 缓存，很快成为长上下文推理的瓶颈。一个实用的补救措施是驱逐不太重要的 KV 条目；然而，现有的驱逐政策在很大程度上是启发式的，很难捕捉到丰富的、依赖于投入的象征性重要性分布。在这项工作中，我们引入了一个可学习的索引器，可以预测 KV 重要性，从而更准确地保留关键词元。同时，天真地驱逐词元会永久丢弃其信息，导致不可逆的遗忘和远距离检索的降级。为了解决这个问题，我们提出了一种轻量级的潜在内存模块，它将被逐出的词元压缩成紧凑的、在线更新的状态，并提供剩余读数以补偿因 KV 逐出而损失的注意力贡献。总的来说，我们的方法能够在有限的 KV 预算下实现准确的长上下文推理，在 Qwen、Mistral 和 Llama 模型上对 RULER (4K/16K) 进行一致的改进（在积极驱逐下最多可达 25 个点），大海捞针检索明显更稳定，并且与现有驱逐策略相比，具有更出色的 LongBench 分数和压缩曲线。