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

Yuhang Wang, Yuling Shi, Shaoqiu Zhang, Jialiang Liang, Shilin He, Siyu Ye, Yuting Chen, Kai Cai, Xiaodong Gu

SWE-Pruner Pro prunes tool outputs by reading relevance signals directly from the agent's own internal hidden states.

Can we prune long-context coding agent outputs by leveraging the agent's own internal hidden states instead of training a separate, external classifier?

Coding agents spend most of their token budget on redundant tool outputs, but existing pruners either use fixed metrics that ignore task intent or require expensive external scoring models to decide what to keep. SWE-Pruner Pro treats the agent's backbone as the scorer: it attaches a lightweight, length-aware head to the backbone's last-layer hidden states to predict keep-or-prune labels in-place during the prefill pass. This approach reduces token consumption by up to 39% while maintaining or improving task resolution rates, all without requiring extra model calls or backbone fine-tuning.

Paper Primer

The core move is to treat the agent's internal representation as a pre-computed relevance map. Because the backbone must process tool outputs to generate the next action, it already encodes which lines are relevant; SWE-Pruner Pro simply reads this signal out using a small, non-linear head trained with a per-sample balanced focal loss.

SWE-Pruner Pro achieves significant token savings across diverse benchmarks without degrading task quality.

On the SWE-QA-Pro benchmark, the method reduced token usage by 39% while maintaining performance, and on SWE-Bench Verified, it improved the resolve rate by +3.8% on the MiMo-V2-Flash backbone. Up to 39% token reduction; +3.8% resolve rate improvement.

The method is designed for efficiency: by colocating the pruning head within the inference engine and reusing the backbone's existing prefill, the overhead is bounded to approximately 15% of total generation time, which is offset by the reduced compute required for all subsequent turns.

Why is this approach more efficient than simply asking the agent to identify what to prune?

Explicitly asking the agent to write a "goal-hint" query or using a separate scoring model adds significant inference overhead and latency. SWE-Pruner Pro avoids this by reading the relevance signal from the backbone's hidden states during the prefill pass that the agent is already performing.

Does this method require retraining the agent's backbone?

No. The backbone remains entirely frozen; only the lightweight pruning head is trained on cached hidden states, making the system compatible with any open-weight model that exposes its internal representations.

Motivation and Core Concept

Why coding agents need internal pruning instead of external classifiers.

Coding agents repeatedly invoke tools (e.g., cat, grep, python) and ingest the raw textual output. Those tool outputs dominate the per‑trajectory token budget, often containing repetitive lines that are never referenced later. Existing pruning pipelines attach a separate classifier and require the agent to emit a goal‑hint query each turn, incurring extra inference cost.

Agents already “read” their tool outputs, so the hidden states they produce contain a latent signal about which lines will be useful later.

**Figure 1.** (Left) Prior work (Wang et al., 2026b) uses a separate Pruner model conditioned on an explicit goal-hint query the agent must write at every turn. (Right) SWE-Pruner Pro prunes directly from the agent's own internal representations, with no external scoring model.

The inefficiency of external classifiers in coding agents makes in‑agent pruning a far more economical solution.

The SWE-Pruner Pro Head

We prune tool-output lines by reading the backbone’s hidden states and scoring them with a learned head.

Tool outputs often contain hundreds of redundant lines, yet the backbone that generates the next turn already attends to them, implicitly encoding which lines matter for the upcoming action.

The head acts like a librarian that scans each line’s hidden representation and decides, based on its content and the total length of the response, which lines to keep and which to discard.

Augment each hidden vector: $\tilde{h}_i = h_i + e(3)$. For token 1, $\tilde{h}_1 = (0.15, -0.02, 0.23, 0.11)$.

Feed $\tilde{h}_i$ through $f_\\theta$ (a two‑layer MLP) to obtain logits $z_i$. Suppose $z_1=0.8$, $z_2=-0.3$, $z_3=0.6$, $z_4=0.2$, $z_5=-0.5$, $z_6=0.9$.

