Efficient Pre-Training with Token Superposition

📄 arXiv: 2605.06546v1 📥 PDF

作者: Bowen Peng, Théo Gigant, Jeffrey Quesnelle

分类: cs.CL

发布日期: 2026-05-07

备注: 25 pages, 11 figures, 28 tables

💡 一句话要点

提出Token叠加训练（TST）方法，通过两阶段训练显著提升大模型预训练效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 预训练效率 Token叠加 多热交叉熵 计算优化 混合专家模型

📋 核心要点

  1. 大模型预训练面临计算成本高昂与数据吞吐量瓶颈，现有优化方案往往需要对模型架构或训练基础设施进行复杂且侵入性的修改。
  2. 提出Token叠加训练（TST），通过将连续Token合并为“包”并采用多热交叉熵损失，在不改变模型架构的前提下提升训练效率。
  3. 实验表明TST在多种模型规模下表现稳健，在10B参数规模下，同等损失水平下预训练时间缩短至原来的40%，即提升了2.5倍效率。

📝 摘要（中文）

大语言模型的预训练过程通常成本高昂且效率低下，往往需要复杂且侵入性的修改才能实现高数据吞吐量。本文提出了Token叠加训练（Token-Superposition Training, TST），这是一种简单的即插即用方法，无需修改并行策略、优化器、分词器、数据或模型架构，即可显著提高预训练期间的单位FLOPs数据吞吐量。TST分为两个阶段：(i) 高效的叠加阶段，将多个连续Token合并为一个“包”，并使用多热交叉熵（MCE）目标进行训练；(ii) 恢复阶段，回归标准训练模式。我们在270M、600M参数规模上进行了广泛评估，并在3B和10B A1B混合专家模型上进行了验证，证明了其在不同设置下的鲁棒性。最终，TST在损失函数和下游任务评估中均优于基线，在等损失设置下，10B A1B规模下的预训练总时间缩短了高达2.5倍。

🔬 方法详解

问题定义：大语言模型预训练的计算开销随规模增长呈指数级上升，现有的并行化或架构优化方法通常需要对底层训练框架进行深度定制，增加了工程实现的复杂度和维护成本。

核心思路：TST的核心思想是利用“叠加”机制增加单次前向传播的信息密度。通过将序列中的多个Token压缩进同一个输入单元，模型在单次计算中处理更多信息，从而在保持模型结构不变的情况下提升FLOPs利用率。

技术框架：TST包含两个阶段：第一阶段为叠加训练阶段，将连续的Token序列打包，利用多热交叉熵（MCE）损失函数引导模型学习叠加后的特征表示；第二阶段为恢复阶段，通过标准训练流程对模型进行微调，以消除叠加带来的噪声并恢复模型对序列顺序的精确感知能力。

关键创新：TST的本质创新在于引入了“叠加”作为一种训练策略，而非模型架构的改变。它通过多热目标函数解决了叠加带来的多义性问题，实现了训练效率与模型性能的解耦。

关键设计：关键技术细节包括多热交叉熵（MCE）损失函数的设计，它允许模型在输出层同时预测多个叠加的Token；此外，叠加比例的选择和两阶段训练的切换时机是影响最终收敛效果的关键超参数。

📊 实验亮点

实验结果显示，TST在270M至10B参数规模下均表现出极高的鲁棒性。在10B A1B混合专家模型（MoE）的基准测试中，TST在达到相同损失值的情况下，将总预训练时间缩短了2.5倍。此外，在下游任务评估中，TST不仅保持了与基线相当的性能，在部分任务上甚至表现出更优的泛化能力。

🎯 应用场景

该方法适用于大规模语言模型的预训练阶段，特别是在计算资源受限或追求极致训练效率的场景下。由于其无需修改模型架构和训练基础设施，TST可直接集成到现有的主流训练框架（如Megatron-LM或DeepSpeed）中，显著降低企业和研究机构开发大模型的算力成本，加速模型迭代周期。

📄 摘要（原文）

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.