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

Lihao Sun; Hang Dong; Bo Qiao; Qingwei Lin; Dongmei Zhang; Saravan Rajmohan

Source-linked AI summary

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

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

arXiv:2604.05655v1Apr 7, 2026cs.CLcs.AIcs.LG

TL;DR

The paper asks whether chain-of-thought reasoning has a structured geometric organization that can reveal correctness and support control. It analyzes step-wise representation trajectories and finds shared, depth-organized subspaces, late-stage correctness signals, and trajectory-based interventions for correction and length control.

Problem
The paper investigates what structured trajectories through representation space reveal about reasoning correctness, training regimes, and control.
Method
The study analyzes activations preceding reasoning steps, derives ideal trajectories from correct solutions, and applies trajectory-based inference-time steering.
Results
Late-step trajectory features predict final-answer correctness before emission, with linear classifiers achieving ROC–AUC 0.87.
Takeaways & Limitations
Reasoning trajectories provide a geometric basis for interpreting, predicting, correcting, and controlling LLM reasoning behavior.
Takeaways & Limitations
The analysis is restricted to math tasks and the Llama 3.1 8B family, leaving generalization to other settings, scales, and architectures open.

Abstract

from arXiv · show

This work characterizes large language models' chain-of-thought generation as a structured trajectory through representation space. We show that mathematical reasoning traverses functionally ordered, step-specific subspaces that become increasingly separable with layer depth. This structure already exists in base models, while reasoning training primarily accelerates convergence toward termination-related subspaces rather than introducing new representational organization. While early reasoning steps follow similar trajectories, correct and incorrect solutions diverge systematically at late stages. This late-stage divergence enables mid-reasoning prediction of final-answer correctness with ROC-AUC up to 0.87. Furthermore, we introduce trajectory-based steering, an inference-time intervention framework that enables reasoning correction and length control based on derived ideal trajectories. Together, these results establish reasoning trajectories as a geometric lens for interpreting, predicting, and controlling LLM reasoning behavior.

1 Introduction

The paper models chain-of-thought reasoning as a trajectory through representation space, asking what step-wise geometry reveals about correctness, training, and control. It finds shared step-specific organization, late-stage correctness divergence, and trajectory-based interventions for correction and length control.

Core framing: Step-preceding activations form linearly separable, step-specific subspaces that organize progressively with layer depth.The analysis extracts activations immediately before explicit Step markers to represent transitions between reasoning states.
Training effects: Step-specific organization is already present in Base models, while reasoning distillation accelerates convergence toward termination-related regions at earlier layers.The underlying linear structure is largely shared across training regimes rather than newly introduced by reasoning training.
Correctness signals: Correct and incorrect solutions follow similar early trajectories but systematically diverge during late-step transitions.This divergence provides a representation-space signal before the final answer is emitted.
Correctness signals: Linear classifiers using late-step features predict final-answer correctness with ROC–AUC 0.87 before final-answer emission.The reported result identifies late trajectory geometry as an actionable mid-reasoning correctness signal.
Interventions: Trajectory-based steering uses ideal correct trajectories to support localized reasoning correction and control reasoning length.Steering toward termination-related regions shortens reasoning, while steering away prolongs intermediate computation.
Implications: The findings support reasoning trajectories as a geometric lens for interpreting, predicting, and influencing LLM reasoning behavior.The proposed applications include mid-reasoning failure detection, correction toward productive computation, and length modulation.

2 Related Work

Prior work studies chain-of-thought behavior, representation subspaces, and inference-time control, but unconditional interventions can degrade correct reasoning. This paper connects these lines through error-prediction gating and adaptive trajectory-based interventions.

