Research map / Large language models / Reasoning and chain of thought

Semantic Tube Prediction: Beating LLM Data Efficiency with JEPA

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Routing, pruning and distillation for reasoning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Semantic Tube Prediction (STP), an identity-predictor JEPA-style auxiliary loss that aligns hidden-state difference vectors across random token triples, improving LLM fine-tuning accuracy and data efficiency.

Problem

Next-token prediction fine-tuning is data-hungry and its hidden-state trajectories drift off error-free semantic paths, causing noisy representations and mode collapse; the paper seeks to beat data-efficiency bounds implied by scaling laws.

Approach

Introduce the Geodesic Hypothesis (token-sequence hidden-state trajectories are locally linear) and a Semantic Tube loss that penalizes the angle between h_t - h_r and h_r - h_s for randomly sampled indices s<r<t, acting as a JEPA-style regularizer with an identity predictor and no extra views or forward passes.

Key results

• STP with 1/16 of NL-RX-SYNTH training data matches full-data regular fine-tuning accuracy on Llama-3.2-1B, per the paper's claim.
• On the .* vs .*.* suffix test, Semantic Tube reaches 88.5% and 68.0% versus LLM-JEPA's 68.9% and 32.0% and regular fine-tuning's 29.9% and 28.0%.
• Optimal lambda values are small, typically 0.0025 to 0.04 across datasets and models, supporting lambda << 1.

Surprising findings

• Accuracy shows a negligible drop when the training set is halved, while regular fine-tuning degrades significantly immediately.
• Replacing the identity predictor with a learned linear projector (Pred variant) degrades performance in all configurations.

Limitations

• Geodesic/local linearity is a hypothesis; theory relies on approximations (e.g., small epsilon, infinite-width Gaussian assumptions) and proofs are sketches.
• Lambda must be tuned per setting, and large lambda causes accuracy collapse with increased variance.
• Evaluated mainly on fine-tuning with small models (1B-8B) and mostly on a single synthetic regex dataset for data-efficiency claims.

Experimental setup

Datasets: NL-RX-SYNTH NL-RX-TURK GSM8K Spider NQ-Open HellaSwag

Benchmarks: NL-RX-SYNTH NL-RX-TURK GSM8K Spider NQ-Open HellaSwag

Metrics: accuracy data efficiency (fraction of training data to match baseline standard deviation over 5 seeds paired one-tailed t-test p-value

Baselines: regular fine-tuning (NTP) LLM-JEPA

Builds on: LLM-JEPA JEPA next token prediction

Closest papers by meaning

• Reinforced Fast Weights via Next-Sequence Prediction (NeurIPS 2026 · similarity 0.87)
• EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism (ICML 2024 · similarity 0.87)
• Compute-Optimal LLMs Provably Generalize Better with Scale (ICLR 2025 · similarity 0.87)
• On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.87)
• Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.86)
• The Role of Sparsity for Length Generalization in LLMs (ICML 2025 · similarity 0.86)