PaperlayerView exact revision on arXiv

Source-linked AI summary

NITP: Next Implicit Token Prediction for LLM Pre-training

Xiangdong Zhang, Debing Zhang, Shaofeng Zhang, Xiaohan Qin, Yu Cheng, Junchi Yan

arXiv:2605.24956v3cs.CL

TL;DR

Standard next-token prediction weakly constrains hidden-state geometry, motivating NITP, which adds continuous semantic supervision from shallow-layer representations. Across dense and mixture-of-experts models, NITP consistently improves downstream performance with minimal computational overhead.

  • Problem

    Standard next-token prediction weakly constrains latent representations, allowing hidden states to drift toward degenerate, anisotropic configurations.

  • Method

    NITP augments next-token prediction with continuous supervision that aligns final hidden states with shallow-layer representations of the next token's implicit semantic content.

  • Results

    Across dense and mixture-of-experts models, NITP consistently improves downstream performance with minimal computational overhead.

  • Takeaways & Limitations

    NITP provides a representation-level pretraining objective that improves representation geometry and downstream performance without changing model architecture or deployment settings.

  • Takeaways & Limitations

    NITP adds hyperparameters for the target layer and loss weight, and further validation across more models is needed to confirm robustness.

Abstract

from arXiv · show

Standard next-token prediction (NTP) supervises language models 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.

1. Introduction

The introduction argues that standard NTP weakly constrains hidden-state geometry, permitting degenerate representations that may harm downstream generalization. It proposes NITP, which adds dense, continuous, self-supervised representation-space targets and reports theoretical and empirical improvements in geometry and downstream performance.

  • Motivation: Standard NTP supervises hidden representations primarily along the target logit direction, leaving many latent-space degrees of freedom weakly constrained.This under-constrained geometry motivates questioning whether next-token prediction alone sufficiently supervises hidden representations.
  • Motivation: Likelihood-based training can collapse representations into a narrow, anisotropic cone, limiting expressive capacity and degrading downstream generalization.The introduction identifies representation degeneration as a geometric failure mode associated with likelihood-based training.
  • Analysis: The authors track representation geometry using Effective Rank and Average Cosine Similarity between paired training-batch tokens.Effective Rank measures utilized subspace dimensionality, while Average Cosine Similarity proxies global anisotropy.
  • Method: NITP augments discrete next-token prediction with dense, continuous supervision directly in the latent representation space.The objective is designed to address the under-constrained nature of hidden representations under NTP.
  • Method and Results: NITP uses shallow-layer representations as implicit tokens and provides theoretical and empirical evidence of improved representation geometry and downstream performance.The shallow-layer targets enable semantically rich self-supervised supervision.

2. Related Work

Related work centers on next-token prediction as the dominant, scalable pre-training objective and on approaches that supervise continuous representations, including layer-wise distillation.

  • Pre-training objectives for large language models: Next-token prediction is the dominant pre-training paradigm for modern large language models and scales effectively with data and model size.NTP optimizes the likelihood of the next discrete token given its context, supporting general-purpose representations.
  • Representation-level supervision: Prior work has explored direct supervision in continuous representation spaces, including layer-wise distillation that aligns intermediate representations between models.These methods aim to transfer structural or task-relevant knowledge through intermediate representations.

3. Methodology

