摘要:针对大模型预训练算力成本高、迭代速度慢的行业痛点,Nous Research提出TST(Token Superposition Training)词元叠加训练方案。本文从技术原理、实现逻辑、实验数据、消融实验、落地优劣全方位拆解,带你看懂这款零侵入、高收益的预训练提效新方案。
关键词:大模型预训练、TST、Token叠加、预训练提速、算力降本、MCE损失
**1.**前言
大模型预训练是AI研发的核心成本瓶颈。传统Next-Token Prediction(NTP)范式全程采用逐Token精细化训练,在训练早期存在严重的算力冗余。现有主流提效方案(MoE、稀疏注意力、MTP)均存在工程复杂度高、推理适配成本大的问题,难以快速落地。
为此,Nous Research团队提出TST 词元叠加训练法,仅修改训练循环逻辑,不改动模型架构、分词器、优化器与训练数据,实现单位FLOPs下数据吞吐最大化,在10B MoE模型上实现2.5倍预训练提速,算力成本降低60%以上。
2. TST****核心技术原理
TST采用粗训 + 精训两阶段渐进式训练,核心是输入词元压缩+多热交叉熵监督,全程无新增复杂模块。
2.1****阶段一:词元叠加粗训阶段
将原始序列长度为L的文本,按固定包大小s,分割为若干个包含s个连续Token的词元包。对每个词元包内的Token嵌入向量做均值处理,完成输入表征压缩:
核心代码逻辑(PyTorch极简实现):
|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| python # TST输入侧词元叠加核心逻辑 if superposition_bag_size > 1: bs, seq = inputs.shape # 重塑维度,拆分词元包 inputs = inputs.reshape(bs, seq // superposition_bag_size, superposition_bag_size) # 对包内嵌入取均值,生成叠加词元 embed = embed.mean(dim=-2) |
输出侧放弃传统单Token交叉熵,采用MCE 多热交叉熵损失函数,对词元包内所有Token做平均损失计算,让模型学习局部Token分布特征,而非单一词序,大幅提升监督信号密度。
2.2****阶段二:标准恢复精训阶段
当训练步数达到总步数的60%-80%(最优区间:r=0.2~0.4),关闭词元叠加机制,回归原生NTP自回归训练。修复粗训阶段丢失的局部词序信息,补齐模型生成精度,保证最终推理模型与传统训练模型完全一致。
**3.**核心优势:区别于传统提效方案
|-------------|-----------|---------|-------|--------------|
| 优化方案 | 改造范围 | 推理成本 | 工程复杂度 | 核心原理 |
| MoE稀疏架构 | 模型架构、并行逻辑 | 增加推理复杂度 | 极高 | 减少单Token激活参数 |
| MTP多Token预测 | 输出头、损失逻辑 | 轻微增加负担 | 中高 | 单位置预测多Token |
| TST词元叠加训练 | 仅训练循环 | 零成本 | 极低 | 粗粒度学习+精细化修正 |
**4.**实验数据验证
团队覆盖多尺度模型完成验证,核心结论如下:
-
小模型(270M/600M):同等算力下,TST训练损失更低,下游分类、理解任务指标全面超越基线;
-
3B稠密模型:20k步TST训练效果,匹配36k步基线训练精度,算力利用率大幅提升;
-
10B-A1B MoE模型(最优收益):算力耗时从12311 B200-h降至4768 B200-h,算力成本降低 61.3%,四大零样本评测指标全部正向提升。
**5.**消融实验与机制分析
消融实验证明,TST的收益来自输入、输出双机制叠加:
-
仅输入叠加:压缩表征、降低训练噪声,提升数据吞吐;
-
仅输出MCE损失:强化局部语义监督,优化梯度更新;
-
完整TST方案:双机制协同,收益近似叠加,效果最优。
**6.**局限性与落地建议
-
适用边界:算力受限场景收益显著,数据稀缺场景不推荐使用(加速数据消耗);
-
超参建议:常规模型包大小s取6-8,百亿级大模型最优s=16,叠加训练比例r固定0.2-0.4即可稳定收益;
-
落地场景:垂直行业小模型、稠密模型、轻量化MoE模型预训练迭代。
**7.**总结
TST是当前性价比最高的大模型预训练优化方案之一,以极低的工程改造成本,实现了算力成本断崖式下降。无需重构架构、无需适配推理、无需海量调参,完美适配中小团队快速迭代、垂直模型批量训练的需求,将成为未来大模型预训练的标配优化策略。