1 Introduction↩︎

Large language models (LLMs) have achieved remarkable success through large-scale pre-training with the Next-Token Prediction (NTP) [1]–[5]. By maximizing the likelihood of the next token over massive corpora, this paradigm learns general-purpose representations that support a wide range of downstream tasks [6]. Although subsequent fine-tuning stages further adapt these models, their effectiveness critically depends on the representations learned during base pre-training [7]. As a result, the pre-training objective plays a central role in determining the capability of LLMs.

Despite its empirical success, standard NTP provides supervision only through discrete one-hot targets in the output token space. While gradients propagate to hidden states via the output projection, the objective primarily constrains representations along the target logit direction, leaving many degrees of freedom in the latent space that remain weakly constrained. This raises a fundamental question: Does next-token prediction alone sufficiently supervise the geometry of hidden representations?

Prior work has identified a phenomenon known as representation degeneration, where likelihood-based training drives learned embeddings to collapse into a narrow, anisotropic cone [8]–[11]. Such geometric collapse limits the expressive capacity of representations and has been linked to degraded generalization on downstream tasks [8], [12]. To examine this issue, we track the geometric evolution of last hidden states throughout training. We utilize two established metrics: (1) Effective Rank [13], which quantifies the effective dimensionality of the utilized subspace, and (2) Average Cosine Similarity between paired tokens within training batches, serving as a proxy for the global anisotropy of the representation space. As illustrated in 1, for NTP, the effective rank diminishes rapidly while cosine similarity increases, indicating that hidden representations drift toward degenerate, anisotropic configurations (detailed analysis in 8). This confirms that standard NTP is geometrically under-constrained: it permits representations to sacrifice expressiveness and distinctiveness without penalty, as long as the next token is correctly predicted.

We attribute this behavior to the under-constrained nature of hidden representations under NTP, where the lack of geometric supervision allows representations to drift into low-expressivity configurations that may harm performance [8]. To bridge this gap, we propose Next Implicit Token Prediction (NITP), an auxiliary pre-training objective that augments discrete next-token prediction with continuous supervision in the representation space. Instead of predicting only the token identity, NITP tasks the model with predicting an implicit token, defined as a latent semantic representation of the next token. We construct this supervision signal from the model’s own shallow layers, which empirically preserve richer lexical and local semantic content than deeper, task-specialized layers [14], [15]. This design enables scalable, self-supervised targets without introducing external encoders or additional annotations. As evidenced by 1, NITP effectively mitigates representation degeneration, maintaining higher effective rank and preventing excessive anisotropy compared to standard NTP. Our main contributions are summarized as follows:

1) We analyze the training objective of NTP and show it provides weak supervision for hidden representations, allowing them to drift toward degenerate space.

2) We propose Next Implicit Token Prediction (NITP), a novel objective that augments next-token prediction with dense, continuous supervision in the latent space.

3) We introduce a self-supervised design for NITP by using shallow-layer representations as implicit tokens to serve as prediction targets, enabling semantically rich supervision.

4) We provide theoretical analysis and empirical evidence showing that NITP improves representation geometry and consistently enhances downstream performance.

2 Related Work↩︎

Pre-training objectives for large language models. The next-token prediction (NTP) objective is the dominant pre-training paradigm for modern large language models [1], [2], and has been shown to scale effectively with data and model size [16]–[18]. By optimizing the likelihood of the next discrete token given its context, NTP enables models to learn general-purpose representations that support a wide range of downstream tasks. To extend the predictive horizon beyond a single token, several works have explored multi-token prediction (MTP) [1], [19], [20] and related objectives [21]–[24]. These methods introduce auxiliary output heads to predict multiple future tokens or compressed summaries of future context, with the goal of improving planning or long-horizon modeling [21]. Despite their differences, these approaches operate primarily in the discrete token space and rely on token-level supervision derived from one-hot targets. In contrast, our proposed Next Implicit Token Prediction departs from prior works by augmenting standard NTP with dense, continuous supervision at the representation level.

Representation-level supervision. Beyond discrete token-level objectives, prior work has explored supervising models directly in continuous representation spaces [25]–[27]. A prominent direction is layer-wise distillation, which aligns intermediate representations between models to transfer structural or task-relevant knowledge [27]–[30]. This paradigm has recently been extended to generative LLMs [31]–[33]. Notably, Sheared LLaMA [33] aligns hidden states to accelerate knowledge recovery after structured pruning. Another line of research focuses on self-supervised consistency, encouraging agreement between latent representations across time or views without discrete labels, e.g., Contrastive Predictive Coding and BYOL [26], [34], as well as self-distillation methods [35], [36]. NITP differs from these approaches in both motivation and formulation: motivated by the under-constrained hidden states during LLM pre-training, NITP adopts an autoregressive objective in the latent space that requires the model to predict the semantic representation of the next token.

3 Methodology↩︎

In this section, we present Next Implicit Token Prediction (NITP), a representation-level objective that explicitly supervises hidden states to complement standard Next-Token Prediction (NTP), as illustrated in 2. By encouraging the model to anticipate the semantic content of the next token in latent space, NITP regularizes representation geometry beyond discrete token-level supervision.

We first revisit the standard NTP paradigm and its limitations in supervising latent representations. We then introduce NITP, which trains the last hidden state \(h_t\) to predict an implicit token—the shallow-layer representation of the \((t\!+\!1)\)-th token with stop-gradient—using a cosine similarity loss jointly optimized with the NTP objective. As the implicit targets are self-generated during the forward pass, NITP incurs negligible computational overhead and scales naturally with model size (see 4 for results).

3.1 Preliminary: Next-Token Prediction↩︎

Given a token sequence \(x = \{x_1, \dots, x_T\}\), let \(h_t \in \mathbb{R}^d\) denote the final hidden state at step \(t\), serving as the contextual representation. Given an unembedding matrix \(W \in \mathbb{R}^{|\mathcal{V}| \times d}\), the standard next-token prediction (NTP) objective minimizes the negative log-likelihood of the target token \(x_{t+1}\): \[\mathcal{L}_{\mathrm{NTP}} = - \mathbb{E}_{x} \left[ \log \frac{\exp(h_t^\top w_{x_{t+1}})}{\sum_{j \in \mathcal{V}} \exp(h_t^\top w_j)} \right],\] where \(w_j\) denotes the embedding vector for token \(j\). We omit layer normalization terms for brevity. While the NTP loss theoretically involves all vocabulary items via the normalization term, the gradient signal is empirically dominated by the target token \(x_{t+1}\) and a few high-probability contenders. Thus, the objective effectively constrains \(h_t\) primarily along the direction of the target embedding \(w_{x_{t+1}}\).

3.2 The Geometric Blind Spot of Next-Token Prediction↩︎

