PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2607.18213v1cs.CLcs.SE

TL;DR

Long-context coding agents accumulate redundant tool-output tokens, while existing pruners obtain pruning signals outside the agent. SWE-Pruner Pro reads line-level relevance from the backbone’s representations and consistently reduces token use across four benchmarks while preserving task quality, with savings reaching 39%.

  • Problem

    Coding agents accumulate redundant tool outputs, while existing pruners rely on external signals despite the backbone already processing those outputs.

  • Method

    SWE-Pruner Pro uses a lightweight head with length-aware embeddings to read line-level keep-or-prune signals directly from the agent backbone.

  • Results

    Across two open-weight backbones and four multi-turn benchmarks, SWE-Pruner Pro consistently reduces token use while keeping task quality close to the unpruned agent, reaching 39% savings.

  • Takeaways & Limitations

    The results indicate that backbone representations already encode tool-output relevance, enabling pruning without a separate scoring model or explicit query.

  • Takeaways & Limitations

    The evaluation covers only open-weight models, and broader coverage across programming languages is left to future work.

Abstract

from arXiv · show

Pruning long context for coding agents has been a vital technology for efficient context management. While existing context pruning methods such as SWE-Pruner realize this by attaching a separate code classifier, we find the agent itself encodes internal representations indicating the relevance of code context when reading tool output. Based on this finding, we propose SWE-Pruner Pro, which prunes tool outputs directly inside the agent. Concretely, a small head turns the agent's own internal representations into a keep-or-prune label for each line, with a length-aware embedding keyed to each tool output's line count. 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, with bounded inference overhead. Notably, on MiMo-V2-Flash SWE-Pruner Pro additionally raises the SWE-Bench Verified resolve rate by +3.8% and the long-context Oolong accuracy by +2.2 points.

1 Introduction

SWE-Pruner Pro reads line-level pruning signals from the coding agent’s own internal representations, avoiding an external scoring model or explicit goal-hint query. Across two open-weight backbones and four multi-turn benchmarks, it consistently reduces token use while keeping task quality close to the unpruned agent.

  • Introduction: Coding agents accumulate redundant tool outputs during multi-turn repository tasks, while general-purpose compression cannot adapt to evolving focus and task-specific pruning requires extra machinery.Prior task-specific pruning conditions on agent intent but uses a second scoring model and an explicit goal-hint query.
  • Introduction: 39% token savings on SWE-QA-Pro and 30% on long-context Oolong accompany task quality close to the unpruned agent, with MiMo-V2-Flash Oolong accuracy improving by +2.2 points.Across seven evaluated methods, SWE-Pruner Pro is the most consistent pruner and the only method reducing end-to-end token use in every SWE-QA and Oolong setting while preserving quality.
  • Introduction: SWE-Pruner Pro shows that the agent backbone’s internal representations encode line-level importance, enabling pruning without a separate scoring model or explicit query.A probing study found kept and pruned lines distinguishable in representation space, even with a simple linear probe.
  • Introduction: A lightweight head reuses frozen-backbone token representations during normal prefill to predict which tool-output lines to keep or prune.The method adds length-aware embeddings so predictions can vary with output length and uses a per-sample balanced focal loss to rebalance keep and prune tokens.

2 Motivation

Tool outputs dominate multi-turn coding agents’ token budgets, while existing compressors reconstruct the agent’s information need externally. The motivation for SWE-Pruner Pro is that this keep-or-prune signal may already be encoded in the agent’s last-layer hidden states, where a linear probe achieves strong held-out discrimination.

  • Motivation: Over 70% of Mini-SWE-Agent’s SWE-Bench Verified tokens are consumed by file-reading commands, and this cost compounds as earlier reads persist across turns.A comparable pattern is observed on GLM-4.6.
  • Motivation: Existing compressors either score tokens with fixed surrogates or use a separate model conditioned on a goal-hint query, reconstructing the agent’s information need externally.These approaches add machinery on top of the agent rather than reading its current information need directly from the backbone.
  • Motivation: The hypothesis is tested on approximately 2,260 multi-turn tool responses comprising approximately 155,000 lines, with each line labeled keep or prune using Claude Sonnet 4.6.The responses come from publicly released SWE-Bench-style and terminal-task datasets.
  • Motivation: AUC 0.83 and best-F1 0.63 show that last-layer hidden states predict whether tool-output lines should be kept or pruned, exceeding the majority-class F1 upper bound of 0.46.The probe’s class distributions have offset means but substantial overlap in the middle band.

