PARTREP：学习在仅解码器LLM中重复什么

# PartRep：学习为解码器专用大语言模型重复什么

**来源：** https://arxiv.org/html/2607.01792

*Andikawati P Widjaja♡ · Yongjun Kim♠ · Hyounghun Kim♠ · Jaeho Lee♠*

♡ 万隆理工学院 · ♠ 浦项科技大学

###### 摘要

尽管解码器专用大语言模型（LLM）在大量自然语言任务中表现出色，但它们因因果注意力机制而受到信息流不对称的困扰：后出现的token比先出现的token拥有更丰富的上下文基础。一个简单而有效的解决方法是重复提示——只需在生成前附加提示的第二份副本，即可重新分配各位置的上下文基础并提升推理性能。然而，完整重复原始提示会使KV缓存占用翻倍，并在预填充阶段将注意力计算成本提升四倍，这使得该方法在长上下文场景中不实用。我们提出 **PartRep**，一种选择性增强方法，仅附加信息最丰富的token，而非整个提示。我们使用逐token的负对数似然（NLL）作为选择信号，其动机是：可预测性较低的token更难从其周围上下文中恢复，因此更受益于后部位置的重复。为避免因评分而进行一次完整前向传递带来的高昂成本，我们训练了一个轻量级门控网络，从早期层的隐藏状态中预测高NLL的token，从而能够在预填充中期通过提前退出实现token选择。在八个基准测试（包括MMLU、GSM8K和RULER）和三个模型系列（Qwen2.5、Llama3.2、Gemma4）上，PartRep 保留了完整重复的大部分收益，同时仅使用了其 59.4% 的 KV 缓存和 79.0% 的预填充FLOPs。

---

## 1 引言

解码器专用大语言模型已成为自然语言处理领域的主流架构，推动了一系列生成和分析任务的重大进展（huang2023advancing; achiam2023gpt）。这些模型的一个决定性特征是因果注意力机制，它允许每个token关注所有前面的token，同时屏蔽未来的token。尽管这种设计实现了有效的自回归生成，但它也导致了信息流的不对称：后面的token比前面的token拥有更丰富的上下文基础（springer2024repetition; behnamghader2024llm2veclargelanguagemodels）。这种不对称性已被认为与一些已知的局限性有关，包括推理中的前提顺序敏感性（chen2024premise）、选项排序偏差（pezeshkpour2024large; wei2024unveiling; ok2026lost）以及"迷失在中间"现象（liu2024lost）。这就引出了一个自然的问题：我们如何才能更均匀地重新分配上下文基础？

