Technical Summary: SWE-Pruner Pro

Problem Statement

Coding agents operating in multi-turn environments accumulate extensive tool outputs (e.g., cat, grep, ls, python execution logs) during their interaction with codebases. These outputs often contain significant redundancy, consuming the majority of the per-trajectory token budget and triggering long-context degradation. Existing context pruning solutions fall into two categories:

1. General-purpose compression: Methods like LLMLingua2 or Selective Context score tokens based on fixed metrics (e.g., perplexity, syntactic structure). These fail to adapt to the agent's evolving task focus.
2. Task-specific pruning: Prior work like SWE-Pruner conditions pruning on the agent's intent but requires an external scoring model and an explicit "goal-hint" query that the agent must generate every turn. This introduces additional overhead and latency.

Both approaches treat the pruning signal as an external quantity to be reconstructed, ignoring the fact that the agent's backbone model has already processed the tool output and likely encoded relevance information within its internal representations.

Methodology: SWE-Pruner Pro

SWE-Pruner Pro proposes reading the pruning signal directly from the agent backbone's internal hidden states, eliminating the need for external scoring models or explicit goal-hint queries. The method operates directly inside the agent's inference loop.

Core Mechanism

1. Internal Representation Utilization: When the agent reads a tool response ($r_t$), the frozen backbone performs a standard prefill, generating last-layer hidden states ($h_i$) for each token. SWE-Pruner Pro attaches a lightweight pruning head to these existing states.
2. Pruning Head Architecture:
   • Length-Aware Embedding: A learned embedding $e(N)$, conditioned on the total line count $N$ of the response, is broadcast-added to each hidden state. This allows the model to adjust its pruning strategy based on response length (e.g., being more conservative with short outputs where every line is critical).
   • Non-Linear Classifier: A small feed-forward network (LayerNorm + two Linear-GELU-Dropout blocks) maps the augmented hidden states to a keep-or-prune logit. This non-linearity is necessary to resolve the overlap in representation space observed in linear probing experiments.
   • Decision Aggregation: Token-level logits are aggregated via majority vote within each line to produce a binary keep/prune decision for the entire line.
3. Training Strategy:
   • Data: The head is trained on 22,609 multi-turn trajectories annotated by Claude Sonnet 4.6 with per-line keep/prune labels.
   • Loss Function: The authors employ a per-sample balanced focal loss. Unlike standard focal loss which balances globally, this approach computes the loss separately for "keep" and "prune" tokens within each sample and averages them with equal weight. This ensures that samples with extreme keep-rates (e.g., only 3% of lines kept) contribute equally to the objective, preventing the model from overfitting to the average keep-rate (~30%).
4. Inference Integration: The head runs "in-server" during the prefill phase. It reuses the hidden states generated for the tool response, adding only a single forward pass of the small head. The pruned response ($\tilde{r}_t$) replaces the raw response in the context for the next turn, while the current turn's generation still attends to the full raw response.

Key Contributions

1. Signal Discovery: The paper demonstrates that an agent backbone's internal representations already encode line-level importance of tool outputs. A simple linear probe achieves an AUC of 0.83, proving that the "keep-or-prune" signal exists without external reconstruction.
2. Architecture: Introduction of SWE-Pruner Pro, a lightweight head that utilizes a length-aware embedding and a per-sample balanced focal loss to make precise pruning decisions directly from the backbone's frozen states.
3. Efficiency: The method avoids extra model calls and explicit goal-hint queries, adding bounded inference overhead while achieving significant token reduction.

Experimental Results

The method was evaluated on two open-weight backbones (MiMo-V2-Flash and Qwen3-Coder-Next) across four multi-turn benchmarks: SWE-Bench Verified, SWE-QA, SWE-QA-Pro, and Oolong.

• Token Savings: SWE-Pruner Pro is the only method among seven evaluated that reduces end-to-end token usage in every setting.
  • Up to 39% reduction on SWE-QA-Pro (Qwen3-Coder-Next).
  • 30% reduction on the long-context Oolong benchmark (MiMo-V2-Flash).
  • In contrast, other methods (e.g., LLMLingua2) sometimes increased token usage by up to 190% on long-context tasks due to overhead.
• Task Quality:
  • SWE-QA/Pro: Maintains task quality within a narrow band of the unpruned baseline (e.g., +0.02 score change on SWE-QA).
  • Oolong: Improves accuracy by +2.2 points on MiMo-V2-Flash.
  • SWE-Bench Verified: On MiMo-V2-Flash, it improves the resolve rate by +3.8% while reducing input tokens by 7.4% compared to the unpruned baseline. On Qwen3-Coder-Next, it achieves the largest input-token reduction (-13.5%) with the smallest degradation in resolve rate (-1.2%).
• Latency: The in-server head adds approximately 15.0% aggregate wall time to the total generation time, a cost offset by the reduced context size for subsequent turns.

Significance and Claims

The paper claims that the "relevance judgment" required for context pruning is not an external signal that must be elicited via additional queries or models, but is already present in the representations formed by the backbone while passively reading an observation.

• Paradigm Shift: Rather than building pruning around the agent (requiring external models or explicit queries), SWE-Pruner Pro simply reads the signal the agent has already formed.
• Practical Impact: This approach resolves the quality-compression trade-off observed in prior work, allowing for aggressive pruning (up to 39% token savings) without degrading task performance, and in some cases, improving it by mitigating long-context degradation.
• Generalizability: While evaluated on Python-centric coding tasks, the pipeline is language-agnostic and has been shown to preserve token savings and quality on natural-language long-context tasks (Oolong), suggesting the paradigm transfers across domains.

The authors conclude that for coding agents, the most efficient path to context management is to leverage the internal state of the backbone itself, avoiding the latency and complexity of external scoring mechanisms.