PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Pre-Training with Token Superposition

Bowen Peng, Théo Gigant, Jeffrey Quesnelle

arXiv:2605.06546v2cs.CL

TL;DR

Large-language-model pre-training must consume rapidly growing amounts of data efficiently under fixed compute. Token-Superposition Training increases token throughput without changing core training components, achieving the same loss at around half the computational cost.

  • Problem

    The paper asks whether pre-training efficiency can improve through higher token throughput without modifying the final model architecture or inference dynamics.

  • Method

    Token-Superposition Training averages contiguous token embeddings into s-tokens, uses a semi-causal multi-hot objective, then returns to standard causal training.

  • Results

    At the same loss, TST achieves the target at around half the computational cost while outperforming equal-FLOPs baseline pre-training.

  • Takeaways & Limitations

    TST provides a high-throughput pre-training paradigm that improves efficiency at the same computational cost without changing architecture, parallelism, tokenizer, or data.

  • Takeaways & Limitations

    TST assumes compute-bound rather than data-bound pre-training, and the paper leaves alternative data-consumption settings for future work.

Abstract

from arXiv · show

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.

1 Introduction

Pre-training efficiency is increasingly important as LLMs scale data and often overtrain beyond compute-optimal estimates. TST addresses this challenge by increasing training-time token throughput without changing the final model architecture or inference dynamics.

  • Motivation: Data scaling and overtraining make efficient raw-text consumption a central pre-training concern under fixed compute.Recent training regimes often overtrain beyond compute-optimal estimates to maximize inference-time performance.
  • Prior Approaches: Prior efficiency methods primarily maximize information, reduce per-token computation, or compress internal representations.These categories include richer training signals, sparse mixture-of-experts or attention, and models that reduce the number of internal representations.
  • Motivation: Training-time efficiency should be separated from inference-time efficiency because inference compute can independently improve downstream performance.The authors therefore aim to leave model architecture and expressivity untouched for inference relative to the baseline.
  • Contribution: TST increases token throughput during training while priming the model for autoregressive prediction and preserving its final inference architecture and dynamics.This perspective motivates asking whether higher training throughput can improve pre-training efficiency without modifying the final model architecture.

2 Related works

Prior work improves pre-training efficiency by predicting multiple or future-token targets and by changing input granularity. Token Superposition Training combines future-token signal with coarser training-time granularity, then returns to baseline granularity and loss.

  • Auxiliary training objectives: Multi-token prediction uses k independent heads to predict the next k tokens simultaneously, improving sample efficiency in some cases but offering limited benefit to smaller models.MTP has also appeared in major state-of-the-art LLM pre-training runs and requires tuning of an unspecified component in the supplied passage.
  • Auxiliary training objectives: Other methods predict future-token representations through cascaded MTP modules, relative token order, or concepts spanning segments.DeepseekV3 uses additional cascaded MTP modules; Zuhri et al. use one additional head for relative-order prediction; Liu et al. propose next concept prediction.
  • Input granularity: Token Superposition Training leverages future-token representations and coarser input granularity during training, then returns to baseline granularity and loss in its second phase.This positions the method as combining two efficiency-oriented changes while restoring the standard setup during recovery.
  • Closest related method: Patch-level training averages consecutive token embeddings into patches, predicts all tokens in the next patch with one output head, and later reverts to standard token-level training.The supplied passage identifies this procedure as algorithmically identical to Token Superposition Training, while noting differences in background theory and execution.

3 Methodology

Token Superposition Training (TST) averages embeddings of contiguous non-overlapping token bags into latent s-tokens and predicts the next bag with a multi-hot cross-entropy objective. It preserves equal FLOPs per step by processing sequences s times longer, increasing token throughput during the superposition phase.

  • Token superposition: TST averages embeddings within non-overlapping s-token bags to form latent s-tokens, each producing one prediction for the next bag of tokens.The latent sequence has length l, while each bag contains s contiguous tokens.
  • Compute and throughput: s times more data tokens are processed per FLOPs on latent s-tokens because the model computes over a coarser-grained input representation.During superposition, the data sequence length L is increased by s to keep each TST step equal-FLOPs to baseline training.
  • Multi-hot objective: TST replaces one-hot cross-entropy with multi-hot cross-entropy, assigning probability 1/s to each of the s valid next-bag labels.The multi-hot objective targets equal probability across the labels in each bag.
  • Causal labeling: Labels shift left by s −1 before bagging, so each input bag at [t, t + s −1] predicts the subsequent bag at [t + s, t + 2s −1].This alignment preserves causality while predicting non-overlapping bags.
  • Loss design: The chosen MCE loss outperformed tested Hinge and BCE bag-loss variants and baseline training without TST, balancing soundness and simplicity.Its simplified form also works with optimized compiled CE-loss kernels without major training-code modifications.

4 Experiments

