STAPLE：通过嵌入模块对预训练检查点进行静态稀疏升级

arXiv:2610.12599v1 公告类型：新论文 摘要：稀疏计算可以扩展大语言模型的参数容量，而无需按比例增加实际计算量，但稀疏架构的预训练仍然成本高昂，且需要大量数据。现有的稀疏升级方法主要针对混合专家模型，其中专家同质性、样本效率低下和系统开销可能限制实际可扩展性。静态稀疏记忆具有可预测的访问模式和词元特定的更新机制，有助于应对这些挑战，但将其直接整合到预训练计算中收益不佳。在本研究中，我们提出 STAPLE，这是一种基于 STEM、将静态稀疏记忆整合到预训练模型中的方案。STEM 用按词元索引的嵌入替换 FFN 的升维投影，而 STAPLE 保留这些投影，并将嵌入作为残差路径加入，从而在保留预训练计算的同时，以可扩展的方式学习互补的词元特定信息。在 OLMo-2 1B 上，相比继续进行稠密训练，STAPLE 在 HumanEval、MBPP、GSM8K 和 DROP 上分别提升了 4.27、2.80、2.27 和 1.65 分。平均性能随静态容量的增加而提升，且在更大的 OLMo-2 7B 骨干模型上仍有改善。STAPLE 还优于 STEM 和 Engram 风格模型等静态稀疏基线；在我们面向吞吐量的解码评估中，其平均得分高于 MoE 升级基线，解码速度达到该基线的 2 倍，而 GPU 模型显存占用和实测每词元 GPU 能耗均不到该基线的一半。重要的是，尽管在训练后期引入了数十亿参数，我们采用主机卸载的实现仅需在 GPU 上缓存使用频率最高的约 5% 的嵌入，就能使生成吞吐量与稠密基线的差距保持在 1.5% 以内，训练吞吐量的差距保持在 3% 以内。这些结果表明，静态稀疏记忆可以为扩展预训练模型提供一条样本高效、可扩展且可持续的路径，而无需从头进行稀疏预训练。