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 mainly constrains hidden representations along target-logit directions, leaving their geometry under-supervised. NITP adds continuous prediction of the next token’s implicit semantic representation, and experiments report consistent downstream gains across dense and MoE models with minimal computational overhead.

  • Problem

    Standard NTP mainly constrains hidden representations along target-logit directions, leaving open whether it sufficiently supervises their geometry.

    Source: §1 · p. 1

  • Method

    NITP adds continuous prediction of the next token’s implicit semantic representation, using the model’s shallow-layer states as self-supervised targets.

    Source: §3.4 · p. 4, §5 · p. 8

  • Results

    Across dense and MoE models, NITP consistently improves downstream performance with minimal computational overhead.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The findings support complementing token-level supervision with representation-level supervision to improve hidden-representation geometry.

    Source: §1 · p. 2, §5 · p. 9

    Further validation on more models is needed to confirm the robustness of NITP’s target-layer and loss-weight choices.

    Source: §5 · p. 9

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

Because pre-training objectives shape the representations that underpin LLM capabilities, the introduction argues that next-token prediction leaves hidden-state geometry weakly constrained and presents NITP as a self-supervised way to improve it.

1. Introduction

  • Standard NTP supervises through discrete one-hot output targets, leaving many latent-space degrees of freedom weakly constrained and allowing hidden representations to lose expressiveness and distinctiveness.The introduction links this geometric degeneration to potentially degraded downstream generalization.Source: §1 · pp. 1–2
  • During training, NTP’s effective rank rapidly diminishes as average cosine similarity increases, while NITP maintains higher effective rank and prevents excessive anisotropy.Effective rank measures utilized subspace dimensionality, and average cosine similarity between paired batch tokens proxies global anisotropy.Source: Figure 1 · p. 1, §1 · p. 2
  • NITP adds dense continuous supervision for the next token’s latent semantic representation, using the model’s own shallow-layer representations as targets.This design supplies self-supervised targets without external encoders or additional annotations.Source: §1 · p. 2
  • The paper analyzes NTP’s weak representation supervision and reports theoretical and empirical evidence that NITP improves representation geometry and consistently enhances downstream performance.The introduction attributes the geometry problem to NTP’s under-constrained hidden representations.Source: §1 · p. 2

2. Related Work

Prior pre-training objectives mostly supervise discrete token predictions, while representation-level methods provide continuous supervision through distillation or self-supervised consistency. NITP is positioned as an autoregressive latent-space objective that predicts the next token’s semantic representation.

Pre-training objectives for large language models

  • NTP is the dominant LLM pre-training paradigm, while MTP extends prediction to multiple future tokens or compressed context summaries; both rely primarily on discrete token-level supervision.The paper contrasts these objectives with NITP’s dense, continuous representation-level supervision.Source: §2 · p. 2

Representation-level supervision

  • Prior representation-level approaches include layer-wise distillation between models and self-supervised consistency across time or views, including contrastive and self-distillation methods.Examples include Sheared LLaMA, which aligns hidden states to aid knowledge recovery after structured pruning.Source: §2 · p. 2
  • NITP differs by addressing under-constrained hidden states during LLM pre-training with an autoregressive latent-space objective that predicts the next token’s semantic representation.Unlike the cited approaches, its formulation specifically predicts the semantic representation of the next token.Source: §2 · p. 2

3. Methodology

NITP complements discrete next-token prediction with continuous supervision that makes the final hidden state predict the next token’s contextualized shallow-layer representation. The method is motivated by NTP’s weak constraints on representation geometry and uses shallow states as semantic anchors to regularize it.

3.1. Preliminary: Next-Token Prediction

  • NTP gradients are dominated by the target token and a few high-probability contenders, primarily constraining the hidden state along the target-embedding direction.This leaves other directions weakly supervised.Source: §3.1 · p. 3

3.2. The Geometric Blind Spot of Next-Token Prediction

  • NTP leaves a broad orthogonal subspace nearly unconstrained, allowing degenerate, anisotropic representations; NITP addresses this by explicitly supervising latent representations.Under NTP, effective rank falls as global cosine similarity rises, even while token-level accuracy improves.Source: §3.2 · p. 3

