A deeper look at depth pruning of LLMs

TL;DR

This paper investigates various importance metrics for pruning large language models, especially focusing on self-attention layers, and explores simple methods to recover performance after pruning, aiming to reduce resource costs without significant accuracy loss.

Contribution

It introduces adaptive importance metrics like Shapley value for pruning LLMs, analyzes layer-specific pruning effects, and proposes lightweight performance recovery techniques.

Findings

Self-attention layers are more amenable to pruning, allowing up to 33% removal without performance loss.

Adaptive importance metrics show a trade-off between task performances.

Simple additive bias or low-rank adapters can effectively recover pruned model performance.

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.