一句话总结

上海交通大学 LLM4SE 实验室与抖音集团提出 SWE-Pruner Pro，该方法通过一个小型头与长度感知嵌入，将 agent 自身的内部表示转换为逐行保留或裁剪的标签，从而直接裁剪 agent 内部的工具输出，在两个开源权重骨干模型和四个多轮基准上实现了高达 39% 的 token 节省，同时保持任务质量，推理开销有限，并且在使用 MiMo-V2-Flash 时，还将 SWE-Bench Verified 的解决率提高了 +3.8%，将长上下文 Oolong 准确率提高了 +2.2 个百分点。

核心贡献

• 编码 agent 在读取工具输出时形成的内部表示，编码了逐行的相关性，从而无需外部评分模型。
• SWE-Pruner Pro 通过附加在骨干模型上的轻量级头读取这一信号，使用长度感知嵌入和逐样本平衡焦点损失来预测保留或裁剪的标签。
• 在两个开源权重骨干模型和四个多轮基准上，SWE-Pruner Pro 节省了高达 39% 的提示和生成 token，同时保持任务质量，并且在 MiMo-V2-Flash 上，将 SWE-Bench Verified 的解决率提高了 +3.8%，将 Oolong 准确率提高了 +2.2 个百分点。

引言

解决代码仓库级任务的编码 agent 在多轮交互中会积累冗长、冗余的工具输出，从而推高 token 成本并降低模型性能。先前的裁剪方法要么采用固定的、与任务无关的指标（如困惑度），忽略了 agent 意图的变化，要么依赖一个独立的评分模型并在每一轮都需要显式的目标提示查询，从而增加了推理开销。作者表明，agent 自身的骨干表示已经编码了逐行的重要性，并提出了 SWE-Pruner Pro，这是一个轻量级头，直接从骨干模型的预填充过程中读取这一裁剪信号。通过引入学习到的长度感知嵌入和逐样本平衡焦点损失，SWE-Pruner Pro 避免了额外的模型调用，节省了高达 39% 的 token，并在多个基准上保持了任务质量。

方法

作者设计了 SWE-Pruner Pro，以直接从 agent 骨干模型的最后一层隐藏状态中对工具响应做出逐行保留或裁剪的决策。在每一轮 $t$t，agent 发出一个工具调用 $c_t$ct​ 并接收到一个原始工具响应 $r_t$rt​。在 agent 生成下一步动作之前，骨干模型将上下文历史 $H_{t-1}$Ht−1​、工具调用 $c_t$ct​ 和原始响应 $r_t$rt​ 预填充到其 KV 缓存中。由于前缀已经缓存，只有新的 $r_t$rt​ token 被前向传播，从而在响应范围内产生最后一层隐藏状态 $h_1, \ldots, h_L$h1​,…,hL​。

SWE-Pruner Pro 在此预填充阶段接入。裁剪头将每个隐藏状态 $h_i$hi​ 转换为一个逐 token 的保留或裁剪 logit，并将其聚合为逐行决策。这种行级粒度保留了被保留代码的语法结构。裁剪在轮次之间应用：agent 在第 $t$t 轮的生成仍然关注完整的 $r_t$rt​，而当轨迹继续到第 $t+1$t+1 轮时，裁剪后的响应 $\tilde{r}_t$r~t​ 替换了 $r_t$rt​。该头从骨干模型已经执行的预填充中读取隐藏状态，这意味着唯一增加的骨干模型工作量是在下一轮对通常短得多的 $\tilde{r}_t$r~t​ 进行一次重新前向传播。

裁剪头由两个主要组件组成：一个长度感知嵌入和一个逐 token 的前馈分类器。设 $N$N 为 $r_t$rt​ 中的行数，$d$d 表示 $h_i$hi​ 的维度。长度感知嵌入 $\mathbf{e}(N) \in \mathbb{R}^d$e(N)∈Rd 是行数 $N$N 的一个学习函数，在分类器之前被广播加到每个隐藏状态上，并且初始化为零，使得分类器在训练开始时退化为其与长度无关的极限情况：

