💬 NLP

Efficient Pre-Training with Token Superposition

This paper introduces Token-Superposition Training (TST), a drop-in method that significantly improves pre-training data throughput and reduces total training time by up to 2.5x without modifying model architecture or parallelism, by initially combining tokens into bags for efficient multi-hot cross-entropy training before reverting to standard training.

Original authors: Bowen Peng, Théo Gigant, Jeffrey Quesnelle

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bowen Peng, Théo Gigant, Jeffrey Quesnelle

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Training AI is Like Running a Marathon in Sand

Imagine you are trying to teach a student (a Large Language Model) to read and write. To do this, you have to show them millions of books. However, the current way of teaching is incredibly slow and expensive. It's like asking the student to read every single word of a book, one by one, while carrying a heavy backpack. They get tired (computationally expensive) and it takes forever to finish the library.

Researchers want to make this process faster without changing the student's brain (the model architecture) or the books they read (the data).

The Solution: "Token Superposition Training" (TST)

The authors propose a clever two-phase training method called Token Superposition Training (TST). Think of it as a "speed reading" phase followed by a "fine-tuning" phase.

Phase 1: The "Smoothie" Phase (Superposition)

In standard training, the AI learns by looking at one word at a time (e.g., "The," then "cat," then "sat").

In the Superposition Phase, the AI is taught to look at a whole "bag" of words at once.

The Analogy: Imagine instead of reading the sentence "The cat sat on the mat," the teacher hands the student a blender full of those words. The student doesn't see the individual words; they see a "smoothie" made of "The + cat + sat."
How it works: The computer takes 8 words (a "bag"), mixes their meanings together into a single "super-word," and asks the AI to predict what the next bag of 8 words will be.
The Benefit: Because the AI is processing 8 words as if they were 1, it can "read" 8 times faster through the data without using any more computer power. It's like the student is now reading 8 pages in the time it used to take to read 1.

Phase 2: The "Fine-tooth Comb" Phase (Recovery)

If you only trained the AI on "smoothies," it would be terrible at writing normal sentences. It wouldn't know the order of words, and its output would be gibberish.

So, after a certain amount of time (usually about 30% of the total training), the researchers stop the "smoothie" phase.

The Analogy: They switch back to the standard method. They take the AI's brain (which has now learned the general structure of language very quickly) and put it back in a normal classroom. They stop using the blender and start showing it individual words again.
The Result: Because the AI already learned the "big picture" so efficiently in Phase 1, it recovers very quickly. It catches up to and often surpasses the performance of models that were trained the "slow way" the whole time.

Why This is a Big Deal

The paper claims this method is a "drop-in" solution. You don't need to change the AI's brain, the computer chips, or the books. You just change how you feed the data during the first part of training.

The Speedup: In their tests with a large model (10 billion parameters), this method allowed them to achieve the same level of intelligence in half the time (a 2.5x speedup) compared to standard training.
The Trade-off: They are trading "data consumption" for "speed." The AI reads more data in the same amount of time, but because it reads it in "bags," it learns the patterns faster.

What the Paper Does NOT Claim

It is important to stick to what the authors actually said:

It doesn't change the final product: Once the training is done, the AI is exactly the same as a normal AI. It can still write coherent sentences, code, and stories. The "smoothie" trick is only used during the learning process.
It's not about "thinking" faster: This doesn't make the AI smarter at reasoning or solving complex math problems during the training; it just makes the training process itself more efficient.
It's not a magic bullet for everything: The authors tested this on models ranging from small (270 million parameters) to large (10 billion). It worked well across the board, but they note that if you have infinite data and time, the benefits might look different.

Summary

Think of TST as a speed-reading course for AI.

First, you teach it to skim: You show it chunks of text mixed together so it can absorb the general vibe of the language very quickly.
Then, you teach it to read normally: You switch back to word-by-word reading to refine its skills.

The result? The AI finishes its "education" in record time, using the same amount of energy, and ends up just as smart (or smarter) than those who took the slow, traditional route.

Technical Summary: Efficient Pre-Training with Token Superposition

Problem Statement

The pre-training of Large Language Models (LLMs) is increasingly prohibitively expensive and inefficient at scale. Current paradigms often require complex, invasive modifications to parallelism strategies, optimizers, tokenizers, or model architectures to achieve high data throughput. Furthermore, while recent research has focused on improving inference-time efficiency or decoupling training/inference efficiency, there is a lack of methods that significantly boost training-time throughput without altering the final model's architecture or inference dynamics. Existing approaches often conflate training and inference efficiency or rely on auxiliary losses that add parameters and hyperparameters without necessarily increasing the raw token processing rate per FLOP.

Methodology: Token Superposition Training (TST)

The authors propose Token Superposition Training (TST), a "drop-in" method designed to improve data throughput per FLOP during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST operates in two distinct phases:

Phase 1: Superposition Phase