3.3. Shallow Layers as Intrinsic Semantic Anchors

  • Shallow-layer representations provide contextualized semantic targets that resolve token ambiguity, retain semantic detail, stabilize earlier, and add negligible computation.Using these representations as anchors encourages deep states to preserve semantic detail and avoid degeneration.Source: §3.3 · p. 4

3.4. Next Implicit Token Prediction

  • NITP forms a stop-gradient target from the next position’s shallow-layer state, projects the final hidden state to it, and uses cosine alignment alongside NTP.The combined objective weights implicit prediction with λ; cosine similarity provides scale-invariant alignment.Source: §3.4 · p. 4

3.5. Theoretical Analysis: Regularizing the Semantic Manifold

  • NITP leaves radial changes unconstrained but adds positive curvature to angular semantic deviations, lifting NTP’s semantic null directions when the shallow target is well aligned.The paper reports NITP loss below 0.1, including 0.04 for the 9B MoE, as evidence that optimization reaches the high-alignment regime.Source: §3.5 · p. 5

4. Experiments

Experiments compare NITP with NTP on dense and MoE models across knowledge, language, and reasoning benchmarks. NITP improves benchmark performance, while ablations identify temporal prediction, shallow targets, and cosine loss as important design choices.

4.1. Experiment Setups

  • The evaluation trains dense and MoE models from scratch across scales and uses few-shot knowledge, language, reasoning, mathematics, and code benchmarks.Within each comparison, NTP and NITP use identical token budgets.Source: §4 · p. 5, §4.1 · p. 6

4.2. Main Results

  • MoE results: NITP beats NTP at all scales, with average gains of 0.8, 2.1, and 2.7 points at 1.9B, 3B, and 9B; 9B MMLU-Pro rises 5.71% (15.29% to 21.00%).On the 9B model, C3 rises 6.36% (56.65% to 63.01%) and CommonsenseQA rises 4.26% (45.70% to 49.96%).Source: Table 1 · p. 6, §4.2 · p. 6
  • On dense models, NITP consistently beats NTP across scales; the 2B average score improves by 1.79 points (31.91% to 33.70%), with C3 +4.16 and MMLU +2.18.At 0.5B, C3 improves by 3.89 points and the average rises from 24.42% to 25.44%; at 3B, the average improves by 1.35 points.Source: §4.2 · p. 6, Table 2 · p. 7

4.3. Ablation Study

  • Performance consistently peaks around λ = 1.0, which is a reliable loss-weight choice with low sensitivity across that range.The exact optimum may vary slightly across model configurations.Source: Figure 4 · p. 7, §4.3 · p. 8
  • Next-token prediction beats same-position alignment; despite alignment loss as low as 0.03, same-position alignment degrades downstream performance.The comparison indicates that temporal prediction, rather than arbitrary cross-layer alignment or alignment-loss minimization alone, supplies the useful predictive signal.Source: Figure 3 · p. 7, §4.3 · pp. 7–8, Table 3 · p. 8
  • Shallow-layer targets outperform middle- and deep-layer alternatives, with an early layer around 20% of model depth performing best.The paper attributes this to shallow layers preserving richer lexical and local semantic structure.Source: Table 3 · p. 8, §4.3 · p. 8
  • Cosine similarity is the most stable loss and consistently performs best; MSE diverges, Smooth-ℓ1 remains uncertain at larger scales, and KL underperforms.MSE’s quadratic penalty amplifies scale mismatch between layers and can trigger destructive gradient spikes.Source: Table 3 · p. 8, §4.3 · p. 8
  • Unlike NITP’s prediction-aligned supervision, a cosine hidden-state regularization baseline yields no improvements, indicating that generic representation regularization alone is insufficient.NITP trains the final hidden state to predict the context-dependent implicit representation of the next token.Source: Table 3 · p. 8, §4.3 · p. 8

5. Conclusion

NITP augments next-token prediction with continuous representation-space supervision to constrain hidden-state degrees of freedom and prevent representation degeneration. The conclusion reports theoretical regularization of representation geometry and consistent downstream gains with minimal computational overhead, while noting that robustness needs validation on more models.