The gradient dominance described above creates a structural “blind spot” in the optimization landscape: the loss becomes near-invariant to perturbations within the vast subspace orthogonal to the target direction. Theoretically, this implies that a wide range of geometrically distinct latent states can yield identical token-level likelihoods.

In practice, this structural freedom does not yield diverse features but instead leads to representation degeneration [8]–[10], where hidden states drift into a low-dimensional, anisotropic cone without penalty. We empirically confirm this collapse by tracking the geometric evolution of hidden states in 1. Under standard NTP, we observe a rapid deterioration in Effective Rank [13] coinciding with a sharp rise in global Cosine Similarity. These trends indicate that the model sacrifices semantic richness for discriminative efficiency, compressing representations even as token-level accuracy improves.

From the perspective of the NTP objective, these degenerate configurations are valid as long as they preserve the correct logit ranking, yet they severely limit the expressive capacity of the latent space. Crucially, such a drop in representation expressiveness has been linked to degraded generalization on downstream tasks [8], [12]. This observation underscores a critical limitation: NTP defines what to predict, but fails to supervise how predictions are represented, necessitating explicit guidance to enforce a semantically structured geometry. To this end, we propose Next Implicit Token Prediction (NITP), complementing NTP by explicitly supervising the latent representation space. Instead of constraining only discrete token identities, NITP requires the model to predict the implicit semantic representation of the next token, thereby directly supervising the geometry and evolution of hidden states.

3.3 Shallow Layers as Intrinsic Semantic Anchors↩︎

To construct an effective supervision target for NITP, we first clarify the semantic role of the last hidden state \(h_t\). Prior works [14], [37] suggest that \(h_t\) contains recoverable semantic information about the next token. Building on this view, we consider an idealized regime that the model predicts the correct next token \(x_{t+1}\) with high confidence. In this setting, \(h_t\) serves as the intermediate representation through which the prediction is realized, and thus is expected to encode semantic information relevant to \(x_{t+1}\) in a form that supports the prediction task.

This observation naturally raises the question: what constitutes an appropriate supervision target for \(h_t\)? A naive choice would be the static embedding of \(x_{t+1}\). However, static embeddings suffer from polysemy—the same token (e.g., “bank”) can correspond to vastly different meanings depending on the context. To account for this ambiguity, a suitable supervision target should be a contextualized representation that captures the precise semantics of the token conditioned on the context \(x_{< t+1}\).

While external encoders [38], [39] could provide such contextualized targets, they introduce prohibitive computational overhead and domain shifts. Instead, we propose an efficient, self-supervised solution: utilizing the model’s own shallow layers as intrinsic semantic anchors. Our rationale is threefold:

• Semantic richness: Recent analysis [15] reveals that shallow layers act as powerful contextual encoders, resolving lexical ambiguity while preserving fine-grained semantic details. Studies [14] explicitly demonstrate that semantic richness peaks in these early layers and degrades in deeper layers, as the model progressively discards details to focus solely on the sparse discriminative features required for NTP.

• Stability: Transformer training typically exhibits a bottom-up convergence pattern [15], where layers closer to the input stabilize much faster than deeper layers. This ensures that the generated implicit tokens are stable relative to the deep predictive state \(h_t\).

• Efficiency: Utilizing shallow layers incurs negligible computational overhead. Unlike external models, the targets are extracted directly from intermediate activations computed during the standard forward pass.

By treating the semantic-rich, contextualized shallow representation of \(x_{t+1}\) as the next implicit token, we impose a fidelity constraint that forces the deep state \(h_t\) to maintain a structured geometry capable of containing full semantic details, thereby preventing representation degeneration.

3.4 Next Implicit Token Prediction↩︎

We introduce Next Implicit Token Prediction (NITP), a continuous representation-level supervision objective designed to complement the discrete next-token prediction. While NTP predicts the discrete identity of the next token, NITP predicts its semantic essence in the latent space.

We formally distinguish between the predictive state \(h_t\) and the prediction target \(z_{t+1}\). We define the implicit token \(z_{t+1}\) as the dense, contextualized semantic representation of the next token, derived from the model’s own shallow layers. The NITP objective forces the final hidden state to align with this implicit token, thereby enforcing a structured and semantic representation geometry.

3.4.0.1 Implicit token construction.

For a given input sequence \(x\), let \(E_{\mathrm{shallow}}(\cdot)\) denote the forward pass through the first \(k\) layers of the model (see 4.3 for details on the choice of \(k\)). We construct the implicit token \(z_{t+1}\), which serves as the optimization target, by extracting the hidden representation at the future time step \(t+1\): \[\label{eq:stop95gradient} z_{t+1} = \mathbf{sg}\left[ E_{\mathrm{shallow}}(x_{\leq t+1})^{(t+1)} \right] \in \mathbb{R}^d,\tag{1}\] where \((\cdot)^{(t+1)}\) denotes selecting the vector at position \(t+1\), and \(\mathbf{sg}[\cdot]\) indicates the stop-gradient operator. Crucially, by sourcing \(z_{t+1}\) from shallow layers, we capture the contextualized semantics of the next token while avoiding the representation collapse often observed in deeper layers, making it a stable and semantically rich anchor.

3.4.0.2 The prediction objective.

NITP tasks the last hidden state \(h_t\) with predicting the implicit token \(z_{t+1}\). To bridge the potential distributional gap between the deep predictive state and the shallow target, we employ a projection head \(\mathcal{P}(\cdot)\) (e.g., an MLP). The objective minimizes the geometric misalignment using cosine similarity: \[\mathcal{L}_{\mathrm{NITP}}(\theta) = 1 - \frac{\mathcal{P}(h_t)^\top z_{t+1}}{\|\mathcal{P}(h_t)\|_2 \cdot \|z_{t+1}\|_2}.\] We favor cosine similarity over other loss functions to enforce scale invariance as demonstrated in 4.3.

3.4.0.3 Combined training.

The overall pre-training framework acts as a dual-supervision mechanism, combining the discrete NTP loss with the continuous NITP supervision: \[\label{eq:loss95all} \mathcal{L}_{\mathrm{total}}(\theta) = \mathcal{L}_{\mathrm{NTP}}(\theta) + \lambda\,\mathcal{L}_{\mathrm{NITP}}(\theta),\tag{2}\] where \(\lambda > 0\) is a hyperparameter controlling the weight of the implicit prediction. By requiring the model to predict both the discrete identity and the continuous implicit token, NITP explicitly regularizes the hidden representations toward a more semantically structured feature space.

3.5 Theoretical Analysis: Regularizing the Semantic Manifold↩︎

