NITP: Next Implicit Token Prediction for LLM Pre-training

Abstract: Standard next-token prediction (NTP) supervises LLMs solely through discrete labels in the output logit space. We argue that this sparse one-hot supervision leaves the latent representation space under-constrained, allowing hidden states to drift into degenerate and anisotropic configurations that can limit generalization. To address this issue, we propose Next Implicit Token Prediction (NITP), which augments discrete prediction with dense continuous supervision directly in the representation space. NITP trains the model to predict the implicit semantic content of the next token, using shallow-layer representations from the same model as stable self-supervised targets. We provide theoretical analysis showing that NITP regularizes the optimization landscape by mitigating under-constrained degrees of freedom and encouraging a compact, structured representation geometry. Empirically, across dense and MoE models ranging from 0.5B to 9B parameters, NITP consistently improves downstream performance with negligible computational overhead. On a 9B MoE model, NITP achieves a 5.7% absolute improvement on MMLU-Pro, along with gains of 6.4% on C3 and 4.3% on CommonsenseQA, with approximately 2% additional training FLOPs and no additional inference cost. Our implementation is available at https://github.com/aHapBean/NITP.

• The paper introduces NITP, a method that augments traditional token prediction with continuous representation-level supervision to counter hidden state degeneration.
• It employs a cosine similarity loss on temporally shifted shallow-layer targets, yielding significant empirical improvements with minimal computational overhead.
• The theoretical analysis shows that NITP regularizes the angular geometry of hidden states, enhancing robustness and transferability across large-scale autoregressive models.

NITP: Next Implicit Token Prediction for LLM Pre-training

Motivation and Problem Statement

Standard LLM pre-training employs Next-Token Prediction (NTP), which supervises models via discrete token labels projected into the output logit space. While effective at learning generative capabilities, NTP only constrains hidden representations along the target logit direction, leaving many orthogonal degrees of freedom under-constrained. As a result, hidden states can drift, yielding degenerate and highly anisotropic representations—phenomena known as representation degeneration—that limit the model's expressiveness and transferability to downstream tasks.

Figure 1: Effective rank dynamics reveal representation collapse under NTP but preservation under NITP.

This geometric degeneration manifests as sharply decreasing effective rank and increased average cosine similarity between hidden states during training, as visualized in Figure 1. Such collapse indicates a loss of semantic richness, with representations compressed into a narrow subspace.

NITP: Method and Representation-Level Supervision

To address the under-constrained nature of NTP, the authors propose Next Implicit Token Prediction (NITP), which augments token-level supervision with continuous representation-level supervision. Rather than predicting only the discrete next token, NITP tasks the model with predicting the latent semantic representation of the next token, termed the "implicit token." These targets are sourced from shallow-layer representations of the same model at future time steps (i.e., $t+1$), acting as semantic anchors.

Figure 2: NITP architecture overview: hidden states are supervised by predicting temporally shifted shallow-layer representations and optimized jointly with NTP.

NITP employs a stop-gradient operation to prevent co-adaptation of targets and introduces a lightweight projector (e.g., MLP) between deep predictive states and shallow semantic targets. The supervision is applied via a cosine similarity loss, enforcing angular alignment while maintaining scale invariance.

Theoretical Foundations

The manuscript provides a rigorous analysis of NTP and NITP's optimization landscapes. The standard NTP loss exhibits high invariance to perturbations orthogonal to the target token embedding, leading to a rank-deficient Hessian and permitting semantic drift. NITP, by contrast, regularizes the angular geometry of hidden representations.

The core findings show that:

• The NITP Hessian is strictly positive-definite in the angular subspace when alignment is high, performing spectral lifting and mitigating null space directions that NTP leaves unconstrained.
• NITP imposes zero curvature on the norm (radial direction), preserving compatibility with NTP and softmax scaling.
• Empirical results confirm rapid and stable convergence of the NITP loss to low values, validating the theoretical curvature guarantee.

Empirical Results

Comprehensive experiments are conducted on both Mixture-of-Experts (MoE) and dense Transformer models, from 0.5B to 45B parameters, across diverse benchmarks: knowledge understanding (MMLU, C-Eval, Xiezhi), reasoning (ARC-Challenge, CommonsenseQA, C3), mathematics (GSM8k), code generation (LCBench), reading comprehension, and more. NITP consistently improves downstream performance with negligible computational overhead (~2% extra FLOPs; no inference cost).

Strong numerical improvements include:

• 9B MoE: +5.7% absolute improvement on MMLU-Pro, +6.4% on C3, +4.3% on CommonsenseQA.
• Dense 2B: +4.2 points on C3, +2.2 on MMLU.
• 45B MoE scaling: +1.7 points average improvement, validating effectiveness at large scale.

  Figure 3: Average performance under varying NITP loss weights confirms low sensitivity and robust improvements.

Ablation studies reveal:

• Temporal shift (predicting next-token latent representations) is critical; same-step alignment fails to improve performance or loss (Figure 4).
• Cosine similarity loss is superior to MSE, smooth-$\ell_1$, or KL-divergence due to scale invariance and training stability.
• Shallow-layer targets (≈20% model depth) maximize semantic richness and stability; deeper layer targets degrade performance.
• NITP supervision must be applied from early training; delayed activation yields lower scores.

  Figure 4: Loss comparison: temporal shift to next-token is essential for meaningful predictive supervision.

  

  Figure 5: Training dynamics of NITP loss, exhibiting initial collapse, transient hump, and stable convergence.

Practical and Theoretical Implications

Practically, NITP offers a low-overhead auxiliary objective that transforms representation geometry, preventing collapse and enhancing generalizability for downstream tasks. Its design is scalable, architecture-agnostic, and incurs no inference-time penalty. Theoretically, NITP provides a mechanism for direct regularization of latent states, mitigating angular null spaces and aligning deep predictive representations with semantically rich anchors.

Implications include:

• Enhanced robustness and transferability of hidden states, as validated by improvements on frozen representation benchmarks such as MTEB.
• Potential for broader exploration of representation-level supervision objectives in autoregressive pre-training beyond NTP and multi-token prediction.
• Structured geometry induced by NITP may augment fine-tuning efficacy and task adaptation in future LLMs.

Future Directions

• Automated target-layer selection and loss weight adaptation strategies to further reduce hyperparameter tuning.
• Integration of NITP with multi-token and concept-level prediction objectives for richer semantic supervision.
• Application to multimodal pre-training, where latent semantic anchors may come from vision or audio streams.
• Investigation of NITP's effects on prompt-based and retrieval-augmented LLMs, where representation quality is critical.

Conclusion

NITP establishes a principled framework for augmenting LLM pre-training with predictive supervision in latent semantic space, formally regularizing representation geometry and preventing degeneracy. Theoretical analysis and empirical validation demonstrate consistent improvements across architectures and scales, with negligible overhead, marking NITP as an effective objective for improving the quality and transferability of LLM representations.