跳转至

Efficient Pre-Training with Token Superposition

会议: NeurIPS 2026
arXiv: 2605.06546
OpenReview: https://openreview.net/forum?id=8td6EGze5A
论文: https://arxiv.org/abs/2605.06546
代码: https://github.com/KellerJordan/modded-nanogpt/pull/178
领域: 优化/理论
关键词: 预训练效率、Token叠加、多标签交叉熵、两阶段训练、表征对齐

一句话总结

提出两阶段Token叠加训练（TST），通过词嵌入平均粗粒化输入与多标签交叉熵预测未来词包，在不改动模型推理架构的前提下大幅提升训练数据吞吐量，在10B MoE规模下实现同等Loss下2.5倍预训练加速。

研究背景与动机

现代通用大语言模型的快速发展在很大程度上依赖激进的数据扩展定律，在实际预训练中往往需要大幅过度训练（overtraining）以换取极致的推理性能。然而，在算力预算严格受限的条件下，如何最大化利用有限FLOPs高效吞吐原始文本成为了亟待解决的核心瓶颈。现有提升预训练效率的方案主要集中在三大路径：其一是信息最大化，例如优化分词策略（BPE、SuperBPE）或引入多Token预测（MTP）、词序预测等辅助多任务损失；其二是计算稀疏化，通过激活子集参数（MoE）或稀疏注意力降低单位Token的FLOPs；其三是压缩建模，在网络内部压缩序列表示（如Byte-Latent Transformer、Bolmo）。

然而，上述方法往往将训练时效率与推理时效率严重混淆。许多方法要么引入额外的网络层、分支预测头或复杂的动态调度，改变了模型原本的推理语义与计算流；要么因引入额外参数使得计算收益被抵消。最近关于分词机制的研究表明，子词模型相比字节级模型的核心性能优势很大程度上源于训练阶段更高的样本吞吐速率。这引发了一个根本性的核心思考：我们能否在完全不破坏推理架构、分词器、优化器和并行机制的前提下，纯粹通过提高训练期的Token吞吐量来提升计算效率？

本文正是基于这一全新视角，提出Token叠加训练（Token-Superposition Training, TST）。其核心洞察在于大语言模型具备极强的结构弹性，预训练可以被解耦为一个高吞吐的粗粒度叠加阶段与一个标准精细自回归恢复阶段。核心 idea：在预训练前期的叠加阶段将相邻的连续Token词嵌入取均值进行输入粗粒化，配合无需新增参数的多标签交叉熵损失监督未来词包预测，使单步等FLOPs吞吐提升 \(s\) 倍；后半程无缝切回标准因果自回归训练迅速恢复有序生成能力，实现零额外开销的大幅预训练加速。

方法详解

整体框架

Token叠加训练（TST）将整个预训练流程划分为两个互补阶段：高吞吐的Token叠加阶段（Superposition Phase，占总步数比例为 \(r\)）和标准自回归恢复阶段（Recovery Phase，占剩余 \(1-r\) 步数）。在叠加阶段，序列按大小为 \(s\) 的非重叠窗口切分为Token词包（bags），嵌入层直接对窗口内的词嵌入做平均压缩为潜变量 \(s\)-token，将数据序列长度扩展为 \(s\) 倍以严格保持每步训练FLOPs恒定；在输出端，模型使用原有的单个分类头输出预测分布，并通过多标签交叉熵（MCE）损失拟合下一个包含 \(s\) 个Token的词包。在恢复阶段，完全移除叠加逻辑，恢复单Token输入与标准的单标签下一词预测，凭借共享的嵌入矩阵与LM Head实现零架构调整的无损对齐。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    subgraph Phase1["第一阶段：高吞吐Token叠加训练（前 r 步）"]
        direction TB
        A["输入原始序列<br/>B×L×V 划分为 s 大小词包"] --> B["输入嵌入平均叠加<br/>非重叠 s-token 均值表征"]
        B --> C["Transformer 主干网络<br/>等 FLOPs 处理粗粒度潜序列"]
        C --> D["单头多标签交叉熵预测<br/>预测下一个未来词包 LMCE"]
    end
    Phase1 -->|无缝切换<br/>保留权重与嵌入| Phase2["第二阶段：标准因果自回归恢复（后 1-r 步）"]
    subgraph Phase2
        direction TB
        E["标准单 Token 输入<br/>移除词包切分与均值操作"] --> F["标准因果自回归监督<br/>单标签交叉熵 LCE 预测下一词"]
        F --> G["最终标准自回归语言模型<br/>推理架构与基线完全一致"]
    end