Standard Next-Token Prediction (NTP) supervises the hidden state \(h_t\) primarily by maximizing its dot product with the target token embedding \(w_{x_{t+1}}\). However, this objective is invariant to perturbations in the subspace orthogonal to the informative directions of the unembedding matrix. Mathematically, this manifests as a rank-deficient Hessian matrix with respect to the semantic geometry, allowing hidden states to drift into degenerate configurations without penalty—a phenomenon closely linked to representation anisotropy [8].

We formally analyze how Next Implicit Token Prediction (NITP) mitigates this by regularizing the angular geometry of the latent space. Unlike NTP, which allows unconstrained drift in null directions, NITP acts as a geometric anchor, introducing positive curvature to the entire semantic subspace.

3.5.0.1 Setup.

Let \(h \in \mathbb{R}^d\) be the last hidden state used for implicit prediction. In practice, the NITP objective utilizes a non-linear projection \(\mathcal{P}(\cdot)\). To characterize the regularization effect, we adopt the Generalized Gauss-Newton (GGN) approximation [40], which models the Hessian with respect to \(h\) as \(J_P^\top H_{\text{NITP}} J_P\), where \(H_{\text{NITP}}\) denotes the Hessian of the NITP loss. Crucially, this quadratic form preserves the spectral properties (specifically positive definiteness) of \(H_{\text{NITP}}\), provided that the projector is locally well-conditioned (i.e., \(J_P\) is full-rank). This implies that the regularization is fundamentally governed by the structure of the loss function. Therefore, we simplify the analysis by examining the NITP objective directly on \(h\) to isolate its intrinsic geometric properties from the architectural details.

We decompose variations in \(h\) into a radial component (change in norm) and an angular component (change in direction). Let \(r = \|h\|_2\) and \(u = h/r\). Let \(z \in \mathbb{R}^d\) be the fixed implicit target derived from shallow layers, with normalized direction \(v = z/\|z\|_2\). The NITP objective is defined as \(\mathcal{L}_{\mathrm{NITP}}(h) = 1 - \cos(h, z) = 1 - u^\top v\).

Proof. See 6.

3.5.0.2 Remark: necessity of semantic targets.

The curvature guarantee relies on high alignment (\(s \to 1\)). This justifies the choice of predictable semantic targets derived from shallow layers over arbitrary vectors, which would yield weak alignment (\(s \approx 0\)). Empirically, we observe that the NITP loss, i.e. \(1 - s\) consistently converges below \(0.1\) (e.g., \(0.04\) for the 9B MoE), confirming that the optimization stably operates in the region where the Hessian is positive definite and geometric regularization is effective.

3.5.0.3 Analysis of spectral properties.

Lemma 1 reveals two critical properties of NITP:

• Radial null space: \(u^\top H_{\mathrm{NITP}} u = 0\). NITP imposes no curvature on the vector norm, allowing the NTP objective to freely adjust the scale (\(r\))—for instance, to optimize softmax temperature—without interference.

• Angular spectral lifting: For any semantic deviation vector \(w\) orthogonal to the current state (\(w \perp u\)), we have strictly positive curvature: \[w^\top H_{\mathrm{NITP}} w \approx \frac{1}{r^2} \|w\|_2^2 > 0.\]

See detailed analysis at 6.

3.5.0.4 Interpretation.

Theorem 1 demonstrates that NITP performs spectral lifting on the angular manifold. While standard NTP leaves vast degrees of freedom orthogonal to the target logit unpenalized, NITP effectively regularizes these under-constrained directions by introducing positive curvature. This mitigates the “flat valleys" in the optimization landscape responsible for semantic drift, forcing the model to maintain robust representations aligned with the semantic anchors derived from shallow layers.

4 Experiments↩︎

In this section, we conduct a comprehensive empirical evaluation to assess the effectiveness of Next Implicit Token Prediction (NITP) in terms of downstream performance. We train both dense and mixture-of-experts models from scratch, covering model scales from hundreds of millions to billions of parameters, and assess performance on a broad suite of knowledge and reasoning benchmarks. Additionally, we conduct ablation studies to examine the effect of key design choices in NITP and to verify its effectiveness.

4.1 Experiment Setups↩︎

Training settings. We evaluate NITP on both Mixture-of-Experts (MoE) [41] and dense language models to assess its generality across model architectures and scales. For MoE models, we adopt the DeepSeek-V2 [42] architecture employing 144 experts with top-\(8\) routing, where each expert is implemented as a SwiGLU-based [43] FFN. We scale this architecture up to a total parameter count of 9B, corresponding to approximately 1B activated parameters per token. For dense models, we conduct experiments on model sizes ranging from 0.5B to 3B parameters. All models are trained from scratch with token budgets scaled according to their parameters. Specifically, the 9B MoE model is trained on 330B high-quality tokens comprising a diverse mix of English, Chinese, code, mathematics, and reasoning data. The context length is fixed at 8192 tokens unless stated otherwise. Comprehensive training details are provided in 7.

Evaluation. We conduct comprehensive few-shot evaluations on a diverse suite of benchmarks to assess the effectiveness of NITP relative to the standard NTP objective. We categorize these benchmarks into two primary domains: 1) Knowledge & Language Understanding: To evaluate broad world knowledge and multi-domain comprehension, we utilize English benchmarks MMLU [44] and MMLU-Pro [45], alongside Chinese datasets C-Eval [46] and Xiezhi [47]. We additionally include LAMBADA [48] to test long-range dependency modeling. 2) Reasoning & Problem Solving: To examine logical deduction and complex problem-solving, we employ ARC-Challenge [49] for scientific reasoning, together with CommonsenseQA [50] and Balanced COPA [51] for commonsense reasoning. We also include BBH [52] and C3 [53] to evaluate challenging multi-step inference and reading comprehension. Furthermore, we assess rigorous application skills using AGIEval [54] on standardized exams, GSM8k [55] on mathematics, and OpenCompass LCBench [56] (a collection of programming problems from Leetcode weekly competitions) on code generation.

4.2 Main Results↩︎

Evaluation on MoE models. We mainly evaluate the performance of the proposed NITP against the NTP baseline across three Mixture-of-Experts (MoE) scales: 1.9B (0.3B activated), 3B (0.5B activated), and 9B (1B activated). The models were pre-trained from scratch with increasing token budgets proportional to their scale. The quantitative results are reported in 1. The results demonstrate that NITP consistently outperforms standard NTP across all model sizes, underscoring the method’s scalability. Specifically, NITP yields average score improvements of 0.8, 2.1, and 2.7 points for the 1.9B, 3B, and 9B models, respectively. For instance, on the 9B model, NITP achieves a 5.71% absolute improvement on MMLU-Pro (rising from 15.29% to 21.00%), a benchmark specifically designed to test robust reasoning beyond simple knowledge retrieval. Similarly, we observe remarkable boosts in reading comprehension and commonsense reasoning, with C3 scores increasing by 6.36% (56.65% to 63.01%) and CommonsenseQA by 4.26% (45.70% to 49.96%). Also, on symbolic and math tasks like ARC-C and GSM8k, NITP maintains a consistent lead.

