🤖 AI Summary
This work proposes Semantic Tube Prediction (STP) to enhance the predictability and geometric structure of hidden-state trajectories in large language models during multi-step reasoning. STP samples at semantic reasoning boundaries and applies geometric regularization to align trajectories with locally linear geodesics. The study reveals, for the first time, the critical influence of sampling positions on regularization efficacy, introduces a novel evaluation metric—multi-step hidden-state prediction mean squared error—and demonstrates that STP-generated trajectories are smooth curves rather than straight lines. Experiments on ProcessBench show that STP improves prediction accuracy by 168× over a frozen baseline; a 3-layer MLP trajectory predictor reduces error by 3–12× compared to linear extrapolation; and decoupling the language modeling loss doubles trajectory predictability.
📝 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.