← 返回 2026-06-02

NITP：面向大语言模型预训练的下一个隐式Token预测 NITP: Next Implicit Token Prediction for LLM Pre-training

Xiangdong Zhang, Debing Zhang, Shaofeng Zhang, Xiaohan Qin, Yu Cheng, Junchi Yan 📅 2026-05-24 👍 36 2026-07-13 08:36
Mixture-of-Experts 几何正则化 自监督学习 表示学习 预训练目标
arXiv HuggingFace PDF

通过在隐空间预测浅层语义表示，正则化表示几何，提升下游任务性能

前置知识

Next-Token Prediction (NTP)

LLM预训练的核心目标函数，要求模型给定前文x<t，预测下一个离散token xt+1。通过最大化似然P(xt+1|x<t) = softmax(ht⊤wxt+1)来优化，其中ht是最后层隐藏状态，wxt+1是token embedding。虽然理论上Softmax归一化涉及所有词表项，但实际梯度主要由目标token和少量高概率候选词主导，导致监督信号稀疏。

本文的核心问题就来自NTP的局限性，理解NTP的工作原理是理解NITP改进动机的基础

Representation Degeneration（表示退化）

语言模型训练中的一种现象，隐藏表示会逐渐聚集到一个低维、各向异性的锥形区域。当不同token的隐藏状态变得越来越相似时，会失去语义区分度。这种现象通过两个指标量化：Effective Rank（有效秩，基于特征值熵的有效维度）和Average Cosine Similarity（平均余弦相似度，衡量全局各向异性）。正常情况下希望有效秩较高、余弦相似度较低。

这是本文要解决的核心问题，NTP会导致表示退化，而NITP旨在缓解这一问题

Effective Rank（有效秩）

衡量表示空间利用的有效维度，定义为exp(H(λ))，其中H(λ) = -∑(λi/∑λi)·log(λi/∑λi)是特征值的熵。如果所有特征值相等，有效秩等于维度d；如果只有少数主导特征值，有效秩会远小于d。这个指标能捕捉表示空间是否真正利用了高维空间的表达能力。

论文用它来定量分析表示退化，Figure 1(a)显示NTP下有效秩快速下降，而NITP能维持更高的有效秩

Anisotropy（各向异性）

在几何学中指不同方向上性质不同的现象。在向量表示的语境下，各向异性指表示向量倾向于沿着某个主导方向对齐，而不是均匀分布在单位球面上。这会限制表示的表达能力，因为垂直于主导方向的变异变得很小。论文用平均余弦相似度来衡量全局各向异性，相似度越高表示各向异性越严重。

Figure 1(b)显示NTP下余弦相似度不断上升，而NITP能显著降低各向异性

Hessian矩阵和曲率

Hessian矩阵是损失函数关于参数的二阶导数矩阵，描述了优化地形在不同方向的弯曲程度（曲率）。如果Hessian在某方向的特征值接近零，说明这个方向是平坦的（曲率小），优化时缺乏约束；如果特征值较大，说明这个方向有强约束。论文通过分析NITP的Hessian证明它在语义零空间中引入正曲率，正则化这些原本不受约束的方向。

这是NITP理论分析的核心，Lemma 3.1和Theorem 3.2证明了NITP的几何正则化作用

Mixture-of-Experts (MoE)

一种稀疏激活的架构，包含多个专家网络，每个token只路由到少数专家（本文用top-8路由）。相比稠密模型，MoE可以在增加总参数量的同时保持每个token的激活参数量较小，从而提高容量而不大幅增加计算成本。论文使用DeepSeek-V2风格的MoE，每层144个路由专家+1个共享专家，每个专家使用SwiGLU FFN。

论文在多个MoE模型规模上验证NITP的有效性（1.9B、3B、9B、45B），证明方法跨架构的泛化能力

Stop-gradient（梯度截断）