Evaluation on dense models. To assess the generality of NITP beyond MoE architectures, we evaluate dense models ranging from 0.5B to 3B parameters on seven representative benchmarks, including MMLU, C-Eval, BBH, ARC-C, C3, AGIEval, and LCBench. The quantitative results are reported in 2. The results demonstrate that NITP is architecture-agnostic, consistently outperforming the NTP baseline across all model scales. On the 0.5B model, NITP induces notable improvements; specifically, C3 scores surge by 3.89 points, and the average score increases from 24.42% to 25.44%. For the 2B model, NITP improves the average score by 1.79 points (31.91% to 33.70%), with notable gains on C3 (+4.16) and MMLU (+2.18). Similarly, on the 3B model, NITP achieves an average improvement of 1.35 points, accompanied by substantial gains on AGIEval and C3. Overall, these results confirm that NITP consistently enhances representation quality across diverse tasks, independent of model architecture such as MoE or dense.

Collectively, these empirical results validate the core design philosophy of NITP. By augmenting discrete prediction with semantic-rich continuous supervision, NITP effectively constrains the latent geometry, preventing hidden states from drifting into degenerate configurations. This translates into consistent performance improvements across diverse architectures (MoE and Dense). Crucially, the substantial gains observed across the evaluated MoE and dense scales further support the scalability of NITP. Beyond downstream benchmark accuracy, we directly evaluate the utility of the learned hidden states as frozen sentence representations in 9, where NITP improves 23 out of 25 MTEB tasks while preserving validation cross-entropy loss.

4.2.0.1 Training efficiency.

We analyze the computational efficiency of NITP by comparing its complexity with the standard NTP baseline. For a MoE model, the training cost is dominated by \(L\) layers of backbone blocks and the unembedding projection, scaling as \(O(L (d^2 + k d d_e) + Vd)\), where \(k\) is the number of activated experts, \(d_e\) is the expert dimension, and \(V\) is the vocabulary size. In contrast, the NITP objective introduces a single projection head and a cosine loss, with a total complexity of \(O(d^2)\). Crucially, since the implicit targets are derived from existing intermediate activations and decoupled via a stop-gradient operator, NITP incurs no additional backbone forward passes or backward propagation through the early layers. Given that \(L \gg 1\), the relative complexity ensures that the extra overhead is inherently marginal. Numerical instantiation on our 9B MoE model (\(d=1280, L=24\)) confirms this advantage. The additional training FLOPs from both forward and backward passes account for approximately 2% of the total computation. This negligible cost ensures that NITP serves as a highly scalable supervision objective without sacrificing training throughput. See details in Appendix 11.

4.3 Ablation Study↩︎

We conduct ablation studies to understand how each component contributes to its effectiveness. In particular, we focus on the following aspects: 1) how to construct the implicit prediction target, including the choice of target layer and temporal shift. 2) how to design a stable and effective auxiliary objective, including the loss function and its role beyond generic regularization. Unless otherwise specified, all ablation experiments are conducted on a 3B-parameter MoE model in 4, trained on 200B tokens.

Importance of temporal shift. To distinguish NITP from static layer-wise alignment, we ablate the temporal structure of the implicit prediction objective. Specifically, we compare predicting the semantic representation of the next token, i.e. the next implicit token, with aligning semantic representations at the current time step. As shown in 3 and 3, predicting the next implicit token significantly outperforms same-position alignment, while the latter leads to terrible loss and performance degradation. Notably, same-position alignment can achieve a very low alignment loss (e.g., 0.03), yet results in substantially worse downstream performance. This indicates that simply regularizing representations or minimizing alignment loss is insufficient. Without the autoregressive temporal shift, the objective fails to provide a meaningful predictive signal of the next token and can collapse into trivial solutions. These results highlight that the effectiveness of NITP crucially depends on predicting the future semantic representation, rather than arbitrarily aligning hidden states across layers.

Loss function. We compare the Cosine Similarity against regression-based losses (MSE, Smooth-\(\ell_1\)) and KL Divergence in 3. The results highlight significant stability differences. MSE leads to catastrophic optimization divergence that loss is observed to diverge temporarily during training, as its quadratic penalty amplifies the inherent scale mismatch between layers, triggering destructive gradient spikes. Though Smooth-\(\ell_1\) maintains stability via implicit gradient clipping, it still remains unknown for larger-scale training due to its hard alignment form. KL Divergence underperforms as it treats continuous feature vectors as distributions, causing geometric distortion. Ultimately, Cosine Similarity proves to be the superior choice. Empirically, it is the most stable objective without any divergence issues and consistently delivers superior performance across tasks.

Target layer selection. A key design choice in NITP is the use of shallow-layer representations as prediction targets (implicit tokens) under the temporal-shifted prediction objective. To validate this choice, we vary the source layer of the implicit tokens across shallow, middle, and deep layers of the model. The results in 3 show that using shallow-layer targets consistently outperforms the alternatives. We attribute this behavior to the fact that shallow layers preserve richer lexical and local semantic structure, whereas deeper layers are more specialized for token discrimination and exhibit worse semantics [15]. Among shallow layers, we further find that selecting an early layer around 20% of the total depth yields the best performance (see [tab:comprehensive_ablation,tab:target_layer_sweep]), which is consistent with prior findings that semantic richness peaks in specific early layers [14], [15] and aligns with our analysis in 3.3.

Is NITP merely a hidden state regularization? To examine whether NITP reduces to generic hidden-state regularization, we compare it with a cosine-based baseline inspired by [11]. Specifically, this baseline applies a cosine similarity penalty between hidden states of different tokens within sequences, encouraging them to be less aligned. While this regularization constrains the geometry of hidden representations, it does not provide a prediction target or token-level supervision. In contrast, NITP introduces a prediction-aligned auxiliary objective by training the final hidden state to predict the next implicit token, which encodes context-dependent information about the upcoming token. As shown in 3, we find that the cosine regularization baseline yields no improvements, whereas NITP consistently improves performance, indicating that the gains stem from prediction-aligned semantic supervision rather than merely representation regularization.

NITP loss weight. We study the effect of the NITP loss weight by sweeping \(\lambda\) in 2 . As shown in 4, performance consistently peaks around \(\lambda = 1.0\), which emerges as the most reliable choice across settings. While the exact optimum may vary slightly across model configurations, strong performance is consistently observed for values around \(1.0\), indicating low sensitivity within this range. The specific loss weights used for different models are summarized in 4.

5 Conclusion↩︎

