及时止损：学习早期剪枝推理路径以实现高效并行推理

用可学习的内部token早期剪掉低质量推理路径，并行准确率提升且token消耗砍七成。

前置知识

并行推理 (Parallel Reasoning / Self-Consistency)

对同一查询让大模型独立采样 $N$ 条完整推理轨迹 $\{\tau_i\}$，再多数投票得 $\hat{y}=\text{vote}(\{\tau_i\})$。算力随 $N$ 线性增长 $C\propto N$，但很多路径从一开始就错且污染投票。

STOP 要解决的核心问题就来自这一范式：很多路径从一开始就是错的，但仍然要消耗完整算力并污染聚合结果。理解"采样—完成—投票"的流程是看懂为什么需要在 prefix 阶段就剪枝的前提。

KV Cache 与 prefix-level 推理

自回归 Transformer 把每个 prefix $p_i$ 的 Key/Value 张量缓存在显存里。STOP 复用这些缓存：先生成短 prefix、缓存 KV，再让几枚 [STOP] token 一次前向就读完全部隐层表征。

本文的"Launch-Check-Resume"三阶段管线能以极低开销 (0.20 s, 0.59%) 评分 prefix，正是建立在 KV Cache 可重用的基础上。不了解 KV Cache 的复用机制，就无法理解 STOP 为什么几乎不增加推理延迟。

LoRA (Low-Rank Adaptation) 与参数高效微调 (PEFT)

LoRA 在原权重旁并联低秩矩阵 $\Delta W=BA$，$r\ll d$，只训 $B$、$A$。它可被条件激活——只在遇到 [STOP] 时才参与计算，从而与冻结基座共存且不损伤通用能力。

STOP 的核心组件"Critique Adapter"就是一个只在遇到 [STOP] token 时才激活的 LoRA，它在不改写 LRM 通用推理能力的前提下，让网络学会把内部隐藏状态映射到"这条 prefix 是否能跑通"的标量评分。理解 LoRA 的条件激活机制是看懂这种"非侵入式接入"的关键。

Monte Carlo 软标签 (Soft Label from MC Estimation)

对 prefix $p_i$ 高温续写 $K$ 次，统计答对比例 $\hat{s}^{mc}=\frac{1}{K}\sum_k\mathbb{1}[\tau_k\text{ correct}]\in[0,1]$。$K=32$ 比 $K=1$ 硬标签方差更小、与 prefix 质量更相关。

STOP 用这种软标签训练分类头，使评分头学到"prefix 越像能跑通的，最终能跑通的概率越大"的连续映射。论文 Table 3 显示这一设计在 AIME24 上把 Cons@N 从 46.67% 提到 53.33%，是模块能稳定工作的关键。

Large Reasoning Model (LRM) 与推理 token 消耗

LRM 指 DeepSeek-R1、GPT-OSS 等长链思维模型，单次推理常需上万 token。并行 $N=64$ 条时单查询可达 $10^5\sim10^6$ token，单次成本近 6 美元——这是 STOP 想压减的对象。

STOP 的全部优化目标——"保留正确率、压缩 token"——都是围绕 LRM 串成长链的过程设计的。理解 LRM 一次推理动辄上 K token 的特性，才能体会到 prefix 阶段剪掉 50%~75% 路径意味着什么量级的算力节省。

研究动机

并行推理已经成为解复杂任务的标准范式：给同一个查询采样 $N$ 条独立推理轨迹再多数投票。代价同样惊人——NVIDIA 的工程报告指出在 NVIDIA 平台跑一遍 LRM 并行采样，单次查询成本逼近 6 美元；Jin 等人也量化了"推理的能耗账本"。更要命的是 LRM 一旦在 prefix 阶段走进错误分支，几乎不可能自我纠偏（Luo et al., 2025），整条轨迹变成"无用算力"且把多数投票搅浑。经验上，一个 $N=64$ 的查询在 AIME24/25、HMMT25 这类基准上常常消耗 50 万~140 万 token，但其中可能只有 30%~50% 的路径最终答对。也就是说，传统 self-consistency 有一半以上算力是"陪跑"的。然而现有剪枝方法各做各的：SlimSC 用表层相似度、DeepPrune/LaBoR/ThinkPRM 用外部奖励模型、DeepConf/AdaDec 用 token 概率或置信度——它们既没有统一比较框架，也没有标准化评测协议，导致整个领域"碎片化"。

