IDEA Prune：生成式语言模型预训练中的集成扩展与剪枝流水线
计算与语言
2025-03-11 v1 人工智能
机器学习
摘要
大型语言模型的最新进展加剧了在有限推理预算内对高效且可部署模型的需求。与从零开始训练目标尺寸模型相比，结构化剪枝流水线在 token 效率方面展现出了潜力。在本文中，我们主张将先前工作中常被忽略的扩展模型预训练纳入剪枝流程。我们将扩展与剪枝流水线作为一个集成系统进行研究，以解决两个关键问题：即使模型永远不会被部署，预训练一个扩展模型是否值得；以及如何优化整个流水线以获得更好的剪枝模型。我们提出了一种集成的扩展与剪枝流水线，它在单一的余弦退火学习率调度下结合了扩展模型训练、剪枝和恢复。该方法进一步辅以一种新颖的迭代结构化剪枝方法，用于逐步移除参数。所提出的方法有助于缓解朴素扩展与剪枝流水线中学习率上升所导致的知识损失，并能在存留神经元之间实现模型容量的有效重新分配，从而促进平滑压缩并提升性能。我们在将 2.8B 模型压缩至 1.3B 且预训练使用多达 2T tokens 的设置上进行了全面实验。结果表明，该集成方法不仅为扩展模型预训练的 token 效率提供了洞见，而且使剪枝模型获得了更优的性能。
引用
@article{arxiv.2503.05920,
title = {IDEA Prune: An Integrated Enlarge-and-Prune Pipeline in Generative Language Model Pretraining},
author = {Yixiao Li and Xianzhi Du and Ajay Jaiswal and Tao Lei and Tuo Zhao and Chong Wang and Jianyu Wang},
journal= {arXiv preprint arXiv:2503.05920},
year = {2025}
}