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

TL;DR

This paper introduces a semantic step sampling method for LLM reasoning trajectories that significantly improves multi-step latent prediction accuracy by focusing on semantic boundaries, revealing a tradeoff between generation quality and geometric regularity.

Contribution

It demonstrates that sampling at semantic reasoning step boundaries enhances geometric regularization and multi-step prediction accuracy in LLMs, surpassing random token sampling methods.

Findings

168x more accurate multi-step latent prediction with semantic boundary sampling

Trajectory shapes are smooth curves, not straight lines, improving predictability

Removing language modeling loss increases trajectory predictability, indicating a tradeoff

Abstract

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 achieved 168x more accurate multi-step latent prediction than frozen baselines on ProcessBench (3,400 samples), compared to only 4x for 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 further 3-12x over linear extrapolation on step-boundary models. Removing the language modeling loss yields trajectories that are 2x more 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.