Apply sigmoid: $p_i = \sigma(z_i)$. This yields $p_1 \approx 0.69$, $p_2 \approx 0.43$, $p_3 \approx 0.65$, $p_4 \approx 0.55$, $p_5 \approx 0.38$, $p_6 \approx 0.71$.

Binarise at $\tau=0.5$: tokens 1, 3, 4, 6 are kept; tokens 2, 5 are pruned.

Vote per line (two tokens per line): Line 1 (tokens 1 & 2) keeps 1/2 → discarded; Line 2 (tokens 3 & 4) keeps 2/2 → kept; Line 3 (tokens 5 & 6) keeps 1/2 → discarded.

The embedding $e(N)$ lets the head be stricter on short responses (Line 1) while being more permissive on longer ones, matching the intuition that losing a line in a brief reply is more damaging.

How does this head differ from a simple linear classifier that directly maps $h_i$ to a keep logit?

The linear classifier lacks the length‑aware embedding $e(N)$ and the non‑linear two‑block feed‑forward network. Without $e(N)$ the model cannot modulate its aggressiveness based on response length, and without non‑linearity it cannot capture the richer decision boundary needed to separate keep versus prune tokens, especially when the hidden states are overlapping.

Per‑turn pruning pipeline (Algorithm 1)

**Figure 3.** Overview of SWE-Pruner Pro. Top (Agent Trajectory): at each turn the agent backbone prefills [history, $c_t$, $r_t$]; only the new tool-response tokens are forwarded, and SWE-Pruner Pro reads their last-layer hidden states {$h_i$} off this prefill at no extra forward on $r_t$. Bottom (SWE-Pruner Pro head): the head scores each $h_i$, aggregates the binarised scores into per-line keep-or-prune decisions, and the pruned response $\tilde{r}_t$ replaces $r_t$ in the next turn.

**Figure 4** Head architecture. The frozen backbone's hidden states $h_i$ receive an additive length-aware embedding indexed by the line count $N$, then pass through a two-block feed-forward stack with LayerNorm and dropout to a per-token keep logit. Line decisions are produced by majority vote within each line at inference.

Performance Benchmarks

Quantitative evaluation of SWE‑Pruner Pro across four benchmarks.

SWE‑Pruner Pro matches baseline quality (within 0.04 score) while cutting token usage by 6.9 % on SWE‑QA.

Table 1 shows a score drop from 8.02 to 7.98 and token count reduction from 590 K to 299 K.

**Table 1.** End-to-end quality and token consumption on the read-only multi-turn benchmarks.

**Table 2.** SWE-Bench Verified results across two agent backbones.

The per‑sample loss balances keep and prune tokens equally, preventing minority‑class recall collapse; the backbone remains frozen, so the head can be trained solely on cached hidden states.

SWE‑Pruner Pro maintains agent quality while significantly reducing token consumption.

Agent Backbones and Ablations

Ablations isolate each design choice, showing token savings, quality impact, and latency trade‑offs.

We evaluate SWE‑Pruner Pro on two Mixture‑of‑Experts (MoE) backbones—MiMo‑V2‑Flash (309 B parameters, 15 B active per token) and Qwen3‑Coder‑Next (80 B parameters, 3 B active). All baselines share identical decoding and hardware, so pruning is the sole variable.

**Figure 2** Linear probe on the agent backbone's frozen last-layer hidden states. (Left) Two-dimensional projection along the LDA discriminant axis and an orthogonal direction. (Right) Distribution of probe scores. The two classes have visibly different means but overlap in the middle band.

SWE‑Pruner Pro cuts input tokens by up to 39 % on the heaviest cell (Qwen3‑Coder‑Next SWE‑QA‑Pro) while keeping judge scores within ±0.24 points.

Token reduction 39 % with a +0.02 judge gain on SWE‑QA‑Pro; comparable reductions on other benchmarks.