$$\tilde{h}_i = h_i + \mathbf{e}(N).$$h~i​=hi​+e(N).

该嵌入使头显式地获知 $N$N，从而使保留或裁剪的映射能够随响应长度变化。这一点至关重要，因为错误裁剪的代价高度不均匀：从一个短响应中移除几行是灾难性的，而在长响应中这样做则影响微乎其微。

逐 token 分类器 $f_\theta$fθ​ 是一个小型非线性网络，包含 LayerNorm、两个隐藏 Linear-GELU-Dropout 块（隐藏宽度为 $d$d），以及一个最终 Linear 投影到单个 logit。仅使用线性分类器无法解决线性探针实验中观察到的中间分数区域的重叠问题，也无法有效地与长度感知嵌入进行交互。

应用于每个增强的隐藏状态，分类器产生保留概率：

$$z_i = f_\theta(\tilde{h}_i), \qquad p_i = \sigma(z_i).$$zi​=fθ​(h~i​),pi​=σ(zi​).

使用逐 token 的 sigmoid 阈值 $\tau = 0.5$τ=0.5，通过对每行 $\ell$ℓ 内二值化后的 token 决策进行多数投票，得到逐行决策 $\hat{y}_\ell$y^​ℓ​：

$$\hat{y}_\ell = \mathbb{1} \left[ \frac{1}{|\ell|} \sum_{i \in \ell} \mathbb{1}[p_i > \tau] > \frac{1}{2} \right].$$y^​ℓ​=1[∣ℓ∣1​i∈ℓ∑​1[pi​>τ]>21​].

在将工具响应追加到 agent 历史之前，$\hat{y}_\ell = 0$y^​ℓ​=0 的行将从工具响应中移除。

作者在 22,609 个来自真实多轮 agent 轨迹的样本上训练该头，这些样本带有逐行保留或裁剪的标注。行标签被扩展为逐 token 标签以用于损失函数。损失函数是逐样本平衡的焦点损失。来自 LLM 标注器的逐行标签本质上是模糊的，但标注器所确定的比率是有信息量的。标准的交叉熵和批次级焦点损失平等地对待每个 token，这可能会过拟合典型的保留率并稀释极端比率的样本。为了保护每个样本中少数类的召回率，作者在一个样本内分别计算保留和裁剪 token 的损失，并以相等的权重取平均。

对于样本 $s$s 中的每个 token $i$i，设 $p_{t,i} = p_i$pt,i​=pi​ 如果 $y_i = 1$yi​=1，否则 $p_{t,i} = 1 - p_i$pt,i​=1−pi​。逐 token 损失为：

$$\mathcal{L}_i^{\mathrm{tok}} = (1 - p_{t,i})^\gamma \cdot \mathrm{BCE}(p_i, y_i), \quad \gamma = 2.$$Litok​=(1−pt,i​)γ⋅BCE(pi​,yi​),γ=2.

损失在每个样本内的保留和裁剪 token 上分别取平均：

$$\mathcal{L}_s^{\mathrm{keep}} = \frac{\sum_i y_i \mathcal{L}_i^{\mathrm{tok}}}{\sum_i y_i}, \quad \mathcal{L}_s^{\mathrm{prune}} = \frac{\sum_i (1 - y_i) \mathcal{L}_i^{\mathrm{tok}}}{\sum_i (1 - y_i)},$$Lskeep​=∑i​yi​∑i​yi​Litok​​,Lsprune​=∑i​(1−yi​)∑i​(1−yi​)Litok​​,

并以相等的权重组合：

$$\mathcal{L}_s = \frac{1}{2} \mathcal{L}_s^{\mathrm{keep}} + \frac{1}{2} \mathcal{L}_s^{\mathrm{prune}}.$$Ls​=21​Lskeep​+21​Lsprune​.

最终目标是 $\mathcal{L} = \frac{1}{S} \sum_s \mathcal{L}_s$L=S1​∑s​Ls​。骨干模型保持完全冻结，因此添加该头不会改变 agent 的通用行为，也不需要重新训练骨干模型本身。该头从缓存的特征中进行训练。

