On the Efficiency of Structured Pruning in Small Language Model Pretraining
Yixiao Li ⋅ Xianzhi Du ⋅ AJAY JAISWAL ⋅ Tao Lei ⋅ Tuo Zhao ⋅ Chong Wang ⋅ Jianyu Wang
Abstract
Companies that deploy generative language models on edge devices often need to train small models from scratch due to legal, data-governance, or product-customization requirements. An alternative approach is to train an enlarged model and compress it to the target size via structured pruning, but this advantage becomes unclear once the cost of pretraining the enlarged model is included. We introduce the enlarge-and-prune framework to study this question under full-cost accounting, and show that conventional enlarge-and-prune pipelines---which train the enlarged model, prune it, and recover in separate stages---do not consistently outperform direct small-model training under the same total token budget. To overcome this limitation, we propose IDEA Prune, an integrated enlarge-and-prune pipeline that unifies enlarged model pretraining, iterative structured pruning, and recovery under a single cosine annealing learning rate schedule. In controlled experiments compressing 2.8B models to 1.3B with up to 2T tokens on the same dataset, IDEA Prune consistently outperforms conventional pipelines and modestly improves overall performance relative to the scratch baseline, improving MMLU accuracy to 46.4\%, compared with 31.4--33.4\% for conventional pipelines. Our analysis further reveals that intermediate checkpoints provide more favorable pruning starting points than fully converged ones, and that an enlargement factor of approximately 2.6$\times$ yields the best pruned models in the tested regime.
Chat is not available.
Successful Page Load