LLMLingua2 inflates token counts by +190 % on Oolong, demonstrating that external scoring overhead can outweigh any compression.

Oolong token count rises from the baseline to 190 % higher when LLMLingua2 is applied.

**Figure 5.** Per-call pruning overhead relative to the immediately following generation step, broken down by trajectory

Enabling the in‑engine pruning head increases wall‑time by 15 % per call, but the resulting token reductions make the net trajectory time faster.

Measured aggregate overhead 15.0 % versus a 34–39 % token reduction across benchmarks.

Related Work

We situate SWE‑Pruner Pro among prior token‑pruning and context‑management approaches.

Prompt compression for code has followed three trajectories: token‑level scoring, retrieval‑based shortening, and structure‑preserving code‑aware methods.

Score individual tokens by self‑information or perplexity and drop the lowest‑ranking ones.

Replace verbatim code fragments with retrieved or aggregated snippets from a database.

Preserve program structure (e.g., AST boundaries) while pruning, typically using fixed heuristics.

Line‑level pruning conditioned on an explicit goal‑hint query; uses a separate scoring model.

Inherits line‑level granularity from SWE‑Pruner but reads importance directly from the agent backbone’s hidden states, removing the external scorer and query.

These baselines attach an auxiliary classifier to the backbone, predict a keep‑probability for each token, and prune tokens based on that score.

Per-turn Execution Loop

How SWE‑Pruner Pro trims redundant tool output line‑by‑line during each interaction.

Tool responses in long‑context tasks are often bloated with lines that add no new information, inflating latency and memory use. SWE‑Pruner Pro tackles this by inspecting the agent’s own hidden states and discarding lines that the model itself deems unimportant.

Think of the loop as a librarian who, after each page of a manuscript arrives, quickly glances at a summary note (the hidden state) and decides whether to keep the page or toss it—so the final bound volume contains only the pages that actually contribute to the story.

Augment each hidden state: $\tilde{h}_1=(0.3,0.2)$, $\tilde{h}_2=(0.9,0.8)$, $\tilde{h}_3=(0.5,0.4)$, $\tilde{h}_4=(0.2,0.15)$.

Apply the head (a single linear projection) to obtain logits $z_i$: $z_1=0.1$, $z_2=0.9$, $z_3=0.4$, $z_4=0.05$.

Sigmoid to get keep probabilities $p_i=\sigma(z_i)$: $p_1\approx0.525$, $p_2\approx0.711$, $p_3\approx0.598$, $p_4\approx0.512$.

Compare to $\tau$: tokens 1–4 all exceed $0.5$, so each line’s majority vote $\hat{y}_{\ell}=1$ (keep).

Now lower the threshold to $\tau=0.6$: tokens 1 and 4 fall below, giving line‑wise votes $\hat{y}_1=0$, $\hat{y}_2=1$, $\hat{y}_3=1$, $\hat{y}_4=0$. Lines 1 and 4 are removed, yielding $\tilde{r}_t$ with only lines 2 and 3.

Changing $\tau$ directly trades off recall of potentially useful lines against the amount of redundancy removed; the loop makes this trade‑off explicit at inference time.

Prefill hidden states $\{h_i\}_{i=1}^{L}$ from the raw response $r_t$ while reusing the cached prefix $H_{t-1}$.

Compute a length‑aware bias vector $\mathbf{b}= \mathbf{e}(N)$ where $N$ is the number of lines.

For each token $i$: form $\tilde{h}_i = h_i + \mathbf{b}$ and obtain keep probability $p_i = \sigma\!\bigl(f_{\theta}(\tilde{h}_i)\bigr)$.

Aggregate token‑level keep flags per line $\ell$ and set $\hat{y}_{\ell}=1$ if more than half of its tokens satisfy $p_i>\tau$.

Remove all lines with $\hat{y}_{\ell}=0$ to produce the pruned response $\tilde{r}_t$.

