← 返回 2026-07-21

SWE-Pruner Pro:编码智能体自身已知该剪枝什么 SWE-Pruner Pro: The Coder LLM Already Knows What to Prune

Yuhang Wang, Yuling Shi, Shaoqiu Zhang, Jialiang Liang, Shilin He, Siyu Ye, Yuting Chen, Kai Cai, Xiaodong Gu 📅 2026-07-20 👍 75 2026-07-25 18:30
上下文压缩 提示压缩 智能体上下文管理 编码智能体 隐状态探针

轻量级头从冻结骨干隐状态直接读取逐行剪枝信号,省去外部评分模型。

前置知识

多轮编码智能体(Coding Agent)

在仓库级编程任务中,智能体把推理与环境交互交错进行:调用 cat、grep、ls、python 等工具,读取原始文本输出,再继续推理。在 SWE-Bench Verified 上,Mini-SWE-Agent 配合 Claude Sonnet 4.5 时仅读文件命令就消耗了超过 70% 的 token 预算,且这些输出会留在后续每轮的上下文窗口里不断累积。

理解本文必须先明白:工具输出是 token 成本的真正大头,且会跨轮次累积、触发长上下文退化,这正是剪枝要解决的问题场景。

骨干模型隐状态(Hidden States)

Transformer 在 prefill 阶段对每个输入 token 做注意力加权的前向传播,最后一层隐状态 $h_i$ 是模型对该 token 的内部表征。本文发现智能体骨干在被动读取工具输出时,其最后一层隐状态已经编码了「这一行是否对下一步动作重要」的判别信号。

本文的核心前提是:剪枝信号已存在于骨干自身表征中,无需外部模型重建。读懂这一点才能理解为什么可以「就地读取」而不是「外部评分」。

线性探针(Linear Probing)

冻结骨干模型,对其最后一层隐状态(按行做均值池化)拟合一个简单的逻辑回归,观察两类标签(保留 vs 剪枝)在该表征空间中是否线性可分。本文探针达到 AUC 0.83、best-F1 0.63,远高于多数类 F1 上界 0.46,证明信号已经存在于骨干表征里。

探针研究是本文动机的关键证据,它量化地说明了「信号已经在里面」,直接支撑了把剪枝头挂在骨干上的设计决策。

焦点损失(Focal Loss)与逐样本平衡

焦点损失 $\mathcal{L}_{tok_i}=(1-p_{t,i})^\gamma\cdot \mathrm{BCE}(p_i,y_i)$($\gamma=2$)让难分样本获得更大权重。本文进一步在每个样本内对保留类和剪枝类分别求平均再等权相加:$\mathcal{L}_s= frac12\mathcal{L}_{keep_s}+ frac12\mathcal{L}_{prune_s}$,使每个样本无论自身保留率高低都贡献相等的两类梯度。

逐样本平衡焦点损失是本文两大技术贡献之一,消融显示它比 BCE 高 +1.13 judge / +0.16 F1,是理解方法新颖性的关键。

KV 缓存与 Prefill

推理时骨干把前缀 [历史 $H_{t-1}$, 工具调用 $c_t$, 工具响应 $r_t$] 写入 KV 缓存;前缀已缓存,只有新的 $r_t$ token 需要前向。SWE-Pruner Pro 的头就挂在这个已有的 prefill 上,读取 $r_t$ span 的最后一层隐状态 $\{h_i\}$,不产生对 $r_t$ 的额外前向。

理解「复用已有 prefill」才能明白为什么该方法几乎没有额外开销(仅 15% 墙钟时间),这也是它击败外部评分模型路线的根本原因。

研究动机

多轮编码智能体在解决仓库级任务时会累积大量冗长的工具输出。在 SWE-Bench Verified 上,Mini-SWE-Agent 配合 Claude Sonnet 4.5 时仅文件读取命令就消耗了超过 70% 的 token 预算(GLM-4.6 上有类似规律),且这些输出会跨轮次保留在上下文窗口中,推高成本并触发已被广泛记录的长上下文退化。现有剪枝分两条路线:一是通用压缩器(LLMLingua2、Selective Context)用困惑度等固定代理指标打分,无法适应智能体不断演化的意图;二是任务专用剪枝器 SWE-Pruner 虽然条件化于智能体意图,却需要额外的评分模型和智能体每轮必须手写的显式 goal-hint 查询。两者的共同症结在于:剪枝信号都从智能体外部获取,而骨干模型本身早已处理过该工具输出。更严重的是,作者实测发现六种已有剪枝器中有四种在至少一个基准单元上反而增加了 token,LLMLingua2 在 Oolong 上甚至膨胀 +190%,即额外开销抵消了压缩收益。