3 Method

SWE-Pruner Pro makes line-level keep-or-prune decisions directly from the agent backbone’s hidden states during tool-response prefill. A length-aware nonlinear head aggregates token predictions into line decisions, while pruning occurs between turns and the frozen backbone is trained only through the head.

  • 3.1 Per-turn pipeline: SWE-Pruner Pro scores tokens with a keep-or-prune head and uses majority voting to produce line-level decisions, preserving code syntax while enabling fine-grained pruning.The head adds a learned embedding indexed by response line count, then applies LayerNorm and two Linear-GELU-Dropout blocks to produce token logits; lines predicted as prune are removed before the next turn.
  • 3.1 Per-turn pipeline: The method reads last-layer hidden states from the backbone’s existing prefill of each new tool response, avoiding an extra forward pass on that response.Only the new response tokens are forwarded because prior history and the tool call are already cached; the agent’s current generation still attends to the full response.
  • 3.2 Pruning head: The length-aware embedding conditions pruning on response length because the cost of mis-pruning differs sharply between short and long responses.The embedding is broadcast-added to every hidden state and zero-initialized, so the classifier begins at its length-agnostic limit.
  • 3.3 Training: Training uses 22,609 Claude Sonnet 4.6-annotated multi-turn samples, per-token labels expanded from line annotations, and a per-sample balanced focal loss.The loss averages keep and prune-token losses equally within each sample, protecting the minority class regardless of the sample’s keep rate; the backbone remains fully frozen and the head trains from cached features.

4 Experiments

The experiments evaluate SWE-Pruner Pro across four multi-turn benchmarks using two long-context open-weight MoE coding agents, with matched pruning comparisons, task-specific metrics, and controlled inference settings.

  • Benchmarks: Experiments span SWE-Bench Verified, SWE-QA, SWE-QA-Pro, and Oolong, covering patch generation, multi-turn question answering, executable environments, and long-context aggregation.The benchmarks contain 500, 144, 260, and 280 instances, respectively, and use either the standard Mini-SWE-Agent harness or a minimal bash-only agent.
  • Agent backbones: Two open-weight MoE backbones support the experiments: 309B-parameter MiMo-V2-Flash with 15B active parameters and 80B-parameter Qwen3-Coder-Next with 3B active parameters.Both models have 256K context windows and are designed or specialized for long-horizon coding-agent workloads.
  • Baselines: The study compares against No Pruning and six prior pruners, including SWE-Pruner as the closest prior task-specific method.Prior methods include LLMLingua2, Selective Context, sliding-window retrieval with bge-reranker-v2-m32, Self-Prune, and LongCodeZip.
  • Metrics and judges: Evaluation uses resolve rate, 1–10 LLM-judge scores, and 0–100 exact-match accuracy, while also tracking total prompt and completion tokens.GPT-5.4-mini judges SWE-QA and SWE-QA-Pro at temperature 0; Oolong uses a rule-based exact-match scorer.
  • Inference configuration: All pruners on a given backbone share decoding, harness, hardware, and rollout limits, isolating pruning as the varying factor.Trajectories are capped at 250 turns, and each backbone follows its official model-card decoding settings.

5 Results

Across read-only and code-modification benchmarks, SWE-Pruner Pro delivers substantial token savings while preserving or improving quality, with backbone-dependent tradeoffs in resolve rate and API calls. Ablations and replay analysis attribute these results to the length-aware embedding, balanced focal loss, and low-overhead in-engine implementation.

  • Read-only multi-turn benchmarks: SWE-Pruner Pro is the only pruner reducing tokens on every read-only benchmark cell, saving up to 39% on SWE-QA-Pro and 30% on Oolong.Four of six prior pruners inflate tokens on at least one cell, while LLMLingua2 reaches +190% on Oolong.
  • Read-only multi-turn benchmarks: 34.7%, 39.4%, and 13.9% token reductions coincide with quality changes of +0.02, +0.24, and −1.4 pp, making SWE-Pruner Pro the only method to preserve quality while reducing tokens substantially.With Qwen3-Coder-Next, other pruners except RAG degrade judge scores by 0.14–0.65, while RAG saves at most 6.9%.
  • Code-modification benchmark: +3.8% resolve rate on MiMo-V2-Flash SWE-Bench Verified accompanies roughly half the token overhead of SWE-Pruner’s +4.2% gain, while Qwen3-Coder-Next loses 1.2 points but saves 13.5% input tokens.On Qwen3-Coder-Next, SWE-Pruner Pro loses only 6 solves and has the most favorable resolve-rate/input-token degradation profile among pruners.
  • Ablations: +1.13 judge and +0.16 F1 make per-sample balanced focal the strongest loss choice over BCE, while Dice and Tversky match F1 but collapse on judge scores.The ablation uses a held-out judge set with n=100 and GPT-5.4-mini scoring on a 1–10 rubric.
  • Ablations: 6.86 to 7.08 judge improvement from the length-aware embedding occurs at essentially identical F1 by redistributing mistakes toward longer responses.The embedding preserves overall line-decision accuracy while reflecting the length-conditioned asymmetry targeted by Eq. 1.
  • Inference overhead: 15.0% aggregate wall-time overhead, with p50 = 14.7% and p95 = 34.8%, is achieved by reusing prefill and colocating the pruning head inside the inference engine.These measurements come from a 16-trajectory MiMo-V2-Flash replay with the Appendix E.3 in-engine head enabled.