本文的目标是本文目标有三层：第一，给路径剪枝这个碎片化领域画一张系统化的"分类地图"，把所有方法按信号来源 (Internal/External) 与可学习性 (Learnable/Non-learnable) 分到 4 个象限。第二，在分类揭示的"研究空白 (Type IV: 内部 + 可学习)"上做出第一个高效实例，命名为 STOP (Super TOken for Pruning)，证明这一象限既保留了内部信号的丰富度，又获得了可学习参数对任务的自适应能力。第三，验证 STOP 在 1.5B~20B 的多种 LRM、5 个推理基准以及 ZebraLogic/AIMO3 等非纯数学场景下都能稳健扩展，并由此蒸馏出可在真实部署中直接套用的"最优保留率"经验公式。

与已有工作不同的是，作者的核心切入角度是："内部信号 + 可学习参数"这一象限长期无人涉足并非偶然——前人尝试往 LRM 内部插分类头时，会因为隐藏表示是为"预测下一个 token"而优化的，难以直接被解码成"这条路能不能走通"。STOP 用一枚专门的 [STOP] 词表 token + 只在它身上激活的 LoRA adapter 解决了这个表征错位问题：token 充当聚合 prefix 全局信息的查询向量，LoRA 把"生成式特征空间"重映射到"价值判别空间"，分类头再读出标量分数。配合 Launch-Check-Resume 三阶段管线，整套方案把额外延迟压到 0.20 秒 (0.59% 开销)，使得"内部可学习"这一象限从理论可能变成工程可部署的范式，这是与 Type II 外部奖励模型（如 ThinkPRM、LaBoR）以及 Type III 朴素置信度方法最大的差异化所在。

核心方法

STOP 思路是"冻结 LRM + 外挂价值判别器"。直觉上，与其花算力生成完整答案再判断，不如在前 1024~4096 个 token 的 prefix 阶段就预测这条路能不能走通；不能就停止补完。技术路线：把新词 [STOP] 加入词表，给 LRM 配一个只在遇到 [STOP] 时激活的 LoRA 适配器，再接一个轻量分类头。训练时先用冻结 LRM 算 prefix 的 KV Cache $C_{p_i}=\text{LRM}(p_i;\Theta)$，再 append [STOP] 序列 $T_s$，把隐藏态送进分类头 $W_\text{cls}$，用 MC 软标签 $\hat{s}^{mc}$ 监督二元交叉熵 $\mathcal{L}=\text{BCE}(\hat{s}^{mc},\sigma(W_\text{cls}h_i))$。推理用三阶段管线：Launch (采 $N$ 条短 prefix 并缓存 KV)、Check (STOP 打分)、Resume (只对 top-$k$ 续写并投票)。整个过程额外开销几乎为零。

STOP 的核心创新是把路径剪枝建模成"用一枚虚拟 token 在 prefix 末尾做一次价值判别"。和已有方法的本质区别有三：第一，与 Type I 的表层相似度 (SlimSC) 相比，STOP 看的是 LRM 内部隐层而非表面文字，能捕捉到"推理是否走向了逻辑死胡同"这种深层信号；第二，与 Type II 的外部奖励模型 (LaBoR/ThinkPRM/DeepPrune) 相比，STOP 复用已生成的 KV Cache，零额外编码开销——Table 7 显示其单次评分 0.20 s，而外部判官 1.13 s，差距接近 5.6 倍；第三，与 Type III 的朴素置信度 (DeepConf/AdaDec) 相比，STOP 的 LoRA adapter 把"next-token 预测"用的特征空间专门重映射到"价值判别"空间，Table 4 显示没有 adapter 时 AIME24 上 avg@8|64 会从 36.67% 掉到 31.67%，说明这个 adaptation 步骤不可省略。

