Semantic Tube Prediction: Beating LLM Data Efficiency with JEPA

Abstract: LLMs obey consistent scaling laws -- empirical power-law fits that predict how loss decreases with compute, data, and parameters. While predictive, these laws are descriptive rather than prescriptive: they characterize typical training, not optimal training. Surprisingly few works have successfully challenged the data-efficiency bounds implied by these laws -- which is our primary focus. To that end, we introduce the Geodesic Hypothesis, positing that token sequences trace geodesics on a smooth semantic manifold and are therefore locally linear. Building on this principle, we propose a novel Semantic Tube Prediction (STP) task, a JEPA-style regularizer that confines hidden-state trajectories to a tubular neighborhood of the geodesic. STP generalizes JEPA to language without requiring explicit multi-view augmentations. We show this constraint improves signal-to-noise ratio, and consequently preserves diversity by preventing trajectory collisions during inference. Empirically, STP allows LLMs to match baseline accuracy with 16$\times$ less training data on the NL-RX-SYNTH dataset, directly violating the data term of Chinchilla-style scaling laws and demonstrating that principled geometric priors can surpass brute-force scaling. Code is available at https://github.com/galilai-group/LLM-jepa#stp.

• The paper introduces Semantic Tube Prediction, a method that augments next token prediction with a regularization enforcing geodesic trajectories in latent space.
• It demonstrates that integrating a cosine-based loss enables models to match baseline accuracy with only 1/16 of the typical training data.
• The study combines theoretical ODE-based justification with extensive empirical validation, highlighting improved signal-to-noise ratios and reduced hallucinations.

Semantic Tube Prediction: Challenging LLM Data Scaling Laws with JEPA-Based Regularization

Introduction and Theoretical Foundation

This work introduces Semantic Tube Prediction (STP), a training strategy for LLMs motivated by geometric structure in representation space, and investigates whether principled inductive biases can surpass the conventional data efficiency ceilings suggested by empirical scaling laws. These scaling laws, describing the relationship between model size, dataset size, and performance, have historically set expectations for how data requirements decrease as models scale. The authors contend that such laws appear to be a consequence of standard objectives—namely, Next Token Prediction (NTP)—rather than properties fundamental to the underlying data or model capacity.

They advance the Geodesic Hypothesis, which asserts that error-free token sequence trajectories correspond to geodesics—locally linear paths—on a smooth "semantic manifold." This is formalized through an ODE perspective on LLM sequence evolution, leveraging the Picard-Lindelöf theorem to argue for uniqueness and non-intersection of trajectory paths under smooth conditions. The local linearity assumption leads to the postulation of a "^^^^1^^^^" around these geodesics, within which valid hidden state trajectories are confined. Deviations perpendicular to the tube correspond to semantic noise, while the parallel component tracks the meaningful signal.

Figure 1: Illustration of the "semantic tube", exemplifying the distinction between signal (along the geodesic) and noise (perpendicular deviations).

Semantic Tube Prediction Loss: Formulation and Dynamics

The Semantic Tube Prediction objective is built as a regularization term, added to the standard NTP cross-entropy objective. For three hidden states $h_s$, $h_r$, $h_t$ at indices $s < r < t$, the loss isolates the component of $h_r-h_s$ perpendicular to $h_t-h_s$ and penalizes it:

$$\mathcal{L}_{\rm STP} = 1 - \cos(h_t - h_r, h_r - h_s)$$

where the cosine term measures alignment, ensuring local collinearity of trajectory increments. The loss is integrated into the overall training objective:

$$\mathcal{L} = \mathcal{L}_{\rm NTP} + \lambda \cdot \mathcal{L}_{\rm STP}$$

with $\lambda$ a tunable hyperparameter. This approach is inspired by JEPA-style architectures but dispenses with explicit two-view scaffolding or learnable predictors, due to the symmetry and local linearity induced by the geodesic assumption. Empirical analysis supports the claim that $\mathcal{L}_{\rm NTP}$ alone cannot prevent inference-time mode collapse and hallucinations, whereas augmenting with $\mathcal{L}_{\rm STP}$ suppresses sample drift and improves signal-to-noise characteristics.

Empirical Validation

Training, Loss Landscape, and Data Efficiency

