IDEA Prune: Integrated Enlarge-and-Prune for LLM Pretraining

Recent work in language model compression has turned to structured pruning as a way to produce deployable models within fixed inference budgets, and pruning pipelines often beat training target-size models from scratch on token efficiency. This paper from researchers at Georgia Institute of Technology, University of Texas at Austin, and Apple pushes that direction further by folding the enlarged model pretraining step into the pruning system instead of treating it as a separate, expendable phase.

The authors study two questions: whether it is worth pretraining an enlarged model even when that model is never deployed, and how to jointly optimize the entire pipeline so the final pruned model is as strong as possible. Their core proposal is an integrated enlarge-and-prune pipeline that runs enlarge model training, pruning, and recovery under a single cosine annealing learning rate schedule. Complementing the schedule is a novel iterative structured pruning method that removes parameters gradually rather than all at once.

The motivation for unifying these stages is that in naive enlarge-and-prune pipelines the learning rate rises again when pruning starts, which causes knowledge loss. The integrated schedule and iterative removal are designed to mitigate that loss and to redistribute model capacity among the surviving neurons, supporting smoother compression and better final performance.

On the evidence side, the authors run comprehensive experiments compressing 2.8B models down to 1.3B with up to 2T tokens of pretraining. They report that the integrated approach offers insights into the token efficiency of enlarged model pretraining and achieves superior performance for the pruned models relative to alternatives.

The scope is bounded by the experimental setting: results are demonstrated at the 2.8B-to-1.3B compression range, and the findings are framed in terms of pretraining cost versus pruned-model quality rather than broader deployment benchmarks. Authors include Yixiao Li and Tuo Zhao (Georgia Tech), Ajay Jaiswal (UT Austin), plus Xianzhi Du, Tao Lei, Chong Wang, and Jianyu Wang, with Li and Jaiswal noting their work was done while at Apple.

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

View Original