6 Related Work

Related work frames code-context reduction around token-level pruning, retrieval-based shortening, and code-aware compression, while noting persistent context-length challenges for coding agents. SWE-Pruner Pro most closely follows SWE-Pruner’s line-level, benchmark-oriented setting but removes its separate scorer and explicit per-turn goal query.

  • Context Reduction Methods: Code-context reduction spans token-level pruners, retrieval-based shorteners, and code-aware methods that preserve program structure.Token-level methods rank tokens using self-information or perplexity, while retrieval-based methods replace verbatim content with retrieved or aggregated representations.
  • Limitations: Code-aware compression preserves program structure but is largely evaluated on single-round proxy tasks and may apply fixed policies regardless of agent intent.The cited work includes Zhang et al. (2022), Wang et al. (2024b), Shi et al. (2025a, 2026b), Hu et al. (2026), Zeng et al. (2025), and Yang et al. (2024a).
  • Closest Work: SWE-Pruner is the closest end-to-end coding-agent work, using line-level pruning conditioned on an explicit goal-hint query at each turn.SWE-Pruner Pro inherits SWE-Pruner’s line-level granularity and benchmark setting while discarding the separate scoring model and explicit query.
  • Motivation: Coding agents remain vulnerable to context overflow and quality degradation as codebase context grows, making context length a first-class engineering problem.These challenges persist despite current coding LLMs offering 128k-plus context windows.

7 Conclusion

SWE-Pruner Pro shows that coding agents’ internal representations encode line-level tool-output importance, enabling lightweight in-agent pruning without a separate scoring model or explicit goal-hint query. Its learned length-aware embedding and per-sample balanced focal loss support consistent results across two open-weight backbones and four multi-turn benchmarks.

  • SWE-Pruner Pro extracts line-level importance from a coding agent’s internal representations while reading tool outputs, eliminating separate scoring and goal-hint queries.A lightweight head shares the backbone’s prefill pass.
  • A learned length-aware embedding and per-sample balanced focal loss provide the main empirical lift.
  • Results are consistent across two open-weight backbones and four multi-turn benchmarks, with reduced token consumption.

Limitations

The evaluation covers only open-weight models and is Python-centric, although the method is designed to extend across backbones and tool-output modalities with limited retraining.

  • Model and language scope: Evaluation is limited to open-weight models, while applying the recipe to a new backbone requires retraining only its pruning head.SWE-Pruner Pro relies on the agent backbone’s exposed hidden states.
  • Model and language scope: Although the agent-task benchmarks are Python-centric, the pipeline is language-agnostic and quality preservation across SWE-QA and Oolong suggests transfer across code and natural-language tool outputs.The stated transfer requires only per-backbone head retraining for a new model.

Ethical Considerations · A TRAINING DATA SWE-Pruner Pro

SWE-Pruner Pro uses publicly released training trajectories without private repositories or proprietary codebases. Its pruning head only compresses redundant tool output, but deployment requires per-backbone validation because aggressive pruning may harm quality beyond benchmark coverage.

  • Ethical Considerations: Training trajectories come exclusively from publicly released datasets, with no private repositories or proprietary codebases used.
  • Ethical Considerations: The pruning head compresses redundant tool output without altering the agent’s generation or reasoning.
  • Ethical Considerations: Deployment should be validated separately for each backbone before safety-critical use, because aggressive pruning may degrade task quality beyond benchmark detection.

