Technical Summary: Context Pruning for Coding Agents via Multi-Rubric Latent Reasoning

Problem Statement

LLM-powered coding agents (e.g., SWE-agent, OpenHands) face a critical "context wall." While state-of-the-art frameworks automate complex repository-level tasks, agents spend approximately 67–76% of their token budget simply reading files. This context bloat incurs prohibitive API costs and triggers "attention dilution" and the "lost-in-the-middle" phenomenon, where irrelevant code segments degrade reasoning precision.

Existing learned pruning methods, such as SWE-Pruner, frame context reduction as a structured sequence labeling task using a single-objective Conditional Random Field (CRF). However, this formulation creates a modeling bottleneck: it attempts to model all facets of code relevance (e.g., contiguous semantic spans and sparse structural support lines like imports or control-flow pairs) with a single transition matrix. Consequently, the model biases toward dense semantic blocks and frequently discards structurally necessary lines (e.g., imports, class headers, try/except blocks) that have low semantic scores but are essential for syntactic validity and agent reasoning. This leads to broken syntax, repeated agent retries, and ultimately higher token consumption than unpruned baselines.

Methodology: LaMR (Latent Multi-Rubric)

The authors propose LaMR, a structured pruning framework that decomposes code relevance into interpretable quality dimensions, modeled by dedicated CRF heads and fused via a query-adaptive gating network.

1. Latent Multi-Rubric Formulation

Instead of a single relevance score, LaMR decomposes retention reasons into $K$ latent rubrics. In practice, the authors use two dimensions:

• Semantic Evidence: Lines directly relevant to the query (topical relevance).
• Dependency Support: Lines structurally necessary to interpret the evidence (imports, scope headers, control-flow siblings like try/except).

Each dimension $k$ is modeled by a dedicated CRF head with its own learnable transition matrix $T^{(k)}$. This allows the model to learn distinct sequential dynamics: semantic relevance favors long contiguous blocks (high self-transition), while dependency support tolerates frequent alternation (scattered lines and paired structures).

2. Query-Adaptive Gating

A lightweight Mixture-of-Experts (MoE) gating network dynamically weights the per-rubric emissions based on the agent's query. For each token $t$, the gate produces weights $g_t \in \Delta^{K-1}$, and the fused emission is computed as:
$$e^{fused}_t = \sum_{k=1}^K g_{t,k} E_{t,k}$$
This allows the model to shift focus between semantic and structural signals depending on the specific task intent.

3. Rubric-Guided Labeling (AST-Based Denoising)

To supervise these dimensions without additional human annotation, the authors introduce Rubric-Guided Labeling. Using the existing binary keep/prune masks from the teacher LLM (SWE-Pruner's training corpus) and static program analysis (AST), they derive fine-grained rubric labels:

• Semantic Labels: Directly map the teacher's document-level relevance score to retained lines.
• Dependency Labels: Propagate scores from retained lines to their structural dependencies (imports, definitions, enclosing headers, and control-flow siblings) via AST edges.
  This process effectively "denoises" the teacher's labels, recovering structurally necessary lines that the teacher might have pruned based on low semantic scores.

4. Architecture and Inference

• Backbone: A 0.6B-parameter reranker (Qwen3-Reranker) with multi-layer feature fusion (concatenating hidden states from early, middle, and final layers).
• Training: The model is trained with a weighted sum of the main CRF loss (on fused emissions) and rubric-specific CRF losses.
• Inference: LaMR acts as a lightweight middleware. It processes code chunks, performs Viterbi decoding on the fused emissions to generate binary keep/prune decisions, and applies a final AST-aware post-processing pass. This pass repairs the pruned code by re-inserting missing scope headers, imports, and control-flow companions to ensure syntactic validity.

Key Contributions

1. Latent Multi-Rubric Formulation: Identifies and resolves the transition conflict in single-objective pruning by modeling code relevance through separate CRF heads for semantic and dependency dimensions.
2. Query-Adaptive Gating: Introduces an MoE gating network to dynamically balance semantic and structural signals based on the agent's query.
3. AST-Based Label Denoising: Develops a zero-cost pipeline to extract multi-dimensional labels from binary masks, correcting missing structural tokens in teacher data without additional annotation.
4. Superior Efficiency-Performance Tradeoff: Demonstrates that preserving self-contained evidence-support units (task-relevant lines + structural context) allows the pruned context to match or exceed full-context baselines while significantly reducing token usage.

Experimental Results

The authors evaluated LaMR on four benchmarks (SWE-Bench Verified, SWE-QA, LCC, LongCodeQA) using two backbone models (Claude Sonnet 4.5 and Opus 4.6).

• Multi-Turn Tasks (SWE-Bench, SWE-QA):

  • LaMR wins 12 out of 16 head-to-head comparisons against the single-objective SWE-Pruner.
  • It saves up to 31% more tokens on multi-turn tasks compared to unpruned baselines.
  • Crucially, while SWE-Pruner sometimes increases token usage with stronger models (due to structural failures causing retries), LaMR consistently reduces or limits overhead.
  • Performance drops are marginal (e.g., ≤0.6 percentage points on success rate), while efficiency gains are substantial.

• Single-Turn Tasks (LCC, LongCodeQA):

  • LaMR improves task quality metrics (Exact Match, Edit Similarity, Accuracy) across all compression constraints (4× and 8×).
  • Notable gains include +3.5 Exact Match on LCC at 8× compression and +1.80 Accuracy on LongCodeQA at 4×.
  • Token-level baselines (LLMLingua-2, Selective-Context) degrade sharply at high compression, whereas LaMR maintains stability through line-level granularity and structural preservation.

• Ablation Studies:

  • Removing the dependency rubric or the AST-aware repair significantly degrades performance, confirming that structural support is distinct from and complementary to semantic relevance.
  • Coarse retrieval (e.g., LongCodeZip) combined with LaMR's fine-grained pruning yields the best results, whereas stacking multiple line compressors hurts performance.

Significance and Claims

The paper claims that LaMR represents a shift from "deleting tokens" to "preserving self-contained evidence-support units." By explicitly separating semantic evidence from dependency support, the framework addresses the inherent tension in code relevance: a line may be irrelevant semantically but critical structurally.

The authors assert that LaMR provides a strictly better tradeoff between context compression and downstream task quality compared to prior methods. It enables coding agents to operate more efficiently without sacrificing the structural integrity required for complex reasoning, effectively filtering distracting noise while maintaining the "signal" necessary for successful task completion. The work suggests that future context management for agents should move beyond single-scalar relevance scores toward structured, multi-dimensional reasoning about code dependencies.