一种技术手段，在前向传播时正常计算，但在反向传播时切断梯度流动，不更新某部分参数或中间激活。在NITP中，对隐式token zt+1 = sg Eshallow(x≤t+1)(t+1)应用stop-gradient，防止目标表示被更新，使其保持作为稳定的语义锚点。如果不加stop-gradient，目标会随预测状态一起漂移，导致正则化效果失效。

论文Table 8显示，禁用stop-gradient会导致性能显著下降，从平均23.58分降到18.23分

研究动机

标准的Next-Token Prediction (NTP)预训练目标存在根本性的几何约束不足问题。NTP只在离散的token输出空间提供稀疏的one-hot监督信号，虽然梯度会传播到隐藏状态，但目标函数主要约束沿着目标logit方向的表示，而对大量与目标方向正交的自由度缺乏约束。在实际训练中，这种欠约束导致隐藏表示漂移到退化、各向异性的配置中。论文通过追踪训练过程中的几何演化发现，NTP下Effective Rank从初始的高值迅速下降，而Average Cosine Similarity持续上升（Figure 1）。以9B MoE模型为例，训练结束时表示空间严重坍缩，虽然token预测准确率可能很高，但表示的表达能力被严重削弱。这种现象被称为Representation Degeneration，已被多篇工作（Ethayarajh 2019; Wang et al. 2020; Barbero et al. 2024）观察到，并与下游任务泛化能力下降相关。

本文的目标是本文的目标是提出一个辅助预训练目标，在表示空间提供密集的、连续的监督信号，补充NTP的离散token级监督。具体来说，要求模型不仅预测下一个token的身份（离散ID），还要预测下一个token的隐式语义表示（连续向量）。通过在隐空间引入额外的监督，约束那些NTP忽略的自由度，防止表示向低维、各向异性的锥形漂移，从而学习到更丰富、更具表达能力的表示，提升下游任务性能。

与已有工作不同的是，本文的独特切入角度在于从几何视角重新审视预训练目标。现有工作主要关注扩展预测范围（如Multi-Token Prediction预测多个未来token）或改变预测内容（如预测未来摘要），但它们仍在离散token空间操作，依赖one-hot目标的token级监督。NITP则是在连续表示空间引入监督，要求模型预测语义表示而非token ID。另一个关键区别是自监督设计：利用模型自身浅层的隐藏表示作为预测目标，无需引入外部编码器或额外标注。这种设计既高效又避免了领域偏移。此外，本文还提供了详细的理论分析（Hessian分析）证明NITP如何通过在语义零空间引入正曲率来正则化表示几何，这与现有工作缺乏理论分析形成对比。

核心方法

NITP（Next Implicit Token Prediction）的整体思路是：在标准的Next-Token Prediction基础上，增加一个辅助目标，要求模型在隐空间预测下一个token的语义表示。直觉上，NTP只要求模型知道"下一个词是什么"，而NITP要求模型理解"下一个词意味着什么"。技术路线上，首先定义"隐式token"为下一个token在浅层网络的上下文化表示，然后训练深层的预测状态ht来预测这个隐式表示。核心创新在于：(1)利用浅层作为稳定的语义锚点，(2)引入时间偏移预测（t→t+1），(3)使用余弦相似度损失实现尺度不变约束。NITP与NTP联合优化，总损失Ltotal = LNTP + λ·LNITP，其中LNITP = 1 - cos(P(ht), zt+1)，P是投影头，zt+1是隐式token。由于隐式token直接从中间激活提取且用stop-gradient截断梯度，NITP的计算开销极小（约2%额外FLOPs）。

