Research map / Architectures and efficiency / RNNs and normalization

Efficient Pre-Training with Token Superposition

NeurIPS 2026 · Architectures and efficiency › RNNs and normalization · approach: Neural scaling laws · 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 Token-Superposition Training (TST), a drop-in two-phase method that averages embeddings of contiguous token bags and predicts next bags with a multi-hot cross-entropy, then reverts to standard next-token training, yielding up to 2.5x pre-training time reduction at 10B-A1B scale.

Problem

Pre-training LLMs is expensive; the goal is to increase data throughput per FLOP during pre-training without changing the parallelism, optimizer, tokenizer, data, or inference-time architecture.

Approach

During a superposition phase, the model consumes bags of s contiguous token embeddings averaged into one latent token and predicts the next bag with a multi-hot cross-entropy (uniform average of CE over bag tokens); a recovery phase then switches to standard next-token training with the same embeddings and LM head, so inference architecture is unchanged.

Key results

• At 10B A1B MoE scale, TST reaches baseline final loss (2.252 vs 2.236) with 4768 vs 12311 B200-hours, up to 2.5x reduction in total pre-training time at equal loss.
• Dense 3B TST (6000 superposition steps, s=6) reaches final loss 2.676 with 247 B200-hours, matching the 36000-step baseline (2.677) which costs 443 B200-hours.
• Dense 600M TST improves HellaSwag from 43.5 to 48.2 and final loss from 3.019 to 2.943 at equal FLOPs.

Surprising findings

• Randomly re-initializing input embeddings and LM head at the start of recovery eliminates TST gains and makes the 3B run worse than baseline (2.938 vs 2.808).
• Using binary cross-entropy or hinge losses for the multi-hot bag was significantly worse than even the baseline without TST.

Limitations

• Gains come from trading more data consumption for loss; assumes compute-bound rather than data-bound regime.
• No statistical significance testing or repeated runs due to compute limits.
• Learning rate for 3B and 10B runs taken from recommendation rather than swept; bag size sensitivity shows U-shaped loss.

Experimental setup

Datasets: DCLM FineWeb-Edu

Benchmarks: HellaSwag ARC-Easy ARC-Challenge MMLU BoolQ OpenBookQA PIQA Winogrande

Metrics: final training loss accuracy (0-shot) B200-hours speedup at equal loss

Baselines: Dense Baseline (standard next-token training) Multi-token prediction Patch-level training

Builds on: patch-level training multi-token prediction SmolLM2 Qwen3 Warmup-Stable-Decay

Closest papers by meaning

• Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.88)
• Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.87)
• The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws (ICLR 2025 · similarity 0.87)
• On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.87)
• Not All Tokens Are What You Need for Pretraining (NeurIPS 2024 · similarity 0.87)
• Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · similarity 0.87)