A deeper look at depth pruning of LLMs

TL;DR

The paper investigates depth pruning in transformer-based LLMs as a route to reduce inference cost. It compares static cosine-based metrics with adaptive Shapley-value metrics and extends pruning down to individual self-attention and FFN layers. Key findings include that adaptive metrics can improve some objectives but harm others, self-attention layers are particularly amenable to pruning, and pruning up to around 33% of self-attention layers can be achieved without MMLU degradation in Mistral 7b. The study also evaluates simple performance-recovery methods, finding that an emulated average update can rival or outperform learned low-rank adapters. These results inform practical deployment of LLMs by guiding metric choice and recovery strategies for depth pruning.

Abstract

Large Language Models (LLMs) are not only resource-intensive to train but even more costly to deploy in production. Therefore, recent work has attempted to prune blocks of LLMs based on cheap proxies for estimating block importance, effectively removing 10% of blocks in well-trained LLaMa-2 and Mistral 7b models without any significant degradation of downstream metrics. In this paper, we explore different block importance metrics by considering adaptive metrics such as Shapley value in addition to static ones explored in prior work. We show that adaptive metrics exhibit a trade-off in performance between tasks i.e., improvement on one task may degrade performance on the other due to differences in the computed block influences. Furthermore, we extend this analysis from a complete block to individual self-attention and feed-forward layers, highlighting the propensity of the self-attention layers to be more amendable to pruning, even allowing removal of upto 33% of the self-attention layers without incurring any performance degradation on MMLU for Mistral 7b (significant reduction in costly maintenance of KV-cache). Finally, we look at simple performance recovery techniques to emulate the pruned layers by training lightweight additive bias or low-rank linear adapters. Performance recovery using emulated updates avoids performance degradation for the initial blocks (up to 5% absolute improvement on MMLU), which is either competitive or superior to the learning-based technique.

Figures (24)

• Figure 1: Our self-attention pruning, adaptive metrics, as well as emulated updates in comparison to cosine block pruning visualized w.r.t. compression ratios on Mistral 7b (see Fig. \ref{['fig:main_llama']} for results on LLaMa-2 7b).
• Figure 2: Comparison of different block influence metrics used for block pruning and their impact on downstream performance in terms of MMLU accuracy on LLaMa-2 7b (top) and Mistral 7b (bottom).
• Figure 3: Comparison of pruning the self-attention and feed-forward layers of a network using cosine metric on LLaMa-2 7b (top) and Mistral 7b (bottom), highlighting a higher propensity for the self-attention layers to be pruned.
• Figure 4: Impact on using performance recovery techniques with cosine block influence on LLaMa-2 7b (left) and Mistral 7b (right). The linear adapter was trained using logit-distillation with a rank of 8.
• Figure 5: Our self-attention pruning, adaptive metrics, as well as emulated updates in comparison to cosine block pruning visualized w.r.t. compression ratios on LLaMa-2 7b (see Fig. \ref{['fig:main_mistral']} for results on Mistral 7b).