NITP augments standard next-token prediction with continuous supervision that aligns final hidden states to contextualized shallow-layer representations of the next token. This objective addresses NTP’s under-constrained representation geometry by adding semantic constraints and positive curvature in directions that NTP leaves largely unpenalized.

  • NITP objective: NITP jointly trains the final hidden state to predict both the discrete next-token identity and the continuous implicit token, aligning representations toward structured semantic geometry.The implicit target is self-generated during the forward pass, and the final state is matched using a cosine similarity loss weighted by λ > 0.
  • NTP limitations: NTP primarily constrains the final hidden state along the target-token embedding direction, leaving a broad orthogonal subspace near-invariant to the loss.This creates geometrically distinct latent states with identical token-level likelihoods.
  • NTP limitations: NTP’s unconstrained degrees of freedom can drive hidden states into a low-dimensional, anisotropic cone, reducing latent expressiveness and limiting downstream generalization.The paper reports that standard NTP exhibits rapid Effective Rank deterioration alongside rising global anisotropy.
  • Implicit targets: NITP treats the next implicit token as a dense, contextualized semantic representation of the next token produced by the model’s shallow layers.Shallow-layer targets capture context-dependent meaning while avoiding the overhead and domain shifts of external encoders.
  • Theoretical effect: NITP performs spectral lifting by introducing strict positive curvature in NTP’s semantic null space, mitigating flat valleys and semantic drift.Theorem 3.2 characterizes the total Hessian as Htotal = HNTP + λHNITP and states that null-space directions acquire strict curvature.

4. Experiments

Experiments show that NITP consistently outperforms standard NTP across MoE and dense models at multiple scales, while adding a low-complexity auxiliary objective. Ablations identify temporal shifting, shallow-layer targets, cosine similarity, and λ around 1.0 as important design choices.

  • Evaluation on MoE models: NITP consistently outperforms standard NTP across MoE models from 1.9B to 9B total parameters, demonstrating scalability.The corresponding average score improvements are 0.8, 2.1, and 2.7 points for the evaluated MoE scales.
  • Evaluation on dense models: NITP consistently outperforms the NTP baseline across dense models ranging from 0.5B to 3B parameters, supporting architecture-agnostic effectiveness.Evaluation covers MMLU, C-Eval, BBH, ARC-C, C3, AGIEval, and LCBench.
  • Overall results: NITP constrains latent geometry and improves performance across diverse MoE and dense architectures, including frozen sentence representations.The experiments associate these gains with preventing hidden states from drifting into degenerate configurations.
  • Training efficiency: The NITP auxiliary objective has O(d2) complexity for a MoE model, compared with the standard training cost dominated by backbone blocks and unembedding.NITP introduces a single projection head and a cosine loss.
  • Importance of temporal shift: Predicting the next implicit token significantly outperforms same-position layer alignment, which produces poor loss and performance.This ablation distinguishes NITP’s temporal prediction objective from static layer-wise alignment.
  • Ablation studies: Cosine similarity is more stable than MSE, Smooth-ℓ1, and KL Divergence, while shallow-layer targets consistently outperform middle- and deep-layer alternatives.Performance peaks around λ = 1.0 and remains strong nearby, indicating low sensitivity in that range.

5. Conclusion

NITP augments next-token prediction with continuous representation-space supervision to constrain hidden-state freedom and prevent geometric degeneration. Across dense and MoE models, it improves downstream performance with minimal overhead, while introducing hyperparameters requiring broader validation.

  • Method: NITP augments standard next-token prediction with continuous supervision, predicting the next token’s implicit semantic representation from shallow-layer self-supervised targets.This objective addresses hidden representations that standard token-level supervision leaves under-constrained and prone to geometric degeneration.
  • Theory: NITP constrains free degrees of freedom, mitigates the optimization null space, and regularizes representation geometry to prevent representation degeneration.The theoretical analysis links these effects to improved geometric structure in the learned representations.
  • Empirical results: NITP consistently improves downstream performance across dense and MoE models with minimal computational overhead.The experiments support the method’s effectiveness across both model types while maintaining low added cost.
  • Limitations: NITP introduces additional hyperparameters for the target layer and loss weight, requiring further validation on more models to confirm robustness.Experiments indicate these choices are highly stable across models, but broader evaluation remains necessary.

Impact Statement … A.2. Gradient Derivation

