🤖 AI

Repair the Amplifier, Not the Symptom: Stable World-Model Correction for Agent Rollouts

This paper introduces WM-SAR, a world-model corrector that outperforms traditional symptom-scanning approaches by identifying and repairing the specific causal subgraph responsible for amplifying errors in large agent planning graphs, thereby enabling efficient in-place stabilization without the computational cost of full-graph replanning.

Original authors: Xinyuan Song, Zekun Cai

Published 2026-07-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xinyuan Song, Zekun Cai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the captain of a massive spaceship (an AI agent) trying to navigate a complex galaxy. Your ship doesn't just move in a straight line; it builds a giant, branching map of its journey, with thousands of stops, tool uses, and checks along the way. This map is called a planning graph.

Sometimes, the ship makes a mistake. Maybe it picks the wrong fuel, misreads a star chart, or forgets a rule. In the past, when the ship messed up, the standard fix was to throw away the entire map and try to redraw the whole journey from scratch.

The authors of this paper say: "Stop doing that!"

They argue that redrawing the whole map is a waste of fuel (computational power) and time. It's like throwing away a whole book because you misspelled one word on page 50. Instead, you should find the specific broken part and fix just that.

Here is the problem: Where exactly is the broken part?

The Trap of "Visible Symptoms"

Most current repair methods act like a doctor who only treats the fever.

The Scenario: The ship's engine overheats (the visible symptom) because a tiny valve in the fuel line was installed backward three hours ago (the root cause).
The Old Way: The "Engineering" repair method scans the map, sees the hot engine, and says, "Ah, the engine is broken! Let's fix the engine." It patches the engine.
The Result: The engine cools down for a moment, but because the backward valve is still there, the engine overheats again on the next leg of the trip. The ship is still broken; you just treated the symptom, not the disease.

The New Solution: WM-SAR (The "Amplifier" Detector)

The authors propose a new method called WM-SAR. Instead of looking for the thing that looks the most broken, WM-SAR looks for the thing that makes the whole system unstable.

Think of it like a sound system with a microphone too close to a speaker.

The Symptom: The speakers are screeching loudly (high error).
The Old Fix: You turn down the volume on the speakers. The screech stops for a second, but as soon as you turn it back up, it screeches again because the microphone is still too close.
The WM-SAR Fix: This method asks, "What is amplifying the noise?" It identifies the feedback loop between the mic and the speaker. It doesn't just turn down the volume; it moves the microphone. Once the feedback loop is broken, the system becomes stable, even if the speakers were the loudest part.

How WM-SAR Works (The "Backward" Search)

Look for the Amplifier: Instead of scanning the map from start to finish, WM-SAR works backward from the failure. It asks: "Which small group of connected steps is responsible for repeating this error?"
Find the "Causal Subgraph": It isolates a tiny, connected chain of events (like the backward valve and the fuel pump) that, if fixed, stops the error from growing.
Send a Tiny Note to the AI: Instead of showing the AI the entire 10,000-step map (which confuses it and wastes memory), WM-SAR sends only that tiny, broken chain to the AI.
The AI Fixes It: The AI looks at this clean, focused slice of the problem, figures out the root cause, and fixes it.

Why This Matters

The paper tested this on simulated agent tasks (like software debugging and web navigation). They found that:

Old methods often fixed the wrong thing or had to look at the whole map to get lucky.
WM-SAR fixed the problem using much less information (fewer "tokens" or words).
Stability: After WM-SAR fixed the problem, the agent didn't just survive the next step; it stayed stable for hundreds of future steps. The error didn't come back.

The Bottom Line

The paper's main message is: Don't just repair the symptom; repair the amplifier.

If you want to fix a complex AI agent, don't just patch the part that is currently screaming. Find the small, hidden connection that is making the whole system scream, fix that, and the rest of the system will calm down on its own. This saves time, saves money, and makes the AI much more reliable.

Technical Summary: Repair the Amplifier, Not the Symptom

Problem Formulation
As agent planning evolves from short tool chains to persistent workflows involving thousands of steps, failures increasingly occur within large planning graphs rather than isolated predictions. Replanning the entire graph after every error is computationally prohibitive and inefficient; it consumes excessive context budgets, exposes Large Language Models (LLMs) to irrelevant symptoms, and risks degrading long-context retrieval performance.

The core problem addressed is World-Model Correction: identifying a minimal, connected subgraph within a failed planning graph whose repair stabilizes the remaining system. The challenge lies in distinguishing between visible symptoms (nodes with high immediate error) and the underlying "amplifiers" (subgraphs that re-amplify errors in subsequent rollouts). Engineering approaches that simply scan for high-error nodes or expand local neighborhoods often fail because they patch symptoms while leaving the causal amplification path active, leading to error recurrence.