实验

评估覆盖了四个编码 agent 基准和两个开源权重混合专家 LLM，将 SWE-Pruner Pro 与六种先前裁剪方法进行了比较。在多轮任务上，SWE-Pruner Pro 是唯一在保持质量的同时实现显著 token 节省的方法，而其他方法要么降低性能，要么增加 token。消融实验证实，逐样本平衡焦点损失和长度感知嵌入对这一效率至关重要，并且延迟分析表明裁剪开销相对于总体 token 减少是适中的。

在只读多轮基准上，SWE-Pruner Pro 是唯一一种持续减少总 token 消耗同时保持或略微提高答案质量的裁剪方法。竞争性的裁剪器要么在长上下文任务上增加 token，要么降低质量，而 SWE-Pruner Pro 利用了原位的骨干信号，以微小的质量影响实现了大量节省。在 Qwen3-Coder-Next 上，SWE-Pruner Pro 在 SWE-QA 和 SWE-QA-Pro 上减少了 35–39% 的 token 使用量，同时评分保持在未裁剪基线的 +0.24 以内，并且将 Oolong token 减少了 14%，准确率仅下降 1.4 个百分点。LLMLingua2 和 Selective Context 在 Oolong 上将 token 消耗增加了一倍以上（最高 +233%），而 Self-Prune 在 SWE-QA/Pro 上损失了 0.55–0.64 个评分点，表明先前裁剪器无法同时实现压缩和保持质量。

在 SWE-Bench Verified 上，裁剪的效果强烈依赖于骨干模型。使用 MiMo-V2-Flash 时，所有裁剪器都提高了解决率，但 SWE-Pruner Pro 实现了大幅提升，同时比次优方法增加了更少的输入 token。在 Qwen3-Coder-Next 上，所有裁剪器都损失了解决数，但 SWE-Pruner Pro 显示出最小的质量下降，同时实现了最大的输入 token 减少。在 MiMo-V2-Flash 上，SWE-Pruner Pro 将解决率提高了 3.8%，token 增加了 7.4%，而解决率最高的裁剪器获得了 4.2% 的提升，但增加了 14.9% 的 token。在 Qwen3-Coder-Next 上，所有裁剪器都降低了解决率；SWE-Pruner Pro 仅损失了 1.2 个百分点，同时将输入 token 减少了 13.5%，在这两个指标上均优于其他裁剪器。在不同的骨干模型上，API 调用和输入 token 的权衡有所不同：在 MiMo-V2-Flash 上，SWE-Pruner Pro 使用了最多的调用次数，而在 Qwen3-Coder-Next 上，它略微增加了调用次数，但仍实现了最大的 token 减少。

对损失函数的消融表明，逐样本平衡焦点损失实现了最高的评分和 F1，而 Dice 和 Tversky 的 F1 与之相当，但评分却低得多，这表明逐行 F1 可能具有误导性。添加长度感知嵌入在不影响 F1 的情况下提高了评分，因为它将裁剪错误转移到更长的响应中，在那里错误裁剪的危害较小。逐样本平衡焦点损失在评分和 F1 上均优于 BCE、Dice 和 Tversky；Dice 和 Tversky 的 F1 几乎持平，但它们的评分崩溃了，表明仅凭逐行 F1 无法捕捉响应的可用性。启用长度感知嵌入将评分从 6.86 提高到 7.08，同时保持逐行 F1 不变，因为它将错误转移到更长的响应中，在那里单行错误裁剪的破坏性要小得多。

SWE-Pruner Pro 在只读多轮基准和使用不同骨干模型的 SWE-Bench Verified 上，与多种裁剪基线进行了评估。它是唯一一种持续减少 token 消耗同时保持或略微提高答案质量的方法，而其他裁剪器要么增加 token，要么导致显著的质量下降。消融实验证实，逐样本平衡焦点损失和长度感知嵌入对于保持响应可用性至关重要，因为仅优化逐行 F1 可能具有误导性。