5. Conclusion

  • NITP predicts the next token’s implicit semantic representation from shallow-layer states, adding continuous supervision to constrain hidden-state degrees of freedom and prevent representation degeneration.The shallow-layer states serve as self-supervised targets.Source: §5 · pp. 8–9
  • NITP mitigates the optimization null space to regularize representation geometry, and experiments across dense and MoE models show consistent downstream gains with minimal computational overhead.Source: §5 · p. 9
  • NITP adds target-layer and loss-weight hyperparameters; although experiments find these choices highly stable across models, validation on more models is needed to confirm robustness.Source: §5 · p. 9

Impact Statement … A.2. Gradient Derivation

NITP changes the pre-training objective without changing data, architecture, or deployment, so its expected societal and ethical impacts align with standard next-token-prediction models. The appendix formalizes its geometric analysis by defining the implicit-prediction notation and deriving gradients and Hessian properties.

Impact Statement

  • NITP changes only the training objective, so its expected societal and ethical impacts align with those of standard next-token-prediction language models.The work uses no new data sources, architectures, or deployment settings; improved downstream performance does not fundamentally change existing use scenarios or risk profiles.Source: Impact Statement · p. 9

A. Detailed Proofs for NITP Geometry

  • The appendix derives the NITP gradient and Hessian and analyzes Hessian spectra to formally prove Lemma 3.1 and Theorem 3.2 on geometric regularization.The spectral analysis is presented as validation of the geometric regularization claims in Section 3.5.Source: A. Detailed Proofs for NITP Geometry · p. 13

A.2. Gradient Derivation

  • The gradient derivation differentiates the normalization map, applies the chain rule to cosine similarity f(h) = u⊤v, and then differentiates L(h) = 1 − f(h).These steps form part of the appendix’s detailed gradient and Hessian derivation for NITP.Source: A. Detailed Proofs for NITP Geometry · p. 13, §A.2 · p. 13

A.1. Notation and Preliminaries

  • For the implicit-prediction loss, h is the projected representation state and z is a fixed implicit target derived from shallow layers.The setup treats h as the projected state to isolate the objective’s geometric incentives.Source: §A.1 · p. 13
  • Assuming r > 0, the tangential difference vector is defined as A := v − su.This definition is used in the appendix’s geometric setup.Source: §A.1 · p. 13
  • A is orthogonal to u, since u⊤A = u⊤v − s(u⊤u) = s − s = 0.The appendix gives this equality as Equation 9.Source: §A.1 · p. 13

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 proofs derive NITP’s Hessian and show that it is flat along the radial direction but positively curved along semantic directions under alignment. Adding NITP therefore mitigates semantic null directions left unconstrained by NTP.

A.3. Hessian Derivation (Proof of Lemma 3.1)

  • The Hessian is derived by differentiating the gradient −r⁻¹A and applying the product rule to r⁻¹ and A.The derivation uses ∇h r = u and A = v − su, with v constant.Source: §A.3 · pp. 13–14
  • Near convergence, s → 1 and ∥A∥ → 0, so the Hessian’s terms involving uA⊤ + Au⊤ vanish and it simplifies to Lemma 3.1’s form.This is the stated condition under which the exact Hessian simplifies.Source: §A.3 · p. 14

A.4. Analysis of Spectral Properties

  • The NITP loss has zero curvature along the radial direction u, so it is locally flat with respect to changes in vector norm.The proof uses that I − uu⊤ projects orthogonally to u and that A ⊥ u.Source: §A.4 · p. 14
  • For perturbations orthogonal to u, NITP has strictly positive curvature when s = cos(h, z) > 0.The proof identifies these tangential perturbations as purely semantic changes.Source: §A.4 · pp. 14–15

A.5. Proof of Theorem 3.2 (Null Space Mitigation)

  • Adding NITP gives strictly positive curvature in semantic null directions of NTP when alignment is high and λ > 0.NTP curvature is concentrated in a low-dimensional subspace, leaving a vast orthogonal complement effectively unconstrained.Source: §A.5 · p. 15

B. Experiment Details

