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

TL;DR

This work tackles the challenge of deploying efficient generative language models under strict inference budgets by integrating enlarged pretraining, structured FFN pruning, and recovery into a single training run called IDEA Prune. It replaces the conventional multi-stage, multi-LR approach with a unified cosine annealing learning rate and an iterative FFN-width pruning strategy, aiming to preserve knowledge while compressing to target sizes. Experiments show that pruning a 2.8B model down to 1.3B with 2T tokens yields strong token efficiency and improved task performance (e.g., MMLU), outperforming naive enlarge-and-prune baselines and benefiting from intermediate pruning seeds and KD when used. The results advocate a more efficient, integrated paradigm for pretraining-based model compression and demonstrate practical gains for deploying powerful decoders under real-world hardware constraints.

Abstract

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.

Figures (4)

• Figure 1: Naive pipeline: 1T for enlarged model training and 1T for pruned model recovery, each with separate learning rate decay. Integrated pipeline: single learning rate schedule matching target-size model training, pruning started at 1T. Left: Perplexity on OpenWebText of 300M model trained with multiple training token budgets from 0.25T to 2T. Right: Training loss curves across enlarge-and-prune pipelines. Training 300M models with 2T total token budget.
• Figure 2: Learning rate schedule for separate (Figure \ref{['fig:lr_minitron']}, \ref{['fig:lr_sheared']}) and integrated (Figure \ref{['fig:lr_integrated']}) enlarge-and-prune pipelines. Minitron uses single shot pruning, resulting in a narrow pruning stage. Our integrated pipeline uses one single learning rate decay schedule for all stages. We use iterative structured pruning to further integrate pruning and recovery training.
• Figure 3: Ablation of the enlarged model size. All model sizes outperforms the training from scratch baseline. Trade-off between the model capacity and pruning degradation happens in the enlarged model size.
• Figure 4: Hyperparameter study. Figure \ref{['fig:pruning_schedule']} shows the a wide range of pipeline schedule works for the best performance. Figure \ref{['fig:impt_score_comb']} implies "mean-max" is the best choice of $f_1$ and $f_2$, though the others do not show significant degradation. Figure \ref{['fig:lambda']} shows the exponential moving average coefficient $\lambda$ is robust in a wide range from 0.3 to 0.8.