Nous Research 提出 Token Superposition Training (TST)，通过在预训练前半程将连续 token 的 embedding 平均并预测下一组 token，在不改模型架构和推理方式的前提下，实现 10B MoE 模型同等 loss 下最高 2.5 倍的预训练提速。
📝 详细摘要
本文详细介绍了 Nous Research 提出的 Token Superposition Training (TST) 方法。TST 的核心思想是在预训练的前半程（叠加阶段），将连续 s 个 token 的 embedding 取平均作为输入，并让模型预测下一组 s 个 token，从而在相同计算量下处理更多原始 token。后半程（恢复阶段）再切回标准逐 token 预测。该方法不改模型架构、并行策略、优化器或分词器，训练结束后模型仍是标准自回归结构。在 10B-A1B MoE 实验中，TST 用不到 40% 的训练时间达到基线模型的最终损失水平，对应约 2.5 倍的预训练提速。文章还详细解释了输入端的张量折叠、输出端的多热交叉熵损失 (MCE)、超参数影响（step ratio 和 bag size 的 U 型趋势）、以及消融实验证明输入与输出叠加机制的独立性。最后指出，跨阶段保持输入输出表示对齐是 TST 收益延续的关键。
💡 主要观点
- TST 通过叠加 token 表示和预测目标，在不改架构下提升预训练吞吐。 在叠加阶段，将连续 s 个 token 的 embedding 平均后输入模型，并预测下一组 s 个 token，使模型在相同 FLOPs 下处理更多原始 token。
💬 文章金句
- TST 不改模型架构、并行策略、优化器、分词器和训练数据，只在预训练前半程把连续 token 的 embedding 临时平均成一个新表示，并让模型预测下一组 token，后半程再切回标准逐 token 预测训练。
- 在 10B-A1B MoE 实验中，TST 用不到 40% 的训练时间达到基线模型的最终损失水平，对应该设置下同等损失约 2.5 倍的预训练提速。
- TST 本质上是用更多数据消耗，换同等计算下更低的训练 loss。
- 跨阶段保持输入输出表示对齐，很可能是 TST 收益能够延续到标准训练阶段的重要条件。
📊 文章信息
AI 初评：86
来源：PaperWeekly
作者：PaperWeekly
分类：人工智能
语言：中文
阅读时间：12 分钟
字数：2918
标签： Token Superposition Training, 预训练加速, Nous Research, MoE, 多热交叉熵