LLM CoT and interpretability: Chain-of-thought research has examined accuracy, faithfulness, and behavioral limitations, while interpretability work identifies subspaces linked to functions and predicts behavior from internal signals.Linear decodability is used as a criterion for representation-level analysis.
Inference-time intervention: Inference-time interventions control model behavior without retraining through methods such as test-time scaling and activation steering.These approaches commonly elicit additional reasoning or modify activations during generation.
Motivation: Unconditional interventions can unintentionally degrade originally correct reasoning, motivating more targeted control.The paper frames error-prediction gating as a response to the overthinking drawback of applying interventions indiscriminately.
Contribution: The paper explores error-prediction gating and adaptive trajectory interventions that operate on evolving representations rather than injecting fixed tokens.This connects internal correctness signals with selective inference-time control.

3 Step-specific Representation Subspaces

The study tests whether reasoning steps occupy distinct representation regions and whether those regions persist across models and prompting formats. Step-specific geometry becomes more separable with depth, is largely shared across training regimes, and transfers beyond explicit Step formatting.

Datasets and models: The experiments use GSM8K and MATH-500 with explicit step markers, comparing Base, Instruct, and R1-Distill Llama 3.1 8B models under controlled generation.The setup holds dataset, prompting format, and decoding strategy constant when comparing training regimes.
Measurement: Figure 1 maps activations immediately before Step markers and tests step identity with t-SNE visualization and binary linear probes.These activations represent the state after one reasoning step and before transition to the next.
Step-specific separability: Step 1 probe accuracy exceeds 0.99 at every layer for all models, while Step 2 reaches this level by layer 2.Steps 3–5 can also approach ceiling accuracy for Instruct and Base after sufficient depth.
Step-specific separability: Step-specific activations become progressively less entangled with depth; for Step 5, layer-0 accuracy ranges from 0.58 to 0.77 and exceeds 0.90 only after layer 6.The pattern indicates that deeper layers disentangle later reasoning steps into more stable regions.
Training regimes reshape geometry: Reasoning training accelerates termination-subspace convergence without uniformly strengthening all step-specific structure.R1-Distill reaches final-answer-marker accuracy above 0.99 at layer 0, while some intermediate steps remain less separable but above 0.90.
Training regimes reshape geometry: Cross-model probe transfer accuracy consistently exceeds 0.90 for nearly all model pairs, indicating largely shared step-specific linear structure.Post-training primarily changes the depth at which the common organization becomes most salient.
Prompt-format robustness: Without formatting constraints, the Instruct model spontaneously uses Step X: formatting in 64.5% of responses.The remaining responses provide natural structural boundaries for activation extraction.
Prompt-format robustness: Probes trained on fixed-form activations transfer to freeform responses, reaching best-layer accuracies from 0.83 to 0.93 across steps and answer markers.Randomly shuffled step labels achieve only 0.59 ± 0.04 average accuracy, supporting a reasoning-progress interpretation rather than a formatting artifact.

4 Correctness in Trajectory Geometry

Correct and incorrect reasoning follow similar early trajectories but diverge systematically at late steps, enabling mid-reasoning correctness prediction and targeted interventions. The trajectory signal generalizes structurally across tasks, while correctness prediction remains task-sensitive and interventions have limited repair coverage.

Late-step divergence: Early reasoning trajectories are largely correctness-invariant, with no significant Step 1→2 distance differences under Euclidean or cosine metrics.Step 1→2 differences remain near zero and within the 95% confidence interval.
Late-step divergence: Late transitions diverge by correctness, reaching Euclidean ∆(I–C)=−13.39 and cosine ∆(I–C)=−0.06 from the last step to the answer marker.The preceding second-last-to-last transition shows Euclidean ∆(I–C)=−4.26 and cosine ∆(I–C)=−0.02.
Correctness prediction: Late-step trajectory features predict final-answer correctness with average AUC 0.83 across layers and peak AUC 0.87 at layer 29, versus 0.63 for early-step geometry.Step-difference features outperform raw activation features, indicating that how the model moves through representation space carries correctness information.
Correctness prediction: Trajectory features achieve AUC 0.852±0.039, outperforming step-count-only AUC 0.649±0.021 and logit-level best AUC 0.765 ± 0.027.Length-balanced resampling preserves AUC at 0.847± 0.006, indicating the signal is not driven by reasoning length.
Error-targeted intervention: Predictor-gated interventions act on only 12.3% of examples and yield gains of up to +35.4% relative to always-on interventions, which can reduce accuracy by up to 36.0%.The selective strategy limits collateral damage from unconditional intervention, though current methods repair only some flagged failures.
Cross-task generalization: Step-specific structure transfers across tasks, but correctness prediction drops from AUC 0.87 in-distribution to 0.73 on MATH-500 and 0.64 on MMLU.The findings distinguish robust structural geometry from task-sensitive correctness geometry shaped by error modes and solution-path diversity.

