---
title: Token-Superposition Training (TST)
url: https://www.emergentmind.com/topics/token-superposition-training-tst
type: topic
---

# Token-Superposition Training (TST)

Token-Superposition Training (TST) is a methodology for increasing data-throughput efficiency in large-scale language model pre-training and chain-of-thought (CoT) reasoning by representing and processing multiple tokens in superposed forms. TST, as formalized in recent works by Peng et al. (“Efficient Pre-Training with Token Superposition”) and extended in the SuperThoughts framework, operates by “fusing” groups of tokens into single embeddings or latent representations, training the model either to predict multiple subsequent tokens jointly or to compress multi-token reasoning steps into fewer forward passes. This strategy achieves substantial reductions in training and inference FLOPs, with robust empirical performance across model scales and tasks [2605.06546, 2606.13862].

## 1. Methodological Foundations

TST is implemented through two broad paradigms:
- **Pre-training throughput optimization** [2605.06546]: In this setting, TST is applied as a two-phase regimen. The model first undergoes a superposition phase in which input tokens are grouped into non-overlapping “bags” of size $s$, each represented by the arithmetic mean of their embeddings. The LM is then tasked with predicting the next bag of $s$ tokens, using a multi-hot cross-entropy (MCE) objective that generalizes the usual one-hot cross-entropy loss to multi-target settings. After a tunable fraction $r$ of training steps, the process reverts to standard autoregressive next-token prediction for the remainder (“recovery phase”) with no modification to model, optimizer, or tokenizer.

- **Token-superposition in CoT reasoning** [2606.13862]: Here, the discrete CoT token sequence $c_{1:L_c}$ is partitioned into pairs $(c_{2i-1}, c_{2i})$, each compressed into a single latent $z_i$ by the Compressor module (either a linear projection or a shallow Transformer). The LLM backbone, acting as a sequence model over $z_i$, alternates with a lightweight Multi-Token Prediction (MTP) head to predict two CoT tokens per inference step. An adaptive mechanism reverts to single-token decoding when the MTP’s prediction confidence is low.

## 2. Formal Algorithms and Training Phases

### TST Pre-Training Workflow [2605.06546]

**Superposition Phase:**  
Given batch shape $B \times L$ and bag size $s>1$:
- Partition each input into contiguous “bags” $\mathbf t = [t_{i+1}, ..., t_{i+s}]$
- Replace $\{e(t_j): j=1..s\}$ with averaged embedding $e'(\mathbf t) = \frac{1}{s}\sum_{j=1}^s e(t_{i+j})$
- Increase either sequence length or batch size by $s$ to maintain constant per-step FLOPs
- For targets, form multi-hot vectors $\mathbf m \in \{0,1\}^V$ for each bag $\mathbf y$ of $s$ next tokens, and minimize
  $$
  \mathcal L_{\rm MCE}(\mathbf z, \mathbf y) = -\frac{1}{s}\sum_{i \in \mathbf y}\log p_i
  $$

**Recovery Phase:**  
After $r$ fraction of total steps, revert to plain next-token prediction.

**Parameter Constraints:**  
Effective settings:
- $s \in \{4,8\}$, $r \approx 0.2\text{–}0.4$

### Token Superposition in Reasoning ("SuperThoughts") [2606.13862]

**Compression:**  
For each token pair $(c_{2i-1}, c_{2i})$, produce latent $z_i = Comp(\Emb(c_{2i-1}), \Emb(c_{2i}))$, with two $Comp$ options:
- Linear: $z_i = P[\Emb(c_{2i-1}) \Vert \Emb(c_{2i})]$
- Tiny Transformer: $z_i = TF(\Emb(c_{2i-1}), \Emb(c_{2i}))_2$

**Inference:**  
- Main LLM processes the $z_i$; at each step $i$,
  - Predicts odd-indexed token $\hat c_{2i+1}$ via shared $\mathrm{LMHead}$
  - MTP head predicts even-indexed token $\hat c_{2i+2}$ from RMS-normed embeddings and hidden states:
    $$
    z'_i = P'[\mathrm{RMSNorm}(\Emb(c_{2i})),\,\mathrm{RMSNorm}(\Emb(\hat c_{2i+1})),\,\mathrm{RMSNorm}(h_i)]
    $$
    $$
    h'_i = TF_{1\text{-layer}}(z'_i)
    $$

**Training Objectives:**
- Compressor distillation (teacher-student): align hidden states using smoothed L1 loss
- Full cross-entropy on all CoT and answer tokens

**Adaptive Decoding:**  
On low MTP confidence ($p^\mathrm{MTP}_i < \tau$), inference reverts to single-token step.

## 3. Empirical Performance

### Pre-Training (TST) [2605.06546]

Comprehensive evaluation up to 10B A1B scale demonstrates:
- On equal-FLOPs, TST enables throughput gains of up to $2.5\times$ at 10B with MoE, as quantified by B200-GPU-hours.
- For dense models (270M, 600M, 3B), TST matches or outperforms baselines in final validation loss and 0-shot downstream tasks.