Insert $\tilde{r}_t$ into the history $H_t$ for the next turn.

Per‑turn pruning loop (simplified)

**Algorithm 1** Per-turn execution of SWE-Pruner Pro **Require:** History $H_{t-1}$, tool call $c_t$, raw response $r_t$ with $N$ lines, head $f_{\theta}$, threshold $\tau$ **Ensure:** Pruned response $\tilde{r}_t$ 1: $\{h_i\}_{i=1}^L \leftarrow \text{Prefill}(r_t \mid H_{t-1}, c_t)$\hfill$\triangleright$ only $r_t$ forwarded; prefix cached 2: $\mathbf{b} \leftarrow \mathbf{e}(N)$\hfill$\triangleright$ length-aware embedding, Eq. 1 3: **for** $i = 1, \dots, L$ **do** 4: $\quad \tilde{h}_i \leftarrow h_i + \mathbf{b}$ 5: $\quad p_i \leftarrow \sigma(f_{\theta}(\tilde{h}_i))$\hfill$\triangleright$ Eq. 2 6: **end for** 7: **for** each line $\ell$ **do** 8: $\quad \hat{y}_{\ell} \leftarrow \mathbb{I}\left[\frac{1}{|\ell|} \sum_{i \in \ell} \mathbb{I}[p_i > \tau] > \frac{1}{2}\right]$\hfill$\triangleright$ Eq. 3 9: **end for** 10: $\tilde{r}_t \leftarrow r_t$ with all lines $\ell$ satisfying $\hat{y}_{\ell} = 0$ removed 11: Substitute $\tilde{r}_t$ for $r_t$ when forming $H_t$ for turn $t+1$ 12: **return** $\tilde{r}_t$

How does this differ from a naïve “keep‑if‑any‑token‑above‑$\tau$” filter?

The naïve filter would discard a line as soon as a single token falls below $\tau$, which can erase useful context when only a few tokens are low‑confidence. SWE‑Pruner Pro instead requires a majority of tokens in the line to be above $\tau$, preserving lines that contain at least one strong signal while still removing clearly irrelevant lines.

Hidden-State Extraction

We extract and prune hidden states in‑server, fixing correctness gaps and slashing payload size.

The pruning head needs the backbone’s per‑token hidden states, but the original SGLang path omitted several safety guards and shipped tensors via a JSON field that balloons to gigabytes.

We pull the last‑layer hidden vector for each token directly from the backbone’s forward pass, then hand it to the pruner without any extra model‑side computation.

Step 1: The matrix is $[[h_{11},\dots,h_{18}],\dots,[h_{41},\dots,h_{48}]]$.

Step 2: The pruner computes a keep‑probability $p_i$ for each token (e.g., $[0.9, 0.2, 0.8, 0.1]$).

Step 3: Tokens with $p_i<0.5$ are dropped, leaving tokens 1 and 3.

Step 4: The retained hidden vectors are concatenated into a $2\times8$ matrix for downstream layers.

The extraction works on the raw tensor; the pruning decision is made solely from these vectors, so any token‑level redundancy is eliminated before the next model component sees it.

How does hidden‑state extraction differ from the log‑prob path that already returns token scores?

The log‑prob path only returns scalar scores per token, while hidden‑state extraction returns the full $[T, H]$ vector. The vector encodes richer contextual information that the pruner can use to decide which tokens are redundant, something a scalar score cannot capture.

Three correctness gaps prevented the hidden‑state path from mirroring the log‑prob path: batch alignment, chunked‑prefill accumulation, and prefix‑cache exemption. We fix each by pre‑allocating slots, concatenating per‑chunk tensors, and threading a `hidden_states_start_len` cap through the scheduler. After these patches the hidden‑state tensor length always matches the request batch and covers the full token span.

**Table 6.** Validation of patched SGLang hidden states against a pure-transformers `flash_attention_2` reference.

