早期剪枝学习！高效并行推理的路径剪枝方法

# 截断亏损！学习及时剪枝路径以实现高效并行推理
源：https://arxiv.org/html/2604.16029

Jiaxi Bi1,3 Tongxu Luo1,2∗ Wenyu Du4 Zhengyang Tang1 Benyou Wang1,2
1香港中文大学（深圳）2深圳环路地区研究院 3北京科技大学 4DualityRL
[email protected] [email protected] [email protected]

等量贡献；按姓氏字母顺序排列。工作完成于在香港中文大学（深圳）实习期间。通信作者。

###### 摘要

并行推理增强了大规模推理模型（LRMs），但由于早期错误导致的无用路径而产生了令人禁止的成本。为了缓解这一问题，在前缀级别的路径剪枝至关重要，但现有研究仍然缺乏标准化框架。在本工作中，我们提出了首个系统化的路径剪枝分类法，根据信号源（内部与外部）和可学习性（可学习与不可学习）对方法进行分类。这一分类揭示了可学习内部方法的未被探索的潜力，促使我们提出STOP（超级剪枝标记）。在参数规模从1.5B到20B的LRMs上的广泛评估证明，STOP相比现有基线方法实现了优越的有效性和效率。此外，我们严格验证了STOP在不同计算预算下的可扩展性——例如，在固定计算预算下，将GPT-OSS-20B在AIME25上的准确率从84%提升到近90%。最后，我们将发现提炼为形式化的经验指南，以便于最优的实际部署。代码、数据和模型可在https://bijiaxihh.github.io/STOP获取。

## 1 引言

并行推理已成为解决复杂问题的标准范例OpenAI（2024）；Wang等人（2025b）。核心原理是采样多条独立的推理路径，随后聚合它们以得出稳健的共识。然而，这种准确性的提升带来了令人禁止的成本。为每个查询生成数十甚至数百条轨迹将计算开销增加数个数量级Jin等人（2025），并将推理成本升至每个查询近6美元NVIDIA公司（2025）。

#### 为何在并行推理中及时剪枝？

至关重要的是，最近的研究Luo等人（2025）；Hassid等人（2025）揭示了这种广泛的计算在很大程度上被浪费了：并非每条路径都对解决方案有贡献。许多轨迹从一开始就存在缺陷，但它们消耗相等的资源来生成，随后污染最终的答案聚合。如图1所示，一旦推理路径以有缺陷的前缀开始，LRM就很难自我纠正，不可避免地陷入无用的轨迹Luo等人（2025）。因此，在前缀级别识别并终止这些不太可能的路径——一种被称为路径剪枝（或前缀拒绝）的技术——是必不可少的。

#### 统一分类法

虽然现有方法尝试使用辅助奖励模型Liao等人（2025）、内部置信度Fu等人（2025）或语义冗余Hong等人（2025）来过滤路径，但它们缺乏标准化的评估协议，导致研究碎片化。因此，我们首先提出了首个系统化的路径剪枝分类法，根据信号的源（内部与外部）和可学习性（可学习与不可学习）对方法进行分类（见图2）。这一分类法揭示了一个重大的研究空白：可学习内部方法的未被探索的潜力。从概念上讲，可学习内部方法提供了独特的优势，因为学习能够实现特定于任务的准确性提升，而内部信号提供了推理失败的早期、细粒度指示，无需产生额外的计算开销。为了弥补这一空白，我们引入STOP（超级剪枝标记），这是该范例的首个高效实现。广泛的评估表明STOP在有效性和效率方面都超过了现有基线方法。

#### 进一步评估与经验分析

尽管路径剪枝前景广阔，但其广泛采用目前受到以下因素的阻碍：在不同计算预算和模型规模上的可扩展性未经验证；以及在实际场景中确定最优剪枝配置的经验指南缺失。为了克服这些问题，我们严格验证了路径剪枝在实际设置中的实用性。我们在不同的模型规模（1.5B到20B）和计算预算上进行了广泛的实验，确认STOP展现出稳健的可扩展性。此外，我们将经验分析提炼为可操作的指南，提供了一种形式化的方法来确定在资源约束条件下的最优保留比例。

#### 贡献

总结而言，本工作做出了四项主要贡献：（1）我们提出了首个系统化的路径剪枝研究和分类法。（2）我们提出了STOP，一种基于可学习内部信号的新型剪枝方法。（3）我们提供了全面的评估，展示了STOP的优越的可扩展性和有效性。（4）我们建立了经验指南以支持路径剪枝的实际实现。

## 2 路径剪枝的统一分类法

### 2.1 问题定义

考虑一个LRM Θ 和一个输入查询 x，并行推理通过生成 N 条独立的轨迹 T = {τᵢ}ᵢ₌₁ᴺ（其中 τᵢ ~ P_Θ(x)）并通过共识策略（如多数投票）对其进行聚合来改进准确性。最终预测 ŷ 通常计算为：

ŷ = vote({τᵢ}ᵢ₌₁ᴺ)  (1)

然而，生成 N 条完整轨迹会产生线性计算成本（C ∝ N）。为了减轻这一成本，路径剪枝旨在在解码过程的早期识别和丢弃不太可能的轨迹。

#### 路径剪枝的形式化

形式上，我们在长度 L_prefix 处定义一个检查点，其中生成被暂停。在这个阶段，模型已生成一组前缀 P = {pᵢ}ᵢ₌₁ᴺ。路径剪枝的核心是一个剪枝信号生成器 S，它将每个前缀映射到一个代表其潜在正确性的标量分数：

sᵢ = S(pᵢ|x, Θ)  (2)