| Model                  | Parameters | Phase I Steps / Total | $s$ | Equiv. Tokens | B200h | Final Loss |
|------------------------|------------|----------------------|-----|---------------|-------|------------|
| Dense Baseline (270M)  | 270M       | –                    | –   | 42B           | 34    | 3.212      |
| Dense TST (270M)       | 270M       | 6,000/20,000         | 6   | 105B          | 34    | 3.142      |
| Dense Baseline (600M)  | 600M       | –                    | –   | 42B           | 61    | 3.019      |
| Dense TST (600M)       | 600M       | 6,000/20,000         | 6   | 105B          | 61    | 2.943      |
| Dense Baseline (3B)    | 3B         | –                    | –   | 42B           | 247   | 2.808      |
| Dense TST (3B)         | 3B         | 6,000/20,000         | 6   | 105B          | 247   | 2.676      |
| MoE Baseline (10BA1B)  | 10BA1B     | –                    | –   | 1.05T         | 12,311| 2.252      |
| MoE TST (10BA1B)       | 10BA1B     | 12,483/49,983        | 16  | 2.00T         | 4,768 | 2.236      |

On standard downstream tasks (HellaSwag, ARC, MMLU, BoolQ, PIQA), TST achieves parity or gains relative to baseline training.

### Chain-of-Thought Reasoning (SuperThoughts) [2606.13862]

Across Qwen2.5-Math models (1.5B, 7B, 14B):

- **CoT length reduction:** 20–35% under adaptive decoding, up to ~50% with fixed 2-token steps.
- **Accuracy drop:** Maintained within 1–2 percentage points for all major benchmarks (MATH500, AMC23, OlympiadBench, GPQA-Diamond).
- **Linear compressor matches Transformer variant**—preferred for efficiency.

Example results for Qwen2.5-Math-7B-Instruct:
- Baseline (MATH500): 83.0% / 538.6 tokens
- SuperThoughts: 80.8% / 357.3 tokens (–34% length, –2.2pp accuracy)

Wall-clock speedups are confirmed (e.g., 32.8% CoT reduction yields 28.3% end-to-end time reduction). Non-adaptive use ($\tau=0$) cuts chain length further but substantially harms accuracy (–10–20pp).

## 4. Mechanistic Insights

TST’s improvements in throughput and data efficiency derive from two complementary mechanisms [2605.06546]:
- **Input Superposition:** The early-phase “averaged” embeddings provide local and global statistical priors cheaply over large corpora at reduced precision, enabling faster initial convergence.
- **Output Superposition:** Predicting bags of tokens via MCE aggregates multiple local targets, increasing the effective data exposure per gradient step.

In the SuperThoughts formulation, latent superposition condenses reasoning steps, but joint supervision and adaptive fallback preserve the fidelity of multi-step reasoning. Critically, all such superposition is restricted to pre-training or intermediate latent space—no inference-time architectural changes are required in TST for language models, and SuperThoughts introduces negligible computational overhead with the lightweight MTP head.

Ablation studies demonstrate that both input and output superposition individually surpass the baseline, but their combination provides maximal benefit. Proper sharing of the token embedding and head between phases is essential; disruption by reinitialization negates the advantage.

## 5. Hyperparameters, Robustness, and Implementation

- **Default settings:** For TST pre-training, $s = 4$ or $8$ and $r \approx 0.2$–$0.4$ provide robust gains. Uniform averaging in the MCE loss is recommended for $s \le 8$, while more complex weighting may help for larger bags.
- **Adaptivity:** In multi-token inference, an MTP confidence threshold $\tau$ controls the fallback policy, optimizing the FLOP–accuracy tradeoff.
- **Implementation:** TST requires no modification to model architecture, optimizer, tokenizer, or parallelization. MCE is realized by repeatedly applying one-hot cross-entropy in a loop. For SuperThoughts, a linear compressor suffices.
- **Ablations:** Varying $s$ and $r$ indicates a U-shaped loss landscape in $s$ (optimal at medium values), and an optimal $r$ band aligning with $0.2$–$0.4$.
- **Inference:** All superposition code is removed after the superposition phase in TST; the model continues with the baseline recipe. In SuperThoughts, the main LLM and MTP modules alternate, with confidence-based fallbacks minimizing quality loss.

## 6. Impact, Limitations, and Relations to Adjacent Methods

TST achieves significant data-throughput scaling without intrusive algorithmic changes or inference penalty. It is compatible with standard dense and mixture-of-experts (MoE) architectures. A core distinction from other compression or acceleration strategies is the explicit preservation of next-token supervision and seamless reversibility: after the superposition phase, the fully expressive autoregressive model is recovered and all efficiency gains are realized in pre-training compute, with no downstream tradeoff.

SuperThoughts extends TST principles to structured reasoning, compressing discrete CoT chains for double-token throughput with minimal degradation—outperforming methods that rely purely on latent-state reasoning without token-level supervision.

A plausible implication is that token-superposition may catalyze further advances in model pre-training regimes, multi-token prediction objectives, and efficient inference strategies across large-scale NLP systems. However, strong accuracy preservation at even higher compression ratios or in settings with highly non-local token dependencies remains to be systematically explored. Constraints on maximal $s$, optimal $r$, and architectural compatibility are empirically delineated but could be model-specific.

## 7. Summary Table: Comparative Results for TST Pre-Training

| Model                  | Scale   | $s$ | Superposition/Total Steps | Equiv. Tokens | GPU Time | Final Loss |
|------------------------|---------|-----|--------------------------|---------------|----------|------------|
| Dense Baseline         | 3B      | –   | –                        | 42B           | 247 h    | 2.808      |
| Dense TST              | 3B      | 6   | 6,000/20,000             | 105B          | 247 h    | 2.676      |
| MoE TST                | 10B A1B | 16  | 12,483/49,983            | 2.00T         | 4,768 h  | 2.236      |

TST and its CoT variant provide a simple, robust, and backward-compatible mechanism for improving large language model efficiency without sacrificing inference quality or requiring complex architectural overhauls [2605.06546, 2606.13862].

Source: https://www.emergentmind.com/topics/token-superposition-training-tst