方法步骤详情

流程分训练和推理。**训练**：(1) 冻结 LRM 采 prefix $p_i$，高温续写 $K=32$ 条得 MC 软标签 $\hat{s}^{mc}=\frac{1}{32}\sum_k\mathbb{1}[\tau_k\text{ correct}]\in[0,1]$；(2) 算 KV Cache；(3) append 4~6 枚 [STOP] 取隐藏态 $h_i$；(4) $s_i=\sigma(W_\text{cls}h_i)$，$\text{BCE}(\hat{s}^{mc}, s_i)$。LRM 全冻结。**推理 (Launch-Check-Resume)**：(1) Launch 并行采 $N$ 条短 prefix 缓存 KV；(2) Check 打 $s_i$ 留 top-$k$；(3) Resume 续写到终止做多数投票。Check 仅几枚 token 且复用 KV Cache，开销几乎为零。

技术新颖性

新颖性有四点。第一，"词表加 [STOP] + 条件激活 LoRA"是非侵入式接入：原 LRM 一次不改写、生成能力零损伤，训练代价仅 LoRA 几百 M 参数 (Table 6 rank 128 时 147.7 M)。第二，把路径剪枝统一到 Internal × Learnable 象限，Table 1 把 9 个已有方法无歧义归类——"先建坐标系再做工作"对领域很有价值。第三，形式化最优保留率幂律公式 $\gamma^{-1}=a\cdot C^{b}\cdot L_\text{prefix}^{c}\cdot L_\text{task}^{d}$，拟合得 $a\approx 1.17\times 10^4$、$b\approx 0.46$、$c\approx 0.40$、$d\approx 4.55$ (Eq. 7)，让"该保留多少"有查表依据。第四，Figure 7 可视化揭示 STOP 机制——高分路径集中注意力到 self-correction 枢纽上，低分路径过早关注终止 token——说明 STOP 学到的是"过程质量"而非"答案形式"。

**Figure 3**: The inference process comprises three stages: caching initial prefixes (Launch), scoring them via the STOP module (Check), and completing only the top-ranked candidates (Resume).

实验结果

实验在 DeepSeek-R1-Distill-Qwen-1.5B/7B、DS-Qwen-3-8B、GPT-OSS-20B 四个 LRM 上跑 5 个推理基准 (AIME24/25、BRUMO25、HMMT25、GPQA-D)，外加 ZebraLogic 和 AIMO3。核心发现五条。F1：STOP 在 4 模型 × 5 基准的 20 组对比里几乎全部最优——如 GPT-OSS-20B AIME25 从 70.99% 提到 75.42% (+4.43%) 同时省 71.62% token；GPQA-D 从 65.55% 提到 77.46% (+11.91%)。F2：Type IV 在所有 compute budget 下稳定领先并跑赢不剪枝基线，Type III 在 1.5B AIME25 上反跌到 23.75% < 22.76%。F3：AIME24 上 MC 软标签把 Cons@N 从 46.67% 提到 53.33%。F4：无 LoRA 时 AIME24 上 36.67% → 31.67%。F5：Table 7 显示单次评分 0.20 s / 0.59% 开销。

**Table 1**: A Unified Taxonomy of Path Pruning Methods. We categorize methods based on the pruning signal source and learnability. Type IV satisfies both Desideratum 1 (Internal) and Desideratum 2 (Learnable).

**Table 2**: Results of avg@k (avg@m|k) across various models and benchmarks. The best result in each row is bolded and the second best is underlined.

**Table 3**: Performance comparison between hard labels (K = 1) and MC-estimated soft labels (K = 32).

**Table 4**: Comparing the STOP module with a simple linear classifier confirms that raw internal states require adaptation to perform effective self-evaluation.