其中 sᵢ ∈ [0,1] 表示剪枝信号。基于这些信号，我们仅保留前 k 条最有前景的路径（其中 k ≪ N）以完成，丢弃其余的。最终聚合的答案然后仅从这个剪枝子集导出：

ŷ_pruned = vote({finish(pᵢ) | sᵢ ∈ {sⱼ}ⱼ₌₁ᵏ})  (3)

因此，路径剪枝的目标是设计一个 S，在最小化计算成本（生成的标记数量）的同时最大化 ŷ_pruned 的准确性。因此，S 的设计决定了整个框架的有效性。

### 2.2 剪枝信号生成器的统一分类法

表1：路径剪枝方法的统一分类法。我们根据剪枝信号源和可学习性对方法进行分类。

第IV类满足期望1（内部）和期望2（可学习）。

如第2.1节所定义的，路径剪枝的有效性完全取决于剪枝信号生成器 S 的质量。虽然 S 的功能是一致的——对前缀评分——但现有方法在如何产生这个信号上从根本上有所不同。为了系统地评估这些方法，我们根据两个关键维度对它们进行分类：信号的源（外部与内部）和生成器的可学习性（可学习与不可学习），如表1所总结。

#### 信号生成器的两项期望

在对具体方法进行分类之前，我们为理想的信号生成器建立了两项期望：

###### 期望1：内部源

理想的 S 应该利用 LRM 的丰富、高维内部状态。

内部信号包含关于不确定性和推理动态的细粒度信息，这些信息在外部方法使用的最终文本输出中经常丢失。

###### 期望2：可学习性

理想的 S 应该能够训练以适应特定的数据分布。

可学习的参数允许生成器捕获复杂的、非线性的错误模式，这是刚性的、预定义的启发式方法无法建模的。

基于这些轴线，我们将现有工作分为四种不同的类型。

#### 外部信号源

这一类别中的方法从生成的文本输出或通过查询单独模型来导出剪枝信号。它们未能满足期望1。

###### 第I类：表面启发式

这些方法依赖于应用于生成文本表面形式的人工设计规则（例如相似性）。

虽然计算成本低廉，但这些启发式方法是刚性的，对模型的实际置信度一无所知。为了克服这一点，下一类型在外部评估过程中引入了可学习性。

###### 第II类：外部评判者

这些方法采用单独的、经过训练的模型来评估推理路径。

虽然它们满足期望2，但由于需要额外的模型推理而产生了显著的计算开销，并且未能访问LRM的内部确定性。为了克服这种刚性，下一类别在外部评估过程中引入了可学习性。

#### 内部信号源

这一类别中的方法直接从LRM的内部状态提取信号，可以访问更丰富的信息（满足期望1）。

###### 第III类：原始置信度

这一范例利用直接从解码过程导出的内在指标，如困惑度或标记概率。

然而，这些方法依赖于固定的置信度定义，违反了期望2；原始概率不一定与推理正确性相关。

###### 第IV类：学习直觉

最后一类代表了两个期望的交集：一个可训练的模块插入到LRM中以处理内部状态。

这种方法可以利用丰富的隐藏表征（内部）同时适应特定任务的错误模式（可学习）。

## 3 方法：超级剪枝标记

如我们的分类法所示，第IV类代表理想的剪枝范例但仍未被探索。在本节中，我们介绍STOP（超级剪枝标记），这是该范例的首个高效实现。我们在第3.1节阐述了动机，随后是第3.2节中的架构设计和工作流程。

### 3.1 第IV类剪枝的动机

如图2所示，之前的方法在信息丰富性或适应性之间进行了权衡。第II类受到高延迟的困扰，而第III类缺乏对复杂错误模式进行建模的能力。第IV类代表了理想的最优点：它结合了访问内部状态的效率与可学习参数的适应性。然而，这一类型仍未被探索，因为设计一个能够提取这些信号而不破坏LRM生成能力的模块具有挑战性。

### 3.2 第IV类剪枝的实现：STOP

为了实现这一类型，我们设计STOP为一个轻量级的、非侵入式的模块，与主干LRM无缝集成。

#### 组件

我们使用三个可学习的组件来增强固定的LRM Θ：（1）一个添加到词汇中的超级标记[STOP]，充当聚合信息的专门查询向量；（2）一个批评适配器LoRA（θ_LoRA），仅在处理[STOP]标记时激活，以提取特定于错误的特征，而不改变LRM的一般推理能力；（3）一个分类头（W_cls），将[STOP]标记的隐藏状态投影到标量概率。

这种设计确保了模块性：原始参数Θ保持冻结，保留了LRM的生成能力，同时实现了高效的参数高效微调（PEFT）。

#### 训练：学习使用内部信息

训练的目标很简单：教会模型区分有前景的前缀和无用的前缀。形式上，对于前缀pᵢ，我们通过蒙特卡洛估计导出一个软标签 sᵢ^mc ∈ [0,1]（详见附录B）。训练过程涉及两个步骤：首先，我们使用冻结的LRM计算前缀的KV缓存：C_pᵢ = LRM(pᵢ; Θ)。其次，我们追加一个可学习[STOP]标记的序列，表示为T_s，并使用LoRA增强的模型处理它们。最终隐藏状态hᵢ被馈入分类器以最小化软二元交叉熵损失：

ℒ = −[sᵢ^mc log σ(W_cls hᵢ) + (1 − sᵢ^mc) log(1 − σ(W_cls hᵢ))]  (4)

其中 hᵢ = LRM(T_s | C_pᵢ; Θ, θ_LoRA)₋₁。

#### 训练成本

构建MC监督需要对每个前缀采样多个延续以估计 sᵢ^mc（例如，K=32），这在数据构建期间引入了前期计算成本。然而，这一成本仅产生一次，并且