摘要：针对大模型预训练算力成本高、迭代速度慢的行业痛点，Nous Research提出TST（Token Superposition Training）词元叠加训练方案。本文从技术原理、实现逻辑、实验数据、消融实验、落地优劣全方位拆解，带你看懂这款零侵入、高收益的预训练提效新方案。

关键词：大模型预训练、TST、Token叠加、预训练提速、算力降本、MCE损失

1. 前言

大模型预训练是AI研发的核心成本瓶颈。传统Next-Token Prediction（NTP）范式全程采用逐Token精细化训练，在训练早期存在严重的算力冗余。现有主流提效方案（MoE、稀疏注意力、MTP）均存在工程复杂度高、推理适配成本大的问题，难以快速落地。

为此，Nous Research团队提出TST词元叠加训练法，仅修改训练循环逻辑，不改动模型架构、分词器、优化器与训练数据，实现单位FLOPs下数据吞吐最大化，在10B MoE模型上实现2.5倍预训练提速，算力成本降低60%以上。

2. TST核心技术原理

TST采用粗训+精训两阶段渐进式训练，核心是输入词元压缩+多热交叉熵监督，全程无新增复杂模块。

2.1 阶段一：词元叠加粗训阶段

将原始序列长度为L的文本，按固定包大小s，分割为若干个包含s个连续Token的词元包。对每个词元包内的Token嵌入向量做均值处理，完成输入表征压缩：

核心代码逻辑（PyTorch极简实现）：

python # TST输入侧词元叠加核心逻辑 if superposition_bag_size > 1:     bs, seq = inputs.shape     # 重塑维度，拆分词元包     inputs = inputs.reshape(bs, seq // superposition_bag_size, superposition_bag_size)     # 对包内嵌入取均值，生成叠加词元     embed = embed.mean(dim=-2)

输出侧放弃传统单Token交叉熵，采用MCE多热交叉熵损失函数，对词元包内所有Token做平均损失计算，让模型学习局部Token分布特征，而非单一词序，大幅提升监督信号密度。

2.2 阶段二：标准恢复精训阶段

当训练步数达到总步数的60%-80%（最优区间：r=0.2~0.4），关闭词元叠加机制，回归原生NTP自回归训练。修复粗训阶段丢失的局部词序信息，补齐模型生成精度，保证最终推理模型与传统训练模型完全一致。

3. 核心优势：区别于传统提效方案

优化方案	改造范围	推理成本	工程复杂度	核心原理
MoE稀疏架构	模型架构、并行逻辑	增加推理复杂度	极高	减少单Token激活参数
MTP多Token预测	输出头、损失逻辑	轻微增加负担	中高	单位置预测多Token
TST词元叠加训练	仅训练循环	零成本	极低	粗粒度学习+精细化修正

4. 实验数据验证

团队覆盖多尺度模型完成验证，核心结论如下：

1. 小模型（270M/600M）：同等算力下，TST训练损失更低，下游分类、理解任务指标全面超越基线；

2. 3B稠密模型：20k步TST训练效果，匹配36k步基线训练精度，算力利用率大幅提升；

3. 10B-A1B MoE模型（最优收益）：算力耗时从12311 B200-h降至4768 B200-h，算力成本降低61.3%，四大零样本评测指标全部正向提升。

5. 消融实验与机制分析

消融实验证明，TST的收益来自输入、输出双机制叠加：

- 仅输入叠加：压缩表征、降低训练噪声，提升数据吞吐；

- 仅输出MCE损失：强化局部语义监督，优化梯度更新；

- 完整TST方案：双机制协同，收益近似叠加，效果最优。

6. 局限性与落地建议

1. 适用边界：算力受限场景收益显著，数据稀缺场景不推荐使用（加速数据消耗）；

2. 超参建议：常规模型包大小s取6-8，百亿级大模型最优s=16，叠加训练比例r固定0.2-0.4即可稳定收益；

3. 落地场景：垂直行业小模型、稠密模型、轻量化MoE模型预训练迭代。

7. 总结

TST是当前性价比最高的大模型预训练优化方案之一，以极低的工程改造成本，实现了算力成本断崖式下降。无需重构架构、无需适配推理、无需海量调参，完美适配中小团队快速迭代、垂直模型批量训练的需求，将成为未来大模型预训练的标配优化策略。