💬 NLP

LLM Reasoning as Trajectories: Step-Specific Representation Geometry and Correctness Signals

This paper characterizes LLM chain-of-thought reasoning as structured trajectories through representation space, revealing that correct and incorrect solutions diverge systematically at late stages to enable accuracy prediction and trajectory-based steering for inference-time correction.

Original authors: Lihao Sun, Hang Dong, Bo Qiao, Qingwei Lin, Dongmei Zhang, Saravan Rajmohan

Published 2026-04-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Lihao Sun, Hang Dong, Bo Qiao, Qingwei Lin, Dongmei Zhang, Saravan Rajmohan

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 a Large Language Model (like the one powering this chat) isn't just a magic text generator, but a hiker climbing a mountain.

For a long time, we thought this hiker just stumbled along a path, hoping to reach the summit (the correct answer) by luck. But this new research from Microsoft suggests something much more structured: The hiker is following a very specific, invisible trail through a landscape of "thoughts."

Here is the breakdown of their discovery, using simple analogies:

1. The "Step-by-Step" Trail Map

When a model solves a math problem using "Chain of Thought" (thinking out loud), it doesn't just jump to the answer. It takes steps: Step 1, Step 2, Step 3...

The researchers found that inside the model's brain, each of these steps happens in its own special "room" or "zone."

The Analogy: Imagine the model's mind is a giant hotel with many floors (layers).
- Step 1 always happens in the "Lobby."
- Step 2 happens in the "Kitchen."
- Step 3 happens in the "Gym."
- And so on.
The Discovery: As the model gets deeper into its thinking (going up the hotel floors), these rooms become more distinct and separated. You can clearly tell the model is in the "Kitchen" (Step 2) versus the "Gym" (Step 3). This structure exists even in models that haven't been trained to reason yet; training just makes the model move between these rooms faster and more efficiently.

2. The "Wrong Turn" Detector

Here is the most exciting part: We can tell if the hiker is going to get lost before they reach the bottom of the mountain.

The Analogy: Imagine two hikers starting a hike. For the first few miles, they walk side-by-side on the exact same path. They look identical.
The Discovery: As they get closer to the end, the "Correct Hiker" and the "Wrong Hiker" start to drift apart. The wrong hiker starts taking weird, jagged paths that don't look like the standard trail.
The Power: By watching how the hiker moves in the last few steps, the researchers built a detector that can predict if the final answer will be right or wrong with 87% accuracy. It's like a GPS that says, "You're still on the right path," or "Uh oh, you're veering off the trail, you're going to get lost."

3. The "Steering Wheel" for Thoughts

Because we can see these paths, we can now steer the model. This is called "Trajectory-Based Steering."

Fixing Mistakes: If the model starts to wander off the "Correct Trail," the researchers can gently nudge it back.
- Analogy: Imagine a self-driving car that sees it drifting toward a cliff. Instead of crashing, it gently turns the wheel back toward the center of the road. This works best on long, difficult problems where the model is most likely to get confused.
Controlling Speed: We can also tell the model to "hurry up" or "take your time."
- Analogy: There is a special "Finish Line Zone" in the model's brain. If you push the model's thoughts toward this zone, it finishes the problem quickly. If you push it away from that zone, it keeps thinking and writing more steps. This lets us control how long the model thinks without retraining it.

4. Why This Matters

Before this, trying to fix a model's reasoning was like trying to fix a broken watch by hitting it with a hammer (unconditional scaling). It often made things worse.

This new approach is like surgery.

Diagnosis: We watch the internal "GPS" to see if the model is about to make a mistake.
Intervention: Only if it's about to fail, we gently nudge it back onto the right path.
Result: We get better answers without breaking the things that were already working.

Summary

The paper argues that reasoning is a geometric journey.

Steps are distinct rooms in a building.
Correctness is determined by the path taken in the final stretch.
Control is possible by gently steering the model back to the "ideal path" when it starts to wander.

This gives us a new "lens" to understand, predict, and fix how AI thinks, moving us from guessing to precise navigation.

1. Problem Statement

Large Language Models (LLMs) generate text autoregressively, updating high-dimensional representations at each timestep. While Chain-of-Thought (CoT) prompting has improved reasoning capabilities, the internal geometric structure of these reasoning processes remains poorly understood. Specifically, it is unclear whether:

Reasoning steps form structured trajectories through the model's representation space.
These trajectories contain geometric signals that distinguish correct from incorrect reasoning before the final answer is emitted.
This geometric understanding can be leveraged to control reasoning behavior (e.g., correcting errors or adjusting length) without retraining.

2. Methodology

The authors characterize reasoning as a sequential geometric process, analyzing the trajectory of hidden states in representation space.