本文的目标是本文的目标是构建一个能直接从智能体骨干自身内部表征读取「保留或剪枝」信号的剪枝器,彻底去掉外部评分模型与每轮 goal-hint 查询。该头应当复用骨干在工具响应上已经执行的 prefill,从而只增加很小且有界的开销。在效果上设定了三个明确且苛刻的指标:第一,要在每一个基准单元(而非平均水平)上都减少 token,最高省 39%;第二,任务质量要落在未剪枝基线很窄的邻域内;第三,每次调用的额外墙钟开销控制在约 15%,且不产生额外的模型前向调用。作者还希望验证一个更广泛的原理:骨干在被动读取观察时形成的表征,已经编码了此前方法需要额外一次模型调用才能恢复的相关性判断。

与已有工作不同的是,本文独特的切入角度是把剪枝信号视作「已经在那里」而非「需要被外部重建」。动机章节的探针研究给出了量化证据:在冻结的 Qwen3-Coder-Next 最后一层隐状态上拟合逻辑回归,在留出轨迹上达到 AUC 0.83、best-F1 0.63,远高于在经验正例率约 30% 下的多数类 F1 上界 0.46。沿 LDA 判别轴两类均值明显错开。这一发现直接颠覆了 SWE-Pruner「需要专门外部模型 + 显式查询」的前提:既然骨干在做注意力加权读取时已经知道哪些行重要,那就把它读出来即可,把智能体本身变成剪枝器,从源头消除外部模型与查询带来的开销。

核心方法

直觉:智能体读取工具响应本身是一次注意力加权前向,故骨干已编码哪些行对下一步重要。SWE-Pruner Pro 据此在骨干已有 prefill 上挂轻量级头。第 $t$ 轮智能体发起 $c_t$、环境返回 $r_t$(往往数百行);骨干把 $[H_{t-1},c_t,r_t]$ 写入 KV 缓存,前缀已缓存故仅新 $r_t$ token 被前向,产生该 span 最后一层隐状态 $h_1,\ldots,h_L$($L=|r_t|$)。剪枝头把每个 $h_i$ 转成逐 token 保留/剪枝 logit 再聚合成逐行决策,$\tilde{r}_t$ 在轨迹进入第 $t+1$ 轮时替换 $r_t$。第 $t$ 轮智能体自身生成仍注意力完整 $r_t$,剪枝发生在轮间。头从骨干已执行的 prefill 读取 $\{h_i\}$,对 $r_t$ 不产生额外前向;唯一新增骨干工作是下一轮对 $\tilde{r}_t$ 重新前向,而 $\tilde{r}_t$ 通常远短于 $r_t$(平均保留率约 30%),其开销被后续生成因上下文变短节省的工作量远远抵消。

与 SWE-Pruner 的本质区别是同时去掉「外部评分模型」与「显式 goal-hint 查询」这两样东西。剪枝信号直接来自骨干在一次本就会执行的前向中形成的表征,因此既没有额外模型调用,也没有每轮查询开销。要让它真正奏效,有两个关键的技术补充。第一是长度感知嵌入 $e(N)$:因为误剪的代价随行数 $N$ 高度不均匀——从 5 行的响应里误删几行是灾难性的,而在 300 行的响应里同样操作几乎无影响,所以头需要显式获知 $N$ 才能让保留/剪枝映射随长度变化。第二是逐样本平衡焦点损失:逐行的 LLM 标注天然模糊,但标注者最终选定的比例本身就有信息——当 100 行里只保留 3 行时,那 3 行承载着「不可替代」的最强信号;当保留 90 行时,被剪的 10 行承载着「可安全移除」的最强信号。把每个样本内保留类与剪枝类分别求平均再等权相加,能保护每个样本少数类的召回。这两个设计共同把线性探针无法解决的中段重叠与长度相关结构处理掉。

方法步骤详情

(1) 第 $t$ 轮智能体发起工具调用 $c_t$,环境返回原始响应 $r_t$。(2) 骨干对 $[H_{t-1},c_t,r_t]$ 做 prefill,仅新 $r_t$ token 前向,得最后一层隐状态 $h_1,\ldots,h_L$。(3) 按行数 $N$ 查分桶长度感知嵌入得 $e(N)\in\mathbb{R}^d$,广播加到每个隐状态:$\tilde{h}_i=h_i+e(N)$。(4) 逐 token 分类器 $f_\theta$(LayerNorm→两个 Linear-GELU-Dropout→Linear)产生 logit $z_i=f_\theta(\tilde{h}_i)$ 与概率 $p_i=\sigma(z_i)$。(5) 逐行决策用 token 二值化多数投票($\tau=0.5$)。(6) 移除被判剪枝的行得 $\tilde{r}_t$,在第 $t+1$ 轮替换 $r_t$。(7) 训练用 22,609 个真实轨迹样本骨干冻结、从缓存特征训练头,损失为逐样本平衡焦点损失,样本内对保留/剪枝 token 分别平均再等权相加。

