🤖 AI

The Shape of Wisdom: Decision Trajectories in Language Models

This paper analyzes 9,000 decision trajectories across three language models to reveal that correctness and stability are distinct properties, demonstrating that attention mechanisms primarily drive stable correct answers while specific text spans critically influence decision margins.

Original authors: Shailesh Rana

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shailesh Rana

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 watching a race where the finish line isn't just a single moment, but a long, winding path. Most people only look at who crosses the finish line first (the final answer). But this paper asks us to watch the entire race to see how the runners actually got there.

Here is the story of the paper, "The Shape of Wisdom," broken down into simple concepts:

1. The Race is a Journey, Not a Snapshot

Usually, we think of an AI model as a machine that reads a question and instantly spits out an answer (A, B, C, or D). This paper argues that's wrong. Inside the model, the answer is actually a journey that happens layer by layer, like climbing a staircase.

At the bottom of the stairs (the early layers), the model might be leaning toward the wrong answer. Halfway up, it might be confused, hovering right on the edge between two choices. By the top, it might finally decide on the right one. Or, it might decide on the right one early, then wobble and almost change its mind, only to snap back at the very end.

The authors call this path a "trajectory." They track how much the model "likes" the correct answer compared to the best wrong answer at every single step of the climb.

2. The Four Types of Runners

By watching these journeys, the researchers found that models don't just give "Right" or "Wrong" answers. They fall into four distinct personality types:

Stable-Correct: The model figures out the right answer early and confidently sticks with it all the way to the top. (The confident winner).
Stable-Wrong: The model picks the wrong answer early and confidently sticks with it. (The confident loser).
Unstable-Correct: The model eventually gets the right answer, but it was a shaky ride. It wobbled, almost changed its mind, and only settled on the right answer at the very last second. (The nervous winner).
Unstable-Wrong: The model picks the wrong answer, wobbles, almost changes its mind, but ends up wrong anyway. (The nervous loser).

The Big Surprise: The most common group in their study was Unstable-Correct. This means that even when the AI gets the answer right, it often hasn't "settled" on it until the very last moment. It's a correct answer that was barely hanging on.

3. What Pushes the Runner? (The Engine Room)

Once they knew how the models moved, they asked what was pushing them. They looked at two main "engines" inside the AI:

Attention: This is like the model's eyes, scanning the question to see what words matter.
MLP (Feed-Forward): This is like the model's internal thinking or memory, processing what it saw.

They found a funny split:

In the Stable-Correct cases, the Attention engine was the one consistently pushing the model toward the right answer, step-by-step.
The MLP engine, surprisingly, often pushed in the wrong direction or didn't help much in these stable cases.

4. The "Delete" Experiment

To prove what was important, the researchers played a game of "delete and see." They took the question text and deleted specific parts:

Deleting the "Evidence": When they removed the part of the text that actually supported the right answer, the model's confidence in the right answer dropped.
Deleting the "Distractors": When they removed the confusing or wrong parts of the text, the model's confidence in the right answer actually went up.

This proved that the model is genuinely reacting to the meaning of the words, not just guessing randomly.

5. The "Time Travel" Test

Finally, they did a weird experiment: they took a "nervous" model that was failing and tried to swap its internal "engine parts" with a "confident" model that was succeeding.

When they swapped the Attention parts, it helped a little.
When they swapped the MLP parts (the internal thinking), it made a huge difference, often turning a failing model into a winning one.

This suggests that while Attention helps the model stay on track step-by-step, the deep internal processing (MLP) is what really locks in the final decision.

The Main Takeaway

The paper concludes that being right isn't the same as being stable.

Just because an AI gives you the correct answer doesn't mean it "knows" it. It might be a fluke, a last-second panic, or a shaky guess. The "wisdom" of the model isn't just the final score; it's the smooth, confident path it took to get there. The best models are the ones that find the right answer early and stay there, rather than the ones that stumble their way to the finish line.

Technical Summary: The Shape of Wisdom: Decision Trajectories in Language Models

Problem Statement

Current evaluations of Large Language Models (LLMs) typically treat the output as a single endpoint: a final token selected at the last layer. This paper argues that this perspective obscures the internal decision-making process. In multiple-choice settings, a model's preference for an answer is not static; it evolves across the depth of the network. A model may prefer the correct answer early, switch to an incorrect one, hover near a decision boundary, or settle late. The paper posits that an answer should be understood as a trajectory through the network's layers rather than a final-layer event. The core challenge is to separate three distinct properties often conflated by endpoint accuracy: the current state of the decision, the direction of its motion, and its proximity to a decision flip.

Methodology