5 Trajectory-Based Inference-Time Interventions

The paper introduces trajectory-based interventions that compare evolving reasoning states with an ideal trajectory to correct errors and control reasoning length. These interventions are most effective for longer reasoning chains and can adjust length with limited accuracy impact at moderate strengths.

Correcting Deviating Reasoning Trajectories: Trajectory-based steering compares projected activations with an ideal path derived from correct GSM8K trajectories, using local and cumulative deviations to trigger low-rank updates.The ideal path is defined by step-wise mean activations and tolerance statistics; interventions move activations toward the corresponding mean along dominant principal directions.
Correcting Deviating Reasoning Trajectories: On six-step GSM8K problems, accuracy improves from 75.44% to 83.04% (+7.60%), while seven-step problems improve from 67.69% to 75.38% (+7.69%).These gains are accompanied by preservation rates of at least 97%.
Correcting Deviating Reasoning Trajectories: Trajectory-based interventions can correct erroneous reasoning while preserving previously correct trajectories through repeated, low-magnitude updates.The method evaluates divergence at every step, allowing it to address abrupt deviations and gradual drift.
Correcting Deviating Reasoning Trajectories: For problems with step count ≤5, the method produces near-zero changes, whereas longer, error-prone chains receive the largest benefits.This pattern is consistent with minimal early-step divergence between correct and incorrect solutions.
Controlling Reasoning Length: Termination-related steering pushes activations toward a termination direction to shorten reasoning or away from it to prolong intermediate computation.For moderate strengths (|α| ≤0.8), reasoning length changes approximately monotonically with around 1% accuracy change for |α| ≤0.4; beyond |α| ≳0.8, behavioral mode collapse occurs.

6 Discussion

The discussion presents reasoning as structured geometric motion whose step-specific organization is already present in Base models and becomes more separable with depth. Late-step geometry supports correctness prediction and inference-time control, while future work concerns faithfulness and finer-grained failure signatures.

Geometric organization: Intermediate reasoning states occupy step-specific regions that become increasingly linearly separable at deeper layers, with this organization already present in Base models.Reasoning distillation primarily reshapes the depth at which convergence occurs rather than introducing new representational structure.
Prediction and control: Late-step geometry provides an actionable signal for predicting final-answer correctness before answer emission and supports selective, error-targeted interventions.The discussion frames these interventions as a way to mitigate overthinking associated with unconditional test-time scaling.
Prediction and control: Reasoning trajectories can be causally manipulated to control reasoning length by steering activations toward or away from termination-related regions.The paper presents trajectories as a unifying abstraction for interpreting, predicting, and influencing reasoning behavior.
Future directions: Future work includes using internal step-activation divergence to study chain-of-thought faithfulness and identifying geometric signatures for specific reasoning failures.These directions are proposed as complements to behavioral faithfulness tests and as a basis for finer-grained failure taxonomies.

Limitations

The study’s conclusions are bounded by its task, model-family, scale, and ideal-trajectory assumptions. Generalization to other reasoning settings, larger or architecturally different models, and cases where correct training examples do not define a reliable reference trajectory remains unverified.

