ReplaceMe: Network Simplification via Depth Pruning and Transformer Block Linearization

TL;DR

This work tackles the high cost of deploying massive transformers by proposing ReplaceMe, a training-free depth pruning method that substitutes a contiguous block of transformer layers with a learned linear transformation estimated from a small calibration dataset. The linear transform is merged into the preceding MLP weights, preserving architecture and avoiding new parameters, while enabling substantial compression with minimal performance loss. The approach supports multiple linear transforms and uses either an analytical $L_2$-based solution or a numerical cosine distance optimization, with regularization to balance accuracy and perplexity. Across LLMs (e.g., Llama, Qwen, Falcon) and vision transformers (e.g., CLIP), ReplaceMe achieves competitive results against state-of-the-art pruning methods, often outperforming training-free baselines and approaching healed baselines at moderate compression, accompanied by an open-source implementation for broader adoption.

Abstract

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 (see Fig.1). 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.

Figures (7)

• Figure 1: Comparison of the proposed LLM compression method with state-of-the-art UIDL. Subplots illustrate (a) compression time, (b) environmental impact ($\text{CO}_{2}$ emissions and energy consumption), and (c) performance accuracy relative to the uncompressed baseline. Our approach attains the shortest compression time, lowest energy use, and reduced emissions, while achieving the highest accuracy, demonstrating superior efficiency, sustainability, and effectiveness over existing methods.
• Figure 2: ReplaceMe compresses and accelerates LLMs by bypassing a contiguous sequence of transformer blocks—illustrated by the red line—while preserving model performance. This is achieved by inserting an estimated linear transformation matrix that maps the MLP output of the $i^{*}$-th block directly to the input space expected by the $(i^{*} + n+1)$-th block, effectively replacing all $n$ blocks in between.
• Figure 3: Comparison between our method and UIDL with different number of pruned layers and different models including, Llama 3 8B, Mistral 3 7B, and Qwen 2.5 7B. We show the average accuracy across the Race, Winogrande, PIQA, BoolQ, OpenBookQA, and SciQ benchmarks.
• Figure 4: Comparison between ReplaceMe and UIDL in terms of computation and environmental impact.
• Figure 5: Comparative analysis of distance metrics and predictive accuracy across layer pruning configurations. Trendlines illustrate the inverse correlation between cosine distance and accuracy, contrasted with the positive correlation between L2 and accuracy degradation . Results are shown for ReplaceMe with LT estimation via Cosine/LS metrics and UIDL baselines.