关键设计

1. 输入嵌入均值叠加：等FLOPs下实现 \(s\) 倍数据吞吐

为解决大模型预训练中单位算力吞吐样本受限的瓶颈，TST在输入端引入非重叠词包均值表征。给定批大小为 \(B\)、序列长度为 \(L\) 的输入Token序列，首先将其重塑为大小为 \(s\) 的连续词包，形状变为 \(B \times (L/s) \times s\)。在嵌入层中，模型直接将每个词包内 \(s\) 个连续Token的嵌入向量取平均值，得到潜序列长度为 \(l = L/s\)、特征维度为 \(d\) 的潜变量表征。为确保公平性并进行严格的等算力扩展，作者在叠加阶段将数据序列长度扩展 \(s\) 倍，使得前向与反向传播在潜序列上的实际计算量（FLOPs per step）与标准基线完全一致。这种输入端粗粒化不仅将单位算力所见到的原始数据吞吐提升了 \(s\) 倍，而且对嵌入空间起到了隐式几何正则化约束，强迫模型在极早期建立能够线性分离并承载局部多Token统计特性的高维语义先验。

2. 单头多标签交叉熵损失：极简未来词包预测目标

以往预测未来多Token的方案（如MTP）通常需要挂载多个预测头或额外的投影模块，导致计算开销与参数膨胀。TST选择保留原模型的单个语言模型头（LM Head），将预测目标从单个离散Token拓展至包含 \(s\) 个Token的未来无序词包。为了使词包内的每个合法目标Token均匀分摊概率质量，作者将标准交叉熵损失拓展为多标签交叉熵（MCE）。通过严谨的KL散度推导，目标分布为词包内的均匀分布 \(t_y = 1/s\)，对应的损失函数形式为：

\[L_{\text{MCE}}(z, y) = \frac{1}{|y|} \sum_{y \in y} \left( -z_y + \log \sum_{i=1}^V \exp(z_i) \right) - \log |y| = \frac{1}{|y|} \sum_{y \in y} L_{\text{CE}}(z, y) - \log |y|\]

由于常数项 \(-\log |y|\) 的梯度恒为0，在实际训练实现中可直接省略，从而转化为对词包内各个目标Token标准交叉熵损失的算术平均值。这种设计巧妙地复用了主流深度学习框架（如PyTorch/TorchTitan）中经过高度编译优化的标准交叉熵算子，无需定制复杂显存算子即可获得高效训练梯度。

3. 零架构调整的跨阶段无缝恢复：保障深层表征连续性

在经历比例为 \(r\) 的叠加预训练后，模型获得了对多Token语义模式的全局理解，但丧失了严格的从左至右自回归采样秩序。为了在不牺牲生成质量的前提下重构自回归链条，TST在第二阶段直接切换回标准的单Token下一词预测，且不冻结任何参数、不引入适配层（Adapter）。消融实验证实，输入嵌入矩阵与输出LM Head在两个阶段之间的权重共享是TST能够成功的决定性支柱——如果人为打乱或重新随机初始化嵌入和输出头，TST前期的所有训练收益将被彻底抹除。正因为没有新加任何投影层，底层神经网络电路在叠加阶段学到的高阶抽象知识与几何流形得以无损迁移至精细自回归阶段，模型仅需在恢复期极短的步数内便迅速反超同等算力基线的收敛曲线。

损失函数 / 训练策略

叠加训练在第一阶段采用等权重多标签交叉熵 \(L_{\text{MCE}}\)，对于大窗口规模（如 \(s \ge 8\)），作者进一步探索了基于词间互信息随距离按幂律衰减规律设计的递减加权交叉熵：