Methodology: WM-SAR
The paper proposes WM-SAR (World-Model Subgraph Amplification Repair), a corrector that works backward from error amplification rather than forward from visible symptoms. The method operates in two phases:

Graph-Based Region Selection (Phase 1):
- Theoretical Basis: The approach is grounded in the observation that rollout error growth is controlled by the spectral radius ( $\rho(B)$ ) of a coupled node-edge operator. If $\rho(B) > 1$ , errors amplify; if $\rho(B) \le 1$ , they decay. The goal is to find a subgraph $R$ such that removing it minimizes the residual spectral radius $\rho(B_{G_f \setminus R})$ .
- Scoring: WM-SAR calculates a Graph Error Amplification Field (GEAF) score for each node, defined as the product of the observed error magnitude and the local spectral radius of the node's neighborhood. It also computes a coupling factor ( $\kappa_v$ ) to measure super-additive interactions between node and edge channels.
- Greedy Growth: Starting from high-scoring seeds, the algorithm greedily expands the repair region by adding boundary nodes that maximize the marginal spectral relief ( $\Delta\rho_{rel}$ ), i.e., the reduction in the residual spectral radius. This process continues until a budget constraint is met or no further relief is possible.
- Pruning: Redundant nodes that do not contribute to spectral relief are removed to ensure the region is compact.
LLM-Based Repair (Phase 2):
- Once the optimal subgraph $R^*$ is identified, only this connected region is serialized and passed to an LLM. The LLM is tasked with identifying the root cause and producing a structured repair decision. This contrasts with full-graph approaches that serialize the entire failed trace, wasting tokens on irrelevant downstream symptoms.

Key Contributions

Problem Formulation: Reframes world-model correction as a constrained optimization problem: finding the smallest connected region that stabilizes the residual world model, rather than merely re-planning the full graph or patching the largest error.
Engineering Baselines: Implements and evaluates a strong family of engineering correctors (including Top-K, sliding windows, k-hop neighborhoods, and cascade scans) to establish a rigorous baseline.
WM-SAR Algorithm: Introduces a spectral-graph approach that uses GEAF and node-edge coupling to identify the "amplifier" subgraph. It provides a theoretical proof that residual spectral radius controls planning regret.
Empirical Evidence: Demonstrates that WM-SAR achieves near-whole-graph stabilization using a fraction of the tokens required by full-graph LLM repair or large neighborhood scans.

Results
Experiments were conducted on a synthetic heterogeneous calling-tree testbed mimicking planner-executor-validator workflows, as well as on topologies inspired by SWE-bench, WebArena, and AgentBench.

Stability vs. Symptom Patching: Engineering correctors (e.g., Greedy-Point, Window-based) successfully reduce immediate error but often fail to flatten the long-horizon error curve because they leave the amplification channel active. WM-SAR, by targeting the residual spectral radius, flattens the error curve, preventing error re-growth.
Token Efficiency: WM-SAR achieves comparable or superior root-cause recovery to full-graph LLM repair while using significantly fewer tokens (e.g., ~855 tokens vs. ~2157+ tokens for full-graph prompts in GPT-4o-mini experiments).
Budget Sensitivity: Under tight repair budgets, WM-SAR reaches high spectral relief with small region sizes (e.g., ~8 nodes), whereas engineering methods require much larger regions (often approaching whole-graph coverage) to achieve similar stability.
Cross-Model Robustness: The performance advantage of WM-SAR holds across different LLM tiers (GPT-4o, GPT-4o-mini, Gemini-2.5-Flash), indicating the benefit stems from better context selection rather than specific model quirks.
Limitations: The method relies on observable per-node error signals. In black-box attribution scenarios (e.g., the "Who&When" benchmark) where only the final answer is known, WM-SAR's performance degrades as the GEAF score cannot be computed, reverting to behavior similar to simple sink-node selection.

Significance and Claims
The paper argues that the primary failure mode of current agent repair strategies is the misalignment between the repair target (visible symptoms) and the stability objective (residual amplification). By shifting the focus to "repairing the amplifier," WM-SAR enables agents to maintain stable world models over long horizons without incurring the prohibitive costs of full-graph replanning.

The authors claim that this approach offers a principled alternative to heuristic scanning, providing a "cleaner" repair target for LLMs that preserves the causal bridge while omitting distracting downstream noise. The work suggests that future agent systems should prioritize identifying and repairing the structural operators that propagate error, rather than simply reacting to the most visible failures.