The paper frames NITP as a representation-level auxiliary objective whose expected societal and ethical impacts align with standard large language models. It also formally derives and analyzes the NITP geometry, beginning with normalized representations, implicit targets, cosine similarity, and the resulting gradient.

  • Impact Statement: NITP modifies only the training objective, without new data sources, model architectures, or deployment settings.Its expected societal and ethical impacts are therefore aligned with those of existing large language models trained using standard next-token prediction.
  • A. Detailed Proofs for NITP Geometry: The appendix derives the NITP objective’s gradient and Hessian and analyzes Hessian spectral properties to formalize its geometric regularization claims.This analysis provides the formal proof for Lemma 3.1 and Theorem 3.2.
  • A.1. Notation and Preliminaries: The geometric analysis represents h ∈ R^d as the projected state for implicit prediction and z ∈ R^d as a fixed target from shallow layers.Treating h as the projected state isolates the objective’s geometric incentives.
  • A.1. Notation and Preliminaries: For r > 0, the tangential difference A := v − su is orthogonal to u, since u⊤A = s − s = 0.The appendix explicitly establishes this orthogonality through equation (9).
  • A.1. Notation and Preliminaries: The appendix then states the NITP objective function before proceeding to its differential derivation.This objective is introduced after defining the representation geometry and tangential difference vector.
  • A.2. Gradient Derivation: The gradient derivation applies the Jacobian of the normalization map u(h) = h/r to cosine similarity f(h) = u⊤v and then to L(h) = 1 − f(h).The sequence uses the chain rule to obtain the loss gradient.

A.3. Hessian Derivation (Proof of Lemma 3.1) · A.4. Analysis of Spectral Properties · A.5. Proof of Theorem 3.2 (Null Space Mitigation)

The appendix derives the exact NITP Hessian, characterizes its radial null space and positive tangential curvature, and proves that adding NITP removes semantic directions left unconstrained by NTP. Under positive alignment and λ > 0, NITP supplies strictly positive curvature in those degenerate directions.

  • A.3. Hessian Derivation (Proof of Lemma 3.1): The Hessian derivation differentiates the gradient −r−1A by separately evaluating the derivatives of r−1 and A before combining the terms.Here A = v − su, with v constant.
  • A.3. Hessian Derivation (Proof of Lemma 3.1): The resulting expression provides the exact Hessian form used to prove Lemma 3.1.Near convergence, s → 1 and ∥A∥ → 0, causing the cross terms involving A to vanish and yielding the lemma’s simplified form.
  • A.4. Analysis of Spectral Properties: The radial quadratic form is zero, confirming zero curvature for changes in the vector norm.This follows because I − uu⊤ projects orthogonally to u while A ⊥ u.
  • A.4. Analysis of Spectral Properties: For perturbations w ⊥ u, the cross terms vanish and the remaining tangential curvature is determined by the NITP Hessian’s angular component.The analysis explicitly evaluates the quadratic form for purely semantic changes orthogonal to u.
  • A.4. Analysis of Spectral Properties: When s = cos(h, z) > 0, the Hessian is strictly positive definite on the subspace orthogonal to h.Thus, reasonable alignment during training produces positive curvature in semantic directions.
  • A.5. Proof of Theorem 3.2 (Null Space Mitigation): The standard NTP Hessian is effectively rank-deficient because its gradient signal concentrates on the target token and a sparse set of high-probability contenders.The resulting low-dimensional curvature leaves a large orthogonal complement unconstrained, defining the semantic null space Nsem among tangential directions.
  • A.5. Proof of Theorem 3.2 (Null Space Mitigation): For w ∈ Nsem, adding λHNITP to HNTP yields strictly positive curvature when s > 0 and λ > 0.Because NITP contributes positive tangential curvature where NTP has approximately zero curvature, it mitigates the semantic null space.

B. Experiment Details

The experiments use consistent optimization and comparison protocols across MoE and Dense models, with architectures and NITP settings adapted to model scale. Downstream averages aggregate results from four representative benchmarks.

  • Hyperparameter settings: All models use AdamW with (β1, β2) = (0.9, 0.95), weight decay 0.1, gradient clipping 1.0, and a WSD learning-rate schedule.The schedule includes 2,000 warmup steps.
  • Model architectures: MoE backbones combine a shallow dense stage with multiple MoE layers, each using 144 routed experts, one shared expert, and top-8 routing.The number of MoE layers increases with model capacity, while expert feed-forward size varies across scales.
  • Model architectures: Dense models use standard Transformers with 24–28 layers, while both model families select intermediate NITP target layers and set the NITP loss weight near 1.0.Larger models receive minor adjustments to the NITP loss weight.
  • Training protocol: Each NTP and NITP pair is pretrained with an identical token budget, using the Qwen2 tokenizer, to ensure fair comparison.The total training-token count scales with model size and architecture.
  • Evaluation: Average downstream performance is computed over MMLU, C3, C-Eval, and ARC-Challenge.These four benchmarks define the average shown in the bottom row of Figure 1.

