解码器大型语言模型重复令牌提升推理性能
速览
PartRep针对解码器-only大型语言模型提出选择性重复方法，仅在低可预测性令牌上重复提示。利用令牌负对数似然作为选择信号，通过轻量级门控网络在早期层隐藏状态预测高NLL令牌，实现中途提前退出。实验在Qwen2.5、Llama3.2、Gemma4等三家模型家族及八个基准上验证，保留全重复的绝大部分收益，但仅需其59.4%的KV缓存和79%的预填充FLOPs。这一创新使长上下文场景下的因果注意力不对称问题得到实用缓解，推动大型语言模型在推理任务中的效率与性能平衡。
AI 深度解读
背景
在解码器-only的大语言模型（LLMs）中，因果注意力机制带来了一种不对称的信息流动：后续位置的标记（tokens）在上下文 grounding 方面比早期位置更丰富。这种结构性限制限制了模型在推理任务上的表现。
一种简单且有效的补救措施是提示词重复（prompt repetition）：在生成前简单追加原提示词的第二份副本，便能将 grounding 分布到更多位置，从而提升推理性能。
然而，全量重复原提示词会使 KV cache footprint 增加一倍，并将 prefill 阶段的注意力成本增加四倍，这在长上下文场景下变得不切实际。
核心内容
原文提出 PARTREP 作为一种选择性增强方法，在生成前仅追加最具信息量的标记，而不是整个提示词。
模型通过标记级别的负对数似然（token-wise negative log-likelihood，NLL）作为选择信号。作者假设，信息密度低的标记（less predictable tokens）从周围上下文中可恢复性较弱，因此从后期位置的重复中获益最大。
为避免在评分阶段进行完整的前向传播，作者训练了一个轻量级门控机制（lightweight gate），该机制从早期层隐藏状态中预测高 NLL 标记。借助这一机制，标记选择可在 mid-prefill 阶段通过提前退出（early exit）实现，显著降低计算开销。
在八个基准测试（包括 MMLU、GSM8K 和 RULER）及三个模型家族（Qwen2.5、Llama3.2、Gemma4）上进行验证，PARTREP 保留了全量重复方法的绝大部分收益，同时仅使用其 59.4% 的 KV cache 和 79.0% 的 prefill FLOPs。
关键要点
- 因果注意力导致后期标记 grounding 更丰富，早期标记相对较弱
- 提示词重复是简单有效的策略，但全量重复在长上下文下因 KV cache 膨胀和注意力开销激增而不实用
- PARTREP 通过 token-wise NLL 筛选最具信息量的标记，仅追加选定部分而非完整提示词
- 轻量级门控从早期隐藏状态预测高 NLL 标记，实现 mid-prefill 阶段的早期退出与低成本选择
- 跨八个基准与三个模型家族，PARTREP 实现收益保留率高（全量重复的绝大部分），同时 KV cache 使用率仅 59.4%，prefill FLOPs 仅 79.0%
意义与影响
PARTREP 为解码器-only LLMs 提供了一个轻量级、可扩展的路径，使模型在推理任务中充分利用后期标记的丰富上下文，而无需支付全量重复带来的计算与内存代价。尤其在长上下文场景下，这一方法显著降低了实际部署门槛，预计将推动更多高效推理优化技术的落地应用。