Layer pruning method improves large language model accuracy by gradual retraining
LOOP: Recovery-Aligned Iterative Layer Pruning via On-Policy Distillation
Machine Learning
Summary
Large language models are often too big to run easily on everyday computers, so parts called layers are removed to make them smaller. But cutting too many layers at once can harm the model's ability to think through problems. The authors propose a new method called LOOP that removes layers one at a time and retrains the model a bit after each removal. This way, the model stays smarter and recovers better. Their tests show this approach keeps accuracy higher than removing many layers all at once.
What this means in practice
- •For machine learning engineers: Build smaller language models that maintain better reasoning accuracy by pruning layers gradually with retraining.
- •For software infrastructure teams: Deploy compact language models that run efficiently on standard hardware without large accuracy loss using recovery-aware pruning.
Authors
Liang Xu, Shufan Shen, Zhongni Hou, Yufei Zhang, Wei Lin, Guojun Yin, Shuhui Wang
Abstract
Layer-level pruning is attractive for large language model (LLM) deployment as it reduces the model size while preserving a dense architecture that can be efficiently executed on standard hardware. Current methods typically follow a one-shot paradigm, which first removes selected layers to reach a target model size and then performs additional training to recover model capabilities. However, pruning directly to the target size can substantially degrade reasoning capabilities, making subsequent recovery more difficult under a limited training budget. Moreover, this paradigm separates layer selection from recovery, limiting the use of recovery information to guide selection and potentially hindering the recovery of reasoning capabilities. To address these limitations, we propose LOOP, a Layer-level Recovery-Oriented On-Policy Pruning framework that interleaves pruning and recovery across multiple shots. In each shot, LOOP removes a single layer and performs short on-policy distillation, mitigating the abrupt degradation in reasoning capabilities associated with one-shot pruning. To guide layer selection in subsequent shots, LOOP reuses on-policy trajectories and optimization signals from recovery to jointly assess each remaining layer's reasoning contribution and recovery sensitivity, thereby making layer selection recovery-aware. Evaluations on several benchmarks show that LOOP consistently outperforms one-shot baselines in average accuracy across different model scales and pruning ratios.