NITP: Next Implicit Token Prediction for LLM Pre-training
TL;DR
Standard next-token prediction (NTP, the usual pre-training objective where the model is scored only on the discrete one-hot label of the next word) supervises a model only in the output logit space. NITP argues this sparse signal leaves the model's internal representation space under-constrained, so hidden states drift into degenerate, anisotropic geometry (vectors collapsing into a narrow cone) that hurts generalization. NITP adds a second, dense, continuous objective directly in representation space: the model also predicts the implicit semantic content of the next token, supervised against the model's own shallow-layer representation of that token, used as a stop-gradient self-supervised target. The paper gives theory that this regularizes the optimization landscape toward a compact, structured geometry. Across dense and MoE models from 0.5B to 9B, NITP improves downstream accuracy for roughly 2% extra training FLOPs and zero inference cost, including +5.7 absolute on MMLU-Pro for a 9B MoE.
┌──────────────────────┐
tokens ───────► │ transformer stack │
└──────────┬───────────┘
│ final hidden
┌──────────────┴──────────────┐
▼ ▼
┌───────────────┐ ┌──────────────────┐
│ next-token │ │ NITP head │
│ head (a) │ │ predict implicit │
│ ►one-hot loss │ │ semantic vector │
└───────────────┘ │ of next token (b)│
└────────┬─────────┘
│ supervised vs
▼ (stop-grad)
shallow-layer rep of that token
total loss = (a)+(b) inference uses ONLY path (a)
Key points
- 9B MoE gains: +5.7% absolute on MMLU-Pro, +6.4% on C3, +4.3% on CommonsenseQA.
- Cost: about 2% additional training FLOPs and zero additional inference cost, since the NITP head is dropped at inference.
- Gains are consistent across the full 0.5B to 9B range, on both dense and MoE architectures.
- The target is self-supervised and stable: it reuses the same model's shallow-layer representations as stop-gradient anchors, so no external teacher or extra data is needed.
- Theory frames the benefit as regularizing under-constrained degrees of freedom and pushing the representation geometry away from anisotropic collapse. Code released at github.com/aHapBean/NITP.
How this relates to prior wiki pages
NITP extends the recurring "the discrete token label is lossy supervision" theme, but from the pre-training and representation-geometry angle rather than the post-training one. It rhymes with TIDE (2026-05-09), which showed every layer carries token-level signal that the final logit loss throws away: NITP turns that observation into an auxiliary pre-training target built from shallow-layer representations. It is also the pre-training counterpart to the auxiliary-head idea in joint MTP + RL (2026-05-28), which adds extra prediction heads during training that are cheap or free at inference. Where on-policy distillation work like TA-OPD (2026-06-01), which asks which token signals a student can actually reach concentrates a post-training signal, NITP injects a richer learnable signal into pre-training itself. The anisotropic-collapse framing connects to representation-geometry concerns tracked in rl-for-llms.md and attention-mechanisms.md rather than contradicting any prior page.
Gaps
The largest reported model is 9B, so whether the 2%-FLOPs-for-+5.7-MMLU-Pro tradeoff survives at frontier scale (tens or hundreds of billions of parameters) is untested, and representation under-constraining may behave differently once the model is already heavily over-parameterized. The choice of which shallow layer supplies the target appears fixed rather than swept across depths. Gains are reported on knowledge and commonsense benchmarks rather than on reasoning or long-context tasks where representation geometry may matter differently.
Source: arXiv 2605.24956 · raw file