NITP的核心创新点是用"隐式token"替代"离散token"作为预测目标。隐式token zt+1 = sg Eshallow(x≤t+1)(t+1)是下一个token在浅层网络的上下文化隐藏表示，用stop-gradient固定，不参与反向传播。这与已有方法的本质区别在于：NTP预测离散token ID（one-hot向量），NITP预测连续语义表示（密集向量）。更关键的是，NITP采用自监督设计，目标来自模型自身浅层，而非外部模型或标注。这个设计的巧妙之处在于：(1)浅层保留了更丰富的词汇和局部语义结构（论文引用Skean et al. 2025，Liu et al. 2024c证明语义丰富度在浅层达到峰值），(2)Transformer训练是自底向上收敛的（Skean et al. 2025），浅层更早稳定，使隐式token成为可靠的目标，(3)无需额外计算开销。这种"预测未来语义表示"的思路与现有所有预训练目标都不同。

方法步骤详情

NITP的完整方法步骤如下：(1)输入处理：给定token序列x = {x1, ..., xT}，标准的前向传播计算每层的隐藏状态。(2)隐式token构造：对于每个位置t，从浅层网络（约总深度的20%，论文实验中第4层/17层或第5层/24层效果最佳）提取位置t+1的隐藏表示，应用stop-gradient得到隐式token zt+1 = sg Eshallow(x≤t+1)(t+1) ∈ ℝ^d。(3)预测状态获取：从最后一层获取位置t的隐藏状态ht。(4)投影头映射：为了处理深层预测状态和浅层目标之间的分布差异，用一个轻量级投影头P（两层MLP，SwiGLU激活，中间维度4d）将ht映射到与隐式token相同的空间：P(ht)。(5)计算NITP损失：使用余弦相似度损失LNITP(ht) = 1 - (P(ht)⊤zt+1)/(∥P(ht)∥2·∥zt+1∥2)。这个损失促使P(ht)与zt+1在方向上对齐，但对尺度不敏感。(6)联合优化：将NITP损失与标准NTP损失加权组合：Ltotal(θ) = LNTP(θ) + λ·LNITP(θ)，论文中λ通常设为1.0或0.8，对大模型可略微降低。(7)推理时丢弃：预训练完成后，丢弃投影头P，推理模型结构与标准NTP完全相同，无额外计算开销。

技术新颖性

NITP的技术新颖性体现在多个方面：首先是"预测隐式token"的全新目标设计，这在预训练文献中首次出现。其次是浅层自监督的巧妙设计，利用模型自身的语义丰富层作为目标，避免了引入外部模型的开销和领域偏移。第三是时间偏移预测（t→t+1）的重要性，论文Table 3和Figure 3证明，如果没有时间偏移（即预测同位置的隐式表示），会导致灾难性性能下降（平均从23.58分降到18.75分），说明NITP的有效性依赖于预测未来而非对齐当前。第四是理论分析的完整性，论文推导了NITP的Hessian矩阵，证明其在与h正交的子空间中提供正曲率（Lemma 3.1），并证明这能缓解NTP的语义零空间问题（Theorem 3.2）。这种理论与实验结合的完整性在预训练目标论文中较为罕见。最后是工程实现的简洁性，NITP只增加一个投影头和一个余弦损失，训练开销仅约2%，且推理零开销，这使得方法极易集成到现有训练流程中。

Overview of NITP. Next Implicit Token Prediction supervises hidden states by predicting temporally shifted implicit tokens, i.e. shallow-layer representations, and is jointly optimized with the standard next-token prediction objective.
Figure 2: Overview of NITP. Next Implicit Token Prediction supervises hidden states by predicting temporally shifted implicit tokens, i.e. shallow-layer representations, and is jointly optimized with the standard next-token prediction objective.
Loss comparison between whether temporal shift or not.
Figure 3: Loss comparison between whether temporal shift or not.
Average performance under different NITP loss.
Figure 4: Average performance under different NITP loss.
Training dynamics of the NITP loss. Evolution of LNITP during pre-training, exhibiting a characteristic three-phase behavior: an initial collapse induced by random initialization, a transient hump caused by the emergence of structured shallow-layer targets, and a long-term stable convergence.
Figure 5: Training dynamics of the NITP loss. Evolution of LNITP during pre-training, exhibiting a characteristic three-phase behavior: an initial collapse induced by random initialization, a transient hump caused by the emergence of structured shallow-layer targets, and a long-term stable convergence.

