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

AI summary

Overview

Research area: Context management for coding agents — specifically, pruning redundant tool outputs from an agent's conversation history so long-horizon software-engineering tasks stay within budget and do not degrade.

Technical level: Intermediate. The paper assumes familiarity with transformer hidden states, KV caching, prefill passes, mixture-of-experts backbones, and agent benchmarks such as SWE-Bench Verified. No new model architecture is proposed; the contribution is a small, trainable classification head bolted onto a frozen backbone.

Scope in one sentence: The paper argues that a coding agent's own last-layer hidden states already encode which lines of a tool output matter, and shows that a lightweight head reading those states can prune tool outputs more effectively than seven competing methods across two backbones and four benchmarks.

What This Paper Is About

Coding agents spend most of their token budget re-reading tool outputs — the raw text returned by commands like cat, grep, and ls — much of which is never referenced again. Existing pruners get their pruning signal from outside the agent: general compressors score tokens with fixed metrics like perplexity, and the task-specific pruner SWE-Pruner runs a separate scoring model conditioned on a goal-hint query the agent must write every turn. The authors ask whether that judgment is already sitting inside the backbone, which has necessarily read the tool output during its normal prefill. They answer yes, and build SWE-Pruner Pro to read it out directly.

Key Contributions

1. Empirical evidence that the signal is already internal. A logistic regression probe on frozen last-layer hidden states of Qwen3-Coder-Next separates kept from pruned lines with AUC 0.83 and best-F1 0.63 on held-out trajectories, against a majority-class F1 upper bound of 0.46 at an empirical positive rate of about 30%.

2. A pruning head attached to the agent's own prefill. SWE-Pruner Pro converts per-token hidden states into a keep-or-prune logit, aggregates token decisions into line-level decisions by majority vote, and requires no separate scoring model and no explicit goal-hint query.

3. Two design choices that carry the empirical lift. A learned length-aware embedding indexed by the tool output's line count, and a per-sample balanced focal loss that averages the loss separately over keep and prune tokens within each sample with equal weight.

4. Evaluation across two open-weight backbones and four multi-turn benchmarks, where the method saves up to 39% of prompt and completion tokens while preserving task quality, and adds 15.0% aggregate wall time on a 16-trajectory replay without an extra model call.

Main Findings

• Token savings without quality collapse. On SWE-QA, SWE-QA-Pro, and Oolong, SWE-Pruner Pro is the only pruner that reduces tokens in every cell, reaching 39.4% savings on Qwen3-Coder-Next / SWE-QA-Pro and 22.6% on MiMo-V2-Flash / SWE-QA-Pro, with quality changes of +0.02, +0.24, and −1.4 points on Qwen3-Coder-Next and −0.04, −0.11, +2.2 on MiMo-V2-Flash.

• Prior pruners often inflate tokens. Four of six prior pruners increase token use on at least one cell; LLMLingua2 reaches +189.8% on Oolong with MiMo-V2-Flash, and +233.3% for Selective Context on Oolong with Qwen3-Coder-Next.

• SWE-Bench Verified is asymmetric across backbones. On MiMo-V2-Flash, every pruner improves the resolve rate; SWE-Pruner Pro reaches +3.8% (345/500 vs 326/500) at roughly half the token overhead of SWE-Pruner (+4.2% resolve at +14.9% tokens). On Qwen3-Coder-Next, every pruner loses solves, but SWE-Pruner Pro gives the best profile: 335/500, losing only 6 solves (−1.2 points) while cutting input tokens by 13.5%.

• API calls move in the opposite direction from tokens. On MiMo-V2-Flash, all pruners increase API calls over No Pruning, with SWE-Pruner Pro using the most (111.8 vs 94.8). On Qwen3-Coder-Next, SWE-Pruner Pro increases calls from 131.9 to 139.8 while still achieving the largest input-token reduction, so the authors report the two efficiency measures separately.

• The loss function is the biggest ablation lever. Per-sample balanced focal loss scores F1 0.635 and judge 7.08 on a held-out set of n=100, beating BCE (0.475 / 5.95) by +1.13 judge and +0.16 F1, and clearly ahead of Focal (0.593 / 6.37), Dice (0.591 / 5.30), and Tversky (0.591 / 3.03).

• The length-aware embedding trades accuracy for harm reduction. Removing it leaves F1 essentially unchanged (0.636 vs 0.635) but drops judge from 7.08 to 6.86, meaning it redistributes errors toward longer responses where mis-pruning a line is less damaging.

• Bounded latency. With the in-engine head on a 16-trajectory MiMo-V2-Flash replay, pruning adds 15.0% aggregate wall time relative to total generation time, with per-trajectory ratios at p50 = 14.7% and p95 = 34.8%.

• The signal is mid-band noisy, motivating a non-linear head. The probe's score distributions have visibly different means but overlap in the middle band, and a linear classifier cannot exploit the length-dependent structure of tool outputs.

Methodology in Plain English

