One-sentence Summary

LLM4SE Lab, Shanghai Jiao Tong University, and Douyin Group propose SWE-Pruner Pro, which directly prunes tool outputs inside the agent by converting the agent's own internal representations into per-line keep-or-prune labels via a small head and length-aware embedding, achieving up to 39% token savings across two open-weight backbones and four multi-turn benchmarks while preserving task quality with bounded inference overhead, and on MiMo-V2-Flash additionally raising the SWE-Bench Verified resolve rate by +3.8% and long-context Oolong accuracy by +2.2 points.

Key Contributions

• The coding agent's internal representations, formed while reading tool output, encode line-level relevance, eliminating the need for an external scoring model.
• SWE-Pruner Pro reads this signal via a lightweight head attached to the backbone, using a length-aware embedding and per-sample balanced focal loss to predict keep-or-prune labels.
• 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, and on MiMo-V2-Flash it raises SWE-Bench Verified resolve rate by +3.8% and Oolong accuracy by +2.2 points.

Introduction

Coding agents that solve repository-level tasks accumulate long, redundant tool outputs across multi-turn interactions, inflating token costs and degrading model performance. Prior pruning methods either apply fixed, task-agnostic metrics like perplexity that ignore the agent’s shifting intent, or rely on a separate scoring model and an explicit goal-hint query at every turn, adding inference overhead. The authors show that the agent’s own backbone representations already encode line-level importance, and they propose SWE-Pruner Pro, a lightweight head that reads this pruning signal directly from the backbone’s prefill pass. By incorporating a learned length-aware embedding and a per-sample balanced focal loss, SWE-Pruner Pro avoids extra model calls, saves up to 39% of tokens, and preserves task quality across multiple benchmarks.

Method

The authors design SWE-Pruner Pro to make per-line keep-or-prune decisions on tool responses directly from the agent backbone's last-layer hidden states. At each turn $t$t, the agent issues a tool call $c_t$ct​ and receives a raw tool response $r_t$rt​. Before the agent generates its next move, the backbone prefills the context history $H_{t-1}$Ht−1​, the tool call $c_t$ct​, and the raw response $r_t$rt​ into its KV cache. Since the prefix is already cached, only the new $r_t$rt​ tokens are forwarded, producing last-layer hidden states $h_1, \ldots, h_L$h1​,…,hL​ over the response span.

SWE-Pruner Pro attaches at this prefill stage. The pruning head converts each hidden state $h_i$hi​ into a per-token keep-or-prune logit and aggregates them into per-line decisions. This line-level granularity preserves the syntactic structure of the kept code. Pruning is applied between turns: the agent's generation at turn $t$t still attends to the full $r_t$rt​, while the pruned response $\tilde{r}_t$r~t​ replaces $r_t$rt​ when the trajectory continues into turn $t+1$t+1. The head reads the hidden states off the prefill the backbone already performs, meaning the only added backbone work is a single re-forward of the typically much shorter $\tilde{r}_t$r~t​ at the next turn.

The pruning head consists of two main components: a length-aware embedding and a per-token feed-forward classifier. Let $N$N be the number of lines in $r_t$rt​ and $d$d denote the dimension of $h_i$hi​. The length-aware embedding $\mathbf{e}(N) \in \mathbb{R}^d$e(N)∈Rd is a learned function of the line count $N$N. It is broadcast-added to every hidden state before the classifier and zero-initialised so the classifier reduces to its length-agnostic limit at the start of training:

$$\tilde{h}_i = h_i + \mathbf{e}(N).$$h~i​=hi​+e(N).

This embedding gives the head explicit access to $N$N, allowing the keep-or-prune mapping to vary with response length. This is crucial because the cost of mis-pruning is highly non-uniform; stripping a few lines from a short response is catastrophic, while the same on a long response is negligible.

The per-token classifier $f_\theta$fθ​ is a small non-linear network comprising LayerNorm, two hidden Linear-GELU-Dropout blocks with hidden width $d$d, and a final Linear projection to a single logit. A linear classifier alone cannot resolve the overlap in the mid-score region observed in linear probe experiments, nor can it interact effectively with the length-aware embedding.

Applied to every augmented hidden state, the classifier produces keep probabilities:

$$z_i = f_\theta(\tilde{h}_i), \qquad p_i = \sigma(z_i).$$zi​=fθ​(h~i​),pi​=σ(zi​).

With a per-token sigmoid threshold $\tau = 0.5$τ=0.5, per-line decisions $\hat{y}_\ell$y^​ℓ​ are obtained by majority vote of the binarized token decisions within each line $\ell$ℓ:

$$\hat{y}_\ell = \mathbb{1} \left[ \frac{1}{|\ell|} \sum_{i \in \ell} \mathbb{1}[p_i > \tau] > \frac{1}{2} \right].$$y^​ℓ​=1[∣ℓ∣1​i∈ℓ∑​1[pi​>τ]>21​].

Lines with $\hat{y}_\ell = 0$y^​ℓ​=0 are removed from the tool response before it is appended to the agent's history.

The authors train the head on 22,609 samples from real multi-turn agent trajectories, annotated with per-line keep-or-prune labels. Line labels are expanded to per-token labels for the loss. The loss function is a per-sample balanced focal loss. Per-line labels from an LLM annotator are inherently fuzzy, but the ratio the annotator settles on is informative. Standard cross-entropy and batch-level focal loss treat every token equally, which can overfit the typical keep rate and dilute extreme-ratio samples. To protect the recall of each sample's minority class, the authors compute the loss separately over keep and prune tokens within a sample and average the two with equal weight.

