ReplaceMe: Network Simplification via Depth Pruning and Transformer Block Linearization
About
We introduce ReplaceMe, a generalized training-free depth pruning method that effectively replaces transformer blocks with a linear operation, while maintaining high performance for low compression ratios. In contrast to conventional pruning approaches that require additional training or fine-tuning, our approach requires only a small calibration dataset that is used to estimate a linear transformation, which approximates the pruned blocks. The estimated linear mapping can be seamlessly merged with the remaining transformer blocks, eliminating the need for any additional network parameters. Our experiments show that ReplaceMe consistently outperforms other training-free approaches and remains highly competitive with state-of-the-art pruning methods that involve extensive retraining/fine-tuning and architectural modifications. Applied to several large language models (LLMs), ReplaceMe achieves up to 25\% pruning while retaining approximately 90\% of the original model's performance on open benchmarks - without any training or healing steps, resulting in minimal computational overhead. We provide an open-source library implementing ReplaceMe alongside several state-of-the-art depth pruning techniques, available at https://github.com/mts-ai/ReplaceMe
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Commonsense Reasoning | HellaSwag | Accuracy65.7 | 1460 | |
| Language Modeling | WikiText | PPL34 | 479 | |
| Physical Interaction Question Answering | PIQA | Accuracy71.1 | 323 | |
| Question Answering | ARC-E | Accuracy65.9 | 242 | |
| Question Answering | SciQ | Accuracy86.4 | 226 | |
| Multitask Language Understanding | MMLU | Accuracy46.4 | 206 | |
| Language Modeling | LAMBADA | Accuracy42.1 | 183 | |
| Reading Comprehension | RACE | Accuracy35.4 | 151 | |
| Science Question Answering | ARC-E | Accuracy65.9 | 138 | |
| Science Question Answering | ARC-C | Accuracy43.7 | 127 |