Experiments use a largely consistent, size-scaled pretraining recipe across MoE and Dense models, with matched token budgets for each NTP/NITP pair.

Hyperparameter settings

  • All models use AdamW, a warmup–stable–decay schedule, 2,000 warmup steps, 8192-token contexts, and size-scaled learning rates and global batch sizes.The shared settings also include β1 = 0.9, β2 = 0.95, weight decay 0.1, global gradient clipping at 1.0, and a fixed decay ratio of 0.2.Source: B. Experiment Details · p. 15

Model architectures

  • MoE models use 144 routed experts and one shared expert per layer with top-8 routing, while their expert FFN sizes vary by scale and dense FFN size stays fixed.The MoE backbone begins with a shallow dense stage; its number of MoE layers grows with capacity, and attention uses 8 query heads and 4 key–value heads.Source: B. Experiment Details · p. 15
  • Dense models use 24–28 Transformer layers, with more attention heads and larger FFNs than their MoE counterparts to match overall parameter scale.For both architectures, NITP targets an intermediate layer and uses a loss weight close to 1.0, with minor adjustments for larger models.Source: B. Experiment Details · p. 15

Pretraining data

  • Pretraining uses Qwen2-tokenized large-scale corpora, with token totals scaled by model size and architecture and identical budgets for each NTP/NITP pair.The model-specific target-layer and loss-weight settings are described alongside this training setup.Source: B. Experiment Details · pp. 15–16

Evaluation

C. Analysis of Representation Geometry Dynamics

NITP changes the training-time representation geometry: it preserves higher, context-responsive effective rank and limits anisotropy compared with NTP, providing a proposed mechanism for improved downstream generalization.

Effective rank dynamics

  • NTP’s effective rank rapidly collapses to a low-variance regime, whereas NITP sustains higher, fluctuating rank by activating context-dependent representation subsets.NITP’s latent-space supervision introduces positive curvature along directions left unconstrained by likelihood-based training; the variance reflects controlled activation, not optimization instability.Source: C. Analysis of Representation Geometry… · pp. 16–17

Cosine similarity and anisotropy

  • NTP’s average cosine similarity rises steadily, while NITP keeps it consistently lower, indicating less anisotropy and greater angular diversity.This lower anisotropy aligns with NITP’s higher effective rank, suggesting a more expressive and better-conditioned representation space.Source: C. Analysis of Representation Geometry… · p. 17

Implications

  • NITP’s latent-space semantic consistency prevents collapse into a static anisotropic cone, offering a geometric explanation for its improved downstream generalization.The analysis characterizes this as a changed geometric equilibrium, not merely a delay in representation collapse.Source: C. Analysis of Representation Geometry… · p. 17

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

NITP improves hidden-state semantic transfer while preserving validation language-model loss; its auxiliary loss evolves as shallow-layer targets become structured, and its FLOPs analysis identifies the added computations.

D. Hidden-State Representation Quality

  • MTEB rose from 39.24 to 41.56 overall, with gains of +2.20 classification, +2.93 STS, and +1.64 retrieval/duplicate-detection points; NITP improved 23/25 tasks.The 3B MoE models were evaluated using frozen, last-layer hidden states that were mean-pooled and ℓ2-normalized, without fine-tuning or task-specific adaptation.Source: Table 5 · p. 17, D. Hidden-State Representation Quality · p. 17

E. Loss Curve of NITP

  • NITP loss first collapses under near-isotropic random states, rises as shallow-layer targets structure (about 40–2k steps), then steadily converges below 0.05 on the 9B model.The transient hump suggests an adaptive predictive task whose difficulty evolves with the model’s representations, rather than a static regularizer.Source: Figure 5 · p. 18, E. Loss Curve of NITP · p. 18

D. Hidden-State Representation Quality

  • Validation CE/PPL remained nearly matched: 3B MoE 2.006 vs. 2.006 (7.43 vs. 7.43), 9B MoE 1.841 vs. 1.840 (6.30 vs. 6.30), and 3B Dense 1.884 vs. 1.882 (6.58 vs. 6.57), NTP vs. NITP.Thus, NITP preserves validation loss while improving hidden-state semantic transferability.Source: D. Hidden-State Representation Quality · p. 18