技术新颖性

技术新颖性体现在三点。第一是范式转换:直接从冻结骨干表征就地读取剪枝信号,既被探针研究(AUC 0.83)验证,又去掉了 SWE-Pruner 的外部模型与查询;这使得相关性的判断不再需要额外模型调用恢复。第二是长度感知嵌入:以行数 $N$ 为键查分桶嵌入表并广播加到每个隐状态(公式 $\tilde{h}_i=h_i+e(N)$,零初始化使其训练初始退化为长度无关极限),它不改变整体逐行准确率(F1 几乎不变)而是把错误重新分配到更长的响应上——那里误删单行的危害小得多,judge 因此从 6.86 升到 7.08。第三是逐样本平衡焦点损失:把保留/剪枝不平衡视为逐样本而非全局性质,比 BCE 高 +1.13 judge / +0.16 F1,并避免了 Dice、Tversky 在逐行匹配 judge 上崩塌到 5.30/3.03 的问题——因为逐行标签匹配无法区分「有用的代码骨架」与「精度正确但不可用」的剪枝。此外还有工程层面的服务器内集成:把头同置于推理引擎内部以避免跨引擎边界传输隐状态,使每次剪枝调用仅增加一次头前向。

Prior work (SWE-Pruner) vs SWE-Pruner Pro 对比示意
Figure 1: Prior work (SWE-Pruner) vs SWE-Pruner Pro 对比示意
Overview of SWE-Pruner Pro pipeline.
Figure 3: Overview of SWE-Pruner Pro pipeline.
Head architecture.
Figure 4: Head architecture.

实验结果

探针研究:冻结 Qwen3-Coder-Next 隐状态上保留/剪枝两类线性可分,留出轨迹 AUC 0.83、best-F1 0.63,高于多数类 F1 上界 0.46。只读基准(表 1):SWE-Pruner Pro 是唯一在每个单元都减 token 的方法。Qwen3-Coder-Next 省 34.7%/39.4%/13.9%(SWE-QA/SWE-QA-Pro/Oolong),质量 +0.02/+0.24/-1.4 pp;MiMo-V2-Flash 省 6.9%/22.6%/30.1%,Oolong +2.2 pp。SWE-Bench Verified(表 2):MiMo-V2-Flash 解决率 +3.8%(345/500)、输入 token 仅 +7.4%,约为次优 SWE-Pruner(+4.2% 但 +14.9% token)的一半;Qwen3-Coder-Next 只丢 6 个解决(-1.2 pp)却取得最大输入 token 削减 -13.5%。消融(表 3)证明逐样本平衡焦点损失与长度感知嵌入均带来增益;延迟(图 5)显示服务器内头仅增加 15.0% 聚合墙钟时间。

End-to-end quality and token consumption on read-only multi-turn benchmarks.
Table 1: End-to-end quality and token consumption on read-only multi-turn benchmarks.
SWE-Bench Verified results across two agent backbones.
Table 2: SWE-Bench Verified results across two agent backbones.
Ablation results on the held-out judge set.
Table 3: Ablation results on the held-out judge set.
Per-call pruning overhead relative to the following generation step, by trajectory.
Figure 5: Per-call pruning overhead relative to the following generation step, by trajectory.
查看结构化数据
任务指标本文基线提升
SWE-QA-Pro(多轮 QA,可执行环境) LLM-judge 分数(1–10,GPT-5.4-mini)+ 总 token Qwen3-Coder-Next:分数 7.84(+0.24)、token 368K(-39.4%) No Pruning:分数 7.60、token 607K token 削减 39.4%,为全文最大节省,同时质量 +0.24
Oolong(长上下文聚合,多轮智能体变体) 规则精确匹配准确率(0–100)+ 总 token MiMo-V2-Flash:准确率 94.6(+2.2 pp)、token 41.2K(-30.1%) No Pruning:准确率 92.4、token 58.9K 省 30.1% token 且准确率 +2.2 pp,唯一兼顾质量与压缩
SWE-Bench Verified(补丁生成,500 题) 解决率 + 平均输入 token MiMo-V2-Flash:345/500(+3.8%)、3,190K(+7.4%) No Pruning:326/500、2,971K 解决率 +3.8%,token 开销仅为次优 SWE-Pruner 的一半
逐行剪枝探针(留出轨迹) AUC / best-F1 AUC 0.83、best-F1 0.63 多数类 F1 上界 0.46(经验正例率约 30%) 证明骨干表征已编码逐行重要性,是全文立论基础