A Training Data

The training corpus combines agent trajectories from five publicly released HuggingFace datasets, covering both SWE-style code modification and heterogeneous CLI/shell tasks. Raw traces are standardized, filtered, diversified, line-labelled, and lightly human-reviewed to train the SWE-Pruner Pro head.

  • Corpus overview: Five HuggingFace datasets provide both SWE-style code-modification rollouts and heterogeneous CLI/shell agent traces for training.The non-code traces span chess, machine learning, cryptography, databases, and shell scheduling, exposing the head to diverse tool outputs.
  • Data construction: The construction pipeline parses trajectories into a uniform per-step schema, light-filters them, selects a diverse 50k pool, labels lines with Claude Sonnet 4.6, and human-reviews malformed labels.The pipeline uses quality-diversity selection before labelling and drops a few hundred malformed labels during review.
  • Pre-filtering and diversity sampling: Each sample is a (history, tool call, tool response) step, while each trajectory is one globally deduplicated instance_id.A per-source adaptor emits one (history, tool_call, tool_response, next_turn) quadruple per tool message and preserves complete assistant–tool pairs in a sliding history window.

B TRAINING DETAILS SWE-Pruner Pro · B Training Details

SWE-Pruner Pro is trained on Claude-labeled tool-response lines using full interaction context, with data distributions targeting roughly 30% retention. Its training details also define a trajectory-level probe split and a length-aware pruning head over frozen backbone states.

  • B TRAINING DETAILS SWE-Pruner Pro: About half of the input pool receives non-empty per-line labels from Claude Sonnet 4.6, which identifies 1-based lines to retain from the full interaction context.The labeller sees the history, triggering tool call, numbered tool response, and next-turn snippet, and emits reasoning plus a confidence flag.
  • B TRAINING DETAILS SWE-Pruner Pro: The labelled corpus targets approximately 30% compression, with a mean keep-ratio of 0.32 and median of 0.23.Tool responses average 76 lines, have a median of 56, span the 10th/90th/99th percentiles at 25/143/294 lines, and are capped at 465 lines.
  • B TRAINING DETAILS SWE-Pruner Pro: The training corpus records language and category composition alongside the tool-response data.Language coverage is recovered from instance_id for Multi-SWE-bench_trajs, whose identifiers encode the programming language.
  • B TRAINING DETAILS SWE-Pruner Pro: The §2 probe uses a random 10% trajectory-level subset containing approximately 625 trajectories, 2,260 tool responses, and 155k lines.Trajectories are split 90/10 into train and evaluation sets, preventing lines from the same trajectory appearing in both splits.
  • B TRAINING DETAILS SWE-Pruner Pro: On held-out trajectories, the linear probe achieves AUC 0.83 and best-F1 0.63, supporting the probe analysis as an existence proof for keep-or-prune information.These metrics are reported in Figure 2 and are computed on the held-out trajectory split.
  • B Training Details: The pruning head is a per-token feed-forward MLP applied to frozen backbone last-layer hidden states, augmented with a length-aware embedding of tool-output line count.The embedding uses 8 log-spaced n_lines buckets: 0–2, 3–5, 6–10, 11–20, 21–50, 51–100, 101–200, and >200.

C OOLONG BENCHMARK CONVERSION SWE-Pruner Pro … E.1 Correctness: closing the hidden-state vs. logprob asymmetry

