基于 Token 叠加的高效 LLM 预训练方法 Efficient Pre-Training with Token Superposition
两阶段预训练叠加 s token 提吞吐，10B 上获 2.5x 加速
前置知识
自回归下一词预测 (Next-Token Prediction)
大语言模型预训练的标准范式：给定前 $t$ 个 token $x_1, x_2, \ldots, x_t$，模型参数化为 $\theta$，输出 logits $z_t \in \mathbb{R}^V$（$V$ 为词表大小），与 one-hot 标签 $y_t$ 通过交叉熵 $\mathcal{L}_{CE}(z, y) = -\log \text{softmax}(z)_y$ 计算损失，再通过反向传播更新参数。这是几乎所有现代 LLM 的基础训练目标，决定了本文方法要与此目标兼容。
TST 的两阶段设计中第二阶段需要完全回到标准的下一词预测，理解它是理解 TST 恢复阶段的前提
Token Embedding 平均 (Embedding Averaging)
在词嵌入层，每个 token $x_i$ 映射为 $d$ 维向量 $e_i \in \mathbb{R}^d$。对连续 $s$ 个 token 的嵌入取算术平均 $e'_j = \frac{1}{s}\sum_{i=js+1}^{(j+1)s} e_i$，得到一个 's-token' 表示。这种平均在数学上等价于把多个 token 折叠成一个向量再送入后续 Transformer 块，能把序列长度从 $L$ 压缩到 $\lfloor L/s \rfloor$，从而显著降低每步的 FLOPs。
这是 TST 输入侧的核心操作，将 s 个连续 token 的嵌入聚合成一个潜在表示，是输入 '叠加' 的实现细节
多热交叉熵 (Multi-Hot Cross-Entropy, MCE)
标准交叉熵只能处理单个正确类别的 one-hot 标签；多热 CE 把标签从单个索引 $y$ 扩展为一组索引 $y = \{y_1, y_2, \ldots, y_s\}$，每个标签赋予相同权重 $1/s$。推导后得到 $\mathcal{L}_{MCE}(z, y) = \frac{1}{|y|}\sum_{y \in y} \mathcal{L}_{CE}(z, y) - \log|y|$，由于 $\log|y|$ 是常数梯度为零，可简化为对 bag 内每个 token 独立计算 CE 再取平均。
这是 TST 输出侧的关键损失函数，使得用一个输出头就能预测下一组 s 个 token，是 '输出叠加' 的数学基础
多 token 预测 (Multi-Token Prediction, MTP)
Gloeckle et al. (2024) 提出的辅助损失方法：用 $k$ 个独立的输出头同时预测接下来的 $k$ 个 token，以增加每步的监督信号密度。DeepSeek-V3 进一步引入级联式的 MTP 模块。MTP 不改变每步的 token 吞吐量（即每步仍处理同样的 token 数），但增加了参数量和计算开销。
本文明确论证 TST 与 MTP 占据不同的设计点：MTP 是 '辅助损失'，TST 是 '吞吐量提升'，两者正交且可叠加
FSDP (Fully Sharded Data Parallel)
PyTorch 推出的数据并行策略，将模型参数、梯度和优化器状态分片到多个 GPU 节点上，每台机器只持有部分分片以降低显存占用。FSDP 在本文中是默认并行方案，作者明确说明 TST 不需要修改 FSDP 设置，因为它的每步 FLOPs 与基线相等。
理解 FSDP 有助于把握为何 TST 能作为 'drop-in' 方法在现有大规模训练栈中无缝应用
MoE 混合专家 (Mixture-of-Experts)
通过路由器将每个 token 路由到部分专家子网络（如 8 选 2），从而在总参数量大幅增加的同时只激活一小部分参数计算。10B-A1B 表示总参 100 亿、激活 10 亿。本文在 10B-A1B MoE 上验证 TST，是论文最大规模的实验。
TST 在 MoE 上同样能取得 2.5x 加速，验证了该方法与稀疏激活架构的兼容性
Warmup-Stable-Decay 学习率调度器
一种三阶段学习率调度：先经过 warmup 阶段线性增长到峰值 LR（如 $2 \times 10^{-4}$），然后保持稳定较长时间，最后在最后 10% 步数中线性或余弦衰减到接近零。WSD 是近期 LLM 预训练（如 MiniCPM、LLaMA3）常用的调度器。
本文所有实验都使用 WSD 调度，warmup 2000 步 + 最后 10% 衰减，是结果可比的统一基线
研究动机
现代 LLM 预训练已经演化为以海量数据为驱动力的范式——不仅模型规模爆炸，token 数据量也常被推到 compute-optimal 估计之上以最大化推理期表现。OpenAI、Anthropic、Meta 等头部实验室的单次预训练已经消耗数千亿乃至数万亿 token。这种 '数据饥渴' 范式下，如何在固定算力下提高每 FLOP 消耗的数据量成了核心瓶颈。现有的训练效率方法大致分三类：(1) 信息最大化，如 SuperBPE 把 BPE 合并为更大的 '超 token'、多 token 预测 (MTP) 增加每步监督信号；(2) 计算稀疏化，如 MoE、稀疏注意力降低每个 token 的 FLOPs；(3) 表征压缩，如 Bolmo、H-Net、Byte-Latent Transformer 在模型内部用更少 token 承载同样信息。但这些方法各自有局限：信息最大化中的 MTP 不增加每步数据吞吐，只是加了辅助损失；SuperBPE 改变了 tokenizer 本身，与下游推理栈不兼容；MoE 修改了推理期架构；压缩建模要么改变模型结构要么需要复杂的对齐适配阶段。
本文的目标是本文的具体目标是在不修改模型架构、不改变 tokenizer、不动 FSDP 并行策略、不动优化器的前提下，仅通过调整训练时的输入和输出表征来成倍提升每步的 token 吞吐，从而显著降低达到目标损失所需的训练时间。作者希望方法能作为 'drop-in' 集成到现有训练栈，并且训练完成后模型可以直接以标准下一词预测方式推理，不需要任何额外修改。
与已有工作不同的是，本文的独特切入角度是把 '压缩式预训练' 重新定位为 'two orthogonal but additive processes'——输入侧通过 embedding 平均降低每 token 的 FLOPs（输入粒度变粗），输出侧通过多热损失让单次预测覆盖更多 token（监督粒度变粗），两个过程独立且可叠加。作者在文中明确指出，与同期工作 Shao et al. (Patch-Level Training) 在算法上几乎相同，但视角不同：Shao 等把方法定位为 '减少总 FLOPs'，而本文定位为 '恒定 FLOPs 下提高 token 吞吐'。同时本文做了更细致的实验设计：每个 TST 步的 FLOPs 与基线严格相等（通过把序列长度 L 扩大 s 倍来补偿），简化了规模化推导；并在最大 10B-A1B MoE 规模上做了 scaling validation（消耗 4768 B200-Hours vs 基线 12311 B200-Hours）。
核心方法
Token-Superposition Training (TST) 是一个两阶段预训练方法。第一阶段（叠加阶段）：把 token 化后的输入序列切成 $s$ 个 token 一组的 'bag'，在 embedding 层对每个 bag 内的 $s$ 个 token 嵌入取算术平均，得到 $l = \lfloor L/s \rfloor$ 个 's-token'；模型照常在这些 s-token 上做 Transformer 计算，但输出头改为预测下一组 $s$ 个 token，使用多热交叉熵损失 $\mathcal{L}_{MCE}$。直觉上等价于把序列压缩 $s$ 倍但把目标扩展 $s$ 倍，从而让单步覆盖 $s$ 倍的原始数据 token。第二阶段（恢复阶段）：完全去除叠加代码，回归标准下一词预测。叠加阶段的步数占比 $r$（如 $r=0.3$）需要调优，太大模型可能学不会恢复，太小则收益不足。整个方法的关键设计是叠加阶段和恢复阶段共用同一个输入 embedding 和输出 LM head，作者通过随机重新初始化实验证明这一 'representation alignment' 是 TST 成功的必要条件。
TST 的核心创新是把 'bag-of-token-embeddings' 输入和 'next bag-of-tokens' 输出作为两个正交机制同时应用，并通过共享 embedding/LM head 跨阶段保持表征一致。技术上和同期 Shao et al. 的 patch-level training 几乎相同（都做 embedding 平均 + bag CE + 切换到标准训练），但视角和定位不同。本文把 TST 定位为 '恒定 FLOPs 下的 token 吞吐提升' 而非 '减少总 FLOPs'，强调它属于训练时效率而非推理时效率，并明确论证与 MTP 等 '辅助损失' 方法占据不同的设计点。与现有方法对比：(1) vs MTP：MTP 不增加每步数据吞吐（仍处理相同 token 数），TST 严格增加；MTP 改架构加参数，TST 不改；(2) vs SuperBPE：SuperBPE 改 tokenizer 永久影响推理，TST 训练期临时改，推理期标准；(3) vs H-Net/BLT：这些方法永久修改模型结构（用动态 chunking），TST 训练后用回标准 next-token；(4) vs Shao et al. 的 patch-level training：算法等价但本文做了更系统的超参数扫描、更大 scale 验证，并发现了 representation alignment 这一关键机制。
方法步骤详情
完整 TST 流程分为以下步骤。Step 1 数据准备：使用 DCLM 数据集（或 10B MoE 实验中 50/50 混合 FineWeb-Edu 和 DCLM），用 Llama3-8B tokenizer 做 token 化。Step 2 叠加阶段 (Phase i)：把形状为 $B \times L \times V$ 的输入张量按 $s$ 个 token 为一组切成 bag，得到 $B \times l \times s \times V$，其中 $l = L/s$。在 embedding 层，对每个 bag 内的 $s$ 个 token 嵌入做算术平均 $e'_j = \frac{1}{s}\sum_{i=js+1}^{(j+1)s} e_i$，得到 $B \times l \times d$ 的 s-token 张量送入 Transformer 块。为保证等 FLOPs 对比，把序列长度 L 乘以 s（即用 $L \cdot s$ 的数据 token 算 L/s 个 s-token）。Step 3 损失计算：标准标签左移 $s-1$ 位后切成 bag，对每个 s-token 位置上的输出 logits $z \in \mathbb{R}^V$ 与 bag $y = \{y_1, \ldots, y_s\}$ 计算多热 CE：$\mathcal{L}_{MCE}(z, y) = \frac{1}{|y|}\sum_{y \in y} \mathcal{L}_{CE}(z, y)$。代码层面直接复用高度优化的 CE kernel（求和在外层 for 循环中，开销可忽略），不修改主训练循环结构。Step 4 恢复阶段 (Phase ii)：从叠加阶段保存的 checkpoint 恢复，加载完整叠加代码并彻底移除，恢复标准 token 级别的下一词预测目标，标签 shift 1，损失用标准 one-hot CE。Step 5 评估：在最终 checkpoint 上用 Eleuther AI LM-Eval harness 做 0-shot 评估（ARC-E、ARC-C、HellaSwag、MMLU、BoolQ、OpenBookQA、PIQA、Winogrande）。Step 6 超参设置：袋大小 $s$ 和步数比 $r$ 通过 sweep 确定最优区间 $s \in [4, 8]$, $r \in [0.2, 0.4]$。
技术新颖性
TST 的技术新颖性主要体现在四点。第一，'正交双机制' 的设计视角：明确把输入叠加和输出叠加解耦为两个独立的 throughput-boosting 机制，消融实验（Figure 6）证明 input-only 和 output-only 各自都能超过基线但都弱于 full superposition，组合起来有叠加效果，这是对 '为什么 TST 有效' 的最直接机理证据。第二，'representation alignment' 假设的实证检验：Table 2 显示在 3B 模型上，如果在恢复阶段开始时随机重新初始化 input embedding 和 LM head，最终损失从 2.676 退化为 2.938（甚至差于基线的 2.808），说明跨阶段的表征对齐是 TST 成功的必要条件，这与现有压缩预训练方法通常需要 adapter 对齐阶段的直觉形成对比——TST 用 '不做任何对齐' 的方式反而成功了。第三，'drop-in' 设计：不修改并行策略、优化器、tokenizer、模型架构，叠加阶段和恢复阶段共享 embedding 和 LM head，是 '不改任何东西就能提速' 的极简方案。第四，超大规模验证：把方法推到 10B-A1B MoE 规模，2T 数据 token 训练最终损失 2.236、B200-Hours 消耗 4768，相比基线 1.05T tokens、12311 B200-Hours 实现 2.5x 加速，把 'bag-of-tokens' 思路从 270M 验证到 10B 工业级规模。
实验结果
实验结果在 270M、600M、3B 密集模型和 10B-A1B MoE 上一致表明 TST 优于基线训练，主要发现可归纳为五点。第一，270M 模型上 TST (s=6, r=0.3) 在 20k 步设置下损失从基线 3.212 降到 3.142（相对改善 2.2%），HellaSwag 从 36.3 提升到 38.6，ARC-E 从 46.7 提升到 47.6，ARC-C 从 24.9 提升到 26.4；100k 步设置下损失从 3.092 降到 3.048。第二，600M 模型上 TST 在 20k 步时损失从 3.019 降到 2.943（改善 2.5%），HellaSwag 从 43.5 提升到 48.2（提升 4.7 个百分点，绝对值），ARC-C 从 25.5 提升到 26.9。第三，3B 模型等 FLOPs 设置下 TST (s=6, r=0.3) 损失从 2.808 降到 2.676（改善 4.7%），且与消耗 36000 步的基线 (loss=2.677) 持平，意味着用更少步数（20000）就达到了几乎相同效果；HellaSwag 从 57.6 提升到 62.4，ARC-C 从 31.9 提升到 36.0，MMLU 从 31.2 提升到 32.8。第四，10B-A1B MoE 上 TST (s=16, r=0.25) 在相同 B200-Hours (4768 vs 12311) 下实现了 2.5x 加速；图 1 显示停止训练匹配相同 loss 时，TST 用 49983 步看 2T tokens 达到基线 1.05T tokens 训练 125000 步的 loss 水平（2.236 vs 2.252），HellaSwag 从 70.1 提升到 71.2，MMLU 从 37.4 提升到 39.0。第五，消融实验（Figure 6）证实 input superposition 和 output superposition 都是贡献源——3B 模型上 input-only 损失 3.200、output-only 损失 3.100、full 损失 3.000，而基线是 3.275；Table 2 的随机化实验显示 3B 模型上 TST (2.676) > TST w/ randomization (2.938) > baseline (2.808)，'representation alignment' 是 TST 成功的关键。超参扫描（Figure 4）显示在 $s \in [4, 8]$、$r \in [0.2, 0.4]$ 范围内都稳定优于基线，U 型曲线说明 s 过大和 r 过大都会损害性能。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 270M 密集模型 (20k 步, 42B tokens) | Final Loss (CE) | 3.142 | 3.212 | -2.2% |
| 270M 密集模型 (20k 步, 42B tokens) | HellaSwag 0-shot | 38.6 | 36.3 | +2.3 pp |
| 600M 密集模型 (20k 步, 42B tokens) | Final Loss (CE) | 2.943 | 3.019 | -2.5% |
| 600M 密集模型 (20k 步, 42B tokens) | HellaSwag 0-shot | 48.2 | 43.5 | +4.7 pp |
| 3B 密集模型 (20k 步等 FLOPs) | Final Loss (CE) | 2.676 | 2.808 | -4.7% |
| 3B 密集模型 (20k 步等 FLOPs) | MMLU 0-shot | 32.8 | 31.2 | +1.6 pp |
| 10B-A1B MoE (等 Loss, 125k 步基线) | B200-Hours | 4768 | 12311 | 2.5x 加速 |
| 10B-A1B MoE (等 Loss) | MMLU 0-shot | 39.0 | 37.4 | +1.6 pp |
| 3B 密集模型 (representation alignment 消融) | Final Loss (CE) | TST: 2.676 / TST w/ rand: 2.938 | 2.808 | TST 改善 4.7%，随机化后恶化 4.6% |
局限与改进
作者明确承认的局限有四点：(1) TST 隐含 '算力受限而非数据受限' 的假设，作者引用 Kim et al. (2025) 指出这一假设在数据增长快于算力增长的趋势下未来可能不成立；output-only 叠加（在不增加数据消耗下提升）是该问题的一种部分解。(2) 由于 s-tokens 折叠原始序列为 1/s，有效上下文长度变为 s 倍，可能对长上下文有正面影响但本文未评估。(3) 受限于 4768 B200-Hours 的总计算预算，没有做多次独立运行的统计显著性检验，也没有做更大规模的 scaling law 推导。(4) 机理层面，本文提出 'pre-pre-training' 和 'embedding 几何正则化' 两个假设，但都未通过可解释性实验验证。本人观察到的额外局限包括：(a) 训练中 s 和 r 都需要在新的训练 setup 上重新调优，没有提供明确的 scaling law 来预测更大规模的最优超参；(b) 文中 Figure 8 提到 power-law 权重方案在 s≥8 时更稳定，但未在 3B/10B 规模上验证；(c) TST 第二阶段需要 70-80% 的标准训练来恢复，这意味着实际的 'TTI (time to inference)' 不一定能从训练时间节约中按比例受益——前向推理速度和标准模型一样；(d) 没有评估对 reasoning、code、multilingual 等非通用 NLP 任务的影响；(e) 评测任务主要是 0-shot 标准基准，没有看 in-context learning 能力是否受到 'superposition 阶段 + 恢复阶段' 这种非平滑训练的影响。
独立分析的弱点
从独立分析视角，本文主要存在以下可改进的弱点。第一，'等 FLOPs' 对比虽然公平但忽略了 '等步数' 场景的实用性。在实际工业界，batch size 经常是显存受限而非算力受限——此时更长的 s-token 序列可能让 OOM 反而更糟。文中提到可以用 '把 micro batch 扩大 s 倍' 替代 '把序列长度扩大 s 倍'，但没做实验验证。第二，r 的选择与总训练 budget 强相关（r=0.3 是在 20k 步里的占比），但 Figure 4 的 x 轴只扫了 (s, r) 组合，没看 '在固定总 tokens 下 r 如何随总 budget 变化' 的趋势，这使得方法的可推广性受限。第三，'representation alignment' 是 TST 成功的关键机制，但作者没有深入分析 '为什么共享 embedding 就足够'。一个可能的解释是 bag 内 token 的平均嵌入恰好落在 embedding 流形的稠密区域，但本文没有用任何 embedding 空间分析（如 PCA、t-SNE、representational similarity）来支撑这一假设。第四，TST 训练完成后模型的标准推理速度没有变，这意味着 2.5x 加速只换来了 '更短训练时间'，但 '推理 cost 一样'；对以 serving cost 为主的应用价值有限。第五，文中没有实验 's 和 r 是否可以动态调度'（如前 30% 步用 s=8、中 30% 步用 s=4、最后 40% 步恢复），这种 schedule 可能比固定 (s, r) 更优。改进方向包括：推导 (s, r) 的 scaling law；在多语言/代码等非英文任务上验证；用 representation geometry 分析 alignment 假设；尝试 curriculum 调度。
未来方向
作者明确提出的未来工作有四点：(1) 与多 token 预测 (MTP) 等辅助损失方法组合，因为 TST 与 MTP 在设计空间正交 (TST 提 throughput，MTP 加监督信号)；(2) long-context 评估，预期 TST 阶段产生的 s 倍有效上下文可能对长上下文任务有正面影响；(3) TST 的 scaling law 推导，预测在更大规模工业级训练中的最优 (s, r) 设置；(4) 解释性研究，理解 '为什么输入叠加有效'（pre-pre-training vs embedding 几何正则化两个假设）。本人基于成果可延伸的方向包括：(a) 把 TST 思路推广到多模态训练（视觉 patch 本身就是叠加的天然候选）；(b) 研究 's-tokens 是否可以作为一种隐式的数据增强'，因为 bag 平均相当于多种 token 排列的 centroid 估计；(c) 探索 'soft superposition'，即用加权平均替代均匀平均，可能带来更细粒度的控制；(d) 拓展到 RLHF / 后训练阶段，看 TST 风格的中间表示是否能加速偏好学习；(e) 推导 TST 在 '数据受限' 场景下的变体（如去掉输入叠加只保留输出叠加）。
复现评估
复现性方面，本文提供了相对完善的实现细节。(1) 代码：作者在 Appendix A 给出了 input/output superposition 的核心代码片段（PyTorch 风格），MCE 损失计算有显式公式，但完整训练代码未在论文中开源（截至 2026-05）；(2) 框架：使用 TorchTitan + FSDP，标准工业栈；用 Llama3-8B tokenizer（公开），模型形状对照 SmolLM2 / Qwen3（公开配置）；(3) 数据：使用 DCLM（公开下载，50/50 mix with FineWeb-Edu for 10B run）；(4) 训练配置：batch size 2M tokens（小模型）/ 8M tokens（10B），WSD 调度器 (2000 步 warmup + 末 10% 衰减)，AdamW ($\beta_1=0.9, \beta_2=0.95$)，学习率 sweep 在 Figure 7 (Appendix B) 中给出；(5) 评测：用 Eleuther AI LM-Eval harness 做 0-shot 评估，标准任务；(6) 算力需求：270M/600M 8 B200 GPU、3B 64 B200 GPU、10B-A1B 64 B200 GPU + 4768 B200-Hours（最贵 run），对大多数学术实验室成本依然过高；(7) 难度：中等——核心思想简单（两阶段 + embedding 平均 + 多热 CE），但要在 10B scale 复现需 5k+ GPU-hours，难度高。复现建议：小规模 (270M) 完全可复现 (单节点 8 GPU 几天)，3B 也可访问 (64 GPU 一天)，10B 几乎只有头部实验室可行。
论文图表