The experiments span multiple model scales and superposition settings, using a consistent TorchTitan/FSDP training setup. Across standard 0-shot evaluations, TST shows robust benefits over baseline under equal-FLOPs or equal-loss comparisons.

  • Experimental setup: Experiments vary superposition bag size, step ratio, model scale, and total training duration.Training uses TorchTitan with FSDP on 64 NVIDIA B200 GPUs for larger models and 8 B200 GPUs for smaller models.
  • Experimental setup: The 10B mixture-of-experts validation model has 10B total and 1B active parameters, trains for 1.05T tokens, and uses 2T data tokens for TST.The run uses a Qwen3-based architecture, a 50/50 FineWeb-Edu/DCLM mixture, and an 8M-token constant batch size.
  • Evaluation: Final-checkpoint evaluations cover ARC, BoolQ, HellaSwag, MMLU, OpenBookQA, PIQA, and Winogrande in a 0-shot EleutherAI LM-Eval setup.Evaluations are conducted after the recovery phase.
  • Results: TST outperforms the baseline across all equal-FLOPs or equal-loss settings, indicating robust benefits across the tested configurations.Figure 4 varies bag size and step ratio, while each point represents a fully trained model converted to a standard autoregressive language model.

5 Discussion

The discussion positions TST as distinct from auxiliary-loss methods because it increases training tokens per FLOP, and attributes its gains to complementary input and output superposition. It further links TST’s effectiveness to coarse-to-fine training and preserving shared representations across phases.

  • Comparison to Related Methods: TST differs from MTP and related auxiliary-loss methods because it increases training tokens per FLOP, whereas they retain baseline token throughput while adding parameters and auxiliary losses.Those methods may target inference-time speculative-decoding gains, which TST does not address.
  • Superposition Ablations: With s = 4 and r = 0.5, full superposition outperforms input-only and output-only variants, indicating complementary input-granularity and output-supervision mechanisms.All superposition settings outperform the baseline, but combining both sides produces the complete improvement without signs of interference.
  • Output Bag Weighting: For s ≥8, power-law output-bag weighting outperforms uniform weighting and is more stable, while smaller bags favor uniform weighting.Loss varies in a U-shaped pattern with bag size, so overshooting the optimal size can produce slightly suboptimal results.
  • Input Superposition: Input superposition may act as coarse-to-fine pre-pre-training or regularize embedding geometry, while its broader pattern suggests reusable granularity scheduling across modalities.The authors do not have interpretability evidence to distinguish the proposed explanations.
  • Phase Representations: Re-initializing the input embedding and output LM head between phases completely eliminates TST’s gains and makes it worse than baseline training.TST instead shares these representations unchanged across superposition and recovery, avoiding representational mismatch.

6 Conclusion

Token Superposition Training increases sample throughput during pre-training without changing core training components, then returns to standard training in a fast recovery phase that quickly outperforms an equal-FLOPs baseline.

  • Token Superposition Training: s-fold sample throughput increases during token superposition without changing per-step FLOPs, parallelism, architecture, tokenizer, or data.The superposition regime is followed by recovery to standard LLM pre-training.
  • Recovery phase: The recovery phase quickly outperforms the loss of an equal-FLOPs baseline pre-training.The passage characterizes recovery as fast after returning to the standard pre-training regime.

7 Limitations, Future Work and Broader Impacts · A Code · B Learning rate sweeps