We propose Next Implicit Token Prediction (NITP), a novel pre-training objective that augments standard next-token prediction with continuous supervision in the representation space. Motivated by the observation that NTP leaves hidden representations under-constrained and prone to geometric degeneration, NITP requires the model to predict the implicit semantic representation of the next token using shallow-layer states as self-supervised targets, which we term as next implicit token. This objective explicitly constrains the free degrees of freedom ignored by token-level supervision and prevents representation degeneration. We provide a theoretical analysis showing that NITP mitigates the optimization null space, thereby regularizing representation geometry. Experiments across both dense and MoE models demonstrate that NITP consistently improves downstream performance with minimal computational overhead.

Limitations. NITP introduces additional hyperparameters, including the target layer and NITP loss weight. While our experiments indicate that these choices are highly stable across models, further validation on more models is required to fully confirm this robustness.

Acknowledgements↩︎

This work was in part supported by Scientific Research Innovation Capability Support Project for Young Faculty (U40) of the Ministry of Education of China (SRICSPYF-ZY2025019) and Xiaohongshu Inc.

Impact Statement↩︎

This work introduces a representation-level auxiliary objective for language model pre-training, with the goal of improving optimization stability and the geometric structure of hidden representations. The proposed method modifies only the training objective and does not involve new data sources, model architectures, or deployment settings. As a result, the societal and ethical impacts of this work are expected to be aligned with those of existing large language models trained using standard next-token prediction. Although improved representation quality and training efficiency may lead to better downstream performance, this does not fundamentally change the usage scenarios or risk profile of current language modeling systems.

6 Detailed Proofs for NITP Geometry↩︎

In this section, we provide the detailed derivation of the gradient and Hessian matrix for the Next Implicit Token Prediction (NITP) objective. We further analyze the spectral properties of the Hessian to validate the geometric regularization claims made in Section 3.5, providing the formal proof for Lemma 1 and Theorem 1.

6.1 Notation and Preliminaries↩︎

Consistent with the setup in Section 3.5, let \(h \in \mathbb{R}^d\) denote the representation used for the implicit prediction loss. We treat \(h\) as the projected state to isolate the geometric incentives of the objective function. Let \(z \in \mathbb{R}^d\) denote the fixed implicit target derived from shallow layers. We define the following variables: \[\begin{align} r &:= \|h\|_2, \quad & u &:= \frac{h}{r}, \\ v &:= \frac{z}{\|z\|_2}, \quad & s &:= u^\top v = \cos(h, z). \end{align}\] We assume \(r > 0\). The tangential difference vector is defined as \(A := v - s u\). Note that \(A\) is orthogonal to \(u\): \[u^\top A = u^\top (v - s u) = u^\top v - s (u^\top u) = s - s = 0.\]

The NITP objective function is: \[\mathcal{L}(h) = 1 - \cos(h, z) = 1 - \frac{h^\top z}{\|h\|_2 \|z\|_2} = 1 - u^\top v.\]

6.2 Gradient Derivation↩︎

First, recall the Jacobian of the normalization map \(u(h) = h/r\): \[J_u(h) = \frac{\partial u}{\partial h} = \frac{1}{r} (I - u u^\top).\] Let \(f(h) = u^\top v\) be the cosine similarity. By the chain rule: \[\begin{align} \nabla_h f(h) &= J_u(h)^\top v \nonumber \\ &= \frac{1}{r} (I - u u^\top) v \nonumber \\ &= \frac{1}{r} (v - u (u^\top v)) \nonumber \\ &= \frac{1}{r} (v - s u) = \frac{A}{r}. \end{align}\] Thus, the gradient of the loss \(\mathcal{L}(h) = 1 - f(h)\) is: \[\nabla_h \mathcal{L}(h) = - \frac{A}{r}.\]

6.3 Hessian Derivation (Proof of Lemma 1)↩︎

The Hessian is the Jacobian of the gradient vector. We differentiate \(\nabla_h \mathcal{L}(h) = -r^{-1} A\) with respect to \(h\): \[H_{\mathrm{NITP}}(h) = \nabla_h \left( - \frac{A}{r} \right) = - \left( \frac{1}{r} \nabla_h A + A (\nabla_h r^{-1})^\top \right).\] We compute the two terms separately.

Term 1: Derivative of \(r^{-1}\). Using \(\nabla_h r = u\): \[\nabla_h (r^{-1}) = - \frac{1}{r^2} \nabla_h r = - \frac{u}{r^2}.\] Thus, the second part of the product rule is: \[A (\nabla_h r^{-1})^\top = A \left( - \frac{u^\top}{r^2} \right) = - \frac{1}{r^2} A u^\top.\]

Term 2: Derivative of \(A\). Recall \(A = v - s u\). Since \(v\) is constant: \[\nabla_h A = - \nabla_h (s u) = - \left( u (\nabla_h s)^\top + s \nabla_h u \right).\] Substituting \(\nabla_h s = \nabla_h f = A/r\) and \(\nabla_h u = \frac{1}{r}(I - uu^\top)\): \[\begin{align} \nabla_h A &= - \left( u \left(\frac{A}{r}\right)^\top + s \frac{1}{r} (I - u u^\top) \right) \nonumber \\ &= - \frac{1}{r} \left( u A^\top + s(I - u u^\top) \right). \end{align}\]

Combining the Terms. Substituting back into the expression for \(H_{\mathrm{NITP}}\): \[\begin{align} H_{\mathrm{NITP}}(h) &= - \left[ \frac{1}{r} \left( - \frac{1}{r} (u A^\top + s(I - uu^\top)) \right) + \left( - \frac{1}{r^2} A u^\top \right) \right] \nonumber \\ &= \frac{1}{r^2} (u A^\top + s(I - uu^\top)) + \frac{1}{r^2} A u^\top \nonumber \\ &= \frac{1}{r^2} \left[ s(I - u u^\top) + (u A^\top + A u^\top) \right]. \end{align}\] This provides the exact Hessian form.

Approximation near convergence. As the training progresses and the prediction aligns with the target, we have \(s \to 1\). Consequently, the difference vector \(A = v - su\) approaches zero (\(\|A\| \to 0\)). Under this condition, the term \((u A^\top + A u^\top)\) vanishes. The Hessian simplifies to: \[H_{\mathrm{NITP}}(h) \approx \frac{1}{r^2} (I - u u^\top) = \frac{1}{r^2} P_{\perp u}.\] This confirms the simplified form presented in Lemma 1. \(\square\)

6.4 Analysis of Spectral Properties↩︎

We now verify the properties stated in the main text regarding the null space and spectral lifting.

