Nous TST 用 token-bag 阶段加 recovery 把预训练 wall-clock 砍 2.5×

Nous Research 本周发布了 Token Superposition Training(TST)——一种两阶段预训练方法,在 FLOPs 一致的前提下把 wall-clock 训练时间最多砍掉 2.5×,而且不改模型架构、不改优化器、不改 tokenizer、不改并行策略、不改训练数据。头条结果在 10B-A1B 混合专家模型规模上:TST 的最终训练 loss 比 equal-FLOPs baseline 还低,消耗 4,768 B200-GPU 小时,baseline 是 12,311。这套方法在四个规模上做过验证——270M 和 600M dense(SmolLM2 形状,改用 Llama3 的建模代码)、3B dense(SmolLM3 形状),以及 Qwen3 家族的一个 10B-A1B MoE——小规模 runs 用 DCLM,MoE run 用 50/50 的 DCLM 加 FineWeb-Edu。所有 run 都用 AdamW + Warmup-Stable-Decay LR 调度,运行在 TorchTitan + FSDP 之上,8 张或 64 张 NVIDIA B200 GPU。最终模型在架构上和常规预训练产出的完全一致;推理行为不变。

机制干净地拆成两个阶段。Phase 1(superposition 阶段,占总训练步数的 r ∈ [0.2, 0.4])把长度为 L 的输入序列切成不重叠的、每袋 s 个连续 token 的"袋",然后把每个袋里 s 个 embedding 取平均、合成一个 latent "s-token"。transformer 接下来处理一个长度 L/s 的序列。为了让每个 TST 步与标准训练步 FLOPs 一致,Phase 1 期间把数据侧的序列长度增加 s×——所以模型每单位 compute 吞进 s× 倍文本,这是吞吐量提升的来源。输出侧每个 latent 位置预测下一袋 s 个 token,损失改成 multi-hot cross-entropy,对每个目标分配 1/s 的概率质量——可以直接用现有融合 CE kernel 实现,不需要新 kernel,不需要 auxiliary head。Phase 2(recovery)从已保存的 checkpoint 续训,用标准 next-token prediction 跑完剩下的 1-r 步。过渡时会有一个 1 到 2 nat 的 transient loss spike,几千步内消化掉;之后 recover 出来的模型就跌到 equal-FLOPs baseline 之下并保持。

Nous 论文里诚实的 hedge 是最值得拿出来的部分。团队明确给了三种对比视角:equal-FLOPs(TST 赢)、equal-loss(TST 赢)、equal-data(baseline 赢,因为 TST 的有效 compute 每数据 token 更少)。这是判定 TST 适用边界的条件——compute-bound 的预训练受益,data-bound 的预训练不受益。考虑到最近行业关于数据稀缺的讨论,实际 data-bound 的店比想象的多。消融结果同样承重:把 Phase 2 起点的输入 embedding 和 LM head 随机重新初始化后,最终 loss 跳到 2.938(比 TST 的 2.676 和标准 baseline 的 2.808 都差)。Phase 1 的表征不是用完即弃——两个阶段共享的表征,才是让 TST 起作用的东西。输入侧机制(token 平均)和输出侧机制(下一袋预测)各自就能超过 baseline,合起来也没有干扰,说明这是两个正交机制,而不是一个 trick。10B-A1B MoE 规模上的具体 benchmark:HellaSwag 71.2 vs baseline 70.1,ARC-Easy 74.2 vs 73.8,ARC-Challenge 47.3 vs 46.3,MMLU 39.0 vs 37.4。

对在预训练任何东西(从一个小 SLM 到前沿级 MoE)的 builder:现实问题就变成——你的 workload 是 compute-bound(TST 帮你)还是 data-bound(TST 让你更难受,因为它每 FLOP 多吞数据 token)。Nous 给的参考配置——r 取 0.2 到 0.4,s 取 6(3B 规模)到 16(10B-A1B 规模)——是 ablate 的起点。这套技术属于和 multi-token prediction(MTP)同类的"未来信号辅助目标"家族,但是其中最便宜的成员:一个输出 head,只换 target,沿用现有 CE kernel。和 MTP 不同,TST 在所有测试规模上都有一致收益,包括 MTP 已经被证明会拉低性能的小模型规模。论文在 arXiv 2605.06546,实现应该会通过 Nous 的标准渠道发布(和本周早些时候的 Hermes Agent 一样的发布模式)。对正在做预训练 roadmap 的店来说,这是一项值得在一个月内放进自己 pipeline 里实际 ablate 的、不常见的工程贡献。