The paper identifies compute–data tradeoffs, untested long-context and scaling behavior, and unresolved interpretability and broader-impact concerns, while documenting the core implementation and learning-rate choices. Future work targets long-context evaluation, statistical validation, scaling laws, and mechanistic understanding.

  • 7 Limitations, Future Work and Broader Impacts: TST trades greater data consumption for improved loss at fixed computational cost, assuming pre-training is compute-bound rather than data-bound.If future regimes become data-bound, output-only superposition could outperform baseline pre-training without increasing data consumption.
  • 7 Limitations, Future Work and Broader Impacts: Folding sequences into bags of s tokens creates a longer effective context than baseline training, but long-context performance was not evaluated.The authors suggest reduced truncation or splitting of native long-context data could be beneficial.
  • 7 Limitations, Future Work and Broader Impacts: Limited compute prevented larger-scale ablations and multiple identical runs for statistical significance, motivating future scaling-law studies for larger models and industry-scale pre-training.Such studies could predict the best TST settings at larger model sizes.
  • 7 Limitations, Future Work and Broader Impacts: Further interpretability work is needed to clarify TST’s underlying mechanisms and the ramifications of token superposition.The paper presents hypotheses about the phenomena involved but does not establish a complete explanation.
  • 7 Limitations, Future Work and Broader Impacts: Improved pre-training efficiency can reduce computational cost and energy usage and broaden researcher accessibility, while potentially accelerating model risks involving harmful content, bias, fairness, and privacy.The contribution is methodological and does not directly introduce those risks.
  • A Code: The implementation folds inputs to shape (bs, seq // superposition_bag_size, superposition_bag_size) when superposition_bag_size > 1 and sums token embeddings across each bag.The embedding sum is performed in float32 for better numerical precision.
  • A Code: The next-bag-of-words loss applies causal label padding and computes a weighted average of cross-entropy terms across superposition positions.The code accumulates weighted losses and divides by the total weight.
  • B Learning rate sweeps: The learning-rate sweeps select final rates of 2×10−3, 2×10−3, and 1×10−3 from left to right across the evaluated model sizes.The optimal learning rate is then used for all training runs.

C Loss Derivations · C.1 MCE: equal-probability targets

The derivation defines MCE for a bag of valid target tokens as a uniform target distribution and explains why KL divergence is used to make the loss vanish at the optimum. It contrasts this with one-hot cross-entropy, whose minimum is zero because one-hot targets have zero entropy.

  • C Loss Derivations: For a target distribution t over the vocabulary, softmax probabilities are defined as P(i) = exp(z_i)/Z.
  • C Loss Derivations: Cross-entropy reaches zero at its minimum only when the target entropy H(t) is zero, namely for one-hot targets.
  • C Loss Derivations: KL divergence instead vanishes whenever the predicted distribution matches the target, P = t, regardless of the target distribution.
  • C Loss Derivations: The derivation considers a bag y containing s valid target tokens and introduces two multi-hot cross-entropy loss variants.
  • C.1 MCE: equal-probability targets: MCE treats the bag as a multi-hot target, pushing each valid token’s probability toward 1/s and forming a uniform distribution over the bag.
  • C.1 MCE: equal-probability targets: For this target, standard cross-entropy is the corresponding cross-entropy objective under the uniform bag distribution.
  • C.1 MCE: equal-probability targets: The uniform target has entropy H(t) = log s, so plain CE bottoms out at log s rather than zero.
  • C.1 MCE: equal-probability targets: Subtracting the target entropy converts the objective to KL divergence, restoring vanishing loss at the optimum; rearrangement yields the stated equivalent form.

C.2 MCEAlt: sum-to-one probability targets

MCEAlt assigns total probability mass of 1 to all valid labels in a bag, allowing the model to choose their relative weighting. Small-scale experiments found it matched MCE after recovery, but its added complexity and lower efficiency led the authors to defer further study.

  • MCEAlt: sum-to-one probability targets: MCEAlt requires the summed probability of all valid bag labels to equal 1, rather than assigning them equal probabilities.It treats the bag as a single composite label with a one-hot target.
  • MCEAlt: sum-to-one probability targets: The formulation lets the model choose its own weighting across tokens within the bag.The objective only constrains the total probability mass assigned to the bag.
  • MCEAlt: sum-to-one probability targets: MCEAlt produced the same final loss as equal-probability MCE after recovery when all other conditions were identical.This result came from limited small-scale experiments.
  • MCEAlt: sum-to-one probability targets: The authors did not pursue MCEAlt because it requires a custom loss, reduces training speed, increases memory usage, and adds unnecessary complexity.They leave more thorough exploration of the variant for future work.

D Non-uniform Multi-hot Cross-Entropy

The section evaluates non-uniform multi-hot target distributions, including uniform, power-law, exponential, and first-token weighting. Performance depends on the superposition bag size and ratio, while a mutual-information-fitted power law provides a slight improvement over the earlier power-law weighting.

  • Target distributions: Four target weightings are considered: uniform, power law, exponential, and first-token distributions, all decreasing with token position except the position-specific first-token target.The distributions are defined over positions within each superposition bag.
  • Results: s = 16 favors the power-law distribution, whereas s = 4 favors uniform weighting, with no distribution best across all superposition ratios.The preferred weighting changes with the superposition bag size and ratio.
  • Motivation: A mutual-information analysis of DCLM token pairs motivates power-law loss weighting because token mutual information decays with distance.The fitted decay is d 7→C0 + a ∗dk, with C0 ≈3.63, a ≈1.35 and k ≈−1.25.
  • Results: The fitted power law yields a slightly better loss than the earlier tested power-law weighting.The fitted values from token mutual information are used to weight losses by position.

E Additional Results

Additional results expand the evaluation across superposition ratios and settings, reporting final losses and downstream task scores for 270M and 600M models. The results include 0-shot evaluations and experiments spanning 20k and 100k training steps.

  • Sensitivity analyses: Figure 11 examines Hellaswag and ARC-Easy downstream evaluations at varying superposition bag sizes and superposition step ratio r.The displayed ratios range from r = 0 to r = 0.6.
  • Loss sensitivity: Final losses are reported across varying r and s for 270M models trained for 20k or 100k total steps.The 20k-step results are given in Table 4, while the 100k-step results are given in Table 5.
  • Loss sensitivity: Final losses are also reported across varying r and s for 600M models trained for 20k total steps.Table 7 additionally specifies that the TST loss uses the power-law weighting described in Appendix D.
Loading 2605.06546v2…