通过先扩大再剪枝,提升大模型预训练效率与压缩后性能
IDEA Prune: An Integrated Enlarge-and-Prune Pipeline in Generative Language Model Pretraining
- 先训大模型再系统性剪枝,用余弦退火优化训练与剪枝流程
- 2.8B模型压缩至1.3B,使用2万亿词训练,性能优于传统方法
- 适合追求高效率、低资源部署的大模型研究者和工程师
大语言模型的发展对高效可部署模型提出了更高要求。结构化剪枝相比从头训练目标规模模型,在词元效率上表现更优。本文提出将通常被忽略的扩大模型预训练纳入剪枝流程,系统研究了‘扩大-剪枝’一体化管道,回答两个关键问题:是否值得在不部署的情况下预训练扩大模型,以及如何优化全流程以获得更优的剪枝模型。我们设计了一种集成的扩大-剪枝管道,结合扩大模型训练、剪枝与恢复,并采用统一的余弦退火学习率调度。该方法还引入一种新型迭代式结构化剪枝策略,逐步移除参数,缓解了传统方法中学习率上升导致的知识损失,促进剩余神经元间能力的有效重分配,实现平滑压缩与性能提升。我们在压缩2.8B模型至1.3B的实验中,使用高达2T tokens的预训练数据,验证了该方法不仅揭示了扩大模型预训练的词元效率,还实现了更优的剪枝模型性能。
原文摘要 · Abstract (English)
Recent advancements in large language models have intensified the need for efficient and deployable models within limited inference budgets. Structured pruning pipelines have shown promise in token efficiency compared to training target-size models from scratch. In this paper, we advocate incorporating enlarged model pretraining, which is often ignored in previous works, into pruning. We study the enlarge-and-prune pipeline as an integrated system to address two critical questions: whether it is worth pretraining an enlarged model even when the model is never deployed, and how to optimize the entire pipeline for better pruned models. We propose an integrated enlarge-and-prune pipeline, which combines enlarge model training, pruning, and recovery under a single cosine annealing learning rate schedule. This approach is further complemented by a novel iterative structured pruning method for gradual parameter removal. The proposed method helps to mitigate the knowledge loss caused by the rising learning rate in naive enlarge-and-prune pipelines and enable effective redistribution of model capacity among surviving neurons, facilitating smooth compression and enhanced performance. We conduct comprehensive experiments on compressing 2.8B models to 1.3B with up to 2T tokens in pretraining. It demonstrates the integrated approach not only provides insights into the token efficiency of enlarged model pretraining but also achieves superior performance of pruned models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。