SWE-Pruner Pro: Let the Coding Agent Prune Its Own Context

Paper: 2607.18213 Authors: Yuhang Wang, Yuling Shi, Shaoqiu Zhang, Jialiang Liang, Shilin He, Siyu Ye, Yuting Chen, Kai Cai, Xiaodong Gu Categories: cs.CL, cs.SE

The Gap

Coding agents like SWE-Bench solvers call tools (grep, find, cat) and get back long outputs — sometimes thousands of lines. Most of that output is irrelevant to the current task. Feeding it all into the context window wastes tokens, slows inference, and can actually hurt performance due to “lost in the middle” effects.

The prior art here is SWE-Pruner (same group’s earlier work), which trains a separate code classifier — an external model that reads tool output and decides which lines to keep. The problem: you need to train and run a whole extra model, it adds latency, and it can’t see what the agent is actually “thinking” at the moment it reads the output.

The key finding of this paper is observational: when a coding agent processes tool output, its own internal hidden states already encode relevance signals. The agent “knows” which lines matter — you just need to extract that knowledge. This eliminates the need for a separate classifier entirely.

Problem: Tool outputs are huge, context windows are finite
    |
    v
Prior approach (SWE-Pruner): Train separate classifier to predict relevance
    |
    v
Limitation: Extra model = extra cost, can't see agent's internal state
    |
    v
Insight: Agent's own hidden states already encode line-level relevance
    |
    v
Method: Small head on agent representations + length-aware embedding
    |
    v
Evidence: 39% token savings, +3.8% resolve rate on SWE-Bench Verified
    |
    v
Conclusion: The agent itself is the best judge of what to prune

The Increment

One sentence: Before this paper, you needed a separate model to decide what context to cut; after this paper, the coding agent’s own internal representations serve as that judge, making pruning cheaper, faster, and more accurate.

Core Mechanism

SWE-Pruner Pro has three components: (1) a line-level feature extractor that reads hidden states from the agent’s transformer layers for each token in a tool output, (2) a length-aware embedding that encodes how many total lines the tool output has (so the pruning head adapts its behavior to different output sizes), and (3) a lightweight classification head that outputs a keep-or-prune probability for each line.

The data flow is straightforward. The agent reads a tool output as usual, producing hidden states at every token position. The method pools these hidden states to get one vector per line. It concatenates each line vector with the length-aware embedding (which is looked up from a small learned table keyed to the output’s line count, bucketed into ranges). The classification head — a two-layer MLP — then scores each line. Lines below a threshold get pruned from the context before the next agent turn.

Tool output (N lines)
        |
        v
Agent processes normally (hidden states at each token)
        |
        v
Pool hidden states --> one vector per line
        |
        v
Concatenate with length-aware embedding
(length embedding = lookup table, bucketed by N)
        |
        v
Classification head (2-layer MLP)
        |
        v
Keep / Prune label per line
        |
        v
Pruned context fed to next agent turn

Think of it like a building superintendent deciding which boxes to move out of storage. Previously (SWE-Pruner), you’d hire an outside appraiser to walk through the storage unit and tag each box as “keep” or “toss.” The appraiser has never met the tenant and doesn’t know what they’re working on. Now, SWE-Pruner Pro says: skip the appraiser. Instead, watch the tenant’s face as they walk past each box. Their expression — the hidden states — tells you exactly which boxes they care about. The length-aware embedding is like the superintendent noting “this is a small closet” vs. “this is a warehouse” — the decision policy shifts depending on how much stuff there is. The classification head is just the superintendent reading the tenant’s expression and stamping “keep” or “toss” on each box.

Key Concepts

• Internal representations as implicit relevance signals: When a transformer model reads a passage, its hidden layers don’t just “process” text uniformly — attention patterns and residual activations vary by token, reflecting how much each part matters to the model’s current task. This paper’s insight is that you don’t need to build a relevance detector from scratch; you can read the agent’s own “body language.” Concrete example: if the agent is trying to fix a bug in utils.py, its hidden states when reading lines of utils.py will look fundamentally different from when it reads README.md — and a simple linear probe can pick that up.