\[L_{\text{MCE}}(z, y, g) = \frac{1}{\sum_i g(i)} \sum_{y_i \in y} g(i) L_{\text{CE}}(z, y_i), \quad g(i) = \frac{1}{i}\]

第二阶段完全切回单Token的 \(L_{\text{CE}}\)。优化器采用AdamW，学习率调度采用Warmup-Stable-Decay（WSD）策略，前2000步线性预热，在退火期（最后10%步数）执行衰减。

实验关键数据

主实验

在270M、600M、3B密集模型以及10B总参数（1B激活参数）MoE架构上，作者基于TorchTitan框架与FSDP并行在NVIDIA B200集群上进行了详尽评测。所有下游评测均在完全恢复为标准自回归模型后的检查点上执行零样本（0-shot）测试。

模型架构 参数量 训练机制 叠加步数/总步数 词包大小 \(s\) 总吞吐数据量 B200卡时 (↓) 最终Loss (↓) HellaSwag (↑) ARC-E (↑) ARC-C (↑) MMLU (↑)
Dense 270M Baseline 0 / 20k – 42B 34 3.212 36.3 46.7 24.9 –
Dense 270M TST 6k / 20k 6× 105B 34 3.142 38.6 47.6 26.4 –
Dense 270M Baseline 0 / 100k – 209B 170 3.092 40.2 47.5 26.2 –
Dense 270M TST 30k / 100k 6× 524B 170 3.048 42.6 50.3 25.5 –
Dense 600M Baseline 0 / 20k – 42B 61 3.019 43.5 51.7 25.5 –
Dense 600M TST 6k / 20k 6× 105B 61 2.943 48.2 52.5 26.9 –
Dense 3B Baseline 0 / 20k – 42B 247 2.808 57.6 60.6 31.9 31.2
Dense 3B Baseline 0 / 36k – 75B 443 2.677 62.3 65.9 34.9 32.7
Dense 3B TST 6k / 20k 6× 105B 247 2.676 62.4 66.3 36.0 32.8
MoE 10B-A1B Baseline 0 / 125k – 1.05T 12,311 2.252 70.1 73.8 46.3 37.4
MoE 10B-A1B TST 12.5k / 50k 16× 2.00T 4,768 2.236 71.2 74.2 47.3 39.0

在3B密集模型上，TST仅消耗247卡时（20k步）即可达到传统基线需要443卡时（36k步）才能达到的2.676验证集Loss；在10B-A1B MoE大模型上，TST在达到相同甚至更优Loss（2.236 vs 2.252）的前提下，将训练步数从125k急剧减少至49,983步，整体GPU耗时降低达 2.5倍（从12,311 B200卡时降至4,768 B200卡时）。

消融实验

为了剖析输入叠加与输出多标签监督的独立贡献，以及跨阶段权重一致性的影响，作者设计了细致的对比消融：

配置 / 变体 叠加设置说明 最终Loss (↓) 相对基线表现
标准基线（Baseline） 纯标准自回归，无任何叠加 3.212 基准
仅输入叠加（Input-Only） 仅输入词嵌入做平均，仍预测单个下一词 3.178 优于基线，吞吐提升但缺乏多标签梯度
仅输出叠加（Output-Only） 输入为单Token，预测未来词包多标签 3.165 优于基线，提供更密集监督信号
完整 TST（Full Superposition） 输入嵌入均值叠加 + 输出多标签交叉熵 3.142 最佳效果，两者呈现无冲突的加成效应
恢复期随机重置（Randomized） 在恢复期开始时随机重新初始化Embedding与Head 2.938（3B模型） 严重恶化，甚至显著差于基线（2.808）

超参数网格搜索（表4、表5）表明：叠加比例 \(r \in [0.2, 0.4]\) 且词包大小 \(s \in [4, 8]\) 是最稳健的泛化区间；过小的 \(s\) 吞吐增益有限，过大的 \(s\)（如 \(s \ge 16\)）需要配合非均匀的幂律加权损失以缓解远距离不确定性带来的噪声。

