Semantic Step Prediction: Multi-Step Latent Forecasting in LLM Reasoning Trajectories via Step Sampling

# Multi-Step Latent Forecasting in LLM Reasoning Trajectories via Step Sampling
Source: [https://arxiv.org/html/2604.18464](https://arxiv.org/html/2604.18464)
###### Abstract

\[This arXiv paper is entirely derived from work conducted as part of the NUS CS52605260course\.\] Semantic Tube Prediction \(STP\) leverages representation geometric to regularize LLM hidden\-state trajectories toward locally linear geodesics during fine\-tuning, thereby greatly improving data efficiency\. The original STP recipe samples random token sub\-spans, which is compatible with the base large language model \(LLM\) training architecture\. Inspired by STP, we are interested to investigate whether the sampling position can further enhance the semantic structure of multi\-step reasoning, and hence affect its geometric impact\. We applied STP at consecutive semantic reasoning step boundaries and achieved168×168\\timesmore accurate multi\-step latent prediction than frozen baselines on ProcessBench \(3,400 samples\), compared to only4×4\\timesfor the random\-token STP\. Probing the latent manifold with a learned non\-linear predictor reveals that STP\-shaped trajectories are smooth curves, not straight lines: a 3\-layer MLP reduces prediction error by a further33–12×12\\timesover linear extrapolation on step\-boundary models\. Removing the language modeling loss yields trajectories that are2×2\\timesmore MLP\-predictable than the combined loss, revealing a tradeoff between generation quality and geometric purity\. Our results identify sampling position as the critical variable in geometric regularization and establish multi\-step latent prediction MSE as a new evaluation metric for this class of methods\.

## 1Introduction

Large Language models \(LLMs\) reason by generating step\-by\-step solutions, producing many tokens per reasoning step\. Each token requires a full autoregressive forward pass, making multi\-step reasoning computationally expensive\. A growing body of work explores latent reasoning where computation is performed in the model’s continuous hidden\-state space rather than decoding discrete tokens at every step\(Hao et al\.,[2024](https://arxiv.org/html/2604.18464#bib.bib4); PLaT,[2026](https://arxiv.org/html/2604.18464#bib.bib15)\)\. For latent reasoning to be feasible, the model’s hidden\-state trajectory must be predictable, which is given the current reasoning state, future states must be forecastable with reasonable accuracy\.

Semantic Tube Prediction \(STP\)\(Huang et al\.,[2026](https://arxiv.org/html/2604.18464#bib.bib8)\)provides a theoretical and practical framework for this\. Building on the Geodesic Hypothesis, which insists that error\-free token sequences trace locally linear geodesics on a smooth semantic manifold, STP adds a supplementary cosine\-displacement loss during fine\-tuning that forces consecutive hidden\-state displacement vectors to be parallel\. The resulting “Semantic Tube” confines trajectories to a tubular neighborhood of the geodesic, improving signal\-to\-noise ratio and enabling16×16\\timesdata efficiency gains\.

STP samples its loss at random token sub\-spans, where three indicess<r<ts<r<tare drawn uniformly from the token sequence\. While this method is elegant, we would like to explore if the semantic structure of the reasoning process, which is the natural boundaries where the model transitions from one reasoning step to the next, can contribute to the geometric smoothness of the method\. We hypothesize that aligning the geometric regularization with the semantic structure of reasoning at step boundaries rather than at arbitrary token positions will produce more predictable trajectories\. We test whether step\-boundary STP creates the latent\-space conditions that future JEPA\(LeCun,[2022](https://arxiv.org/html/2604.18464#bib.bib9)\)or COCONUT\-style systems would need, as measured via post\-hoc multi\-step prediction accuracy of both linear and learned predictors\.

#### Our key contributions\.

1. 1\.Sampling position is the critical variable\.Applying STP at consecutive step boundaries achieves168×168\\timesprediction improvement vs\.4×4\\timesfor random\-token STP—a40×40\\timesgap from where the loss is sampled, not what the loss computes \(§[5\.1](https://arxiv.org/html/2604.18464#S5.SS1)\)\.
2. 2\.Multi\-step latent prediction MSE as evaluation metric\.We introduce and validate this metric with decoding fidelity tests: predicted embeddings decode to the correct next token 93\.4% of the time \(§[5\.4](https://arxiv.org/html/2604.18464#S5.SS4)\)\.
3. 3\.Trajectories are smooth curves, not straight lines\.A learned MLP predictor reduces error by33–12×12\\timesover linear on STP\-shaped models but finds no systematic structure on baselines, revealing three distinct manifold regimes \(§[6\.1](https://arxiv.org/html/2604.18464#S6.SS1)\)\.
4. 4\.ℒNTP\\mathcal\{L\}\_\{\\text\{NTP\}\}creates a tradeoff\.Removing the language modeling loss yields2×2\\timesbetter MLP prediction at the cost of 3\.7 pp GSM8K accuracy—generation quality vs\. geometric purity \(§[5\.1](https://arxiv.org/html/2604.18464#S5.SS1), §[6\.1](https://arxiv.org/html/2604.18464#S6.SS1)\)\.
5. 5\.Negative finding\.Geometric smoothness does not encode step correctness \(AUC≈0\.5\\approx 0\.5; §[A\.2](https://arxiv.org/html/2604.18464#A1.SS2)\): the Semantic Tube captures organized thinking, not correct thinking\.

## 2Background

### 2\.1Semantic Tube Prediction

Huang et al\. \([2026](https://arxiv.org/html/2604.18464#bib.bib8)\)propose the Geodesic Hypothesis, which describes that token sequences generated by LLMs trace geodesics on a smooth semantic manifold, and deviations from these geodesics represent noise\. They formalize this as the STP loss:

ℒSTP=𝔼s<r<t​\[1−cos⁡\(hr−hs,ht−hr\)\]\\mathcal\{L\}\_\{\\text\{STP\}\}=\\mathbb\{E\}\_\{s<r<t\}\\left\[1\-\\cos\(h\_\{r\}\-h\_\{s\},\\;h\_\{t\}\-h\_\{r\}\)\\right\]\(1\)wherehs,hr,hth\_\{s\},h\_\{r\},h\_\{t\}are hidden states at randomly sampled token positions\. MinimizingℒSTP\\mathcal\{L\}\_\{\\text\{STP\}\}forces consecutive displacement vectors to be parallel, confining trajectories to a tube around the geodesic\. Combined with the standard next\-token prediction loss, the training objective isℒ=ℒNTP\+λ⋅ℒSTP\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{NTP\}\}\+\\lambda\\cdot\\mathcal\{L\}\_\{\\text\{STP\}\}\.

With the predictability of a geometric constraint, STP is able to dramatically reduce the amount of data required for training\. Finetuning with16×16\\timesless training data across multiple model families\(Huang et al\.,[2026](https://arxiv.org/html/2604.18464#bib.bib8)\)resulted in comparable results with full\-data training\.

### 2\.2Joint Embedding Predictive Architectures

We draw motivation from the JEPA paradigm\(LeCun,[2022](https://arxiv.org/html/2604.18464#bib.bib9)\), which advocates predicting in latent space rather than reconstructing inputs\. I\-JEPA\(Assran et al\.,[2023](https://arxiv.org/html/2604.18464#bib.bib1)\)and V\-JEPA\(Bardes et al\.,[2024](https://arxiv.org/html/2604.18464#bib.bib2)\)demonstrate this for images and video; LLM\-JEPA\(Huang et al\.,[2025](https://arxiv.org/html/2604.18464#bib.bib7)\)extends it to language\. Interestinly, STP was interestingly introduced to solve the two\-view problem of LLM\-JEPA, which requires the text data to be paired with the code form of the same information\. We do not implement a JEPA architecture as our training lacks target encoder, learned predictor module, and continuous embedding feedback loop\. Instead, we test the preconditions that the latent space supports accurate multi\-step prediction\.

### 2\.3Geometric Analysis of LLM Hidden States

Several works analyze the geometry of LLM hidden states diagnostically\.Wang et al\. \([2025](https://arxiv.org/html/2604.18464#bib.bib10)\)measure curvature across layers;Zhou et al\. \([2026](https://arxiv.org/html/2604.18464#bib.bib13)\)apply Menger curvature to reasoning trajectories;Herrmann et al\. \([2025](https://arxiv.org/html/2604.18464#bib.bib5)\)introduce PHi loss and show that hidden\-state unpredictability correlates with correct solutions on difficult problems\.Jiang et al\. \([2026](https://arxiv.org/html/2604.18464#bib.bib18)\)decompose reasoning traces into Progress \(displacement magnitude\) and Stability \(curvature\) at the token level\.Sun et al\. \([2026](https://arxiv.org/html/2604.18464#bib.bib16)\)extract activations at step markers and find that correct and incorrect solutions diverge geometrically at late reasoning steps\.Carson & Reisizadeh \([2025](https://arxiv.org/html/2604.18464#bib.bib19)\)model reasoning trajectories as drift\-diffusion systems with regime switching\.Damirchi et al\. \([2026](https://arxiv.org/html/2604.18464#bib.bib3)\)find that smooth layer\-wise trajectories correlate with valid reasoning\.

### 2\.4Latent Reasoning

Latent reasoning to reduce or avoid token\-wise decoding has been explored in several recent works\. COCONUT\(Hao et al\.,[2024](https://arxiv.org/html/2604.18464#bib.bib4)\)feeds the last hidden state back as the next input embedding, enabling latent iteration without token decoding\. PLaT\(PLaT,[2026](https://arxiv.org/html/2604.18464#bib.bib15)\)decouples reasoning from verbalization entirely\. CoLaR\(Tao et al\.,[2025](https://arxiv.org/html/2604.18464#bib.bib22)\)trains a latent head with a next\-compressed\-embedding objective at random token positions\. LightThinker\(Zhang et al\.,[2025](https://arxiv.org/html/2604.18464#bib.bib11)\)compresses intermediate steps into gist tokens at step boundaries\.Teoh et al\. \([2025](https://arxiv.org/html/2604.18464#bib.bib17)\)train a next\-latent prediction head with Smooth L1 loss at every token, supporting multi\-step rollouts for world modeling\. At inference time, The Geometric Reasoner\(Zhuang et al\.,[2026](https://arxiv.org/html/2604.18464#bib.bib20)\)scores candidates via bumpiness penalties at chunk boundaries, and STEP\(Sun et al\.,[2026b](https://arxiv.org/html/2604.18464#bib.bib21)\)trains a classifier on step\-boundary hidden states to detect reasoning errors\.

## 3Method

### 3\.1STP at Semantic Step Boundaries

![Refer to caption](https://arxiv.org/html/2604.18464v1/SSP_concept.png)Figure 1:Semantic Step Prediction concept\. The token\-level trajectory \(pink\) oscillates around the step\-level geodesic \(green\)\. Step boundary positions \(numbered 1, 2, 3\) define the trajectory𝐳=\(z0,z1,…,zK\)\\mathbf\{z\}=\(z\_\{0\},z\_\{1\},\\ldots,z\_\{K\}\)on which the STP loss enforces consecutive displacement parallelism, confining the step\-level path to a smooth tube\.We insert a special<\|step\|\>delimiter token between reasoning steps in the training data \(Figure[1](https://arxiv.org/html/2604.18464#S3.F1)\), yielding sequences of the form:

\[question\]​<\|step\|\>​\[step1\]​<\|step\|\>​\[step2\]​<\|step\|\>​⋯​<\|step\|\>​\[stepK\]​<\|step\|\>\[\\text\{question\}\]\\;\\texttt\{<\|step\|\>\}\\;\[\\text\{step\}\_\{1\}\]\\;\\texttt\{<\|step\|\>\}\\;\[\\text\{step\}\_\{2\}\]\\;\\texttt\{<\|step\|\>\}\\;\\cdots\\;\\texttt\{<\|step\|\>\}\\;\[\\text\{step\}\_\{K\}\]\\;\\texttt\{<\|step\|\>\}\(2\)The hidden states at<\|step\|\>positions form a trajectory𝐳=\(z0,z1,…,zK\)\\mathbf\{z\}=\(z\_\{0\},z\_\{1\},\\ldots,z\_\{K\}\)\. We compute the STP loss on consecutive triples:

ℒSTPstep=1K−1​∑k=1K−1\(1−\(zk−zk−1\)⋅\(zk\+1−zk\)‖zk−zk−1‖⋅‖zk\+1−zk‖\+ϵ\)\\mathcal\{L\}\_\{\\text\{STP\}\}^\{\\text\{step\}\}=\\frac\{1\}\{K\-1\}\\sum\_\{k=1\}^\{K\-1\}\\left\(1\-\\frac\{\(z\_\{k\}\-z\_\{k\-1\}\)\\cdot\(z\_\{k\+1\}\-z\_\{k\}\)\}\{\\\|z\_\{k\}\-z\_\{k\-1\}\\\|\\cdot\\\|z\_\{k\+1\}\-z\_\{k\}\\\|\+\\epsilon\}\\right\)\(3\)This is mathematically identical to STP’s loss but differs in where the indices are sampled\. We sample at semantically meaningful step boundaries rather than at random token positions\. The combined objective isℒ=ℒNTP\+β⋅ℒSTPstep\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{NTP\}\}\+\\beta\\cdot\\mathcal\{L\}\_\{\\text\{STP\}\}^\{\\text\{step\}\}withβ=1\\beta=1\.

Fine\-tuning uses LoRA\(Hu et al\.,[2022](https://arxiv.org/html/2604.18464#bib.bib6)\)with rank 16 onq,k,v,oq,k,v,oprojections \(∼\\sim4\.4M trainable parameters on a 1\.5B base model\)\.

### 3\.2Experimental Grid

Table 1:Six\-model experimental grid\. Each model differs in loss composition and/or STP sampling strategy\.This grid \(Table[1](https://arxiv.org/html/2604.18464#S3.T1)\) isolates three factors: \(i\) token\-level vs\. step\-level sampling \(C vs\. A2\), \(ii\) random vs\. consecutive at the step level \(A2 vs\. A\), and \(iii\) necessity ofℒNTP\\mathcal\{L\}\_\{\\text\{NTP\}\}\(A vs\. A1\)\.

### 3\.3Multi\-Step Prediction MSE

For each trajectory\(z0,z1,…,zK\)\(z\_\{0\},z\_\{1\},\\ldots,z\_\{K\}\)and each valid positionk≥1k\\geq 1, we predictmmsteps ahead via linear extrapolation:

z^k\+m=zk\+m⋅\(zk−zk−1\)\\hat\{z\}\_\{k\+m\}=z\_\{k\}\+m\\cdot\(z\_\{k\}\-z\_\{k\-1\}\)\(4\)and compute the normalized prediction error:

MSEm=1N​∑\(k,sample\)‖z^k\+m−zk\+m‖2‖zk\+m‖2\\text\{MSE\}\_\{m\}=\\frac\{1\}\{N\}\\sum\_\{\(k,\\text\{sample\}\)\}\\frac\{\\\|\\hat\{z\}\_\{k\+m\}\-z\_\{k\+m\}\\\|^\{2\}\}\{\\\|z\_\{k\+m\}\\\|^\{2\}\}\(5\)MSEm≈0\\text\{MSE\}\_\{m\}\\approx 0means the trajectory is perfectly linear;MSEm≈1\\text\{MSE\}\_\{m\}\\approx 1means the prediction error equals the signal magnitude\.

### 3\.4Trajectory Smoothness Scores

We report the cosine score \(matching the training objective\):

cos\_scorek=1−\(zk−zk−1\)⋅\(zk\+1−zk\)‖zk−zk−1‖⋅‖zk\+1−zk‖∈\[0,2\]\\text\{cos\\\_score\}\_\{k\}=1\-\\frac\{\(z\_\{k\}\-z\_\{k\-1\}\)\\cdot\(z\_\{k\+1\}\-z\_\{k\}\)\}\{\\\|z\_\{k\}\-z\_\{k\-1\}\\\|\\cdot\\\|z\_\{k\+1\}\-z\_\{k\}\\\|\}\\in\[0,2\]\(6\)and the perpendicular score \(geometric interpretation\):

perp\_scorek=‖dk−\(dk⋅s^k\)​s^k‖‖dk‖=sin⁡\(θk\)∈\[0,1\]\\text\{perp\\\_score\}\_\{k\}=\\frac\{\\\|d\_\{k\}\-\(d\_\{k\}\\cdot\\hat\{s\}\_\{k\}\)\\hat\{s\}\_\{k\}\\\|\}\{\\\|d\_\{k\}\\\|\}=\\sin\(\\theta\_\{k\}\)\\in\[0,1\]\(7\)wheredk=zk−zk−1d\_\{k\}=z\_\{k\}\-z\_\{k\-1\}andsk=zk\+1−zk−1s\_\{k\}=z\_\{k\+1\}\-z\_\{k\-1\}\. The two scores are monotonically related and produce identical model orderings \(Spearmanρ=1\.0\\rho=1\.0\)\.

## 4Experiments

### 4\.1Setup

models\.Qwen2\.5\-Math\-1\.5B \(math\-specialized\), Qwen2\.5\-1\.5B \(general\-purpose\), Llama\-3\.2\-1B \(different architecture, 2048\-dim hidden states vs\. 1536\)\. All fine\-tuned via LoRA with rank 16\.

Training data\.6,132 MATH competition problems and solutions split at paragraph boundaries with<\|step\|\>delimiters\. 3 epochs \(∼\\sim1,150 optimizer steps per model\)\.

Evaluation data\.ProcessBench\(Zheng et al\.,[2024](https://arxiv.org/html/2604.18464#bib.bib12)\): 3,400 step\-by\-step mathematical solutions across four difficulty levels with human\-annotated error positions\.

## 5Results

### 5\.1Prediction MSE

Table 2:Six\-model ablation on ProcessBench with Qwen2\.5\-Math\-1\.5B\. Mean cosine score \(↓\\downarrow, matches training objective\), mean perpendicular score \(↓\\downarrow\), and multi\-step prediction MSE \(↓\\downarrow\)\. Improvement factor computed atm=1m\{=\}1relative to B1\.![Refer to caption](https://arxiv.org/html/2604.18464v1/mse_and_accuracy.png)Figure 2:Multi\-step prediction MSE \(solid, left axis, log scale\) and prediction accuracy \(dashed, right axis\) vs\. skip distancemm\. Step\-boundary models \(A, A1, A2\) achieve orders\-of\-magnitude better prediction than frozen \(B1\) and random\-token STP \(C\)\. Model A remains usable \(MSE<1<1\) up tom=8m\{=\}8\.Sampling at consecutive step boundaries \(model A\) achieves168×168\\timesmore accurate prediction than the frozen baseline atm=1m\{=\}1, while the original random\-token recipe \(model C\) achieves only4×4\\times\(Table[2](https://arxiv.org/html/2604.18464#S5.T2), Figure[2](https://arxiv.org/html/2604.18464#S5.F2)\)\. Model A’s prediction atm=3m\{=\}3\(MSE 0\.093\) is more accurate than B1’s prediction atm=1m\{=\}1\(MSE 0\.955\)\. The role ofℒNTP\\mathcal\{L\}\_\{\\text\{NTP\}\}sees to be a tradeoff and not a redundancy\. Model A1 \(ℒSTPstep\\mathcal\{L\}\_\{\\text\{STP\}\}^\{\\text\{step\}\}only\) tests whetherℒNTP\\mathcal\{L\}\_\{\\text\{NTP\}\}can be dropped\. Under linear extrapolation, A1 matches A within 12%\. Under a learned MLP predictor, A1 outperforms A by 47–57% across all skip distances \(§[6\.1](https://arxiv.org/html/2604.18464#S6.SS1)\)\. The choice is a tradeoff between generation quality \(ℒNTP\\mathcal\{L\}\_\{\\text\{NTP\}\}: \+3\.7% GSM8K\) and geometric predictability whereℒSTP\\mathcal\{L\}\_\{\\text\{STP\}\}alone:2×2\\timesbetter MLP prediction\.

### 5\.2Cross\-model and Cross\-Dataset Validation

Table 3:B1 vs\. A across three model families on ProcessBench\. Accuracy=1−MSE=1\-\\sqrt\{\\text\{MSE\}\}\.![Refer to caption](https://arxiv.org/html/2604.18464v1/cross_model_and_dataset.png)Figure 3:Robustness of the geometric improvement\.\(a\)Cross\-model: B1 vs\. A prediction MSE atm=1m\{=\}1across three model families \(log scale\)\. The improvement ranges from39×39\\times\(Qwen2\.5\-1\.5B\) to335×335\\times\(Llama\-3\.2\-1B\)\.\(b\)Cross\-dataset: B1 vs\. A on four ProcessBench splits of increasing difficulty\. The improvement is consistent \(145×145\\times–183×183\\times\) and slightly stronger on harder datasets\.To investigate if geometric representation finetuning is transferable to non\-math models, we similarly finetuned two models: Qwen2\.5\-1\.5B, which is a direct comparison to Qwen2\.5\-Math\-1\.5B, as well as Llama\-3\.2\-1B, which represents a different model architecture with different dimensions\. From the results, we observed that the effect transfers across architectures, hidden dimensions, and pretraining regimes \(Table[3](https://arxiv.org/html/2604.18464#S5.T3), Figure[3](https://arxiv.org/html/2604.18464#S5.F3)a\)\. Llama\-3\.2\-1B achieves the strongest improvement \(335×335\\times\) despite having no math\-specific pretraining\.

The improvement also holds across difficulty levels \(Figure[3](https://arxiv.org/html/2604.18464#S5.F3)b\) on four ProcessBench splits, GSM8K \(165×165\\times\), MATH \(145×145\\times\), OlympiadBench \(183×183\\times\), and Omni\-MATH \(180×180\\times\)\. The geometric shaping is consistent and slightly stronger on harder problems, thereby ruling out the concern that STP overfits to easy problem structure\.

### 5\.3Task Accuracy: No Degradation

Table 4:GSM8K accuracy \(greedy decoding\)\.Overall, all models perform somewhat equally on accuracy, confirming that step\-boundary STP does not degrade the model’s reasoning ability\. Model A achieves accuracy identical to vanilla fine\-tuning \(B2\) while simultaneously achieving168×168\\timesbetter latent prediction MSE \(Table[4](https://arxiv.org/html/2604.18464#S5.T4)\)\. Random\-token STP is the only variant that showed a decrease in accuracy relative to the frozen baseline\. This suggests that the semantic structure of reasoning prevents interference with the language model objective\. The accuracy and MSE tests in fact reveal that two losses contribute orthogonally\.ℒ​NTP\\mathcal\{L\}\{\\text\{NTP\}\}seem to impact accuracy improvement whileℒ​STP\\mathcal\{L\}\{\\text\{STP\}\}brings in geometric alignment to shape how models think through latent space\.

### 5\.4Decoding Validation

Table 5:Decoding from predicted embeddings \(200 ProcessBench samples\)\.The decoding validation \(Table[5](https://arxiv.org/html/2604.18464#S5.T5)\) experiment shows that two hidden\-state vectors can be close in Euclidean distance yet differ in the specific subspace that the LM head uses for next\-token prediction\. The information relevant to decoding occupies a lower\-dimensional manifold within the full 1536\-dimensional hidden\-state space\. To test whether our predicted embeddings are functionally equivalent and not merely geometrically close, we pass both the predictedz^k\+m\\hat\{z\}\_\{k\+m\}and the actualzk\+mz\_\{k\+m\}through the model’s final RMSNorm and LM head, producing two probability distributions over the vocabulary\. Atm=1m\{=\}1, the distributions agree on the top\-1 token 93\.4% of the time \(the frozen baseline B1 had 19\.1%\), and with a KL divergence of just 0\.019 nats \(vs\. 4\.581 for B1\)\. The step retrieval metric provides a complementary view: the predicted embedding’s nearest neighbor among all step\-boundary embeddings \(excluding the source pointszkz\_\{k\}andzk−1z\_\{k\-1\}\) is the correct target step 92\.2% of the time, confirming thatz^k\+m\\hat\{z\}\_\{k\+m\}lands at the right reasoning state\. An instructive contrast emerges between model A and B1 atm=1m\{=\}1: B1 achieves relatively high step retrieval \(70\.2%\) but low token agreement \(19\.1%\), meaning the frozen model’s predicted embedding is geometrically near the correct step but encodes different next\-token information—the geometric position and functional content are misaligned\. STP training aligns both: model A’s 92\.2% retrieval and 93\.4% token agreement show that geometric proximity and functional equivalence coincide after geometric shaping\. At longer horizons, model A’s token agreement degrades gracefully—71\.1% atm=2m\{=\}2and 34\.1% atm=3m\{=\}3—while B1 collapses to near\-chance \(10\.1% and 8\.3%\)\. Notably, model A atm=2m\{=\}2\(71\.1%\) still exceeds B1 atm=1m\{=\}1\(19\.1%\), mirroring the MSE result and confirming that the prediction advantage extends well beyond one\-step forecasting\.

## 6Analysis

### 6\.1Manifold Structure

Table 6:MLP vs\. linear extrapolation on ProcessBench \(all six models\)\. Ratio = MLP MSE / Linear MSE; decreasing ratio withmmindicates smooth curvature\.We train a residual multilayer perceptron \(MLP\) post\-hoc predictor probe,z^k\+m=zk\+m​\(zk−zk−1\)\+gϕ​\(zk,zk−1\)\\hat\{z\}\_\{k\+m\}=z\_\{k\}\+m\(z\_\{k\}\-z\_\{k\-1\}\)\+g\_\{\\phi\}\(z\_\{k\},z\_\{k\-1\}\), wheregϕg\_\{\\phi\}is a 3\-layer MLP with zero\-initialized last layer \(architecture details in Appendix[A\.3](https://arxiv.org/html/2604.18464#A1.SS3)\)\. The zero initialization ensures that the MLP starts as an exact copy of the linear predictor and can only improve upon it\. Any reduction in MSE represents a genuine non\-linear structure that linear extrapolation missed\. The MLP probe is trained on 80% of ProcessBench step\-boundary pairs and evaluated on the held\-out 20%\. The probe serves purely as a measurement instrument for the latent geometry of the finetuned model without changing the models\.

To characterize model’s manifold structure, we compared the MLP against a trivial linear probe on the same task of predicting the hidden state at stepk\+mk\+mwhen given hidden state of stepkkand stepk−1k\-1\. The ratio of MLP MSE to and liner MSE reveals how much learnable strcture exists beyond what a striaght\-line approximation captures\. Crucially, a ratio that decreases withmmindicates smooth, compounding carvature, which is the signature of a coherent curving manifold rather than random noise\. \(Table[6](https://arxiv.org/html/2604.18464#S6.T6), Figure[4](https://arxiv.org/html/2604.18464#S6.F4)\)

Models B1, B2 and C show flat or weakly decreasing ratios, while models with step\-boundary finetuning, such as A, A1, and A2, show strongly decreasing ratios\. Model A goes from 0\.285 to 0\.085 \(a 70% drop overm=1m\{=\}1tom=3m\{=\}3\), and A1′goes from 0\.137 to 0\.029 \(a 79% drop\)\. Decreasing ratio withmmis the signature of a smoothly curving path\. curvature compounds quadratically with horizon, so longer extrapolations expose proportionally more non\-linear structure for the MLP to learn\. A1 \(pureℒ​STP\\mathcal\{L\}\{\\text\{STP\}\}, no language modeling loss\) has the steepest decrease and lowest ratios of any cell, producing the cleanest geometric structure\. It’s MLP atm=3m\{=\}3achieves MSE 0\.003407 — a1,098×1\{,\}098\\timesimprovement over B1’s linear prediction atm=1m\{=\}1\(MSE 0\.954\)\. With a learned predictor and pure geometric training, we can forecast three reasoning steps ahead more accurately than a frozen pretrained model can forecast one\.

![Refer to caption](https://arxiv.org/html/2604.18464v1/ratio_vs_m.png)Figure 4:MLP/linear MSE ratio vs\. skip distancemmfor all six models\. Flat ratio \(B1, B2, C\) indicates noise without systematic curvature\. Decreasing ratio \(A2, A, A1\) is the signature of a smoothly curving manifold where curvature compounds withmm\. A1 \(STP only\) has the steepest decrease\.
### 6\.2Why Sampling Position Matters

The40×40\\timesgap between model A and model C arises from optimization efficiency: step\-level STP concentrates all gradient signal on theKKstep\-boundary positions that define the trajectory we measure, while token\-level STP distributesqqtriples acrossTTtoken positions withO​\(1/n\)O\(1/n\)dilution per cross\-step triple\. The effective gradient ratio is∼n\(K−1\)/q≈100×\\sim n\(K\{\-\}1\)/q\\approx 100\\timesfor typical values\. Both models use identical LoRA parameters, training data, epochs, and compute\. Hence, likely the40×40\\timesMSE gap is entirely attributable to sampling position\.

The MLP/linear ratio decreasing withmm\(Table[6](https://arxiv.org/html/2604.18464#S6.T6)\) follows from curvature bounds\. The linear prediction error scales asO​\(m2​ϵ\)O\(m^\{2\}\\epsilon\)while a curved predictor achievesO​\(m2​ϵ2\)O\(m^\{2\}\\epsilon^\{2\}\), yielding ratioO​\(ϵ\)O\(\\epsilon\)\. Model A1 \(ϵ=0\.005\\epsilon=0\.005\) has lower ratios than model A \(ϵ=0\.014\\epsilon=0\.014\) at everymm, consistent with theO​\(ϵ\)O\(\\epsilon\)prediction\. The2×2\\timesMLP gap between A1 and A follows from gradient interference:ℒNTP\\mathcal\{L\}\_\{\\text\{NTP\}\}adds perturbations orthogonal to the geometric direction that linear tolerates but the MLP cannot fit through\.

Our results establish a clear hierarchy of latent\-space support for multi\-step prediction\. Frozen pretrained representations \(model B1\) that are based on token\-based sampling are not suitable for prediction\. The linear MSE≈1\.0\\approx 1\.0atm=1m\{=\}1, means the prediction error equals the signal magnitude, faired worse compared to a simple learned MLP \( reduced to MSE≈0\.39\\approx 0\.39\)\. This confirmed that the frozen trajectory has no coherent geometric structure for a neural predictor to exploit\. Vanilla LM fine\-tuning \(model B2\) reduces linear MSE to 0\.60, but the MLP/linear ratio remained flat at∼\\sim0\.45 across all skip distances\. Random\-token STP \(model C\) provided a further reduction to linear MSE at≈0\.23\\approx 0\.23, however, its MLP/linear ratio \(0\.466\) is the highest of any cell, suggesting the residual structure is less learnable than the frozen baseline\. This may mean that Random\-token STP denoises but it does not create a manifold\. Only step\-boundary STP showed capability in predictable trajectories\. Model A achieved linear MSE of≈0\.006\\approx 0\.006with a strongly decreasing MLP/linear ratio \(0\.285→\\to0\.085\)\. The gap between Cell C \(MSE 0\.226\) and Cell A \(MSE 0\.006\) is40×40\\times, arising entirely from where the loss is sampled, not from what the loss computes\. Model A1, without theℒ​NTP\\mathcal\{L\}\{\\text\{NTP\}\}pushed further to MLP MSE≈0\.0009\\approx 0\.0009, about a1,098×1\{,\}098\\timesimprovement over B1’s linear baseline\.

## 7Limitations and Future Work

Limitations\.\(1\) We characterize prediction potential but do not build a latent reasoning system; utilizing the geometric structure for generation would require COCONUT\-style architectural changes\(Hao et al\.,[2024](https://arxiv.org/html/2604.18464#bib.bib4)\)\. \(2\) Step boundaries are manually marked via<\|step\|\>at paragraph breaks; some MATH solutions have noisy boundaries \(e\.g\., individual lines of aLaTeXalign\*block\)\. Despite this, the168×168\\timesimprovement is robust\. \(3\) Evaluation is on math reasoning only\. \(4\) MSE atm=1m\{=\}1is closely related to the STP training objective; MSE atm≥2m\\geq 2and the MLP predictor provide independent validation\.

Future work\.\(1\) Combining step\-boundary STP with COCONUT curriculum training—STP straightens the trajectory that COCONUT iterates along\. \(2\) Data efficiency evaluation\. \(3\) Automatic step boundary detection\. \(4\) Cross\-domain evaluation\.

## 8Conclusion

In this paper, we investigates the effect of STP’s geometric loss and step sampling on shaping language model’s reasoning trajectories\. We showed that changing STP from random token sub\-spans to semantic step boundaries transforms hidden\-state trajectories from unpredictable walks \(linear MSE≈1\.0\\approx 1\.0\) into smoothly curving manifolds \(linear MSE≈0\.006\\approx 0\.006; MLP MSE≈0\.0009\\approx 0\.0009\)\. And also achieved multi\-step latent prediction MSE168×168\\timesto1,098×1\{,\}098\\timesimprovement validated across three model families and four datasets\.

## Code Availability

The code used to reproduce all experiments and figures in this work is publicly available at[hhttps://github\.com/YYDreamzure/SSP/](https://arxiv.org/html/2604.18464v1/hhttps://github.com/YYDreamzure/SSP/)\. The repository includes training scripts, evaluation pipelines, and pre\-trained model checkpoints\. Any additional data or materials required to reproduce the results reported in this paper are available from the corresponding author upon reasonable request\.

## References

- Assran et al\. \[2023\]M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. Ballas\.Self\-supervised learning from images with a joint\-embedding predictive architecture\.In*CVPR*, 2023\.
- Bardes et al\. \[2024\]A\. Bardes, Q\. Garrido, J\. Ponce, X\. Chen, M\. Rabbat, Y\. LeCun, M\. Assran, and N\. Ballas\.V\-JEPA: Latent video prediction for visual representation learning\.*arXiv:2404\.16930*, 2024\.
- Damirchi et al\. \[2026\]H\. Damirchi et al\.Geometric analysis of transformer layer trajectories in reasoning\.2026\.
- Hao et al\. \[2024\]S\. Hao, B\. Sukhbaatar, D\. Su, X\. Li, Z\. Hu, J\. Weston, and Y\. Tian\.Training large language models to reason in a continuous latent space\.*arXiv:2412\.06769*, 2024\.
- Herrmann et al\. \[2025\]V\. Herrmann et al\.PHi: Hidden\-state prediction loss for evaluating reasoning quality\.In*ICML*, 2025\.
- Hu et al\. \[2022\]E\. J\. Hu, Y\. Shen, P\. Wallis, Z\. Allen\-Zhu, Y\. Li, S\. Wang, L\. Wang, and W\. Chen\.LoRA: Low\-rank adaptation of large language models\.In*ICLR*, 2022\.
- Huang et al\. \[2025\]Y\. Huang et al\.LLM\-JEPA: Joint embedding prediction for language models\.*arXiv:2509\.14252*, 2025\.
- Huang et al\. \[2026\]Y\. Huang et al\.Semantic tube prediction\.*arXiv:2602\.22617*, 2026\.
- LeCun \[2022\]Y\. LeCun\.A path towards autonomous machine intelligence\.*OpenReview*, 2022\.
- Wang et al\. \[2025\]X\. Wang et al\.Latent cosine of expertise: Geometric analysis across transformer layers\.*arXiv:2410\.13640*, 2025\.
- Zhang et al\. \[2025\]K\. Zhang et al\.LightThinker: Thinking step\-by\-step compression for LLM reasoning\.2025\.
- Zheng et al\. \[2024\]C\. Zheng et al\.ProcessBench: Identifying process errors in mathematical reasoning\.*arXiv:2412\.06559*, 2024\.
- Zhou et al\. \[2026\]Y\. Zhou et al\.Geometry of reasoning in large language models\.*arXiv:2510\.09782*, 2026\.
- CoLaR \[2025\]CoLaR\.Compressed latent replay for efficient chain\-of\-thought\.2025\.
- PLaT \[2026\]J\. Wang et al\.Latent Chain\-of\-Thought as Planning: Decoupling Reasoning from Verbalization2026\.
- Sun et al\. \[2026\]X\. Sun, Y\. Dong, et al\.LLM reasoning as trajectories: Representation, verification, and steering\.*arXiv:2604\.05655*, 2026\.
- Teoh et al\. \[2025\]T\. Teoh et al\.NextLat: Next\-latent prediction transformers for multi\-step world modeling and reasoning\.*arXiv:2511\.05963*, 2025\.
- Jiang et al\. \[2026\]Y\. Jiang et al\.TRACED: Beyond scalars—progress and stability in reasoning trajectories\.*arXiv:2603\.10384*, 2026\.
- Carson & Reisizadeh \[2025\]J\. Carson and A\. Reisizadeh\.Statistical physics of language model reasoning\.In*ICML*, 2025\.*arXiv:2506\.04374*\.
- Zhuang et al\. \[2026\]S\. Zhuang et al\.The Geometric Reasoner: Training\-free geometric reasoning via smoothness and diversity penalties\.*arXiv:2601\.18832*, 2026\.
- Sun et al\. \[2026b\]S\. Sun et al\.STEP: Hidden states as early signals for reasoning quality\.*arXiv:2601\.09093*, 2026\.
- Tao et al\. \[2025\]L\. Tao et al\.CoLaR: Compressed latent reasoning via next\-embedding prediction\.In*NeurIPS*, 2025\.*arXiv:2505\.16552*\.

## Appendix ASupplementary Material

### A\.1Perpendicular Score Geometry

![Refer to caption](https://arxiv.org/html/2604.18464v1/SSP_legvssecant.png)Figure 5:Perpendicular score computation\.Left: interior positions use the two\-sided secantrefrk=zk\+1−zk−1\\text\{ref\}\_\{r\_\{k\}\}=z\_\{k\+1\}\-z\_\{k\-1\}as the reference direction\.Right: the last position uses the one\-sided secantrefrK=zK−1−zK−2\\text\{ref\}\_\{r\_\{K\}\}=z\_\{K\-1\}\-z\_\{K\-2\}\(the arriving direction\)\. The perpendicular score issin⁡θk∈\[0,1\]\\sin\\theta\_\{k\}\\in\[0,1\]whereθk\\theta\_\{k\}is the angle between the displacementdkd\_\{k\}and its projection onto the reference direction\.Binary detection AUC is indistinguishable from random \(Table[7](https://arxiv.org/html/2604.18464#A1.T7)\)\. The Semantic Tube captures the*process*of organized reasoning—a structured, predictable flow through latent space—not the*outcome*of correct reasoning\. This is consistent with the PHi loss literature\(Herrmann et al\.,[2025](https://arxiv.org/html/2604.18464#bib.bib5)\): hidden\-state unpredictability correlates with correct solutions on difficult problems\.

This finding resolves an apparent conflict withDamirchi et al\. \([2026](https://arxiv.org/html/2604.18464#bib.bib3)\), who report that smooth trajectories correlate with valid reasoning\. The two analyses operate on different axes: smoothness across*layers*\(vertical\) may reflect well\-conditioned computation, while smoothness across*steps*\(horizontal\) reflects predictable reasoning content—which need not correlate with correctness\.

### A\.2Smoothness≠\\neqCorrectness

Table 7:ProcessBench error detection\. Perpendicular score does not encode correctness\.
### A\.3Skip\-Predictor Architectures

We use two post\-hoc predictors to probe the geometric structure of step\-boundary trajectories\. Both take as input the hidden states at positionskkandk−1k\{\-\}1and predict the hidden state at positionk\+mk\{\+\}m, without modifying the underlying model\.

#### Linear predictor \(zero parameters\)\.

Extrapolates along the current displacement direction:

z^k\+m=zk\+m⋅\(zk−zk−1\)\\hat\{z\}\_\{k\+m\}=z\_\{k\}\+m\\cdot\(z\_\{k\}\-z\_\{k\-1\}\)\(8\)This assumes the trajectory is locally linear\. If consecutive displacements are perfectly parallel \(ℒSTPstep=0\\mathcal\{L\}\_\{\\text\{STP\}\}^\{\\text\{step\}\}=0\), linear prediction is exact\. Any prediction error reflects curvature or noise in the trajectory\.

#### MLP predictor \(learned,∼\\sim6M parameters\)\.

Adds a learned non\-linear correction on top of the linear baseline:

z^k\+m=zk\+m⋅\(zk−zk−1\)⏟linear baseline\+gϕ​\(zk,zk−1\)⏟learned correction\\hat\{z\}\_\{k\+m\}=\\underbrace\{z\_\{k\}\+m\\cdot\(z\_\{k\}\-z\_\{k\-1\}\)\}\_\{\\text\{linear baseline\}\}\+\\underbrace\{g\_\{\\phi\}\(z\_\{k\},\\;z\_\{k\-1\}\)\}\_\{\\text\{learned correction\}\}\(9\)wheregϕg\_\{\\phi\}is a 3\-layer MLP \(input:2​D→2048→2048→D2D\\to 2048\\to 2048\\to D, GELU activations\)\. The last layer is zero\-initialized, so at epoch 0 the MLP prediction equals the linear prediction exactly—it can only improve, never hurt\. Training minimizes MSE on 80% of ProcessBench step\-boundary pairs; evaluation uses the held\-out 20%\.

The ratioMSEmMLP/MSEmlinear\\text\{MSE\}\_\{m\}^\{\\text\{MLP\}\}/\\text\{MSE\}\_\{m\}^\{\\text\{linear\}\}characterizes the manifold structure: ratio≈1\\approx 1means the trajectory is maximally linear \(MLP adds nothing\); ratio≪1\\ll 1means systematic non\-linear structure exists that the MLP exploits\. A ratio that*decreases*withmmindicates smooth curvature that compounds over longer horizons—the signature of a curving tube rather than a straight line or random walk\.