For each token $i$i in sample $s$s, let $p_{t,i} = p_i$pt,i​=pi​ if $y_i = 1$yi​=1 and $p_{t,i} = 1 - p_i$pt,i​=1−pi​ otherwise. The token-level loss is:

$$\mathcal{L}_i^{\mathrm{tok}} = (1 - p_{t,i})^\gamma \cdot \mathrm{BCE}(p_i, y_i), \quad \gamma = 2.$$Litok​=(1−pt,i​)γ⋅BCE(pi​,yi​),γ=2.

The losses are averaged separately over keep and prune tokens within each sample:

$$\mathcal{L}_s^{\mathrm{keep}} = \frac{\sum_i y_i \mathcal{L}_i^{\mathrm{tok}}}{\sum_i y_i}, \quad \mathcal{L}_s^{\mathrm{prune}} = \frac{\sum_i (1 - y_i) \mathcal{L}_i^{\mathrm{tok}}}{\sum_i (1 - y_i)},$$Lskeep​=∑i​yi​∑i​yi​Litok​​,Lsprune​=∑i​(1−yi​)∑i​(1−yi​)Litok​​,

and combined with equal weights:

$$\mathcal{L}_s = \frac{1}{2} \mathcal{L}_s^{\mathrm{keep}} + \frac{1}{2} \mathcal{L}_s^{\mathrm{prune}}.$$Ls​=21​Lskeep​+21​Lsprune​.

The final objective is $\mathcal{L} = \frac{1}{S} \sum_s \mathcal{L}_s$L=S1​∑s​Ls​. The backbone remains fully frozen, so adding the head does not alter the agent's general behaviour and requires no retraining of the backbone itself. The head is trained from cached features.

Experiment

The evaluation spans four coding-agent benchmarks and two open-weight Mixture-of-Experts LLMs, comparing SWE-Pruner Pro against six prior pruning methods. On multi-turn tasks, SWE-Pruner Pro uniquely achieves substantial token savings while preserving quality, whereas other methods either degrade performance or inflate tokens. Ablations confirm that per-sample balanced focal loss and length-aware embedding are critical to this efficiency, and latency analysis shows that pruning overhead is modest relative to the overall token reductions.

Across read-only multi-turn benchmarks, SWE-Pruner Pro was the only pruning method that consistently reduced total token consumption while preserving or slightly improving answer quality. Competing pruners either inflated tokens on long-context tasks or degraded quality, whereas SWE-Pruner Pro leveraged in-place backbone signals to deliver large savings with negligible quality impact. On Qwen3-Coder-Next, SWE-Pruner Pro reduced token usage by 35–39% on SWE-QA and SWE-QA-Pro while judge scores remained within +0.24 of the unpruned baseline, and cut Oolong tokens by 14% with only a 1.4-point accuracy drop. LLMLingua2 and Selective Context more than doubled token consumption on Oolong (up to +233%), while Self-Prune lost 0.55–0.64 judge points on SWE-QA/Pro, showing that prior pruners could not simultaneously compress and maintain quality.

On SWE-Bench Verified, the effect of pruning depends strongly on the backbone. With MiMo-V2-Flash all pruners raise the resolve rate, but SWE-Pruner Pro achieves a large gain while adding far fewer input tokens than the next-best method. On Qwen3-Coder-Next every pruner loses resolves, yet SWE-Pruner Pro shows the smallest quality drop and simultaneously delivers the largest input-token reduction. On MiMo-V2-Flash, SWE-Pruner Pro improves resolve rate by 3.8% with a 7.4% token increase, whereas the best-resolving pruner gains 4.2% but adds 14.9% more tokens. On Qwen3-Coder-Next, all pruners degrade resolve rate; SWE-Pruner Pro loses only 1.2 percentage points while cutting input tokens by 13.5%, outperforming other pruners on both metrics. API calls and input tokens trade off differently across backbones: on MiMo-V2-Flash SWE-Pruner Pro uses the most calls, while on Qwen3-Coder-Next it slightly increases calls yet still achieves the largest token reduction.

Ablating the loss function shows that per-sample balanced focal achieves the highest judge score and F1, while Dice and Tversky match its F1 but suffer much lower judge scores, indicating that per-line F1 can be misleading. Adding a length-aware embedding improves the judge score without affecting F1 by redirecting pruning errors to longer responses where they are less harmful. Per-sample balanced focal loss outperforms BCE, Dice, and Tversky on both judge score and F1; Dice and Tversky nearly match F1 but their judge scores collapse, revealing that per-line F1 alone does not capture response usability. Enabling the length-aware embedding raises the judge score from 6.86 to 7.08 while keeping per-line F1 unchanged, because it shifts mistakes toward longer responses where a single mis-pruned line is far less damaging.

SWE-Pruner Pro was evaluated against several pruning baselines on read-only multi-turn benchmarks and on SWE-Bench Verified with different backbone models. It was the only method that consistently reduced token consumption while preserving or slightly improving answer quality, whereas other pruners either inflated tokens or caused notable quality degradation. Ablation experiments confirmed that per-sample balanced focal loss and a length-aware embedding are critical for maintaining response usability, as optimizing for per-line F1 alone can be misleading.