SWE-Pruner Pro prunes redundant tool-output lines from a coding agent’s context by reading a small classifier head off the frozen backbone’s own hidden states, cutting up to 39% of tokens while matching or beating unpruned quality, because the agent already encodes line-level relevance during its normal prefill.
You’ve shipped a coding agent that runs grep, cat, and test commands across a repo. Each tool call dumps hundreds of lines into the conversation, most of it license headers, blank rows, or files the agent glances at once and never reads again. By turn 30 you’re paying to re-attend to megabytes of stale output, and the model’s answers get worse as the window fills.
The existing task-specific fix is SWE-Pruner, which bolts a second scoring model onto the agent and makes it write an explicit “here’s what I’m looking for” hint every turn. This paper’s move is to drop both of those and read the keep/prune decision straight out of representations the backbone already computed.
The key observation: when the agent reads a tool response, its attention layers already have to figure out which lines matter for the next action. So that judgment should be sitting in the last-layer hidden states, no extra query needed. The authors confirm this with a linear probe on frozen Qwen3-Coder-Next hidden states, which hits AUC 0.83 on distinguishing kept from pruned lines. That’s a signal, not noise.
The deployed system attaches a small MLP head to those hidden states. When a tool returns, say, 200 lines of cat output, the backbone prefills them into the KV cache as it normally would. The head reads each token’s hidden state, adds a length-aware embedding keyed to the line count (mis-pruning one line of a 5-line output is catastrophic; one line of 300 is nothing), pushes through two Linear-GELU-Dropout blocks, and emits a keep logit per token. Lines are decided by majority vote across their tokens. Pruned lines are dropped from the history the next turn attends to; the current turn still sees the full response.
Training uses a Per-sample balanced focal loss: standard Focal loss weights hard tokens more, but here the authors additionally average keep-loss and prune-loss separately within each sample before combining. This stops the loss from being dominated by the typical 30% keep rate and preserves signal on samples where almost everything (or almost nothing) should be kept. Labels come from ~22.6k tool responses annotated by Claude Sonnet 4.6.
# Per turn, after tool call c_t returns response r_t:
h = backbone.prefill(history, c_t, r_t) # already happening
hidden = h[r_t_span] # last-layer states over r_t
N = count_lines(r_t)
logits = head(hidden + length_embedding(N))
keep_token = sigmoid(logits) > 0.5
keep_line = majority_vote_per_line(keep_token)
r_t_pruned = drop_lines(r_t, ~keep_line)
history = history + [c_t, r_t_pruned] # next turn sees pruned version
The backbone is fully frozen. Only the head trains, from cached features, in about 15 minutes on one 8xH200 node.
The prevailing approach to context pruning for agents is to bolt on external machinery: a separate scoring model, a retrieval index, or an explicit goal-hint query the agent writes every turn. This paper shows the opposite. The backbone already computed the relevance judgment while reading the tool output; a small head on its hidden states extracts it for free, and every layer of external machinery you add on top wipes out the compression gains through its own overhead. The clearest evidence is that four of six prior pruners actually inflate end-to-end tokens on at least one benchmark cell, one by +190%, while the in-place head is the only method that reduces tokens on every configuration.
The load-bearing finding is consistency, not peak score. On the read-heavy benchmarks SWE-QA, SWE-QA-Pro, and Oolong, SWE-Pruner Pro is the only method of seven that reduces end-to-end tokens in every backbone-by-benchmark cell. Prior pruners often shrink the tool response but pay it back in extra scoring calls, retrieval queries, or worse trajectories that drag on longer. LLMLingua2 inflates Oolong tokens by +190% on MiMo-V2-Flash; the paper’s method cuts them by 30%.
Secondary numbers, all against the unpruned agent:
•
SWE-QA-Pro on Qwen3-Coder-Next: −39.4% tokens with judge score actually up +0.24.
•
Oolong on MiMo-V2-Flash: −30.1% tokens and +2.2 points accuracy (this is a long-context natural-language benchmark, so it’s an out-of-domain check).
•
SWE-bench Verified on MiMo-V2-Flash: resolve rate +3.8% (326/500 → 345/500) at roughly half the token overhead of the next-best pruner.
•
On SWE-bench Verified with Qwen3-Coder-Next, all pruners lose resolves; this method loses the fewest (−1.2 points) while cutting input tokens the most (−13.5%).
Ablations isolate the two design choices. Swapping the per-sample balanced focal loss for plain BCE drops the LLM-judge score from 7.08 to 5.95; Dice and Tversky match on F1 but collapse to 5.30 and 3.03 on the judge, because F1 rewards precise-but-narrow skeletons that the agent can’t actually use downstream. The length-aware embedding leaves per-line F1 unchanged but lifts the judge from 6.86 to 7.08 by redistributing errors toward long responses where a single mis-pruned line barely matters.
On latency, the in-engine head adds 15% aggregate wall time relative to total generation, with p95 at 34.8%. That’s paid once per turn and offset by shorter contexts on every subsequent turn.
Reach for this pattern when you run an open-weight model you serve yourself and your agents accumulate long tool outputs. Concrete scenario: you’re operating a repo-navigation agent on SGLang or a similar engine, and cat/grep outputs are 70%+ of your token bill. Instead of adding a second model call per turn or writing a retriever, train an ~18M-parameter head on cached hidden states from ~20k labeled trajectories and colocate it in the inference engine. You reuse the prefill the backbone already runs; the only extra compute is one head forward per tool response.
The paper doesn’t mention a released code repository or model weights; the training data is aggregated from five public HuggingFace trajectory datasets (terminal-wrench-trajectories, TIGER-Lab/SWE-Next-SFT-Trajectories, ByteDance-Seed/Multi-SWE-bench_trajs, zai-org/CC-Bench-trajectories, AweAI-Team/Scale-SWE-Distilled) with per-line labels from Claude Sonnet 4.6. The engine work (correctness patches to hidden-state return, binary payload envelope, in-engine head) is described in enough detail to reproduce against SGLang 0.5.10, but no patch or fork link is given. This is closed-source in practice unless you rebuild it.
When your model already read the data, don’t hire a second model to summarize what it thought. The relevance judgment is sitting in the hidden states from the forward pass you already paid for; a small head on top beats an external scorer both on quality and on the only metric that actually matters, end-to-end tokens after the pruner’s own overhead is counted.
•
The whole approach requires access to the backbone’s internal hidden states, so it only works with open-weight models you serve yourself. Closed API models are out.
•
The head is per-backbone. Switching from MiMo-V2-Flash to a different model means re-extracting features and retraining the head, though the paper reports this takes ~15 minutes on one node.
•
Behavior on SWE-bench Verified is asymmetric across backbones: it improves resolve rate on one and degrades it on the other, so “preserves quality” is a benchmark-average claim, not a guarantee for your specific model-and-workload combination. Validate before shipping to production.