• Length-aware embedding: A small learned vector that conditions the pruning head on how long the current tool output is. The intuition: if a tool output has 50 lines, you might keep 30; if it has 5000 lines, you might keep 200. The threshold should be adaptive. The authors implement this by bucketing output lengths into ranges and learning one embedding per bucket, concatenated to each line’s hidden-state vector before classification.

• In-context pruning vs. pre-processing pruning: SWE-Pruner (the predecessor) prunes tool outputs *before they enter the agent’s context, using an external classifier. SWE-Pruner Pro prunes inside the agent — the agent sees the full output, and then the pruning head operates on the agent’s own representations. This is architecturally different: the pruning decision happens after the agent has “read” the output, not before.

Framework Shift

Before (mainstream approach):        After (this paper):

Tool output                          Tool output
      |                                    |
      v                                    v
Separate classifier model            Agent processes output normally
(trained independently)                    |
      |                                    v
      v                              Agent's own hidden states
Keep/Prune labels                          |
      |                                    v
      v                              Small learned head
Pruned context                           |
      |                                    v
      v                              Keep/Prune labels
Agent processes pruned output              |
                                           v
                                     Pruned context
                                     (for next turn)

From “outsourcing the pruning decision” to “the agent judges its own reading,” the core shift is that the model doing the understanding is now the same model doing the pruning — eliminating the gap between comprehension and selection.

Expert Assessment

Problem choice: This is a real, practical gap. Context management for coding agents is a genuine bottleneck — everyone working on SWE-Bench-style systems deals with it. The observation that hidden states encode relevance is not surprising in hindsight (probing studies have shown this for years), but applying it to tool-output pruning for agents is a clean, useful contribution. It sits at the intersection of efficient inference and agent design, which is a hot trajectory.

Method maturity: This is a clever insight executed with minimal machinery — a 2-layer MLP head and a lookup table. That’s a strength, not a weakness. The simpler approach being overlooked? Possibly just truncating tool output by recency or keyword matching, which might get you 60% of the benefit at 0% overhead. But the paper does show meaningful quality improvements, not just savings, which elevates it above simple heuristics.

Experimental integrity: The baselines are reasonable — they compare against SWE-Pruner (the predecessor) and no-pruning baselines across multiple benchmarks (SWE-Bench Verified, Oolong, and others) and two backbones (Qwen2.5-Coder and MiMo-V2-Flash). The +3.8% resolve rate improvement on SWE-Bench Verified with MiMo-V2-Flash is notable — it’s rare for a pruning method to *improve accuracy, which suggests the original context was actually hurting via noise. One concern: the paper tests on a limited set of backbones. Would this generalize to GPT-4-class closed models where you can’t access hidden states? Probably not without API changes, which limits applicability.

Writing quality: The paper is clearly written and well-structured. The related work section is thin — it doesn’t engage deeply with the broader probing/interpretability literature that motivated the approach. The ablation study could be stronger: how much does the length-aware embedding matter vs. just using a fixed threshold? That’s the kind of detail that would elevate the paper from good to solid.

Verdict: weak accept — Clean idea, practical results, but the insight (agents encode relevance internally) is more evolutionary than revolutionary, and the experimental scope could be broader.

Takeaways

Three things you can steal:

1. Probe before you build: Before training a separate model for any selection/filtering task, check whether the main model’s hidden states already encode the signal you need. A quick linear probe can save you months of classifier engineering. This transfers to RAG (retrieval), summarization, multi-document QA — anywhere you’re filtering context.

2. Length-aware conditioning is cheap and effective: If your downstream decision depends on the size of the input, add a bucketed embedding. It’s a 10-line change that gives you adaptive behavior without complex architectures.

3. Pruning can improve quality, not just save cost: The MiMo-V2-Flash results show that removing noisy context *helps the agent reason better. This is a strong argument for aggressive pruning — you’re not just saving tokens, you’re reducing distraction. If you’re building coding agents, test whether shorter context actually beats longer context.