关键发现

  • 双重正交机制的叠加收益：输入叠加改变了模型感知的输入粒度，实质上降低了每单位有效语义信息的FLOPs；输出叠加则提供了紧凑的未来词包局部先验。二者协同运作，完全没有相互抵消。
  • 表征对齐的脆弱性与关键性：多阶段预训练中，以往文献广泛采用适配层（Adapter）冻结主干进行迁移，反而遮蔽了这一机制的潜力。必须保留完全一致的嵌入与Head参数，模型才能在恢复阶段实现跳跃式快速收敛。
  • 算力受限与数据受限的权衡：TST本质上是在算力约束主导的场景下，用更激进的数据吞吐量换取更低的FLOPs消耗。如果未来进入极度缺乏新数据的极限场景，仅输出叠加（Output-Only）可以在不增加数据消费的前提下稳步超越标准基线。

亮点与洞察

  • 极简工程落地的优雅设计：不需要调整模型层级结构、不增加参数量、不重训Tokenizer，仅用数行PyTorch原生代码实现输入均值与交叉熵累加，即可直接插入TorchTitan等工业级预训练框架。
  • 解耦训练效率与推理架构：摆脱了以往效率研究强行绑定推理端加速（如推测解码、投机采样）的思维定式，将所有激进的效率优化局限在预训练期，最终交付的模型完全保留标准的因果自回归架构与完整表达力。
  • 宏观粗粒度到微观精细度的认知范式：证实了大语言模型具备类似人类认知的"先掌握粗粒度概括、再精细推导具体词句"的双阶段自适应弹性，为未来长文本、层次化语言建模提供了重要的理论支点。

局限与展望

  • 数据需求量显著放大：由于等FLOPs下序列长度扩展了 \(s\) 倍，TST在叠加期消耗的原始Token数量增加。在算力受限（compute-bound）时极其高效，但在数据受限（data-bound）场景下可能受限于未标注语料规模。
  • 超长词包下的词序信息缺失：当词包大小 \(s > 8\) 时，无序词包内部的信息熵大幅上升，若不借助复杂的自适应衰减加权，容易出现训练不稳定或退化。
  • 未深入探索长上下文潜在红利：将序列压缩 \(s\) 倍意味着在叠加阶段模型的有效上下文窗口被等比例扩大了 \(s\) 倍，这对于保留原生超长文本的全局连贯性可能具有尚未被量化评估的潜在利好，值得后续研究深挖。

相关工作与启发

  • vs Multi-Token Prediction (MTP / DeepSeek-V3): MTP通过挂载额外的独立前向头预测未来 \(k\) 个Token，增加了模型参数量与训练内存占用，且并不增加单位FLOPs的数据吞吐率；TST仅使用单一原始LM Head监督无序词包，严格以 \(s\) 倍提升原始Token吞吐，两者概念正交且理论上可正交结合。
  • vs Patch-Level Training (Shao et al.): 二者在核心机制上呈现出惊人的概念收敛，但Shao等将重点放在固定数据集下的FLOPs缩减，而TST基于Gigant等人的吞吐假说，从等步FLOPs的Token吞吐放大角度统一了输入与输出双重正交视角，并验证了10B级MoE与2T Token的超大规模鲁棒性。
  • vs Future Summaries (Mahajan et al.): 后者依赖挂载辅助头与二值交叉熵（BCE）损失预测未来概括，带来了额外开销；TST证明直接采用单头多标签交叉熵并在后期无损切回，能够以更低的系统复杂性达成更优的收敛表现。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 突破传统单阶段自回归训练定势，以极简形式打通粗粒度叠加与精细恢复。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖270M到10B MoE规模、从42B到2T Tokens，包含系统级网格搜索与关键表征消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨清晰，数学推导完备，代码完全透明可复现。
  • 价值: ⭐⭐⭐⭐⭐ 工业级即插即用的预训练提效方案，具有显著降低前沿大模型算力成本的实际落地价值。

相关论文