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

Summary

Recent advancements in large language models necessitate efficient deployment options within constrained inference budgets. While structured pruning offers better token efficiency than training target-sized models from scratch, previous work frequently ignores the pretraining of enlarged models. To resolve this gap, researchers developed an integrated enlarge-and-prune pipeline that unifies enlarged model training, pruning, and recovery under a single cosine annealing learning rate schedule. This system incorporates an iterative structured pruning technique to gradually eliminate parameters, redistributing model capacity across surviving neurons and avoiding the knowledge loss seen in naive approaches. Evaluated on compressing 2.8B parameter models to 1.3B across up to 2T pretraining tokens, the approach yields superior downstream model performance.

Takeaways

1. 1

   An integrated pipeline unifies enlarged model pretraining, pruning, and recovery under a single cosine annealing learning rate schedule.

2. 2

   Iterative structured pruning enables gradual parameter removal to redistribute model capacity among surviving neurons and mitigate knowledge loss.

3. 3

   Compressing 2.8B parameter models to 1.3B with up to 2T pretraining tokens demonstrates token efficiency and superior pruned model performance.