In this phase, the model processes a "bag" of contiguous tokens as a single latent unit, referred to as an "s-token."

Input Superposition: A contiguous sequence of $s$ tokens is segmented into non-overlapping bags. The embeddings of these $s$ tokens are averaged to create a single superposed latent token embedding. This reduces the sequence length by a factor of $s$ (from $L$ to $\lfloor L/s \rfloor$ ).
Throughput Mechanism: To ensure a fair comparison with baseline training (equal-FLOPs), the data sequence length $L$ is increased by a factor of $s$ during this phase. Consequently, the model ingests $s$ times more data tokens for the same computational cost per step.
Output Superposition: Instead of predicting a single next token, the model predicts the next "bag" of $s$ tokens. The standard one-hot cross-entropy (CE) loss is replaced with a Multi-Hot Cross-Entropy (MCE) loss. The MCE objective treats the target bag as a multi-hot label, pushing the probability of each valid token in the bag toward $1/s$ . The loss is formulated as the average of the standard CE losses for each token in the bag, minus a constant term ( $\log |y|$ ) which does not affect gradients.
Causality: To preserve causality, standard next-token prediction labels are shifted left by $s-1$ before being split into non-overlapping bags.

Phase 2: Recovery Phase

A model trained exclusively with TST produces nonsensical outputs because it loses token ordering within the bags and lacks standard autoregressive sampling capabilities. To remedy this:

The training reverts to the standard next-token causal prediction objective.
The model resumes training from the checkpoint saved at the end of the superposition phase.
The TST code is fully removed to ensure the final model architecture and inference dynamics are identical to the baseline.
The ratio of steps spent in the superposition phase is denoted as $r$ (e.g., $r=0.3$ means 30% of steps are TST, 70% are standard).

Key Contributions

Novel Training Paradigm: Introduction of TST, which increases token throughput $s$ -fold during a specific training regime without changing the model's final architecture or inference behavior.
Two-Phase Design: A simple mechanism combining input superposition (averaging embeddings) and output superposition (multi-hot prediction) followed by a recovery phase to align the model back to standard autoregressive behavior.
Orthogonality to Existing Methods: The authors position TST as orthogonal to auxiliary-loss methods (like Multi-Token Prediction). While MTP adds parameters and targets future tokens without increasing tokens-per-FLOP, TST strictly increases throughput and leaves the architecture untouched.
Empirical Validation: Extensive evaluation across model scales (270M, 600M, 3B, and 10B Mixture-of-Experts) demonstrating robustness across different hyperparameter settings.

Results

The paper reports the following findings based on experiments using the DCLM and FineWeb-Edu datasets:

Loss Reduction: TST consistently outperforms baseline models in terms of final training loss under equal-FLOPs constraints.
Downstream Performance: Models trained with TST show improved performance on downstream benchmarks (HellaSwag, ARC-E/C, MMLU, etc.) compared to baselines trained for the same number of steps or same computational cost.
Efficiency Gains:
- Under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B parameter scale.
- In the 10B A1B (10B total, 1B active) MoE experiment, the TST model saw 2T data tokens compared to the baseline's 1.05T tokens while achieving a lower final loss (2.236 vs. 2.252) in significantly fewer B200-hours (4768 vs. 12311).
Hyperparameter Sensitivity: The method is robust within a reasonable range of superposition bag sizes ( $s \in [4, 8]$ ) and step ratios ( $r \in [0.2, 0.4]$ ).
Ablation Studies:
- Input vs. Output: Both input-only and output-only superposition improve performance over the baseline, but the combination of both yields the best results, suggesting they are orthogonal mechanisms.
- Representation Alignment: Randomly re-initializing the input embeddings and output LM heads between the superposition and recovery phases completely eliminates the gains of TST, indicating that maintaining representation alignment across phases is critical.

Significance and Claims

The paper claims that TST offers a significant improvement in pre-training efficiency by maximizing training-time throughput independently of the model's inference-time architecture. The authors argue that this approach challenges the monolithic pre-training paradigm by demonstrating that a "coarse-to-fine" granularity schedule is a reusable ingredient for efficient pre-training.

Key claims include:

Cost Reduction: Achieving the same model quality (loss) at roughly half the computational cost.
Simplicity: The method requires no complex architectural changes, new optimizers, or additional parameters, making it a practical "drop-in" solution.
Decoupling Efficiency: It successfully decouples training-time efficiency from inference-time efficiency, addressing a gap in current literature where many methods conflate the two.
Scalability: The method scales effectively from small dense models (270M) to large Mixture-of-Experts models (10B), suggesting broad applicability.

The authors remain modest regarding the underlying mechanisms, noting that while they hypothesize the input superposition acts as a form of "pre-pre-training" or implicit regularization of embedding geometry, further interpretability work is needed. They also acknowledge limitations, such as the assumption that pre-training is currently compute-bound rather than data-bound, and the lack of large-scale ablations due to compute constraints.