Token Superposition Training (TST): 高效 LLM 预训练

论文: Peng, Gigant & Quesnelle (Nous Research, 2026) — arXiv:2605.06546

核心问题

LLM 预训练在大规模下计算成本极高，现有优化方法（MoE、稀疏注意力、压缩建模）通常需要侵入式修改模型架构。能否在不改动模型架构的前提下，仅通过提高训练时 token 吞吐量来提升预训练效率？

方法：Token Superposition Training (TST)

TST 是一个简单的 drop-in 方法，分两阶段：

阶段一：叠加阶段（Superposition Phase）

输入叠加：将连续 s 个 token 的 embedding 取平均，形成单个 "s-token"
输出叠加：预测下一个 bag 的全部 s 个 token，使用 multi-hot-cross-entropy 损失
效果：等 FLOPs 下吞入 s× 更多数据 token

阶段二：恢复阶段（Recovery Phase）

完全回归标准 next-token prediction 训练
不做任何 adapter 或投影层——embedding 和 LM head 保持不变

关键发现

2.5× 加速：在 10B A1B MoE 模型上，等 loss 条件下预训练时间减少 2.5 倍
表示对齐至关重要：叠加和恢复阶段共享 embedding 和 LM head——若在两阶段之间重新初始化，所有增益消失
超参数鲁棒：bag size s ∈ [4, 8]，叠加比例 r ∈ [0.2, 0.4] 内均有效
输入+输出叠加均有贡献，但输入叠加的增益机制仍有待解释

核心洞察

TST 的本质是 粗→细粒度调度（coarse-to-fine-granularity）：先用低分辨率、高吞吐量的数据分布进行"预-预训练"，再切换到标准分辨率。这与 ViT 中的 patch size scheduling 和 byte-level → subword 转移共享同一设计哲学。

概念网络

token-superposition-training — 方法总览
multi-hot-cross-entropy — 核心损失函数
input-superposition — 输入侧的 token 叠加
two-phase-pretraining — 两阶段训练范式
representation-alignment — 跨阶段表示对齐
coarse-to-fine-granularity — 底层设计原则
throughput-hypothesis — 吞吐量假说

2.5 KiB Raw Blame History Unescape Escape