IDEA Prune:将放大与剪枝整合进生成语言模型预训练

大语言模型的推理成本越来越高,如何在有限预算内获得可部署的小模型成为关键。结构化剪枝已被证明比从头训练同尺寸模型更省 token,但以往工作往往忽略了一个环节:是否值得先预训练一个更大的模型,即使它最终永不上线。本文把“放大-剪枝”当作一个整体系统来研究,并给出了一套集成方案。

核心做法是将放大模型训练、剪枝和恢复三个阶段统一在同一个余弦退火学习率调度下,不再各自独立。这避免了朴素放大-剪枝流程中因学习率回升导致的知识丢失,同时配合一种新的迭代式结构化剪枝方法,逐步移除参数。这样做的好处是能把模型容量重新分配到幸存神经元上,让压缩过程更平滑,最终剪枝后的模型性能更佳。

实验部分,作者在预训练阶段将 2.8B 模型压缩到 1.3B,累计训练数据量达到 2T tokens。结果表明,这种一体化方法不仅在 token 效率上优于从头训练,而且剪枝后的模型表现也更好。文章还借此回答了开头提出的两个问题:放大预训练即使不部署也有价值,以及如何优化整个管线来获得更好的剪枝结果。

论文作者来自佐治亚理工学院、德克萨斯大学奥斯汀分校,工作完成于 Apple。目前仅以论文摘要形式公开,尚未披露具体性能数值或详细实验对比。对于关注模型压缩和高效预训练的读者来说,这是一个值得跟踪的方向。

IDEA Prune: An Integrated Enlarge-and-Prune Pipeline in Generative Language Model Pretraining

查看原文