![图1](https://arxiv.org/html/2607.01792#S1.F1)
图1：token 重复策略概览。Vanilla 提示仅使用一次原始提示。完整重复附加整个提示，以将输入长度翻倍至 2L 为代价提高了准确性。我们的方法 **PartRep** 仅附加选定的关键 token，产生一个更短的序列（长度为 (L+τL)），在保持重复带来的准确率提升的同时降低了计算成本。

针对这种不对称性，一个出奇有效的补救措施是提示重复：在生成前呈现相同的提示两次。通过在输入末尾放置提示的第二份副本，每个原始token都获得了一个更靠后的"回响"，它可以关注整个上下文，从而有效地在整个序列中重新分配上下文基础。尽管不需要参数更新或架构改变，这种简单的干预已被证明能够提升各种推理任务的性能（leviathan2025prompt; xu2024re; springer2024repetition）。然而，它也引入了一些限制：附加提示的完整副本会使 KV 缓存占用翻倍，并将自注意力FLOPs提升四倍。随着提示长度的增长，这些开销变得越来越难以承受，使得这种方法在长上下文场景下不可扩展。

为此，我们提出了 **PartRep**，一种仅选择性重复最关键 token（而非复制整个提示）的方法（图1）。PartRep 根据信息重要性对每个输入 token 进行评分，然后在生成前仅附加得分最高的 token。作为评分标准，我们采用每个 token 的负对数似然（NLL）作为信息密度的代理：高 NLL（即高惊奇度）的 token 是模型认为最难以从前面上下文预测的，表明其内容没有被周围 token 冗余编码。正如我们在附录 B.1 节中所示，重复这样的 token 能带来最大的边际效益，因为它们的回响注入了原本会在后面位置的注意力中代表性不足的信息。

由于精确计算 NLL 需要一次完整的前向传递，这将会损害我们的效率目标，我们转而训练一个轻量级门控网络，直接从早期层的隐藏状态预测前 τ 比例的高 NLL token，从而能够在预填充中期进行选择。选中的 token 通过一个简短的自然语言桥接（例如，"请关注这些关键token……"）附加到原始提示之后，然后通过一次前向传递产生最终输出。

我们在8个基准测试（包括 GSM8K 和 RULER）以及一组多样化的解码器专用大语言模型（涵盖 Qwen2.5、Llama3.2 和 Gemma4 系列）上展示了我们方法的鲁棒性。PartRep 保留了完整重复的准确性优势，同时仅需要 59.4% 的 KV 缓存预算和 79.0% 的预填充FLOPs。我们的贡献总结如下：

- 我们提出了 **PartRep**，一种选择性提示增强方法，以远低于完整重复的预填充计算和 KV 缓存内存成本，近似实现了完整重复的收益。
- 我们展示了逐 token 的 NLL 为 token 选择提供了原则性的重要性信号，并设计了一种高效的门控机制，无需完整前向传递即可近似计算该信号。
- 我们在八个基准测试和三个模型系列上实证验证了 PartRep，表明它在各种设置下始终匹配或超越完整重复，同时大幅降低了预填充计算和内存开销。

---

## 2 相关工作

#### 提示重复

多项研究表明，提示重复可以提升大语言模型的性能。Xu等人发现，在链式思维推理之前"重读"问题能提升大语言模型的性能，他们认为重复使因果掩码模型能够部分模拟双向注意力，并缓解解码器专用架构的局限性。Arora等人将类似的想法扩展到循环语言模型，而 Springer 等人通过"回响嵌入"形式化了这一现象，表明重复 token 的表示在检索和相似性任务上显著优于单次传递的嵌入。在更大规模上，Leviathan 等人观察到，在非推理模式下，提示重复在开源和闭源模型上都能带来一致的性能提升。尽管先前的工作主要关注完整提示重复，Ok等人关于部分（选项级别）重复的相关发现进一步表明，恢复被因果注意力阻断的信息路径能够带来益处。然而，选项级别的重复仅限于多项选择设置，无法泛化到现实世界大语言模型使用中普遍存在的开放式生成任务。相比之下，我们引入了一个任务无关的框架，仅选择性重复任意提示中信息最丰富的 token，既避免了完整重复翻倍的 KV 缓存成本，也避免了选项级别重复的多项选择约束。

#### KV缓存驱逐

许多先前的工作建议驱逐"较不重要"token 的 KV 缓存，以减轻长上下文大语言模型推理的高内存 I/O 成本。例如，H2O（zhang2023h2o）保留那些在下一个 token 预测过程中频繁获得高注意力的 token，而 KVzip（kim2025kvzip）则要求模型重建提示并利用此过程中的注意力。最近的工作则侧重于避免在预填充阶段为每个 token 计算完整的注意力矩阵，以减少延迟开销。例如，FastKVzip（kim2026fast）和 KVzap（jegou2026kvzap）引入了轻量级预测器，训练其根据早期 transformer 层的隐藏状态预测 token 重要性，其中 token 重要性由先前的方法（如 KVzip）估计。我们采用了类似的基于预测器的策略，但目标根本不同。驱逐是估计哪些 token 可以安全地丢弃，而部分重复是估计哪些 token 在被重新关注时能产生额外的收益，并且这两种重要性的概念并非彼此的对偶。我们在第 3.1 节中详细说明，例如，提示的最后一个 token 在驱逐下通常至关重要，但重复时却最没有用处。

#### 摘要

另一类工作通过压缩输入提示本身来减少内存成本。LLMLingua（jiang2023llmlingua）使用一个紧凑的语言模型来识别并移除"较不重要"的 token，然后将摘要后的提示输入主模型。在这个意义上，摘要可以被视为 KV 缓存驱逐的文本级别类比：不是在预填充后丢弃键值对，而是在 token 进入模型之前移除冗余的 token。我们的工作同样估计 token 重要性并构建一个缩短的文本表示形式作为摘要。然而，目标根本不同。LLMLingua 假设原始提示包含固有的冗余，并通过压缩来减少内存消耗。相比之下，我们的目标不仅仅是缩短提示，而是在受限的内存预算下，近似实现"冗余的"完整重复所带来的准确率提升。

---

## 3 问题形式化

假设给定一个初始**用户提示**作为输入，其形式为 token 序列：

t = (t₁, t₂, ..., tₜ) (1)

长度为 L。对于解码器专用大语言模型，预填充自注意力的计算复杂度按 O(L²) 的比例二次增长，而存储的键值（KV）缓存按 O(L) 的比例线性增长。

**完整重复**（leviathan2025prompt）是重复整个提示的想法。换句话说，提示被复制成 t_FR = t ⊕ t，其中 ⊕ 表示序列拼接。已知预填充 t_FR 能够为序列启用关键性的伪双向注意力，从而提高模型准确性。另一方面，预填充复杂度增加了四倍，KV 缓存占用翻倍。

扩展这一范式，我们形式化了**部分重复**的问题，其目标是获得重复的好处，而不继承其高昂的内存和延迟成本。目标是分离原始提示中一个**重要的、信息丰富的**子集，使得仅附加这个子集（而非完整提示）就足以提高模型准确性。具体来说，我们考虑如下结构的提示：

t_PR = t ⊕ t_part (3)

其中 t_part 表示 t 的一个**关键词子序列**。给定这种提示结构，我们的目标是找到一个 token 选择器 f(t) = t_part，它能够最大化模型准确性。更具体地说，我们的目标是解决：

max_f acc(t ⊕ f(t)) (4)

其中 acc(Q) 表示基础大语言模型在 Q 提示下的准确性，并满足约束条件：

f(t) ⊂ t， |f(t)| ≤ τ · |t| (5)

其中 ⊂ 表示子序列关系（而非子集），τ ∈ (0, 1) 是为满足给定内存预算而施加的重复阈值。

将问题形式化为**token选择**（即施加 f(t) ⊂ t）的主要优势在于计算效率。预测器 f 可以看作是对每个输入 token 做出二元决策，充当一个接受或拒绝每个 token 的"门控"。由于输出维度简单，我们可以预期 f 能够用一个轻量级模型来实现。实际上，我们在目标模型隐藏状态之上，使用一个带注意力的两层 MLP 来实现它。如果让 f(t) 的输出落在通用词汇空间（例如，t 的文本摘要）或连续 token 空间，则可能无法做到这一点。

#### 效率

通过压缩重复的提示，我们可以预期预填充自注意力的计算与 (1+τ)²L² 成比例，这低于完整重复的 4L²。同样，KV 缓存将与 (1+τ)L 成比例，这小于完整重复的 2L。然而，请注意，部分重复会为 token 选择过程 f(·) 引入额外的推理成本。幸运的是，正如我们将在第 6.2 节中看到的，这个成本很小。

![图2](https://arxiv.org/html/2607.01792#S6.SS2)
图2：所提出的 PartRep 的推理过程。我们首先用原始提示预填充大语言模型，然后将其中早期层的隐藏状态传递给门控模块。接着，我们选择前 τ 比例的 token，并重复所选 token（即将其附加在原始提示之后），然后继续预填充。

### 3.1 与 KV 缓存驱逐的比较

回顾一下，KV 缓存驱逐的任务是关于决定"丢弃什么"（zhang2023h2o）。部分重复的任务，即决定"重复什么"，在需要估计每个 token 的重要性以减少 KV 缓存成本方面是类似的。然而，这两个任务在两个根本方面有所不同。

首先，通过 KV 缓存驱逐来解决部分重复在计算上是次优的。为了说明这一点，考虑将 KV 缓存驱逐方法应用于完全重复的提示，以此作为部分重复的一种手段。在这种情况下，我们仍然必须用 2L 个 token 进行预填充，从而失去了预填充计算的任何优势。

其次，"重要性"的概念在这两个任务中是不同的。为了说明这一点，考虑提示的最后一个 token。在部分重复中，这是最不重要的重复 token，因为原始提示中的最后一个 token 已经关注了所有其他 token。然而，在 KV 缓存驱逐中，（重复提示中的）最后一个 token 通常被视为必不可少的（zhang2023h2o）。

---

## 4 算法

我们开发了 **PartRep**，一种基于学习的方法，通过仅重复高度重要的 token 来获得重复的好处，而无需承受其内存和延迟成本。为了选择这些 token，我们首先论证负对数似然（NLL）作为信息密度的一个原则性代理（第 4.1 节）。由于在运行时计算精确的逐 token NLL 代价高昂，我们引入了一个轻量级门控网络，训练其从早期层的隐藏状态预测这些得分（第 4.2 节）。第 4.3 节描述了用于在多种数据上离线训练该门控网络的流程。