Nous Research releases Token Superposition Training with 2-3× speedup

VIEWS40.9KBOOKMARKS200

Did they just reinvent Patched Training for LLMs ? 🤔

They even have a very similar plot 👀

Today we release Token Superposition Training (TST), a modification to the standard LLM pretraining loop that produces a 2-3× wall-clock speedup at matched FLOPs without changing the model architecture, optimizer, tokenizer, or training data.

During the first third of training, the model reads and predicts contiguous bags of tokens, averaging their embeddings on the input side and predicting the next bag with a modified cross-entropy on the output side. For the remainder of the run, it trains normally on next-token prediction. The inference-time model is identical to one produced by conventional pretraining.

Validated at 270M, 600M, and 3B dense scales, and at 10B-A1B MoE.

The work on TST was led by @bloc97_, @gigant_theo, and @theemozilla.

38d40.9K242200

LIKES316RETWEETS29

elie@eliebakouch

nice pre training work by nous claiming ~2.5x efficiency gains, building on previous research like MTP/SuperBPE. overall intuition is that at each step you want the model to process and predict more tokens

Nous Research@NousResearch

Today we release Token Superposition Training (TST), a modification to the standard LLM pretraining loop that produces a 2-3× wall-clock speedup at matched FLOPs without changing the model architecture, optimizer, tokenizer, or training data.

During the first third of training, the model reads and predicts contiguous bags of tokens, averaging their embeddings on the input side and predicting the next bag with a modified cross-entropy on the output side. For the remainder of the run, it trains normally on next-token prediction. The inference-time model is identical to one produced by conventional pretraining.

Validated at 270M, 600M, and 3B dense scales, and at 10B-A1B MoE.

The work on TST was led by @bloc97_, @gigant_theo, and @theemozilla.

38d25.8K316178

REPLIES10

Teknium 🪽@Teknium

Check out our researchers' latest paper that introduces Superposition, a potential path to multiplying training speed during pre-training.

Nous Research@NousResearch

Today we release Token Superposition Training (TST), a modification to the standard LLM pretraining loop that produces a 2-3× wall-clock speedup at matched FLOPs without changing the model architecture, optimizer, tokenizer, or training data.

During the first third of training, the model reads and predicts contiguous bags of tokens, averaging their embeddings on the input side and predicting the next bag with a modified cross-entropy on the output side. For the remainder of the run, it trains normally on next-token prediction. The inference-time model is identical to one produced by conventional pretraining.

Validated at 270M, 600M, and 3B dense scales, and at 10B-A1B MoE.

The work on TST was led by @bloc97_, @gigant_theo, and @theemozilla.

38d14K21437

Clive Chan@itsclivetime

this is quite remarkable

Nous Research@NousResearch

Today we release Token Superposition Training (TST), a modification to the standard LLM pretraining loop that produces a 2-3× wall-clock speedup at matched FLOPs without changing the model architecture, optimizer, tokenizer, or training data.

During the first third of training, the model reads and predicts contiguous bags of tokens, averaging their embeddings on the input side and predicting the next bag with a modified cross-entropy on the output side. For the remainder of the run, it trains normally on next-token prediction. The inference-time model is identical to one produced by conventional pretraining.

Validated at 270M, 600M, and 3B dense scales, and at 10B-A1B MoE.

The work on TST was led by @bloc97_, @gigant_theo, and @theemozilla.

38d23.2K9664

Anjney Midha@AnjneyMidha

very cool

a 2-3x speed up in training by essentially letting the model learn more flexibly in its early stages than rigid regimes

sort of akin to how homeschooling is much better for some kids than factory education

Nous Research@NousResearch

Today we release Token Superposition Training (TST), a modification to the standard LLM pretraining loop that produces a 2-3× wall-clock speedup at matched FLOPs without changing the model architecture, optimizer, tokenizer, or training data.

During the first third of training, the model reads and predicts contiguous bags of tokens, averaging their embeddings on the input side and predicting the next bag with a modified cross-entropy on the output side. For the remainder of the run, it trains normally on next-token prediction. The inference-time model is identical to one produced by conventional pretraining.

Validated at 270M, 600M, and 3B dense scales, and at 10B-A1B MoE.

The work on TST was led by @bloc97_, @gigant_theo, and @theemozilla.

37d20.9K8962

François Fleuret@francoisfleuret

Awesome. LLM mixup weirdness.

Nous Research@NousResearch

Today we release Token Superposition Training (TST), a modification to the standard LLM pretraining loop that produces a 2-3× wall-clock speedup at matched FLOPs without changing the model architecture, optimizer, tokenizer, or training data.