Property 1: Radial Null Space. We examine the curvature along the radial direction \(u\). We compute the quadratic form \(u^\top H_{\mathrm{NITP}} u\): \[u^\top H_{\mathrm{NITP}} u = \frac{1}{r^2} \left[ s \underbrace{u^\top (I - uu^\top) u}_{0} + \underbrace{u^\top (u A^\top) u}_{u^\top u (A^\top u)} + \underbrace{u^\top (A u^\top) u}_{u^\top A (u^\top u)} \right].\] Since \(I - uu^\top\) projects onto the orthogonal complement of \(u\), the first term is 0. For the remaining terms, recall that \(A \perp u\), so \(A^\top u = 0\) and \(u^\top A = 0\). \[u^\top H_{\mathrm{NITP}} u = \frac{1}{r^2} [0 + 1 \cdot 0 + 0 \cdot 1] = 0.\] This confirms that the loss function imposes zero curvature (is locally linear or flat) with respect to changes in the vector norm.

Property 2: Angular Spectral Lifting. Consider a perturbation vector \(w \in \mathbb{R}^d\) such that \(w \perp u\) (representing a purely semantic change). We compute \(w^\top H_{\mathrm{NITP}} w\): \[w^\top H_{\mathrm{NITP}} w = \frac{1}{r^2} \left[ s w^\top (I - uu^\top) w + w^\top (u A^\top + A u^\top) w \right].\] Since \(w \perp u\): 1. \(w^\top (I - uu^\top) w = w^\top w - (w^\top u)^2 = \|w\|_2^2\). 2. \(w^\top u A^\top w = (w^\top u) (A^\top w) = 0\). 3. \(w^\top A u^\top w = (w^\top A) (u^\top w) = 0\).

Thus, the curvature in the tangential direction is: \[w^\top H_{\mathrm{NITP}} w = \frac{s}{r^2} \|w\|_2^2.\] Assuming reasonable alignment during training (\(s = \cos(h, z) > 0\)), the Hessian is strictly positive definite on the subspace orthogonal to \(h\).

6.5 Proof of Theorem 1 (Null Space Mitigation)↩︎

We consider the total loss Hessian \(H_{\mathrm{total}} = H_{\mathrm{NTP}} + \lambda H_{\mathrm{NITP}}\).

Existence of null space. As analyzed in 3.2, although the NTP objective theoretically involves all vocabulary items via the Softmax normalization, the gradient signal is empirically dominated by the target token and a sparse set of high-probability contenders. Consequently, the curvature of \(H_{\mathrm{NTP}}\) is concentrated within this low-dimensional subspace, leaving the vast orthogonal complement effectively unconstrained (rank-deficient).

Based on this, let \(\mathcal{N}_{\mathrm{sem}}\) be the semantic null space of the standard NTP objective, defined as the set of tangential directions \(w \perp u\) that lie in this unconstrained subspace (\(w^\top H_{\mathrm{NTP}} w \approx 0\)).

For any such degenerate direction \(w \in \mathcal{N}_{\mathrm{sem}}\), we compute the quadratic form of the total Hessian: \[\begin{align} w^\top H_{\mathrm{total}} w &= w^\top (H_{\mathrm{NTP}} + \lambda H_{\mathrm{NITP}}) w \nonumber \\ &= \underbrace{w^\top H_{\mathrm{NTP}} w}_{\approx 0} + \lambda w^\top H_{\mathrm{NITP}} w. \end{align}\] Substituting the result from Property 2 derived above (\(w^\top H_{\mathrm{NITP}} w = \frac{s}{r^2} \|w\|_2^2\)): \[w^\top H_{\mathrm{total}} w \approx 0 + \lambda \left( \frac{s}{r^2} \|w\|_2^2 \right).\] Under the empirical observation that NITP converges to high alignment (\(s \to 1\)), we have \(s > 0\). Given \(\lambda > 0\), it follows that: \[w^\top H_{\mathrm{total}} w > 0 \quad \text{for all } w \in \mathcal{N}_{\mathrm{sem}}, w \neq 0.\] This proves that the addition of NITP mitigates the semantic null space by ensuring strictly positive curvature in these previously unconstrained directions. \(\square\)

7 Experiment Details↩︎

Hyperparameter settings. 4 summarizes the hyperparameter configurations used for pretraining both MoE and Dense models. Across all model scales, we adopt a largely consistent training recipe to ensure fair comparisons between standard NTP and NITP. All models are optimized using AdamW with momentum coefficients \((\beta_1, \beta_2) = (0.9, 0.95)\), weight decay set to \(0.1\), and global gradient clipping at \(1.0\). We employ a warmup–stable–decay (WSD) learning rate schedule [57] with 2,000 warmup steps and a fixed decay ratio of \(0.2\). The learning rate and global batch size are scaled according to model size, while the context length is fixed to 8192 tokens for all experiments.

Model architectures. For MoE models, the backbone consists of a shallow dense stage followed by multiple MoE layers, with the number of MoE layers increasing with model capacity. Each MoE layer contains 144 routed experts and one shared expert, using top-8 routing throughout. To control the overall parameter budget, the expert feed-forward hidden size is adjusted across scales, while the dense FFN size is kept fixed. All MoE models use multi-head attention with 8 query heads and 4 key–value heads, with head dimensions scaled proportionally to the hidden size.

Dense models follow a standard Transformer architecture with 24–28 layers depending on model size. They use a larger number of attention heads and higher-dimensional FFNs to match the overall parameter scale of the MoE counterparts. For both MoE and Dense models, the NITP target layer is chosen from intermediate depths, and the NITP loss weight is set close to \(1.0\), with minor adjustments for larger models. All models are pretrained on large-scale tokenized corpora using the Qwen2 tokenizer. The total number of training tokens ranges from 105B to 330B, depending on model size and architecture.

For the average downstream performance shown in the bottom row of 1, we average over four representative benchmarks: MMLU, C3, C-Eval, and ARC-Challenge.

8 Analysis of Representation Geometry Dynamics↩︎

We provide a detailed analysis of the geometric phenomena observed in 1, focusing on the joint behavior of effective rank and average cosine similarity during training.

8.0.0.1 Effective rank dynamics.

We periodically measure the effective rank of the last-layer hidden states during training by collecting all valid token representations within a mini-batch, applying mean-centering, and computing the entropy-based effective rank from the eigenvalue spectrum of the empirical covariance matrix. Measurements are reported every fixed number of training steps to track long-term geometric trends.

As shown in 1 (a), where the background indicates the variance of effective rank, standard NTP exhibits a rapid and monotonic collapse of effective rank, converging to a low-variance regime early in training. This behavior indicates that the learned hidden states concentrate on a small, fixed subspace, regardless of contextual variation. Such convergence reflects a degenerate equilibrium: once the model discovers a narrow anisotropic cone that suffices for next-token prediction, no gradient signal encourages re-expansion along semantically meaningful but task-irrelevant directions.