局限与改进

作者明确承认两点。其一,方法读取骨干内部隐状态,故当前评估只覆盖开源权重模型;该方案可扩展到任何暴露隐状态的骨干,但换模型需重新训练头。其二,智能体任务基准以 Python 为中心,尽管 Oolong 是自然语言长上下文聚合的域外检验(在两个骨干上都同时保住 token 节省与质量),但更广的编程语言覆盖被留给未来工作。补充我自己观察到的局限:(1) 必须掌握服务控制权(要给 SGLang 打补丁),对闭源 API(GPT、Claude)无法即插即用;(2) 22k 训练样本严重依赖 Claude Sonnet 4.6 的逐行标注,标注噪声被作者承认并由逐样本平衡缓解,但仍构成依赖与成本;(3) 在 Qwen3-Coder-Next 的 SWE-Bench 上仍丢 6 个解决,剪枝对代码修改并非零代价;(4) 输入 token 与 API 调用数在不同骨干上方向相反(Qwen 上 API 调用从 131.9 升到 139.8),成本模型是混合的;(5) 训练集保留率约 30% 带有数据集特性,向其他领域的泛化未被验证。

独立分析的弱点

独立分析的主要弱点与改进方向如下。第一,对隐状态访问的硬依赖:只能用于自托管开源模型,无法服务闭源 API 用户,改进方向是设计不读隐状态的代理信号或蒸馏方案。第二,对 Claude Sonnet 4.6 标注的重度依赖:22k 样本 + 约 155k 行的标注成本与噪声都是隐患,可用自训练或弱监督自蒸馏替代以降低成本。第三,换骨干需重训头:跨骨干迁移能力未验证,可探索表征对齐或共享头。第四,长度嵌入 $e(N)$ 依赖固定桶数 $K$:在极端长度或分布外长度上的行为未测,可改为连续长度编码。第五,SWE-Bench 结果不对称:Qwen3-Coder-Next 实际丢解,「保质量」主张主要成立于只读基准,建议报告按任务类型分层的结果。第六,固定阈值 $\tau=0.5$:可改为动态/自适应阈值。第七,仅做逐行剪枝:可向 AST 感知的结构化剪枝、甚至多模态工具输出(图像/表格)扩展。

未来方向

作者显式提出的方向包括:扩展到任何暴露隐状态的骨干;用同一管线覆盖更广的编程语言;通过更紧的隐状态量化进一步压缩开销;以及附录 E.3 讨论的离引擎变体与剩余优化空间。基于该成果可延伸的方向:(1) 跨骨干迁移或借表征对齐共享头,避免每个模型重训;(2) 自训练/弱监督去掉 LLM 标注器依赖;(3) 与主动式历史折叠方法联合,形成「观察边界压缩 + 历史折叠」双层管理;(4) 超越逐行的结构化(AST 感知)剪枝;(5) 自适应阈值 $\tau$;(6) 在代码之外的智能体基准(网页、数据分析)上测试;(7) 可解释性分析——研究头究竟学到了哪些与保留相关的特征,反哺对骨干注意力机制的理解。

复现评估

代码:论文给出 Project Page https://github.com/Ayanami1314/swe-pruner-pro。模型:两个开源骨干 MiMo-V2-Flash(309B MoE,激活 15B)与 Qwen3-Coder-Next(80B MoE,激活 3B),需在打过补丁的 SGLang 栈上服务,算力门槛极高。数据:训练用 22,609 个真实轨迹样本,探针用约 2,260 个工具响应(约 155k 行),均来自公开 SWE-Bench 风格与终端任务数据集并由 Claude Sonnet 4.6 逐行标注(产生付费 API 成本)。基准 SWE-Bench Verified/SWE-QA/SWE-QA-Pro/Oolong 均公开,超参与服务细节在附录 B/E.1,标注与评判 prompt 在附录 H。难度整体偏高——复现主结果需 80B–309B MoE 服务设施、GPU 资源、打过补丁的推理引擎与标注管线;但头本身训练很便宜(用缓存特征)。综合看中等偏难,关键瓶颈在超大 MoE 服务与标注成本而非算法本身。