Scope: The trajectory structure is established on GSM8K, MATH-500, and MMLU, but its presence in open-ended reasoning, multi-hop QA, or program synthesis remains an open question.The evidence therefore does not establish the same geometric organization across these other settings.
Model scope: The analysis is restricted to the Llama 3.1 8B family, leaving behavior at larger scales and across architecturally distinct model families unverified.The paper notes that larger models could exhibit qualitatively different trajectory dynamics.
Intervention assumptions: Trajectory-based interventions assume that correct training examples provide a reliable ideal trajectory, but this assumption may break down outside the studied settings.The paper reports empirical support for the assumption in its evaluated settings while identifying its broader validity as a limitation.

Ethical Considerations

The study analyzes internal representations and inference-time interventions on standard public benchmarks, focusing on low-risk arithmetic reasoning tasks without human subjects or personal data.

Study scope: The research uses standard, publicly available benchmarks and focuses on low-risk arithmetic reasoning tasks.It does not involve deployment in real-world decision-making contexts.
Data and participants: The study does not involve human subjects or personal data.The paper reports no additional ethical concerns beyond those generally associated with interpretability and model-analysis research.

C Complete Linear Probe Results

Linear probes capture reasoning-progress information across formats, including freeform outputs without explicit Step markers, while Figure 6 shows late-stage geometric divergence for incorrect reasoning.

Figure 6 shows statistically significant geometric divergence between correct and incorrect reasoning in late steps.
Table 7 reports average accuracy across all 32 layers and the best single-layer accuracy with its corresponding layer for each format category.
Best-layer probe accuracies exceed 0.82 across all format categories, with most exceeding 0.85.
Non-Step X: examples achieve at least 0.84 best-layer accuracy for every step, indicating probes capture reasoning progress rather than surface-token identity.

F Experimental Details

The experiments use deterministic generation, hidden-state extraction, linear probes, correctness predictors, and trajectory-based steering evaluated for transfer, stability, cost, and computational overhead.

Setup: Generation uses deterministic greedy decoding, while hidden states are extracted in a second forward pass with 33 vectors per token position.Each vector has dimension 4,096, comprising the embedding output and 32 transformer-layer activations.
Probe evaluation: Linear probes use balanced logistic regression, while correctness predictors use single-layer logistic regression with regularization selected by 5-fold stratified cross-validation.
Probe evaluation: Cross-model step-specific probe transfer is averaged across all layers, and fixed-form probes are evaluated on freeform formats without retraining.
Trajectory steering: Trajectory steering estimates an ideal trajectory from correct activations, projects it with PCA, and applies rank-32 displacement updates toward step-wise means using held-out thresholds.
Cost and stability: Steering operates at step boundaries rather than every token, with experiments reporting three classifier seeds and approximately 32× lower computation than per-token vector addition.
Cost and stability: The prototype incurs a 1.38× wall-clock slowdown because output_hidden_states=True disables KV caching and reprocesses the growing sequence at every token.
Cost and stability: Table 8 reports highly stable linear-probe accuracies but modest correctness-predictor AUC variance, consistent with fewer incorrect examples.

H Step-Count Distributions Across Models

Reasoning-chain lengths differ substantially across model variants, but step-specific geometry remains transferable and is tied to reasoning progress rather than absolute step position.

Step-count distributions: Instruct produces shorter chains with median K=4, compared with median K=6 for Base and K=5 for R1-Distill.Instruct chains are concentrated in 3–5 steps.
Cross-model geometry: Despite different chain lengths, cross-model probe transfer accuracy consistently exceeds 0.90.
Layer-wise structure: t-SNE visualizations show step-structured organization in all models, with earlier and stronger separation in R1-Distilled and Instruct than in Base.
Layer-wise structure: Linear probes separate early reasoning steps from shallow layers, whereas later steps require deeper layers to become separable.