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 的、不常見的工程貢獻。