只重复关键句子,让大模型推理更准还省资源。
PARTREP: Learning What to Repeat for Decoder-only LLMs

- 选重要词重复,而非整段重写,降低计算负担。
- 在8个基准上接近全量重复效果,节省40%显存和21%算力。
- 适合长文本推理、资源受限场景下的大模型优化。
尽管解码器仅有的大语言模型在众多自然语言任务中表现优异,但其因果注意力机制导致信息流不对称:后序标记比前序标记具有更强的上下文支撑。一种简单有效的方法是提示重复——在生成前附加一份原始提示的副本,可重新分配上下文支撑并提升推理性能。然而,完整重复会使KV缓存占用翻倍,预填充阶段的注意力计算成本增加四倍,在长上下文场景下不切实际。本文提出PartRep,一种选择性增强方法,仅重复最具有信息量的标记,而非整个提示。采用逐标记负对数似然(NLL)作为选择信号,基于假设:不可预测的标记更难从上下文中恢复,因此更受益于后期重复。为避免全前向传播带来的高开销,训练了一个轻量级门控模块,从早期层隐藏状态预测高NLL标记,实现在预填充中期通过提前退出完成标记选择。在八个基准(包括MMLU、GSM8K和RULER)及三个模型族(Qwen2.5、Llama3.2、Gemma4)上,PartRep保留了全量重复约90%的性能增益,同时仅使用59.4%的KV缓存和79.0%的预填充浮点运算量。
原文摘要 · Abstract (English)
While decoder-only LLMs excel at a vast array of natural language tasks, it suffers from an asymmetric information flow induced by causal attention: later tokens are richer in contextual grounding than earlier ones. A simple and effective remedy is prompt repetition -- just appending a second copy of prompt before generation can redistribute grounding across positions and improve reasoning performance. However, full repetition of the original prompt doubles the KV cache footprint and quadruples attention cost during prefill, making it impractical for long-context settings. We propose PartRep, a selective augmentation method that appends only the most informative tokens -- rather than the entire prompt. We use token-wise negative log-likelihood (NLL) as a selection signal, motivated by the hypothesis that less predictable tokens are less recoverable from surrounding context and therefore benefit more from late-position repetition. To avoid the heavy cost of a full forward pass for scoring, we train a lightweight gate that predicts high-NLL tokens from early-layer hidden states, enabling token selection during mid-prefill via early exit. Across eight benchmarks (including MMLU, GSM8K, and RULER) and three model families (Qwen2.5, Llama3.2, Gemma4), PartRep retains most of the gains of full repetition while using only 59.4\% of its KV cache and 79.0\% of its prefill FLOPs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。