IDEA Prune：面向生成式语言模型预训练的集成式放大-剪枝流水线

作者Yixiao Li†**、Xianzhi Du、Ajay Jaiswal‡**、Tao Lei、Tuo Zhao†、Chong Wang、Jianyu Wang

查看论文

近年来大语言模型的快速发展，使得在有限推理预算下构建高效、可部署模型的需求愈发迫切。结构化剪枝流水线在 token 效率方面相较于从头训练目标尺寸的模型展现出了明显优势。然而现有工作往往忽视了"先放大再剪枝"中的放大预训练阶段。本文主张将放大模型预训练纳入剪枝流程，并将先放大再剪枝视为一个整体系统来研究两个关键问题：即使放大后的模型最终不会被部署，预先训练它是否依然值得？以及如何优化整条流水线以获得更好的剪枝模型？我们提出了一套集成式的先放大再剪枝流水线，在统一的余弦退火学习率调度下融合了放大模型训练、剪枝与恢复三个阶段，并辅以新颖的迭代式结构化剪枝方法，实现参数的渐进式移除。该方法能有效缓解朴素先放大再剪枝流程中因学习率回升而造成的知识损失，同时促进模型容量在存留神经元之间的合理再分配，从而实现更平滑的压缩效果与更优的性能。我们在将 2.8B 模型压缩至 1.3B、最高使用 2T 预训练 token 的条件下开展了全面实验，结果表明该集成方法不仅揭示了放大模型预训练在 token 效率方面的深层机理，还使剪枝后的模型取得了更出色的表现。

• † 佐治亚理工学院
• ‡ 德克萨斯大学奥斯汀分校
• ** 工作完成于 Apple 任职期间