**Table 7.** Per-request hidden-state payload by serialization format.

Binary envelope construction for hidden‑state transfer.

Finally we colocate the pruning head inside the SGLang scheduler, eliminating the cross‑process tensor transfer. This in‑engine placement cuts the per‑call overhead from 19.3 % to 15.0 % while preserving the same pruning quality.

Qualitative Analysis

We examine how the SWE‑Pruner Pro head scores each line of tool output and what it prunes.

Across the four sampled tool families the head’s per‑line sigmoid scores cluster in a narrow 0.4–0.7 band, producing a fine‑grained ranking rather than a hard binary threshold.

**Figure 7** Read case: view /testbed/`snapshot_dbg_cli`/__init__.py. Legend (shared by all four cases). The left score gutter encodes the head’s predicted keep probability $p_i = \sigma(z_i)$, fading from white (low) to green (high). G is filled when the gold annotator marked the line as keep; P is filled when the head predicts keep at $\tau=0.5$. Predicted-keep rows additionally get a yellow row wash with crisp black text; pruned rows are dimmed. Case. The head assigns the lowest scores in the response (0.25–0.48) to the license/docstring boilerplate—the omitted lines 2–13 and the visible # limitations under the License plus the package docstring on lines 14–19—while concentrating mass at 0.51–0.66 on the imports, the __version__ assignment, and the two function bodies (main, `run_main`), exactly the symbols the agent will reference when wiring its patch. The path-header line 1 is also kept, anchoring the file identity for downstream turns.

**Figure 8** Search case: `cat -n /app/flask_1.1.1/src/flask/templating.py | grep -A 25 'def render_template_string'`. The head allocates the highest scores to the function's executable body—the `_app_ctx_stack` lookup, the `update_template_context` call, and the final `_render(...)` return—which together fully specify how a string template is dispatched. The intervening docstring lines are conservatively pruned even though the annotator keeps them; the head prioritises the call sequence over the prose as the authoritative spec.

**Figure 9** Listing case: `ls -la /home/legacy_admin/configs/{cron,db,secrets}/`. The head keeps the issued command and every actual file (`backup_schedule`.cron, database.conf, `api_keys`.env), which are the only artefacts the agent will need to act on. The remaining eight over-keeps are concentrated on group boundaries rather than scattered noise: the three directory-path headers (/cron/:, /db/:, /secrets/:), two of the three total 12 summaries, two leading . entries, and the closing shell prompt. The head still cleanly prunes the redundant .. parents, blank separators, and the third total 12, so its conservative bias on this listing costs structural boilerplate, not signal.

**Figure 10.** Test case: python `test_fix`.py (lollipop-jsonschema). The head locks onto the outer frames of the stack trace: the Traceback header, the &lt;module&gt; and `test_list` call sites, the public `json_schema` entry, the offending set(`any_of_validators`[0].choices) expression, and the terminal TypeError: unhashable type: 'list' (the highest-scoring line in the response at 0.71). The intervening inner-dispatch frame (js = _json_schema(...) on line 78 of the same file) is conservatively pruned, since it only forwards into the line that already carries the highest keep score. The four passing lines (String OK, Integer OK, Float OK, Boolean OK) are kept too, providing the contrast that localises the bug to the list path. The trailing harness footer (cwd, exit-code banner) prunes cleanly.

When evaluating with F1, a head can achieve high scores by keeping a tiny, high‑confidence subset of lines, yet those lines may be useless for the agent’s next action.

**Figure 11.** F1 can rank a useless head above a useful one. A 52-line view of pdm/resolver/providers.py; the agent is investigating a recursion bug in `register_provider` (L50). Per-sample balanced focal (PSBF, green) keeps 30 lines covering imports, the `TYPE_CHECKING` block, the provider registry, and the bodies of `get_provider`, `provider_arguments`, and `register_provider`; BCE (red) keeps only 4 isolated lines (L31, L34, L38, L50), namely the variable and function signatures, with no bodies.

