SWE-Pruner Pro: The Coder LLM Already Knows What to Prune

🕒 Published (v1): 2026-07-20 17:47 UTC · Source: Arxiv · link

Why this paper was selected

Microsoft Research; coding agent self-prunes context without separate classifier — deploy-ready

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

SWE-Pruner Pro eliminates the need for a separate pruning model in coding agents by reading keep-or-prune decisions directly from the agent backbone's own last-layer hidden states. A lightweight trainable head—with a length-aware embedding and per-sample balanced focal loss—prunes tool-output lines in-server during normal prefill. Across two MoE backbones and four benchmarks, it saves up to 39% of tokens while matching or exceeding unpruned task quality.

Problem

Multi-turn coding agents accumulate massive, redundant tool outputs (file reads via cat, grep, ls, etc.) that consume the bulk of per-trajectory token budgets, causing cost inflation and long-context degradation. Existing pruners either use fixed surrogate metrics (perplexity, syntactic structure) that ignore agent intent, or require a separate scoring model conditioned on an explicit goal-hint query the agent must generate every turn—both approaches reconstruct the pruning signal externally despite the agent backbone having already processed the tool output during prefill.

Method

SWE-Pruner Pro attaches a small trainable head to a frozen agent backbone. At each turn \(t\), the backbone prefills \([H_{t-1}, c_t, r_t]\) into its KV cache; since the prefix is already cached, only the raw tool response \(r_t\) is forwarded, producing last-layer hidden states \(h_1, \ldots, h_L\). The head operates in two components:

1. Length-aware embedding: a learned embedding \(e(N) \in \mathbb{R}^d\) keyed to the line count \(N\) of the tool output, broadcast-added to every hidden state: \(\tilde{h}_i = h_i + e(N)\). Zero-initialized so training starts from the length-agnostic case.
2. Per-token classifier \(f_\theta\): LayerNorm → two Linear-GELU-Dropout blocks → Linear projection to a scalar logit \(z_i\), yielding \(p_i = \sigma(z_i)\). Line-level decisions aggregate token predictions by majority vote within each line \(\ell\):

\[\hat{y}_\ell = \mathbb{1}\!\left[\frac{1}{|\ell|}\sum_{i \in \ell} \mathbb{1}[p_i > \tau] > \tfrac{1}{2}\right]\]

Training uses a per-sample balanced focal loss: focal loss is computed separately over keep and prune tokens within each sample and averaged with equal weight, so samples with extreme keep/prune ratios contribute proportionally to their minority class rather than being dominated by the global 30% keep rate. The backbone is fully frozen; only the head is trained from cached prefill features.

Key Contributions

• Probing study showing that a frozen backbone's last-layer hidden states already encode line-level keep-or-prune signal (linear probe AUC 0.83, best-F1 0.63 on held-out trajectories), motivating in-model pruning without external scoring.
• SWE-Pruner Pro architecture: in-server pruning head reusing existing prefill, with no extra backbone forward pass on the raw response.
• Length-aware embedding that redistributes pruning mistakes toward longer responses where individual line loss is less harmful.
• Per-sample balanced focal loss that protects recall of each sample's minority class regardless of per-sample keep rate imbalance.
• Evaluation on two open-weight MoE backbones (MiMo-V2-Flash 309B, Qwen3-Coder-Next 80B) across four benchmarks with six pruner baselines.

Results

• Token savings: up to 39% on SWE-QA-Pro (Qwen3-Coder-Next) and 30% on Oolong (MiMo-V2-Flash); the only method that reduces tokens on every benchmark cell across both backbones.
• Quality preservation: Qwen3-Coder-Next scores +0.02/+0.24 on SWE-QA/SWE-QA-Pro vs. no pruning, with −1.4 pp on Oolong; MiMo-V2-Flash scores −0.04 on SWE-QA, +0.07 on SWE-QA-Pro, +2.2 pp on Oolong.
• SWE-Bench Verified: MiMo-V2-Flash resolve rate +3.8% (345/500 vs. 326/500) at +7.4% input tokens—roughly half the overhead of SWE-Pruner (+4.2% resolve at +14.9% tokens); Qwen3-Coder-Next loses only 6 resolves (−1.2%) while achieving the largest input-token reduction (−13.5%).
• Latency: 15.0% aggregate wall-time overhead on generation time (p50 = 14.7%, p95 = 34.8%) on a 16-trajectory MiMo-V2-Flash replay.
• Ablation: per-sample balanced focal loss beats BCE by +0.16 F1 and +1.13 judge score; length-aware embedding adds +0.22 judge at near-identical F1.
• Contrast with baselines: LLMLingua2 inflates Oolong tokens by +190%; four of six prior pruners inflate tokens on at least one benchmark cell.

Limitations

• Backbone must be frozen; the head is trained per-backbone, so a new head must be trained when the agent backbone changes.
• Training requires 22,609 multi-turn tool-response samples annotated by Claude Sonnet 4.6—annotation cost is non-trivial and label quality depends on the annotating LLM.
• API call count increases on both backbones for SWE-Bench Verified (up to +18% on MiMo-V2-Flash), which can matter for rate-limited or per-call-billed deployments.
• On Qwen3-Coder-Next, SWE-Bench Verified resolve rate still degrades (−1.2%), suggesting the pruning signal quality is backbone-dependent.
• Latency overhead (p95 = 34.8%) may be material for short trajectories or latency-sensitive deployments; off-engine variants are discussed but not fully characterized.
• Evaluation is limited to two MoE backbones specialized for long-horizon coding; generalization to other model families or non-coding agent tasks is unverified.

Relevance to Agentic AI / LLM Agents

This work directly targets one of the core scalability bottlenecks in multi-turn LLM agents: context window inflation from accumulated tool outputs, which triggers both cost and long-context quality degradation. The key insight—that the agent's own backbone already encodes line-level relevance during normal prefill, making an external scoring model unnecessary—has broad implications for agent efficiency research and suggests that introspective, in-model approaches to context management are practically viable. SWE-Pruner Pro advances the line of work connecting agent scaffold efficiency (context pruning, KV cache management) with the agent's internal information processing, complementing approaches like retrieval-augmented context selection and goal-conditioned pruners. For researchers studying agentic infrastructure, this establishes a concrete design pattern: train lightweight heads on frozen backbone representations rather than adding separate models to the agent's tool loop.

🔗 Related in this collection

• SWE-QA: Can Language Models Answer Repository-level Code Questions?