**Table 5**: Effect of the number of [STOP] tokens (DS-Qwen-2.5-1.5B, AIME 2024, Lprefix = 2048).

**Table 6**: Effect of LoRA rank (DS-Qwen-2.5-1.5B, AIME 2024).

**Table 7**: Inference overhead analysis. STOP achieves near-zero cost by avoiding re-encoding.

**Table 8**: Generalization on ZebraLogic. STOP robustly generalizes beyond math and science tasks.

**Table 9**: Results on the AIMO3 competition setting with tool use (GPT-OSS-120B).

**Figure 4**: Performance vs. compute for four types of S on math and stem benchmarks.

**Figure 5**: Performance comparison under different retention ratios (γ) and prefix lengths (Lprefix).

**Figure 6**: Inverse retention ratio γ−1 vs. compute-to-prefix ratio. The theoretical curves (Eq. 7) closely align with empirical observations across varying reasoning progress levels.

**Figure 7**: Attention Analysis of [STOP] Decision-Making. High-scoring paths prioritize logical pivots (e.g., self-correction markers), whereas low-scoring paths fixate on terminal answer tokens. This contrast confirms that STOP functions as a process-oriented evaluator, rewarding reasoning integrity over premature closure.

查看结构化数据

任务	指标	本文	基线	提升
AIME24 (1.5B 模型)	avg@8\|64 (从 64 条选 8 条的均值准确率)	37.92% / 204.3k tokens (-73.88%)	不剪枝 30.10% / 782.3k tokens；Type I 26.25%；Type II 32.50%；Type III 32.92%	+7.82% 绝对提升，同时省 73.88% token；比最强基线 Type II 高 5.42%
AIME25 (GPT-OSS-20B)	avg@8\|64	75.42% / 191.1k tokens (-71.62%)	不剪枝 70.99% / 673.4k tokens；Type III 70.83% / -70.64%	+4.43% 绝对提升；省 71.62% token；为论文 abstract 重点宣传的成果
GPQA-Diamond (GPT-OSS-20B)	avg@8\|64	77.46% / 143.4k tokens (-48.26%)	不剪枝 65.55% / 277.2k tokens；Type III 68.43% / -47.36%	+11.91% 绝对提升，是表中最大单点涨幅
BRUMO25 (DS-Qwen-3-8B)	avg@8\|64	78.75% / 277.5k tokens (-78.98%)	不剪枝 75.00% / 1320k tokens；Type II 75.83% / -70.10%	+3.75% 绝对提升，省 78.98% token (表中最大节省)
HMMT25 (GPT-OSS-20B)	avg@8\|64	52.92% / 216.1k tokens (-76.27%)	不剪枝 48.13% / 910.8k tokens；Type III 46.25% / -74.97%	+4.79% 绝对提升，省 76.27% token
ZebraLogic (DS-Qwen-2.5-7B, 4×4 网格逻辑谜题)	avg@8\|64 (多选模式)	77.23%	不剪枝 73.73%	+3.50% 绝对提升；证明方法可迁移到非数学/非科学任务
AIMO3 工具调用竞赛 (GPT-OSS-120B, 单卡 H100, 5 小时 50 题)	竞赛得分 (public leaderboard)	STOP(24→8)=42, STOP(16→8)=43 (银牌)	不剪枝 + 工具 39	+3~+4 分，工具增强推理下仍稳定领先
推理延迟开销 (DS-Qwen-2.5-7B, L_prefix=2048, 单 H100)	每次 Check 延迟 / 相对开销	0.20 s / 0.59%	Type I 0.38 s / 0.93%；Type II 1.13 s / 3.37%	延迟降到 Type II 的 17.7%，Type I 的 52.6%

局限与改进