F. Training FLOPs Analysis

  • Training FLOPs are approximated as 6× the number of parameters involved in forward computation, accounting for forward, backward, and parameter-gradient updates.The analysis compares NITP’s added computation with the baseline NTP objective.Source: F. Training FLOPs Analysis · p. 19

F.1. Baseline NTP FLOPs

  • Baseline NTP estimates include approximately 18d² training FLOPs for GQA attention, 18kdde for k activated MoE experts, and 6Vd for unembedding per token.The input embedding lookup is excluded because it incurs 0 FLOPs.Source: §F.1 · p. 19

F.2. Additional FLOPs Introduced by NITP

  • NITP adds a SwiGLU projection head with 12d² parameters and cosine-similarity computation of approximately 6d operations; stop-gradient makes target extraction cost-free.The supplied analysis describes the projection head’s three transformations as d → 4d, d → 4d, and 4d → d.Source: §F.2 · p. 19

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

In the 9B MoE setting, NITP adds about 1.8% measured training time, while inference FLOPs remain identical to NTP. The FLOPs analysis estimates a small pretraining overhead.

F.3. Numerical Instantiation: 9B MoE Model

  • NITP’s estimated overhead is 1.18 × 10^8 FLOPs for the 9B MoE configuration.The estimate uses d = 1280, de = 640, k = 9, L = 24, and V = 152064.Source: §F.3 · p. 19
  • The estimated NTP baseline totals 5.06 × 10^9 FLOPs, comprising 3.89 × 10^9 for backbone layers and 1.17 × 10^9 for unembedding.Source: §F.3 · p. 19
  • The target-layer sweep reports that the best target layer consistently lies in a shallow contextualized region around 20% of total depth.Scores are averaged over MMLU, MMLU-Pro, CSQA, BBH, and LCBench.Source: Table 6 · p. 20

F.4. Empirical Wall-Clock Overhead

  • NITP takes 16h 18m versus 16h 1m for NTP in a 5k-step run, an additional 17 minutes or approximately 1.8% wall-clock overhead.The comparison uses the same number of GPUs and a global batch size of 1024; the analysis concludes the overhead remains marginal (∼2%).Source: §F.4 · p. 20
  • The measured wall-clock increase is slightly below the theoretical FLOPs ratio because MoE communication and system overheads are largely unaffected by NITP’s projection head and cosine loss.Source: §F.4 · p. 20

G. Inference FLOPs Analysis

  • NITP adds zero inference overhead: its projection head is discarded after pretraining, leaving the deployment architecture and per-token FLOPs identical to baseline NTP.The passage states this preserves generation speed and serving costs.Source: G. Inference FLOPs Analysis · p. 20

H. Additional Ablation Studies

The ablations favor shallow but contextualized implicit targets and show that NITP works best with stop-gradient, activation from step 0, and a projection head.

Target-layer sweep

  • The best targets are layer 4 of 17 for the 3B MoE and layer 5 of 24 for the 9B MoE, both near 20% depth.Embedding-layer targets are weaker, and moving the target deeper reduces the gain.Source: H. Additional Ablation Studies · p. 20

Stop-gradient on implicit targets

  • Stop-gradient stabilizes training and improves downstream benchmarks, whereas target gradients cause co-adaptation, unstable optimization, and degraded performance.The study therefore uses stop-gradient to keep implicit targets fixed as semantic anchors.Source: H. Additional Ablation Studies · p. 20, Table 7 · p. 21

Start step of NITP

  • Delaying NITP activation consistently degrades performance, particularly on reasoning benchmarks, while starting at step 0 yields the best overall results.Early application can shape representation geometry before higher-level semantics stabilize.Source: H. Additional Ablation Studies · p. 21, Table 7 · p. 21

Effect of the projection head

  • Removing the two-layer SwiGLU projection head degrades performance across benchmarks, while the projector flexibly bridges the mismatch between deep representations and shallow targets.Direct alignment without the projector imposes an overly restrictive constraint.Source: H. Additional Ablation Studies · p. 21, Table 7 · p. 21

Related papers

Loading 2605.24956v3…