实验结果

论文的主要发现可以总结为四点。第一，NITP在多个模型规模和架构上持续提升下游性能。在MoE模型上，1.9B、3B、9B模型的平均得分分别提升0.8、2.1、2.7个百分点。最显著的提升在9B MoE上，MMLU-Pro从15.29%提升到21.00%（+5.71%绝对提升），C3从56.65%提升到63.01%（+6.36%），CommonsenseQA从45.70%提升到49.96%（+4.26%）。在稠密模型上，0.5B、2B、3B模型的平均得分分别提升1.02、1.79、1.35个百分点。第二，NITP有效缓解表示退化。Figure 1显示，NITP下Effective Rank维持更高值，Average Cosine Similarity更低，证明表示空间更丰富、更少各向异性。第三，NITP提升隐藏状态本身的语义质量，即使不微调，直接用冻结的最后一层隐藏状态作为句子表示，在MTEB的25个任务上，NITP在23个任务上优于NTP，总体得分从39.24提升到41.56（+2.33）。第四，NITP的计算开销极小。在9B MoE上，额外训练FLOPs仅约2.3%，实测wall-clock开销约1.8%（从16小时1分增加到16小时18分），且推理零开销。第四，NITP在更大规模（45B MoE）上依然有效，平均得分从52.36提升到54.02（+1.66）。第五，NITP的语言模型困惑度几乎不变，在Pile验证集上，3B MoE的PPL都是7.43，9B MoE都是6.30，说明NITP在不牺牲原始语言建模目标的情况下提升表示质量。

Main results. Performance comparison between NTP and NITP. We categorize benchmarks into Knowledge & Language Understanding (5 tasks) and Reasoning & Problem Solving (8 tasks). Best results are bolded.
Table 1: Main results. Performance comparison between NTP and NITP. We categorize benchmarks into Knowledge & Language Understanding (5 tasks) and Reasoning & Problem Solving (8 tasks). Best results are bolded.
Results for dense models. Performance comparison between NTP and NITP on seven representative benchmarks. Best results are bolded. Avg. denotes the arithmetic mean across all benchmarks.
Table 2: Results for dense models. Performance comparison between NTP and NITP on seven representative benchmarks. Best results are bolded. Avg. denotes the arithmetic mean across all benchmarks.
Ablation studies. We evaluate the impact of key design choices in NITP using the 3B MoE model trained on 200B tokens, including the target layer selection, temporal shift strategy, loss function, and regularization. The default NITP configuration (shallow L4, next-token target, cosine loss) achieves the best performance.
Table 3: Ablation studies. We evaluate the impact of key design choices in NITP using the 3B MoE model trained on 200B tokens, including the target layer selection, temporal shift strategy, loss function, and regularization. The default NITP configuration (shallow L4, next-token target, cosine loss) achieves the best performance.
Hyper-parameters for MoE and Dense model pretraining.
Table 4: Hyper-parameters for MoE and Dense model pretraining.
Scaling result on a 45B MoE model. We report results on a larger MoE model trained for 240B tokens. Best results are bolded.
Table 5: Scaling result on a 45B MoE model. We report results on a larger MoE model trained for 240B tokens. Best results are bolded.
MTEB evaluation of hidden-state quality. Scores are reported as percentages and macro-averaged over task-level main scores within each group. NITP improves representation utility on 23 out of 25 tasks using frozen last hidden states.
Table 6: MTEB evaluation of hidden-state quality. Scores are reported as percentages and macro-averaged over task-level main scores within each group. NITP improves representation utility on 23 out of 25 tasks using frozen last hidden states.
Target-layer sweep across scales. We report the average score over MMLU, MMLU-Pro, CSQA, BBH, and LCBench. The best target layer consistently lies in a shallow contextualized region around 20% of the total depth.
Table 7: Target-layer sweep across scales. We report the average score over MMLU, MMLU-Pro, CSQA, BBH, and LCBench. The best target layer consistently lies in a shallow contextualized region around 20% of the total depth.
Top: Representation geometry of the last hidden states under NTP and NITP. Bottom: Average downstream performance of 9B MoE and 2B dense models (details in Appendix B).
Figure 1: Top: Representation geometry of the last hidden states under NTP and NITP. Bottom: Average downstream performance of 9B MoE and 2B dense models (details in Appendix B).
查看结构化数据
任务指标本文基线提升
MMLU-Pro（知识推理） 准确率 21.00% 15.29% +5.71%（绝对提升）
C3（阅读理解） 准确率 63.01% 56.65% +6.36%
CommonsenseQA（常识推理） 准确率 49.96% 45.70% +4.26%
ARC-Challenge（科学推理） 准确率 53.95% 51.20% +2.75%
MMLU（多任务语言理解） 准确率 46.14% 43.71% +2.43%
GSM8k（数学推理） 准确率 32.52% 30.09% +2.43%
BBH（链式思维推理） 准确率 28.67% 28.07% +0.60%
MTEB（嵌入质量，25任务平均） 平均得分 41.56% 39.24% +2.33%
45B MoE平均得分 平均准确率 54.02% 52.36% +1.66%