In contrast, NITP maintains a substantially higher effective rank with noticeably larger variance over training. Importantly, this variance should not be interpreted as optimization instability. Rather, it is a direct consequence of mitigating the semantic null space inherent to NTP. By explicitly supervising the prediction of the next token in latent space, NITP introduces positive curvature along directions that are otherwise unconstrained under likelihood-based training. As a result, the model dynamically activates different subsets of the representation space depending on contextual demands, leading to controlled fluctuations in effective rank.

8.0.0.2 Cosine similarity and anisotropy.

To quantify global anisotropy, we report the average cosine similarity between randomly sampled pairs of token representations from the last layer. At each reporting step, a fixed number of token pairs is sampled across sentences, and cosine similarity is computed after \(\ell_2\) normalization.

As shown in 1 (b), under NTP the average cosine similarity increases steadily throughout training, signaling growing anisotropy and alignment toward a common dominant direction. This behavior is consistent with prior observations of representation degeneration, where embeddings become increasingly indistinguishable despite maintaining predictive accuracy. NITP alleviates this effect. Although cosine similarity still increases as training progresses, the overall magnitude is consistently lower, indicating that representations preserve greater angular diversity. This reduced anisotropy aligns with the higher effective rank observed under NITP, jointly suggesting that the representation space remains more expressive and better conditioned.

8.0.0.3 Implications.

Taken together, these results clarify that NITP does not merely slow down representation collapse, but fundamentally alters the geometric equilibrium of training. While NTP permits hidden states to sacrifice expressiveness as long as token prediction remains correct, NITP enforces semantic consistency in latent space, preventing representations from collapsing into a static anisotropic cone. This geometric regularization provides a mechanistic explanation for the improved downstream generalization observed with NITP.

9 Hidden-State Representation Quality↩︎

NITP is designed to improve the representations learned during pre-training, rather than merely to change final benchmark scores through the output head. While 8 characterizes this effect through representation geometry, we further examine whether the learned hidden states encode richer semantic information than the baseline.

9.0.0.1 Hidden-state quality on MTEB.

We evaluate frozen sentence representations on MTEB v2.12.0 [58]. Following standard practice, for each example, we extract the last-layer hidden states from the 3B MoE models, mean-pool over valid tokens, and apply \(\ell_2\) normalization. No fine-tuning, contrastive training, or task-specific adaptation is applied, so the evaluation serves as a direct probe of the representation quality induced by the pre-training objective. We evaluate 25 English MTEB tasks and exclude seven large-scale clustering and reranking tasks due to computational cost. According to the task types used in our evaluation pipeline, the remaining tasks are grouped into Classification (10 tasks), STS/Similarity (9 tasks), and Retrieval/Duplicate-detection (6 tasks).

The evaluated tasks are: Classification: AmazonCounterfactualClassification, AmazonReviewsClassification, Banking77Classification, EmotionClassification, MTOPDomainClassification, MTOPIntentClassification, MassiveIntentClassification, MassiveScenarioClassification, ToxicConversationsClassification, and TweetSentimentExtractionClassification; STS/Similarity: BIOSSES, SICK-R, STS12, STS13, STS14, STS15, STS16, STS17, and STSBenchmark; and Retrieval/Duplicate-detection: AskUbuntuDupQuestions, SciDocsRR, SprintDuplicateQuestions, StackOverflowDupQuestions, TwitterSemEval2015, and TwitterURLCorpus.

As shown in 5, NITP improves the overall MTEB score from 39.24 to 41.56. The gains are consistent across task families, with improvements of +2.20 points on classification, +2.93 points on semantic textual similarity, and +1.64 points on retrieval/duplicate-detection tasks. At the individual task level, NITP improves 23 out of 25 tasks; the only regressions are small drops on AskUbuntuDupQuestions (\(-\)​0.15 points) and ToxicConversationsClassification (\(-\)​0.88 points). These results indicate that NITP improves the transferability of the hidden states themselves, providing direct evidence for its role as a representation-level pre-training signal.

9.0.0.2 Language modeling preservation.

Finally, we verify that the stronger representations do not come from sacrificing the original language modeling objective. On the Pile [59] validation set, NITP achieves nearly identical validation cross-entropy loss to NTP across model configurations: 3B MoE obtains 2.006 vs. 2.006 (PPL 7.43 vs. 7.43), 9B MoE obtains 1.841 vs. 1.840 (PPL 6.30 vs. 6.30), and 3B Dense obtains 1.884 vs. 1.882 (PPL 6.58 vs. 6.57), where each pair reports NTP vs. NITP. This suggests that token-level cross-entropy alone does not fully characterize the quality of the learned hidden states: models with matched next-token likelihood can still differ substantially in representation geometry and transferability, consistent with our core claim. NITP therefore preserves validation loss under the original next-token prediction objective while learning hidden-state representations with better semantic transferability.

10 Loss Curve of NITP↩︎

To further understand the optimization dynamics of Next Implicit Token Prediction, we monitor the trajectory of \(\mathcal{L}_{\text{NITP}}\) throughout pre-training (5). The loss curve exhibits a characteristic three-phase evolution:

• Phase I: random-initialization collapse. In the very early stage of training, the NITP loss drops sharply from its initial value within the first few dozen steps. This behavior is primarily induced by random initialization: both the predictive state \(h_t\) and the implicit target \(z_{t+1}\) are nearly isotropic and weakly differentiated across tokens, causing their cosine similarity to rapidly concentrate around a common value. As a result, the NITP objective collapses quickly before any meaningful representation learning takes place.

• Phase II: emergence of structured targets. Following the initial drop, the loss exhibits a transient increase, forming a local peak (around 40–2k steps, i.e. warm-up phase). This behavior is consistent with the bottom-up convergence pattern observed in language models [15]. As shallow layers converge faster, they begin to extract more structured and contextualized representations, transforming the implicit target \(z_{t+1}\) from a near-constant signal into a more informative prediction target. This transition temporarily increases the difficulty of the NITP objective for the deeper predictive state \(h_t\).

• Phase III: steady convergence. After the peak, the loss enters a prolonged and stable decline, indicating that deeper layers progressively adapt to the representational manifold defined by shallow layers. The final convergence to a low loss value (e.g., below 0.05 for our 9B model) is consistent with NITP providing a stable and learnable auxiliary supervision signal throughout pre-training.

Overall, the observed hump in the loss curve suggests that the NITP objective functions not merely as a static regularizer, but as an adaptive predictive task whose difficulty evolves alongside the model’s internal representations.

11 Training FLOPs Analysis↩︎

This section provides a detailed analysis of the training FLOPs introduced by NITP, in comparison with the baseline Next-Token Prediction (NTP) objective. Following common practice [17], [18], we approximate total training FLOPs as approximately \(6\times\) the number of parameters involved in forward computation, accounting for forward, backward, and parameter gradient updates.

11.1 Baseline NTP FLOPs↩︎