The study analyzes 9,000 model-prompt trajectories across three instruction-tuned 7–8B parameter models: Qwen2.5-7B-Instruct, Llama-3.1-8B-Instruct, and Mistral-7B-Instruct-v0.3. All models were evaluated on the same 3,000 four-choice MMLU prompts.

Observables and Metrics

The authors define three primary quantities tracked at every layer (from Layer 1 to the final layer):

State (Margin): The score of the correct option minus the score of the strongest incorrect competitor. A positive margin indicates the correct option leads; a negative margin indicates a lead by an incorrect option.
Motion (Drift): The change in the margin from one layer to the next.
Boundary Distance: The absolute value of the margin, representing the distance to a decision flip (where the margin crosses zero).

Trajectory Classification

Based on the final answer and the stability of the trajectory's tail (the final layers), the authors classify trajectories into four operational regimes:

Stable-Correct (SC): Final answer is correct, and the tail is stable.
Stable-Wrong (SW): Final answer is incorrect, and the tail is stable.
Unstable-Correct (UC): Final answer is correct, but the trajectory exhibits significant late-layer movement.
Unstable-Wrong (UW): Final answer is incorrect, with significant late-layer movement.

Mechanistic Analysis

To explain the motion within these trajectories, the authors utilized a traced subset of 600 prompts per model (150 per trajectory type).

Scalar Accounting: Stored traces recorded scalar summaries of how much Attention blocks and MLP blocks moved the margin at each layer.
Span Deletion: An intervention where specific spans of the prompt (labeled as evidence, distractor, or neutral) were removed to measure the causal effect on the answer margin.
Counterfactual Replay: Simulated removal and substitution experiments where component contributions (attention or MLP scalars) from a "failing" trajectory were replaced with those from a "stable-correct" source trajectory to observe changes in the final margin.

Key Results

1. Heterogeneity of Correctness

The most significant empirical finding is that Unstable-Correct (UC) is the largest trajectory group, comprising 42.1% of all model-prompt rows (3,790 trajectories). In contrast, Stable-Correct (SC) accounts for only 19.0% (1,711 trajectories).

This implies that a correct final answer does not necessarily mean the model has "settled" on that answer early; many correct answers are reached through persistent uncertainty and late-layer movement.
Endpoint accuracy conflates "settled correctness" with "fragile correctness."

2. Mechanistic Drivers of Motion

Linear Accounting: A linear combination of attention and MLP scalars successfully reconstructs layer-to-layer margin drift, achieving high $R^2$ scores (0.90 for Qwen, 0.76 for Llama, 0.91 for Mistral).
Stable-Correct Dynamics: In stable-correct trajectories, the average attention scalar is positive (0.14), indicating it pushes the correct answer ahead. Conversely, the average MLP scalar is negative (-0.22).
Stable-Wrong Dynamics: Both attention and MLP scalars are negative in stable-wrong trajectories.

3. Input Interventions

Span deletion experiments confirmed causal links between prompt content and margin movement:

Removing evidence spans caused the correct margin to drop (mean effect: +2.92 logit units).
Removing distractor spans caused the correct margin to rise (mean effect: -2.18 logit units).
Neutral spans showed negligible effects. This demonstrates that specific text segments causally drive the margin defined by the trajectory regimes.

4. Counterfactual Substitution

When substituting late-layer components from stable-correct trajectories into failing ones:

MLP substitution produced a larger average shift in the final margin (e.g., +5.75 in the "legacy first" protocol) compared to attention substitution (-0.31).
This suggests that while attention provides a consistent layer-by-layer push toward stability, the MLP carries a larger transferable "boost" for the final margin in this specific replay protocol.

Significance and Claims

The paper explicitly avoids claiming to discover a complete hidden-state circuit or a universal physical phase transition. Instead, its significance lies in providing a reproducible framework to distinguish between different types of model behavior:

Decoupling Accuracy and Stability: The work demonstrates that "getting the answer right" and "having the answer settled" are distinct properties. The prevalence of unstable-correct trajectories suggests that many correct outputs are fragile and remain near decision boundaries until the final layer.
Legibility of Decision Processes: By tracking state, motion, and boundary distance, the authors make the decision process legible without relying on complex circuit discovery.
Identifying Sources of Motion: The study identifies that attention is the clearest positive contributor to stable-correct motion layer-by-layer, while MLP components appear to drive larger simulated shifts in final margins during counterfactual interventions.
Causal Input Sensitivity: The results confirm that removing answer-supporting text or adding distractors moves the decision margin in predictable, interpretable directions.

The authors conclude that the "best behavior" in these models is not merely reaching the correct endpoint, but reaching it early and maintaining that position (settling). The provided metrics and interventions offer a way to measure which answers are settled, which are fragile, and which measured forces (attention vs. MLP, evidence vs. distractors) drive these outcomes.