局限与改进

论文讨论了NITP的几个局限性。首先是引入了额外的超参数：目标层选择和NITP损失权重λ。虽然论文的消融实验表明这些选择在模型间高度稳定（目标层约20%深度，λ≈1.0），但仍需在更多模型上验证鲁棒性。其次，理论分析假设隐式token与预测状态有良好对齐（s = cos(h,z) → 1），论文通过实验验证NITP损失确实收敛到低于0.1（如9B MoE为0.04），但在极早期训练阶段这个假设可能不成立。第三，虽然NITP在0.5B到45B参数规模上都有效，但尚未在更大规模（如100B+）上验证，且论文主要在中文和英文数据上测试，多语言场景的泛化性未知。第四，NITP主要关注表示几何正则化，但它是否对所有类型的下游任务都有帮助尚不清楚。论文在MTEB上有23/25任务提升，但有2个小幅下降（AskUbuntuDupQuestions -0.15%，ToxicConversationsClassification -0.88%），说明某些任务可能不受益或轻微受损。第五，NITP的MTEB测试使用冻结的隐藏状态，这能证明表示质量提升，但实际应用中通常会微调模型，冻结表示场景的实用性可能有限。

独立分析的弱点

独立分析的第一个弱点是NITP的理论假设在实际训练中可能不完全成立。论文分析依赖隐式token与预测状态的高对齐（s→1），但Figure 5显示NITP损失曲线有三阶段行为，Phase II（约40-2k步）会出现一个hump，损失上升，表明在这个阶段对齐度下降。虽然论文解释这是因为浅层开始提供更结构化的目标，增加了难度，但这个阶段的理论正则化效果可能较弱。改进方向是研究如何更早启动NITP或调整损失权重调度，减少Phase II的负面影响。第二个弱点是目标层选择对模型规模敏感。论文Table 7显示，3B MoE的最佳目标层是第4层/17层（约24%），9B MoE是第5层/24层（约21%），45B MoE也是第5层/24层（约21%），虽然都接近20%，但有细微差异。如果这个规律在更大模型上继续漂移，需要更系统的目标层选择策略。改进方向是研究自适应目标层选择，如基于每层表示的Effective Rank或语义丰富度自动选择。第三个弱点是NITP主要针对表示几何，但下游任务性能可能还受其他因素影响，如输出logit的分布、模型对特定模式的记忆等。虽然论文报告了MTEB的提升，但某些任务的小幅下降表明NITP可能不是"万能"的。改进方向是将NITP与其他辅助目标结合，如多token预测或未来摘要预测，形成更全面的预训练目标组合。第四个弱点是NITP在推理时被丢弃，这虽然避免了推理开销，但意味着预训练学到的表示对齐信息没有被利用。某些下游任务可能从这种对齐中受益。改进方向是研究如何在推理阶段保留部分NITP信息，如使用投影头进行表示一致性检查。

