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

预训练模型不再重训:加一张词表记忆,HumanEval涨4.27分,吞吐只掉1.5%

导语

在OLMo-2 1B上,把零初始化的词索引嵌入表作为残差通路接到已预训练的FFN上,继续训练后十任务均分从44.39升到45.59,HumanEval提高4.27分,而生成吞吐仍保持在稠密基线的1.5%以内。

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

正文

预训练稠密模型可以在继续训练阶段通过新增静态稀疏记忆扩容,而不必从零预训练稀疏架构。 此前的稀疏扩容主要面向专家混合模型,克隆出的专家在功能上可能仍然相似,限制了继续扩容的收益。 在OLMo-2 1B上,十任务平均分从继续稠密训练的44.39升到45.59,HumanEval提高4.27分、MBPP提高2.80分、GSM8K提高2.27分、DROP提高1.65分。

新记忆以残差方式叠加在原有前馈网络的上投影之上,而不是替换它,因此预训练计算被保留下来。 STEM把前馈网络的上投影替换为按词元索引的嵌入表,在四层消融中即使加上额外预热也未能追平继续稠密训练。 嵌入表零初始化以恢复种子模型的计算图,预热阶段冻结上投影,让新表先学习互补信号;在OLMo-2 7B上平均分从57.11升到58.01,十个评测任务中九个改善。

接下来

报告覆盖的是OLMo-2 1B与7B骨干上的继续训练设置,1B主对比使用约100B预热词元加100B继续训练词元,7B因算力成本改用约30B预热加70B继续训练。嵌入表放在前部解码层,扩容沿增强层数与嵌入维度两条轴进行,大表留在CPU内存、只把需要的行搬到GPU。作者提出的下一步是考察静态记忆与专家混合联合扩容能否继续推进性能与效率的边界。

移除全部嵌入表后十任务均分从45.59降到32.17,说明模型在继续训练后依赖这些记忆,但完全移除后的模型并不等于原始稠密种子,因为其余权重已与静态记忆共同适应。7B上MMLU提升有限,作者提出这可能与更短的预热有关,但该解释未在7B实验中单独隔离。词元级干预能改变输出,作者同时说明这并未确立通用的编辑可靠性或孤立的事实表示。

来源