The authors start from a simple observation about how agents work. When an agent calls a tool, the backbone must read the returned text to decide what to do next. Reading is an attention-weighted pass over the tokens, so the backbone has already formed some opinion about which lines matter. The authors test this by taking about 2,260 tool responses (roughly 155k lines) from public SWE-Bench-style and terminal-task datasets, having Claude Sonnet 4.6 label each line keep or prune, freezing Qwen3-Coder-Next, mean-pooling the last-layer hidden states per line, and fitting a logistic regression. The two classes turn out to be separable well above chance, which is the paper's central empirical claim.

Given that, the pipeline is: at each turn the agent calls a tool, the environment returns a response, and the backbone prefills the history, the call, and the response. Only the new response tokens need forwarding because the prefix is already in the KV cache. The pruning head reads the hidden states for those tokens straight off that prefill — no extra forward pass on the response. The head adds a learned embedding indexed by the response's line count, then runs a small LayerNorm and two feed-forward blocks to produce a per-token keep logit, and takes a per-line majority vote to decide what to delete.

Training uses 22,609 annotated samples drawn from 6,252 unique trajectories across five public HuggingFace datasets. Labels come from Claude Sonnet 4.6 at the line level, and are expanded to token labels for the loss, with uncertain rows kept as positives. The loss is what the authors consider their key design choice: rather than a global class balance, they compute the focal loss separately over the keep tokens and the prune tokens of each individual sample and average the two branches with equal weight. This protects minority-class recall in extreme-ratio samples — the cases where only 3 of 100 lines matter, or where 90 of 100 lines are safe. The backbone is fully frozen throughout; hidden states are cached to disk so head training is a feature-reading job that completes a 10-epoch pass in about 15 minutes on a single 8×H200 node.

Evaluation compares against No Pruning and six prior pruners (LLMLingua2, Selective Context, RAG with bge-reranker-v2-m3, Self-Prune, LongCodeZip, and SWE-Pruner) on MiMo-V2-Flash (309B-parameter MoE, 15B active, 256K context) and Qwen3-Coder-Next (80B-parameter MoE, 3B active, 256K context), using SWE-Bench Verified (500 issues, Mini-SWE-Agent harness), SWE-QA (144 questions), SWE-QA-Pro (260 questions), and Oolong (280 long-context aggregation instances).

Why This Matters

Impact on research. The paper reframes context pruning as a readout problem rather than a scoring problem. Prior task-specific pruners reconstructed the agent's information need externally through a second model and a per-turn query; this work suggests that the relevance judgment is already computed as a byproduct of the forward pass the backbone must run anyway. If that holds generally, an entire class of auxiliary scoring machinery becomes unnecessary, and the same argument may extend beyond pruning to retrieval, memory management, and observation filtering.

Real-world applications.

• Long-horizon coding agents that run hundreds of tool calls on a repository and hit context limits, where pruning incoming observations before they enter the history defers or avoids truncation and summarization.
• Cost control for agent deployments, since tool outputs dominate the per-trajectory token bill — file-reading commands alone account for over 70% of tokens in the Mini-SWE-Agent + Claude Sonnet 4.5 setting cited by the authors.
• Terminal and CLI agents handling non-code output, since the training corpus deliberately mixes SWE-style rollouts with a domain-mixed CLI corpus covering chess, machine learning, cryptography, databases, and shell scheduling.
• Long-context aggregation and analysis pipelines, where Oolong serves as an out-of-domain natural-language check on both backbones.

Industry relevance. The method works on open-weight backbones, runs inside the serving engine alongside the existing prefill, and adds 15.0% aggregate wall time on the reported replay while cutting tokens on every subsequent turn. That combination — no extra model call, no extra query, bounded overhead — is what makes it deployable rather than merely accurate. The caveat the authors state plainly is that the recipe requires access to hidden states, so the present evaluation covers only open-weight models, though they argue the approach extends to any backbone that exposes them.

Future Directions

• Broader language coverage. The benchmarks are Python-centric; the authors state the pipeline is language-agnostic and that broader coverage across programming languages reuses the same pipeline, leaving it to future work. The training corpus is already 39

Authors’ abstract

Pruning long context for coding agents has been a vital technology for efficient context management. While existing context pruning methods such as SWE-Pruner realize this by attaching a separate code classifier, we find the agent itself encodes internal representations indicating the relevance of code context when reading tool output. Based on this finding, we propose SWE-Pruner Pro, which prunes tool outputs directly inside the agent. Concretely, a small head turns the agent's own internal representations into a keep-or-prune label for each line, with a length-aware embedding keyed to each tool output's line count. Across two open-weight backbones and four multi-turn benchmarks, SWE-Pruner Pro saves up to 39% of prompt and completion tokens while preserving task quality, with bounded inference overhead. Notably, on MiMo-V2-Flash SWE-Pruner Pro additionally raises the SWE-Bench Verified resolve rate by +3.8% and the long-context Oolong accuracy by +2.2 points.