未来方向

作者提出的未来工作方向包括在更多模型规模和架构上验证NITP的鲁棒性，以及探索NITP与其他预训练目标的结合。基于论文成果，可延伸的未来研究方向包括：(1)自适应NITP：研究如何根据训练阶段或数据特性动态调整λ权重。论文Figure 4显示λ=1.0在多数情况下最佳，但更大模型（45B）用λ=0.6更好，未来可设计调度策略，如随着训练进行或模型规模增大自动降低λ。(2)多粒度隐式token：当前NITP只预测单个token的语义表示，未来可扩展到预测n-gram或句子的语义表示，形成层次化的隐式预测目标。(3)跨模态扩展：NITP的思想可以推广到多模态模型，要求视觉模态预测文本模态的隐式表示，或反之，这可能提升跨模态对齐。(4)理论拓展：论文的Hessian分析假设隐式目标固定，未来可研究如果目标随训练缓慢变化（如不加stop-gradient或用软stop-gradient）的几何效应，这可能进一步提升性能。(5)下游微调策略：研究NITP预训练的模型是否需要特殊的微调策略，如在微调时保留部分NITP监督，以维持表示几何。(6)与其他正则化方法结合：将NITP与表示去各向异性方法（如Gao et al. 2019的cosine regularization）结合，形成更强的几何正则化。

复现评估

论文的可复现性良好。代码已在GitHub开源（https://github.com/aHapBean/NITP），这使研究团队能够直接复现实验。论文详细描述了实验设置，包括超参数（Table 4）、模型架构、训练数据（105B-330B tokens，混合英文、中文、代码、数学、推理数据）和评估基准（13个主流基准+MTEB 25任务）。然而，复现仍有挑战：首先，训练9B MoE模型需要大量计算资源（330B tokens，约24层，1280 hidden dim），论文没有报告总训练时间或GPU数量，预估可能需要数百GPU-week。其次，论文使用DeepSeek-V2风格的MoE架构，包含144个专家，这种架构的工程实现较为复杂，对分布式训练框架有较高要求。第三，论文使用自有的训练数据混合，未公开详细的数据组成和分词后token数量，这可能影响精确复现。第四，MTEB的某些任务（如大规模聚类和重排序）因计算成本被排除，这会影响结果的可比性。总体而言，对于有计算资源的团队，复现主要实验是可行的，但完整复现所有实验（包括45B MoE）需要相当大的投入。

论文图表

Table 8: Additional ablations for NITP design choices. We report results on key reasoning benchmarks. Avg. denotes the arithmetic mean of MMLU, MMLU-Pro, CSQA, BBH, and LCBench.

Table 8展示额外的消融实验。Stop-gradient：启用(37.37%)远优于禁用(31.89)，证明固定隐式token作为锚点的重要性。NITP启动步数：从步数0开始(37.37%)最优，延迟启动（3000步: 20.17, 2000步: 21.40, 1000步: 21.22）导致性能下降，尤其在推理任务上。投影头：使用投影头(37.37%)远优于不用(35.40)，说明深层和浅层表示之间存在分布差异，需要投影头桥接。

这张表格补充了Table 3的消融实验，提供了更多设计选择的证据。Stop-gradient的对比（+5.54绝对提升）是最强的消融结果，证明了"固定目标"是NITP的核心。启动步数的对比表明NITP需要在训练早期就激活，以在高层语义稳定前塑造表示几何。投影头的对比证明NITP不是简单地对齐表示，而是需要通过投影来适应深层和浅层之间的分布差异。