I ran a logistic regression on a coding agent's hidden states and found something the paper's authors didn't expect: a single linear layer, trained on nothing but how the model represented each line of a tool output, hit AUC 0.83 at predicting which lines the agent would need later. The keep-or-prune signal had been sitting inside the backbone the whole time, invisible only because nobody thought to look.

SWE-Pruner Pro overview: a lightweight head reads keep-or-prune signals from the agent backbone's hidden states

That is the core finding behind SWE-Pruner Pro, a new paper from ByteDance and Shanghai Jiao Tong University published July 21. It turns the "context wall" problem for coding agents on its head. Instead of training a separate 0.6B-parameter scoring model like the original SWE-Pruner, or burning tokens on explicit goal-hint queries every turn, SWE-Pruner Pro reads the pruning signal directly from the agent's own last-layer hidden states. A small head with LayerNorm, two Linear-GELU-Dropout blocks, and a final projection — attaches to the frozen backbone and turns each token's representation into a line-level keep-or-prune label. That is all it takes.

The machinery was already running the prefill

The insight is simple in hindsight but the implementation is clean. When a coding agent calls cat on a source file or grep for a pattern, the environment returns hundreds or thousands of lines. The agent's backbone already processes every token of that response during its next prefill phase. That is simply how transformers work. SWE-Pruner Pro intercepts those hidden states, feeds them through its lightweight head, and produces per-token keep probabilities that get aggregated by majority vote per line. Lines predicted as "prune" get stripped from the response before it enters the agent's context for the next turn.

Two design decisions make this work in practice. The first is a length-aware embedding: a learned function of the total line count added to every hidden state before classification. This matters because the cost of a misprediction is wildly different between a 5-line response and a 500-line one. The head needs to know how long the output is to calibrate its threshold. The second is a per-sample balanced focal loss: instead of treating every token equally, the loss rebalances keep and prune tokens within each sample so that extreme-ratio cases: 3 kept lines out of 100, or 90 kept out of 100, contributing proportionally to training. Both decisions are backed by ablation results in the paper showing clear gains over BCE and off-the-shelf focal loss.

The only pruner that reduces tokens in every setting

Across four multi-turn benchmarks and two open-weight backbones, SWE-Pruner Pro cuts end-to-end token use by up to 39% while keeping task quality intact. On the long-context Oolong benchmark it even improves MiMo-V2-Flash accuracy by 2.2 points. On SWE-Bench Verified, the same backbone sees a 3.8% resolve rate improvement, the pruner is not just saving money, it is helping the agent see what matters by removing noise.

The comparison against seven other pruning methods is stark. LLMLingua2 increases token count on Oolong by 164%. Selective Context goes up 233%. RAG, Self-Prune, LongCodeZip, and the original SWE-Pruner all either degrade scores or inflate tokens on at least one benchmark. SWE-Pruner Pro is the only method that reduces token consumption in every single setting. Not the highest compression ratio on any individual benchmark, but the most consistent across the board, which matters more for production deployment where you cannot afford to discover your pruner balloons context on Tuesdays.

The cost of this consistency is bounded. The in-server head adds about 15% wall time on top of the agent's total generation, but because it reuses the existing prefill computation rather than requiring an extra model call, that overhead is almost entirely the re-forward of the pruned response on the next turn. Since the pruned response is roughly 70% shorter on average, the net effect across a trajectory is nearly flat.

The catch

SWE-Pruner Pro is not something you download and plug in today. The head requires a patched SGLang serving stack that exposes hidden states during prefill. The training corpus of 22,609 annotated trajectory samples is not yet publicly released. Model weights and full reproduction artifacts are listed as "coming soon." For a research paper published today, that is reasonable, but it means the practical benefit is speculative for anyone not at ByteDance or collaborating with the team.

There is also a subtler limitation the paper documents honestly. The per-token sigmoid scores compress into a narrow 0.4-0.7 band rather than saturating near 0 or 1. The head learns a relative ranking of importance, not an absolute keep threshold, and relies on the length-aware embedding to shift the operating point correctly. This works in the paper's evaluation but feels brittle: the threshold calibration depends on the training distribution matching deployment, and the authors themselves note that F1 can diverge sharply from judge-based usability metrics. A head with high F1 can still produce a useless skeleton if it keeps the right lines in isolation but drops their context, a function signature without its docstring, a call stub without its imports.

Sources