C. Analysis of Representation Geometry Dynamics

The analysis shows that standard NTP drives hidden states toward low-rank, anisotropic configurations, whereas NITP preserves richer, context-dependent geometry. This geometric regularization prevents collapse into a static anisotropic cone and explains NITP’s downstream gains.

  • Effective rank dynamics: Standard NTP rapidly and monotonically collapses effective rank, concentrating hidden states in a small, fixed subspace early in training.The resulting low-variance regime reflects a degenerate equilibrium in which a narrow anisotropic cone suffices for next-token prediction.
  • Effective rank dynamics: NITP maintains substantially higher effective rank with larger variance, reflecting context-dependent activation of different representation subspaces rather than optimization instability.Latent-space supervision introduces positive curvature along directions unconstrained by likelihood-based training.
  • Cosine similarity and anisotropy: Under NTP, average cosine similarity increases steadily during training, signaling growing anisotropy and alignment toward a common dominant direction.NITP alleviates this effect, although cosine similarity still increases as training progresses.
  • Implications: Together, the results indicate that NITP fundamentally alters the geometric equilibrium by enforcing latent semantic consistency and preventing collapse into a static anisotropic cone.This geometric regularization provides a mechanistic explanation for NITP’s improved downstream performance.

D. Hidden-State Representation Quality … F.2. Additional FLOPs Introduced by NITP

NITP improves frozen hidden-state utility across MTEB while preserving language-modeling quality, and its loss evolves from initialization collapse through structured-target emergence to stable convergence. The method adds a projection head and cosine loss, with stop-gradient target extraction avoiding additional target-extraction cost, while training FLOPs are analyzed against a 6× parameter-count approximation.

  • D. Hidden-State Representation Quality: NITP improves representation utility on 23 of 25 MTEB tasks using frozen last hidden states.Evaluation uses mean-pooled, ℓ2-normalized last-layer states from 3B MoE models without fine-tuning or task-specific adaptation.
  • D. Hidden-State Representation Quality: NITP raises the overall MTEB score from 39.24 to 41.56, with gains of +2.20 classification, +2.93 semantic textual similarity, and +1.64 retrieval/duplicate-detection points.The only regressions are AskUbuntuDupQuestions (−0.15 points) and ToxicConversationsClassification (−0.88 points).
  • D. Hidden-State Representation Quality: NITP preserves language-modeling performance, with nearly identical Pile validation cross-entropy and perplexity to NTP across 3B MoE, 9B MoE, and 3B Dense models.Reported NTP versus NITP pairs are 2.006 vs. 2.006, 1.841 vs. 1.840, and 1.884 vs. 1.882 for cross-entropy.
  • E. Loss Curve of NITP: The NITP loss follows three phases: rapid random-initialization collapse, a transient warm-up hump as structured targets emerge, and prolonged stable convergence.The final loss falls below 0.05 for the 9B model, while the evolving difficulty indicates an adaptive predictive task rather than merely a static regularizer.
  • F. Training FLOPs Analysis: Training FLOPs are approximated as approximately 6× the number of parameters involved in forward computation, including forward, backward, and parameter-gradient updates.The analysis compares this estimate for NITP with the baseline Next-Token Prediction objective.
  • F.1. Baseline NTP FLOPs: For baseline NTP, GQA self-attention contributes approximately 18d^2 training FLOPs per token.The estimate assumes 8 query heads and 4 KV heads, with projections contributing approximately 3d^2 parameters.
  • F.1. Baseline NTP FLOPs: For baseline NTP, each MoE layer with k activated SwiGLU experts contributes approximately 18kdde training FLOPs, while unembedding contributes approximately 6Vd.Input embeddings are excluded because lookup operations incur 0 FLOPs.
  • F.2. Additional FLOPs Introduced by NITP: NITP adds a SwiGLU projection head with 12d^2 parameters and a cosine-similarity loss, while stop-gradient target extraction incurs no cost.The projection head uses intermediate dimension 4d; cosine similarity requires approximately 6d operations before the backward-pass accounting.

