Integrated Expansion-Pruning Pipeline for Pre-training Generative Language Models: IDEA Prune

Structured Pruning pipelines are gaining attention for the efficient deployment of Large Language Models (LLMs). This work explores whether integrating pre-training of expanded models, often overlooked in previous research, into the pruning process is valuable, and how the entire pipeline can be optimized.

The proposed IDEA Prune integrates the processes of expanded model training, pruning, and recovery under a single Cosine Annealing learning rate schedule. It also introduces a novel iterative structured pruning method for gradual parameter removal. This approach mitigates knowledge loss caused by learning rate increases in existing methods and effectively redistributes model capacity among surviving neurons, achieving smooth compression and performance improvements.

In experiments compressing a 2.8B model down to a 1.3B model, up to 2T tokens of pre-training data were utilized. The integrated approach provided insights into the token efficiency of expanded model pre-training and demonstrated the superiority of the pruned models.