作者在 Limitations 中明确点出两点。第一是规模上限：所有实验覆盖到 20B 参数 (GPT-OSS-20B) 与 $N\le 64$ 的采样规模，但 70B+ 模型与 $N\ge 1000$ 的大规模并行设置下 STOP 是否仍然稳健，论文没有直接验证——这恰是工业部署最常见的尺寸，留下了空白。第二是结构灵活性：实验全部基于"在固定 token 位置 (L_prefix=512/1024/2048/4096) 做单次剪枝"的设定，作者承认没有探索多阶段级联剪枝 (例如 64→32→16→8) 或动态 checkpoint (让模型自己决定何时该被检查)。这意味着如果任务的"失败征兆"出现在第 100 token 或第 8000 token，单点固定 checkpoint 可能错过最佳剪枝时机。这两点都是后续工作的明确抓手。

独立分析的弱点

独立审视后有四个可改进之处。第一，**单一标量输出**丢失不确定性信息：STOP 只输出 $s_i\in[0,1]$，若 prefix 处于"半对半错"灰色地带，硬 top-$k$ 筛选可能误剪；改进方向是输出 (mean, variance) 或 Dirichlet。第二，**prefix 长度 $L_\text{prefix}$ 是手动调超参**：最优 $\gamma^{-1}=aC^bL_\text{prefix}^cL_\text{task}^d$ 公式要求预知 $L_\text{task}$，而 $L_\text{task}$ 本身就是要预测的；可让模型根据 prefix 内部"推理进展信号"自动决定剪枝。第三，**训练数据查询分布偏窄**：MC 标签主要在数学/科学上构造，对话/代码/多跳 QA 上可能差异大，Table 8 ZebraLogic +3.50% 虽可观但不能简单外推。第四，**Type IV 与 OTV 对比缺失**：作者承认 OTV 同期也是 Type IV，但无 head-to-head，留下孰优孰劣的盲点。

未来方向

作者已指出两个明确方向。**Progressive Multi-Stage Pruning**：把 64→32→16→8 这种漏斗式剪枝串联起来，更短 prefix 先粗筛、长 prefix 再精筛，从而把单次 check 精度要求降低、用更小 $N$ 起步。**STOP 加速 RL 训练**：把 STOP 当成 PPO/GRPO rollout 阶段的在线拒答器，提前终止低质量轨迹以提高单位算力的高质量样本密度。我个人还看到三条可延伸方向：一是把 [STOP] 推广为**多任务多目标 token** (例如一枚剪枝 token + 一枚重排序 token)，在同一前向里同时解决"哪条该剪、哪条该排第一"；二是和**投机解码 (speculative decoding)** 结合，让 STOP 既决定"继续还是停下"又决定"用什么小模型先草拟"，进一步压低单路径开销；三是把 Eq. 7 推广到**在线自适应 $\gamma$** 控制器，根据当前 batch 的 prefix 分数分布动态决定 top-$k$ 阈值，这对 latency-bounded 服务尤其有价值。

复现评估

复现友好度较高。代码、训练数据与 checkpoint 都会开源 (https://bijiaxihh.github.io/STOP)，MC 软标签也承诺 release。门槛有三：(1) **算力中等偏高**——4 LRM × 5 benchmark × 多组超参需若干 A100/H100 跑数日，1.5B/7B 单卡 H100 可复现主结果，20B 与 AIMO3 需多卡；MC 软标签构造对每 prefix 采 32 条续写开销最重但只一次。(2) **词表与 tokenizer 改动**——需把 [STOP] 加进 tokenizer 并随机初始化 embedding + 配套 LoRA，vLLM/HF Transformers 都支持扩展词表。(3) **评测协议标准化**——作者明确"每 query 生成 64 条初始路径、剪到 8 条、prefix 长度 2048"的统一规则，附录 C 列全超参与硬件。综合：核心方法 + 主要数字复现难度 4/5，单点对照 (如 1.5B + AIME24) 复现约 2/5。

及时止损：学习早期剪枝推理路径以实现高效并行推理 Cut Your Losses! Learning to Prune Paths Early for Efficient Parallel Reasoning