A comprehensive set of experiments examines the effectiveness of STP across architectures, dataset domains, data regimes, and parameterizations. Key results demonstrate that for the NL-RX-SYNTH dataset, models trained with STP match baseline accuracy while using only $\frac{1}{16}$ of the training data—a direct contradiction of the predictions made by canonical "Chinchilla-style" scaling laws. Accuracy degradation is negligible when halving the dataset, and only apparent at extreme data reduction. Standard NTP fine-tuning, in contrast, exhibits significant accuracy loss even for moderate reductions in data volume.

Figure 2: Loss curves during fine-tuning, showing divergent long-term behavior of $\mathcal{L}_{\rm NTP}$ and $\mathcal{L}_{\rm STP}$.

Figure 3: STP yields consistent accuracy improvements across a range of tasks and model sizes compared to both standard fine-tuning and JEPA baselines.

STP also yields smooth, stable training curves, with the auxiliary loss contributing a persistent decrease in $\mathcal{L}_{\rm STP}$ even after $\mathcal{L}_{\rm NTP}$ plateaus. Performance is robust across Llama, Gemma, OpenELM, and OLMo model families, and persists across multiple orders of magnitude in parameter count.

Regularization Strength and Ablation

Extensive tuning of $\lambda$ reveals a consistent, concave accuracy-versus-$\lambda$ relationship. Empirically, optimal performance is typically achieved in the range $0.01 \leq \lambda \leq 0.08$, and performance is not overly sensitive within this range, but exceeds this value performance drops sharply, supporting the need for subtle regularization rather than hard enforcement of collinearity.

Figure 4: Impact of $\lambda$ on model accuracy, demonstrating consistently optimal values in the $0.01$–$0.08$ interval across datasets.

Ablation studies on predictor variants (e.g., learned linear projection instead of identity), start/end anchorings, and pooling strategies reveal that the direct STP objective outperforms all alternatives, including variants modeled more closely on classical JEPA-style architectures.

Figure 5: Ablation results. The STP objective outperforms all alternative configurations, corroborating the central claims of the Geodesic Hypothesis.

Representation Geometry and Diversity

Representation analysis via SVD on the difference between latent encodings for pairs of related sequences shows that STP enforces a highly structured, low-dimensional geometry on the direction (normalized differences), while allowing for more complex, polymorphic behavior in the raw, unnormalized space. This supports the idea that STP regularizes the shape of semantic change without forcing collapse onto a trivial subspace, thus preserving semantic diversity and mitigating mode collapse.

Figure 6: Decomposition of hidden state relationships, confirming that STP enforces structure in normalized vector directions but preserves polymorphism without normalization.

Theoretical and Practical Implications

The STP framework unifies multiple lines of inquiry in LLM representation geometry. It makes explicit connections to the Manifold Hypothesis by interpreting principal learning dynamics as local linear geodesics, and it generalizes the Linear Representation Hypothesis to the level of sequence trajectories:

Figure 7: Concept direction alignment when the sentence traces a geodesic path in latent space, per the Linear Representation Hypothesis and its extension.

The formal ODE/SDE treatment shows that inference-time drift can be viewed as a Brownian motion around the geodesic, with the regularization coefficient $\lambda$ controlling the variance. Suppressing noise perpendicular to the geodesic lowers the probability of accidental trajectory collisions—hence the reduction in hallucinations and mode collapse observed empirically.

The practical consequence is an actionable regularization scheme that can be implemented with negligible overhead, requiring only computation of random-index cosines in the last-layer hidden states. This makes Semantic Tube highly amenable to integration into existing LLM pretraining and fine-tuning pipelines. The empirical success in drastically reducing data requirements implies that power-law data scaling is not a fundamental limitation, but rather contingent on the lack of geometric priors in the training objective.

Conclusion

Semantic Tube Prediction provides a geometrically grounded regularization for LLMs, built on the hypothesis that semantic progression follows locally linear geodesics in the high-dimensional latent manifold. By constraining hidden state evolution to a tube centered on these paths, STP enhances signal-to-noise ratio, robustly preserves semantic diversity, and enables accuracy retention under dramatic dataset reduction. These findings refute the inevitability of power-law-limited data efficiency and suggest that embedding geometric inductive biases into training objectives is a critical frontier for next-generation LLMs. The STP framework thus opens new theoretical and practical avenues for reducing the resource intensiveness of large-scale language modeling, and raises new questions regarding the interplay between geometry, representation learning, and sample efficiency in deep sequence models.