Datasets & Models:
- Tasks: GSM8K (math word problems), MATH-500, and MMLU.
- Models: Three variants of Llama 3.1 8B: Base (pretrained), Instruct (instruction-tuned), and R1-Distill (reasoning-distilled).
- Protocol: A fixed CoT format was used to explicitly mark reasoning steps (e.g., "Step 1:", "Step 2:") and the final answer.
Data Extraction:
- Activations were extracted immediately preceding each "Step" token ( $h_{t(\text{Step } k)-1}$ ) and the final answer marker ( $h_{t(\text{term})-1}$ ).
- These snapshots represent the model's internal state after completing a reasoning step but before transitioning to the next.
Analysis Techniques:
- Linear Probes: Trained to classify which step number an activation corresponds to, measuring linear separability of step-specific subspaces.
- Trajectory Distance: Calculated Euclidean and Cosine distances between consecutive steps to measure movement magnitude and direction changes.
- Correctness Prediction: Logistic regression classifiers trained on trajectory features (early vs. late steps) to predict final-answer correctness (ROC-AUC).
- Inference-Time Interventions:
  - Steering: Adding low-rank updates to hidden states to nudge trajectories toward an "ideal" path or a "termination" subspace.
  - Gating: Intervening only when a predictor flags an impending failure (error-targeted).

3. Key Contributions & Findings

A. Step-Specific Representation Geometry

Structured Subspaces: Reasoning steps occupy distinct, linearly separable subspaces in the representation space.
Depth-Dependent Separability: These subspaces become increasingly separable as the model goes deeper into its layers. Early steps are separable in shallow layers, while later steps require deeper layers.
Training Regime Effects:
- The step-specific structure exists in the Base model, indicating it is an inherent property of the architecture.
- Reasoning Distillation (R1-Distill) does not create new structures but accelerates convergence toward the "termination-related" subspace (the state before the final answer) at shallower layers.
Generalization: The linear structure of step-specific regions is largely shared across training regimes and transfers robustly to different datasets (MATH-500, MMLU) and free-form prompting styles, suggesting it captures reasoning progress rather than surface formatting.

B. Correctness Signals via Trajectory Divergence

Early vs. Late Dynamics: Correct and incorrect reasoning follow highly similar trajectories during early steps (Step 1–2).
Late-Stage Divergence: As reasoning progresses, correct and incorrect trajectories systematically diverge, particularly in the transitions from the second-to-last step to the final answer.
Predictive Power:
- Classifiers using late-step trajectory features achieve a ROC-AUC of 0.87 in predicting final correctness before the answer is emitted.
- Early-step features perform near chance (AUC ~0.63).
- Trajectory-based predictors significantly outperform baselines based on step count or logit-level features (e.g., entropy, token rank).

C. Trajectory-Based Interventions

The authors propose an inference-time framework to control reasoning based on these geometric insights:

Error-Targeted Correction:
- Instead of unconditional intervention (which often degrades correct reasoning), the system uses a predictor to flag potential failures.
- When a failure is predicted, low-rank steering nudges the current trajectory back toward an "ideal" trajectory derived from correct examples.
- Result: This yields consistent accuracy gains (up to +7.6% on 6-step problems) with high preservation rates (>97%) for correct trajectories.
Reasoning Length Control:
- By steering activations toward the "termination subspace," the model converges faster (shortening reasoning).
- By steering away, computation is prolonged.
- Result: Reasoning length can be modulated monotonically with minimal impact on accuracy for moderate steering strengths ( $|\alpha| \le 0.8$ ). Excessive steering causes "mode collapse" (repetitive loops).

4. Results Summary

Separability: Step-specific probes achieve >0.99 accuracy for Step 1 across all layers and models.
Prediction: Late-step trajectory features predict correctness with 0.87 AUC, outperforming step-count baselines (0.65) and logit-lens baselines (0.77).
Intervention Efficacy:
- Unconditional token injection (e.g., "Wait") often reduces accuracy by up to 36%.
- Gated interventions (only acting on predicted errors) convert these drops into gains, improving accuracy on long reasoning chains (6+ steps) by ~7.6%.
- Length control allows for monotonic adjustment of reasoning steps without significant accuracy loss.

5. Significance

This work establishes reasoning trajectories as a geometric lens for interpreting LLM behavior. Its significance lies in three areas:

Interpretability: It reveals that reasoning is not a black box but a structured geometric motion where steps occupy distinct, separable regions.
Predictability: It demonstrates that correctness is encoded in the path of the trajectory (how the model gets to the answer), not just the final state, allowing for early error detection.
Controllability: It provides a causal mechanism for inference-time intervention. By manipulating the trajectory geometry, researchers can correct errors and control reasoning length without retraining, offering a practical solution to the "overthinking" problem and improving the reliability of LLM reasoning.

The paper concludes that while the geometric organization of reasoning is inherent to the model architecture, reasoning training refines the speed and precision of convergence, and these geometric signals can be harnessed to build more robust and controllable reasoning systems.