**Figure 12.** The same failure mode on a second case. A 34-line view of pandas/core/internals/managers.py showing the start of quantile(). FOCAL (green) keeps the function signature plus the docstring sections explaining reduction semantics and the comment block at the start of the body; BCE (red) keeps only the parameter list (L3–L10) as a contiguous high-precision block with no docstring and no body.

Training Data Composition

Describes the composition and sourcing of the training corpus for SWE‑Pruner Pro.

The training set aggregates agent trajectories from five public HuggingFace sources, mixing code‑modification rollouts with CLI‑style tasks to reflect real‑world tool‑output diversity.

**Figure.** Data collection pipeline. Raw trajectories from five HuggingFace sources are parsed and filtered, then sub-sampled to a 50k diverse pool by quality-diversity sampling (see Section 4.6), and finally passed through a small human review pass that results in the final training dataset.

Training Configuration

How the training corpus is filtered, labeled, and used to train the pruning head.

The raw trajectory data are first parsed into (history, `tool_call`, `tool_response`, `next_turn`) quadruples following the OpenAI tool‑calling schema. A sliding‑window history preserves complete (assistant, tool) turn pairs; any partial pair that does not fit is dropped rather than truncated, guaranteeing that tool calls are never split.

A light pre‑filter removes responses with fewer than two non‑blank lines, implausibly low code‑line ratios, or pure stack‑trace output, focusing the labeling budget on substantive tool responses.

The remaining pool is far larger than needed, so a greedy facility‑location selection (maximising diversity over language, repository, size bucket, and a log‑step‑length kernel) extracts a 50 k diverse subset; lazy evaluation yields the classic $(1-1/e)$ approximation guarantee at negligible cost.

Table 5 reports the language and category composition of the final training corpus, showing Python (≈ 40 %) and Bash‑style CLI content (≈ 14 %) dominate, together covering about 83 % of the data.

Each retained quadruple is labeled by Claude Sonnet 4.6 via the Anthropic API. The annotator sees the full history, the triggering tool call, the line‑numbered tool response (with internal line numbers stripped), and the next‑turn snippet, and returns a list of 1‑based line indices to keep, a short reasoning, and a confidence flag.

Approximately half of the inputs receive a non‑empty keep set; the rest are rejected as boilerplate. Samples where the annotator cannot identify any keepable lines are marked *skeleton* (≈ 17 % of the corpus), indicating the entire output should be retained unchanged.

During training the kept‑line set is expanded to per‑token binary labels: a token gets label 1 if its character span overlaps any kept line, otherwise 0. Tokens outside the \<`tool_response`\> span are masked out of the loss.

The tool‑response length distribution is heavy‑tailed (mean 76 lines, median 56, 99‑th percentile 294, max 465). The keep‑ratio (kept ÷ total lines) averages 0.32 with a median of 0.23, matching the target compression of ≈ 30 %.

Language recovery uses the `instance_id` for Multi‑SWE‑bench trajectories and source‑dataset metadata for the others; the final breakdown mirrors Table 5, with Python and CLI content together accounting for ≈ 83 % of tokens and six non‑Python languages contributing the remaining ≈ 17 %.

A random 10 % subset of the labeled corpus (≈ 625 trajectories, ≈ 2 260 tool responses, ≈ 155 k lines) is reserved for the linear‑probe analysis in §2. It is split 90/10 into train/eval at the trajectory level, yielding AUC 0.83 and best‑F1 0.63 on the held‑out split.

The pruning head is a per‑token feed‑forward MLP applied to the frozen backbone’s last‑layer hidden states $h_i$. It also receives a continuous length‑aware embedding $e(N)$ realized over eight log‑spaced line‑count buckets (0–2, 3–5, 6–10, 11–20, 21–50, 51–100, 101–200, >200) that is added to each hidden vector before the MLP.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers