MemoSight：统一上下文压缩与多Token预测以加速推理

摘要

虽然思想链 (CoT) 推理使 LLM 能够解决具有挑战性的推理问题，但由于 KV 缓存随着生成的token数量线性增长，CoT 推理面临着速度和内存使用方面的扩展问题。在这项工作中，我们提出了 MemoSight（基于内存预见的推理），这是一个集成上下文压缩和多标记预测的统一框架，以减轻效率问题，同时保持 CoT 推理性能。我们的框架采用相同的简约设计，通过特殊token及其针对每种token类型定制的相应位置布局来进行上下文压缩和多token预测。对四个推理基准的综合实验表明，MemoSight 将 KV 缓存占用空间减少了高达 66%，并将推理速度提高了 1.56 倍，同时优于现有的 CoT 压缩方法。

MemoSight：统一上下文压缩与多Token预测以加速推理 — 图 1：（左）上下文压缩：与 Vanilla CoT 相反，CoT 压缩利用内存token来压缩上下文并减少迭代推理和内存过程中的 KV 缓存占用空间。（右）多token预测：使用多个 LM 头的传统 MTP，与基于特殊token的 MTP 形成对比，后者使用单个 LM 头和交错寄存器token ⟨ r ⟩ \langle\text{r}\rangle 实现并行未来预测（向前 d d 步）。