During the first third of training, the model reads and predicts contiguous bags of tokens, averaging their embeddings on the input side and predicting the next bag with a modified cross-entropy on the output side. For the remainder of the run, it trains normally on next-token prediction. The inference-time model is identical to one produced by conventional pretraining.

Validated at 270M, 600M, and 3B dense scales, and at 10B-A1B MoE.

The work on TST was led by @bloc97_, @gigant_theo, and @theemozilla.

37d15.3K6326

Asking why@a66mike99

@NousResearch @witcheer I had to screenshot this as the perfect encapsulation of AI 2026.

38d8.3K785

Dimitris Papailiopoulos@DimitrisPapail

@NousResearch Hi 😝 https://openreview.net/pdf?id=TZxPjnYDBI

38d3.1K3111

elie@eliebakouch

@eric_alcaide this plot is even better to show similarity imo

38d1.4K399

Eric Alcaide@eric_alcaide

Link: https://arxiv.org/abs/2407.12665v2 with code here: https://github.com/shaochenze/PatchTrain from Chenze Shao who's now at DeepSeek

38d1.6K189

Ajith@ajith_io

@NousResearch I made an explainer video on how this system really works. If you don’t have time right now, feel free to save it for later.

Credits:- created using distilbook tool..

38d98477

Nous Research@NousResearch

TST works in two phases. In phase 1, which covers the first 20-40% of training, the model reads contiguous bags of k tokens, with input embeddings averaged within each bag, and predicts the next bag of k tokens using a sum of standard cross-entropy terms.

In phase 2, it returns to ordinary next-token prediction for the remainder of the run. The loss curve resolves within a few thousand steps and crosses below the matched-FLOPs baseline for the rest of the run.

The entire input-side change is a reshape and a mean over the embedding lookup. The output-side change is a sum of standard cross-entropy terms. No new kernel, no auxiliary head, and no change to the output projection.

38d2.4K471

Nous Research@NousResearch

We validated TST at four model scales: 270M, 600M, and 3B dense, and a 10B-A1B mixture-of-experts. The 3B and 10B runs were using non-cherry-picked standard training hyperparameters. The 10B run beats the matched-FLOPs baseline on HellaSwag (+1.1 pts), ARC-Easy (+0.4), ARC-Challenge (+1.0), and MMLU (+1.6); a stronger result than a loss improvement alone, given the well-known tendency of training-loss reductions to fail to transfer to downstream tasks.

The method has two knobs: bag size k and step ratio r. It is possible that the optimal bag size grows with model size, where bigger models can process a bigger bag of tokens during training. The step ratio is less sensitive; any r between 0.2 and 0.4 works at every scale we tested.

38d2K421

nathan chen@nathancgy4

@NousResearch fascinating paper! would be curious if there's any specific tricks used to handle lr scheduling / batch size at phase transition, or was it basically just wsd throughout regardless of which phase you're in (😳)

38d1.5K263

catid@MrCatid

@NousResearch How does it improve on https://github.com/shaochenze/PatchTrain

38d980132

Clive Chan@itsclivetime

tl;dr:

3x fewer steps iso-data, by pre-pre-training on a new objective: - segment of 8 tokens (mean pooled embeddings) ==> next segment of 8 tokens (multi-hot cross-entropy)

Clive Chan@itsclivetime

this is quite remarkable

38d84944

Nous Research@NousResearch

What we find most useful about TST is the decoupling. The training-time efficiency is fully separated from the inference-time architecture, which makes TST a clean addition on top of other pretraining improvements: sparse attention, MoE routing, alternative tokenizers, optimizer changes. Most pretraining-efficiency interventions don't satisfy this criterion, and we think more should.

TST is the first of several research releases from the Nous pretraining group going up over the next two weeks. If you want to work on problems like this, find us on Discord.

38d1.9K21

Jean-Gabriel BARTHELEMY@JG_Barthelemy

@NousResearch How is it different from patch-level training ? It has been around since July 2024 : https://arxiv.org/abs/2407.12665 Interestingly it doesn't work on hybrid architecture (GDN or Mamba-2)

38d1.7K21

Nous Research@NousResearch

@witcheer The Research part of Nous Research

38d45326

Clive Chan@itsclivetime

someone in the comments pointed out similarity to this prior work, which i've never heard of before:

Jean-Gabriel BARTHELEMY@JG_Barthelemy

@NousResearch How is it different from patch-level training ? It has been around since July 2024 : https://arxiv.org/abs/2407.12665 Interestingly it doesn't work on hybrid architecture (GDN or Mamba-2)

38d61562