SWE-Pruner Pro converts Oolong into a multi-turn CLI exploration benchmark and executes pruning per tool-response turn using cached backbone hidden states. Its implementation freezes the backbone, trains a lightweight line classifier, and restores hidden-state-path correctness guards to match the logprob path.

  • C OOLONG BENCHMARK CONVERSION SWE-Pruner Pro: The classifier is a LayerNorm followed by two Linear-GELU-Dropout blocks and a final single-logit projection, with hidden dimension resized to each backbone.The head consumes hidden states from the tool-response span and is zero-initialized so training begins at the length-agnostic limit.
  • C OOLONG BENCHMARK CONVERSION SWE-Pruner Pro: Training uses AdamW with a 3 × 10−5 peak learning rate, 5% linear warmup, cosine decay to 1.5 × 10−5, gradient clipping at 1.0, and seed 42.No backbone parameters are updated; the schedule keeps the effective learning rate within [1.5 × 10−5, 3 × 10−5].
  • C OOLONG BENCHMARK CONVERSION SWE-Pruner Pro: The objective is a per-sample class-balanced focal loss with γ = 2 and equal 0.5/0.5 keep/prune weighting.Samples containing only one class retain the surviving branch with weight 1 and omit the absent branch from the average.
  • C OOLONG BENCHMARK CONVERSION SWE-Pruner Pro: The backbone remains frozen while cached hidden states support efficient head-only training and loss or length-bias ablations.Features are extracted once per dataset–backbone pair and stored as memory-mapped files; a full 10-epoch pass over 22,609 samples takes approximately 15 minutes on one 8 × H200 node.
  • D Per-turn execution: SWE-Pruner Pro replaces each raw tool response with its pruned version before constructing the next turn’s history.The per-turn pipeline pre-fills the response while caching the prefix, adds a length-aware embedding keyed to the response’s line count, classifies lines, removes predicted-pruned lines, and substitutes the result into subsequent history.
  • E In-server hidden-state extraction: In-server extraction obtains last-layer hidden states from the backbone’s existing tool-response prefill through SGLang’s return-hidden-states path.The implementation identified alignment, chunked-prefill, and prefix-cache correctness gaps, plus impractical JSON transport for tensors far larger than typical text payloads.
  • E.1 Correctness: closing the hidden-state vs. logprob asymmetry: The correctness fixes restore guards from the logprob path for hidden-state extraction, covering mixed-batch alignment and chunked-prefill accumulation.Mixed batches previously risked IndexError because hidden-state outputs lacked per-request slots, while chunked prefill captured only the final chunk rather than the full prompt-crossing span.

E IN-SERVER HIDDEN-STATE EXTRACTION SWE-Pruner Pro … H Prompts

The engineering sections validate reliable hidden-state extraction, reduce serialization overhead, and show that in-engine head colocation lowers pruning cost. Qualitative analyses show length-aware relative scoring, metric disagreement, and prompts designed to preserve an actionable code skeleton.

  • E IN-SERVER HIDDEN-STATE EXTRACTION SWE-Pruner Pro: All 48 validation samples match hidden-state shapes exactly, with median per-token cosine 0.997 against a flash_attention_2 Transformers reference.The minimum cosine is 0.33, attributed to bf16 arithmetic differences across attention backends and GEMM reduction order rather than a residual correctness bug.
  • E.3 In-engine pruning head: 15.0% overhead replaces 19.3% when the 18M-parameter pruning head moves inside SGLang, with residual cost from prefix-cache-exempt prefill.In-engine colocation modifies the request schema and ties deployment to a specific head architecture, unlike the engine-agnostic off-engine path.
  • F.1 Read case; F.2 Search case; F.3 Listing case; F.4 Test case: The read, search, listing, and test cases keep 18/23, 6/12, 12/4, and 15/13 predicted/gold lines, respectively.The listing example preserves the command and actual files but over-keeps directory headers, summaries, dot entries, and the shell prompt.
  • F QUALITATIVE CASES: PER-LINE HEAD SCORES ON TOOL OUTPUTS SWE-Pruner Pro: Across four held-out tool-response cases, the head assigns compressed per-line scores in a ∼0.4–0.7 band and uses length-aware calibration to set response-specific keep rates.The cases cover cat, grep, ls, and test execution, illustrating relative ranking rather than a single global keep threshold.
  • G QUALITATIVE CASES: WHERE F1 DISAGREES WITH THE LLM JUDGE SWE-Pruner Pro: F1 can prefer unusable caller-only or parameter-only skeletons: BCE scores 0.53 versus PSBF 0.49, yet the judge scores them 2/10 versus 8/10 (∆= 6).F1 rewards precision on small kept sets, while the judge evaluates whether bodies, imports, docstrings, and other context support the agent’s next action.
  • H Prompts: The prompts define pruning as a readable skeleton that preserves directly used lines, structural boundaries, key definitions, and error-relevant output while removing unrelated or repetitive content.The trajectory labelling prompt also distinguishes confident selection from skeleton fallback when the agent’s next action is ambiguous, and requires outer line numbers.
  • H Prompts: The evaluation prompts score pruner outputs on recall and precision, while SWE-QA judges final answers on correctness, completeness, relevance, clarity, and reasoning.The architecture-ablation judge uses a 1–10 rubric based on whether the agent could proceed identically with a compact response.
Loading 2607.18213v1…