We consider a transformer-based MoE model with model hidden size \(d\) and vocabulary size \(V\). Each layer consists of a self-attention module and a MoE block. Additionally, the unembedding layer projects the final hidden state to the vocabulary space. For each token:

• Self-attention (GQA) [60]: Adopting Grouped Query Attention (GQA) with 8 query heads and 4 KV heads, the parameter count is reduced compared to standard Multi-Head Attention. The projections (\(W_Q, W_K, W_V, W_O\)) contribute \(\approx 3d^2\) parameters (where \(W_K\) and \(W_V\) are halved), resulting in \(\approx 18d^2\) training FLOPs.

• MoE MLP (SwiGLU [43]): Each expert uses three projections (\(d \times d_e\)). For \(k\) activated experts, the training FLOPs are \(\approx 6 \times (3 k d d_e) = 18 k d d_e\).

• Unembedding Layer: The final projection to the large vocabulary space (\(d \times V\)) contributes significant computation. The training FLOPs are \(\approx 6 V d\). Note that the input embedding layer is excluded as it involves only lookup operations (0 FLOPs).

The total baseline training FLOPs per token (summing over \(L\) layers and the output head) are: \[\mathrm{FLOPs}_{\mathrm{NTP}} \approx L \times (18d^2 + 18 k d d_e) + 6 V d.\]

11.2 Additional FLOPs Introduced by NITP↩︎

NITP introduces a projection head and a cosine loss. Crucially, extracting the target \(z_{t+1}\) incurs no cost due to the stop-gradient.

11.2.0.1 Projection head.

We employ a SwiGLU projection head with an intermediate dimension of \(4d\). It consists of three linear transformations (\(d \to 4d\), \(d \to 4d\), and \(4d \to d\)), resulting in \(12d^2\) parameters. The training FLOPs are: \[\mathrm{FLOPs}_{\mathrm{proj}} \approx 6 \times 12d^2 = 72d^2.\]

11.2.0.2 Cosine similarity loss.

The computation involves dot products and norms (\(\approx 6d\) operations). Including the backward pass: \[\mathrm{FLOPs}_{\mathrm{cos}} \approx 3 \times 6d = 18d.\]

11.3 Numerical Instantiation: 9B MoE Model↩︎

Using the 9B MoE model parameters specified in 4 (\(d=1280\), \(d_e=640\), \(k=9\), \(L=24\), \(V=152064\)), we estimate the FLOPs as follows:

• NITP Overhead (Total): \(72 d^2 + 18d \approx 72 \times 1280^2 \approx 1.18 \times 10^8\) FLOPs.

• Baseline (Backbone Layers): \(24 \times (18 \times 1280^2 + 18 \times 9 \times 1280 \times 640) \approx 24 \times 1.62 \times 10^8 \approx 3.89 \times 10^9\) FLOPs.

• Baseline (Unembedding): \(6 \times 152064 \times 1280 \approx 1.17 \times 10^9\) FLOPs.

• Baseline (Total): \(3.89 \times 10^9 + 1.17 \times 10^9 \approx 5.06 \times 10^9\) FLOPs.

The resulting relative training overhead is: \[\frac{\mathrm{FLOPs}_{\mathrm{NITP}}}{\mathrm{FLOPs}_{\mathrm{NTP}}} \approx \frac{1.18 \times 10^8}{5.06 \times 10^9} \approx 2.3\%.\]

11.4 Empirical Wall-Clock Overhead↩︎

We also measure the practical training-time overhead in the same 9B MoE setting. Using the same number of GPUs and a global batch size of 1024, we compare NTP and NITP over a 5k-step run. The NTP baseline takes 16h 1m, while NITP takes 16h 18m, corresponding to an additional 17 minutes, or approximately 1.8% wall-clock overhead.

This empirical overhead is consistent with the FLOPs estimate above. The measured wall-clock increase is slightly smaller than the theoretical FLOPs ratio because large-scale MoE training also involves communication and system overheads that are largely unaffected by the NITP projection head and cosine loss.

11.4.0.1 Conclusion.

The analysis confirms that the overhead introduced by NITP remains marginal (\(\sim 2\%\)), ensuring scalability.

12 Inference FLOPs Analysis↩︎

Crucially, NITP introduces zero additional computational overhead during inference. Since the NITP objective serves as an auxiliary supervision signal, the projection head is discarded after pre-training. The model architecture used for deployment remains identical to the standard transformer backbone. Therefore, the inference FLOPs per token are exactly the same as the baseline NTP model, ensuring that NITP improves model performance without compromising generation speed or increasing serving costs.

13 Additional Ablation Studies↩︎

Target-layer sweep. We further study the choice of the layer used to construct the implicit target. As shown in 6, the best target consistently appears in a shallow but contextualized region: layer 4 out of 17 for the 3B MoE model and layer 5 out of 24 for the 9B MoE model, both close to 20% of the total depth. Using the embedding layer as the target is weaker, suggesting that purely lexical embeddings lack sufficient contextual information. Moving the target deeper also reduces the gain, likely because deeper representations become increasingly specialized for token discrimination and provide less complementary semantic supervision.

Stop-gradient on implicit targets. We ablate whether gradients are allowed to propagate into the implicit target representations in 1 . Allowing gradients to flow into the targets leads to co-adaptation between the predicted and target representations, which frequently results in unstable optimization and degraded performance. In contrast, applying a stop-gradient operation consistently stabilizes training and yields substantial improvements across downstream benchmarks. The results in 7 highlight that treating the implicit target as a fixed semantic anchor is critical for effective latent-space prediction, whereas disabling stop-gradient performs significantly worse, undermining the supervisory signal by allowing the target to drift during training. Accordingly, we adopt stop-gradient on implicit targets in all experiments.

Start step of NITP. Motivated by the three-phase dynamics observed in the NITP loss curve, we study when NITP supervision should be activated during training. We compare enabling NITP from the beginning with introducing it at later stages. As shown in Table 7, delaying NITP activation leads to consistent performance degradation, particularly on reasoning-oriented benchmarks, while applying NITP from step 0 yields the best overall results. These results indicate that NITP is most effective when applied early, where it can shape representation geometry before higher-level semantics stabilize, thereby promoting more coherent alignment between shallow and deep representations throughout training.

Effect of the projection head. We study the role of the projection head in NITP by comparing our default design, which applies a two-layer MLP projector using SwiGLU on the last hidden states, with a projector-free variant that directly aligns the last hidden states with the implicit targets. Removing the projector consistently results in notable performance degradation across all benchmarks as shown in 7. We attribute this performance drop to a distributional mismatch between deep representations and shallow implicit targets. Directly enforcing alignment imposes an overly restrictive constraint, while a lightweight MLP projection head provides sufficient flexibility to bridge this mismatch without limiting the backbone capacity.

References↩︎