Iterative Layer-wise Distillation for Efficient Compression of Large Language Models
Grigory Kovalev, Mikhail Tikhomirov
TL;DR
The paper tackles the challenge of deploying large language models in resource-constrained settings by proposing an iterative, data-driven layer pruning framework coupled with targeted fine-tuning. It blends an objective layer-importance assessment with a joint loss that combines $L_{KL}$-based logits alignment and $L_{MSE}$-based hidden-state alignment, enabling substantial reductions in transformer depth while preserving accuracy. Experiments on the Qwen2.5-3B model demonstrate that pruning from 36 to 28 layers yields only about $9.7\%$ quality loss (to a $0.574$ aggregate score from $0.636$) and that further pruning to 24 layers incurs ~ $18\%$ loss, with the middle layers (roughly 17–24) identified as least critical for inference. This approach enables efficient deployment in resource-limited environments and provides a reproducible pipeline, with code and a distilled model released for public use.
Abstract
This work investigates distillation methods for large language models (LLMs) with the goal of developing compact models that preserve high performance. Several existing approaches are reviewed, with a discussion of their respective strengths and limitations. An improved method based on the ShortGPT approach has been developed, building upon the idea of incorporating iterative evaluation of layer importance. At each step, importance is assessed by measuring performance degradation when individual layers are removed, using a set of representative datasets. This process is combined with further training using a joint loss function based on KL divergence and mean squared error. Experiments on the Qwen2.5-3B model show that the number of layers can be reduced from 36 to 28 (resulting in a 2.47 billion parameter model) with only a 9.7% quality loss, and to 24 layers with an 18% loss. The findings suggest that the middle transformer layers contribute less to inference, underscoring the potential of the proposed method for creating efficient models. The results demonstrate the effectiveness of iterative distillation and fine-tuning, making the approach suitable for deployment in resource-limited settings.