跳到主要内容
返回时间线
arXiv来源发表:

预训练模型不再重训稀疏架构:加一条零初始化记忆通路,十项基准平均涨1.20分

核心概要

预训练稠密模型可以在继续训练阶段扩容,而不必从零训练稀疏架构。STAPLE 保留原有的前馈网络上投影,只把零初始化的词元索引嵌入表作为残差通路接入,让新参数学习互补的词元信息。在 OLMo-2 1B 上,它相对继续稠密训练把十项基准平均分从 44.39 提到 45.59,HumanEval 涨 4.27 分、MBPP 涨 2.80 分、GSM8K 涨 2.27 分、DROP 涨 1.65 分;增益在 OLMo-2 7B 上仍然存在。

Source-provided article image: STAPLE: Static Sparse Upcycling of Pretrained Checkpoints via Embedding Modules
Figure 3 · arXiv

深度剖析

预训练稠密模型扩容时,可以保留原有计算通路,只把新增的静态稀疏记忆当作补充,而不是替换掉已有计算。 此前的静态稀疏方案以 STEM 为代表,用词元索引嵌入表替换前馈网络的上投影;在四层消融里,这种替换式做法即使加上额外预热,也追不上继续稠密训练。 在相同训练预算下,保留上投影并把嵌入作为残差通路接入,下游表现明显更好;在 OLMo-2 1B 上,十项基准平均分从继续稠密训练的 44.39 升到 45.59。

新增的嵌入表用零权重初始化,先恢复种子模型原本的计算图,再通过一段短预热学习通用记忆。 替换式做法缺少能保住原有功能的初始化方式,随机、词元均值、主成分以及变换后的输入嵌入初始化都没有改善长期表现。 预热阶段在预训练语料的一个子集上训练,并冻结前馈网络上投影;预热在 MMLU 和常识推理平均分上持续优于直接扩容,但在高等数学与编程任务上没有一致增益。

静态记忆的容量可以沿两条轴扩大,即增加接入嵌入表的层数和增大嵌入维度,而新增容量对吞吐影响很小。 接入层数从 4 增到 12,平均下游分数从 45.10 升到 45.59;低秩变体在 1B 上把词元索引表缩小到原来的四分之一,平均分为 44.97,仍高于继续稠密训练的 44.39。 在 OLMo-2 7B 上,接入前 16 层后平均分从 57.11 升到 58.01,十项任务中九项改善,数学平均提升 4.16%,代码平均提升 3.68%。

大表放在主机内存、只把当前需要的行搬到 GPU,可以让新增的数十亿参数不拖慢训练和生成。 稠密基线把全部权重常驻 GPU;MoE 扩容基线训练慢 13%,权重与优化器状态占用约 2.9 倍 GPU 内存。 STAPLE-12L 有 91.9 GiB 的表与优化器矩留在主机 DRAM,训练吞吐为 113.9K 对稠密的 117.0K 词元每秒,生成吞吐在批大小 64 时与稠密相差 1.5% 以内,GPU 每词元能耗比所评估的 MoE 实现低 58.3%。

启示与展望

报告的范围是 OLMo-2 1B 与 7B 两个稠密骨干上的继续训练实验,训练数据为 Dolma 3 Dolmino Mix 高质量混合数据,预热取自 DCLM Baseline 1.0 的约 100B 词元。1B 主对比使用约 100B 预热词元加 100B 继续训练词元,7B 因算力成本改用约 30B 预热加 70B 继续训练。系统效率结论来自面向吞吐的解码评测,以及一次在单张空闲 B200 上、批大小 64、上下文长度 256 的能耗测量。

7B 上 MMLU 增益较小,一种解释是预热更短,与 1B 的预热消融一致,但该解释在 7B 实验中未被单独隔离。移除全部嵌入表后的模型并不等价于原始稠密种子,因为其余权重已与静态记忆共同适应。词元行替换实验显示静态嵌入能因果地改变输出,但一致上下文会逐步压过它,这一现象与上下文相关门控控制记忆贡献的说法相符。附录中的数值追踪把词元可寻址的记忆干预与连贯的下游推理联系起来,但未确立通用的编辑可靠性,也未分离出孤立的事实表示。