F.3. Numerical Instantiation: 9B MoE Model · F.4. Empirical Wall-Clock Overhead · G. Inference FLOPs Analysis

For the 9B MoE model, NITP adds approximately 1.18 × 10^8 training FLOPs against a 5.06 × 10^9 baseline, yielding modest practical overhead. NITP adds no inference FLOPs because its projection head is discarded after pre-training.

  • F.3. Numerical Instantiation: 9B MoE Model: 1.18 × 10^8 FLOPs constitute the estimated total NITP overhead for the 9B MoE configuration.The estimate uses d = 1280 in 72d^2 + 18d.
  • F.3. Numerical Instantiation: 9B MoE Model: 5.06 × 10^9 FLOPs constitute the estimated total baseline cost, combining backbone layers and unembedding.The backbone contributes 3.89 × 10^9 FLOPs, while unembedding contributes 1.17 × 10^9 FLOPs.
  • F.3. Numerical Instantiation: 9B MoE Model: The best target layer consistently lies in a shallow contextualized region around 20% of total model depth.This finding comes from a target-layer sweep across scales using average scores over MMLU, MMLU-Pro, CSQA, BBH, and LCBench.
  • F.4. Empirical Wall-Clock Overhead: 17 minutes, or approximately 1.8% wall-clock overhead, separates NITP from NTP in a 5k-step 9B MoE training run.Using the same GPUs and global batch size of 1024, NTP takes 16h 1m and NITP takes 16h 18m.
  • F.4. Empirical Wall-Clock Overhead: The measured wall-clock increase is slightly smaller than the theoretical FLOPs ratio because MoE communication and system overheads are largely unaffected by NITP.The unaffected components include communication and system overheads associated with large-scale MoE training.
  • F.4. Empirical Wall-Clock Overhead: Approximately 2% overhead remains marginal in the 9B MoE setting, supporting NITP scalability.This conclusion summarizes the practical training-time overhead analysis.
  • G. Inference FLOPs Analysis: Zero additional computational overhead occurs during inference because the auxiliary NITP projection head is discarded after pre-training.The deployed architecture remains identical to the standard transformer backbone.
  • G. Inference FLOPs Analysis: Inference FLOPs per token are exactly the same for NITP and baseline NTP models.The analysis therefore preserves generation speed and does not increase serving cost.

H. Additional Ablation Studies

Additional ablations show that NITP performs best with shallow contextualized targets, stop-gradient targets, activation from step 0, and a two-layer SwiGLU projection head. These choices improve stability and preserve performance across downstream and reasoning benchmarks.

  • Target-layer sweep: Layer 4 of 17 for the 3B MoE model and layer 5 of 24 for the 9B MoE model give the best implicit targets.Both layers are near 20% of total depth; embedding targets are weaker, and deeper targets reduce the gain.
  • Stop-gradient on implicit targets: Stop-gradient targets stabilize training and substantially improve downstream performance, whereas gradient flow causes co-adaptation, instability, and degradation.The ablation directly compares allowing gradients into implicit target representations against applying stop-gradient.
  • Start step of NITP: Activating NITP from step 0 yields the best overall results, while delaying activation consistently degrades performance, especially on reasoning benchmarks.Early activation is presented as shaping representation geometry before higher-level semantics stabilize.
  • Effect of the projection head: Removing the two-layer SwiGLU projection head causes notable performance degradation across all benchmarks.The default projector operates on last hidden states, while the projector-free variant directly aligns them with implicit targets.
Loading 2605.24956v3…