论文
PARTREP:学习仅解码器 LLM 的重复内容
PARTREP: Learning What to Repeat for Decoder-only LLMs
摘要
虽然仅解码器的 LLM 在大量自然语言任务中表现出色,但它受到因果注意引起的不对称信息流的影响:后来的标记比早期的标记具有更丰富的上下文依据。一个简单而有效的补救措施是提示重复——只需在生成之前附加提示的第二个副本就可以重新分配各个位置的上下文依据并提高推理性能。然而,完全重复原始提示会使 KV 缓存占用空间增加一倍,并使预填充期间的注意力成本增加四倍,这对于长上下文设置来说是不切实际的。我们提出 PartRep,一种选择性增强方法,仅附加信息最丰富的标记,而不是整个提示。我们使用 token-wise 负对数似然(NLL)作为选择信号,其动机是假设不可预测的 token 不太容易从周围的上下文中恢复,因此从后期位置重复中受益更多。为了避免完整前向传递进行评分的高昂成本,我们训练了一个轻量级门,该门可以从早期层隐藏状态预测高 NLL 词元,从而可以通过提前退出在中期预填充期间进行词元选择。在八个基准测试(包括 MMLU、GSM8K 和 RULER)和三个模型系列(Qwen2.5、Llama3.2、Gemma4)中,PartRep 保留了完全重复的大部分收益,同时仅使用 59.4% 的 KV 缓存和 79.0% 的预填充 FLOP。