S hadow LLM : Predictor-based Contextual Sparsity for Large Language Models (2024)

Abstract: The high power consumption and latency-sensitive deployments of large language models (LLMs) have motivated efficiency techniques like quantization and sparsity. Contextual sparsity, where the sparsity pattern is input-dependent, is crucial in LLMs because the permanent removal of attention heads or neurons from LLMs can significantly degrade accuracy. Prior work has attempted to model contextual sparsity using neural networks trained to predict activation magnitudes, which can be used to dynamically prune structures with low predicted activation magnitude. In this paper, we look beyond magnitude-based pruning criteria to assess attention head and neuron importance in LLMs. We develop a novel predictor called ShadowLLM, which can shadow the LLM behavior and enforce better sparsity patterns, resulting in over 15% improvement in end-to-end accuracy compared to prior methods. In addition, ShadowLLM achieves up to a 20% speed-up over the state-of-the-art DejaVu framework. These enhancements are validated on Llama-2 and OPT models with up to 30 billion parameters. Our code is available at https://github.com/abdelfattah-lab/shadow_llm/

## introduction large language models (llms) are emerging as a core component of many computing applications. their ability to perform in-context learning, i.e., to perform a task by conditioning on examples without any gradient updates @xcite @xcite , make them broadly applicable to numerous applications. yet, their large size combined with the latency-sensitivity of llm-based applications make them expensive to deploy @xcite . a key optimization in llm deployment is sparsification, where weights or activations are pruned to reduce the computation and memory requirements at run time. sparsification can either be static, which permanently removes an attention head or neuron, or contextual, which prunes based on the current input and context. while some works investigate task-specific static pruning methods @xcite , they typically have a large impact on in-context learning, reducing downstream task accuracy compared to contextual sparsity. contextual sparsity can be leveraged at run time to dynamically prune llms, yet it requires making fast and accurate predictions based on predetermined pruning criteria. these criteria can have large effects on the overall accuracy and performance of the model, as shown in figure 1 . our method shadowllm uses more accurate pruning criteria and a unified predictor at the beginning of the model, which leads to a stronger accuracyperformance tradeoff compared to the recent work dejavu @xcite . 0 100 200 300 400 500 600 700 800 head 10 0 10 1 10 2 10 3 10 4 rank variance random activation power-law activation head activation rank variance of opt-1.3b 0 5 10 15 20 layer figure 3: heads with higher rank variance, calculated using gradnorm, indicate greater context dependence. ## pruning criteria contextual sparsity requires dynamically understanding which neurons to prune (i.e. assessing the neurons importance relative to an input) and ranking the neurons relative to each other. figure 4 depicts how we can use information about the activations and gradients to prune a llm for this contextual sparsity. consider a model m and dataset d, containing prompts (inputs) along with the target output sequence. we then wish to define performance on the dataset as p m (d). now suppose a subset of the model c ⊂ m is pruned out. ideally, we would like to be able to estimate p m (d) -p m\c (d) @xcite . the optimal pruning strategy is found in equation 1. if we look at aggressive attention head pruning of even small transformers (prune 56 out of 64 heads in each layer), exhaustive search in a single layer would require 64 c 8 evaluations, and this would have to be repeated for every layer, making the problem intractable. we can feed a subset of the dataset @xmath0, and calculate the loss l. further, we can also get access to the activations (a), as well as the parameters of the up-projection ffn of transformer at layer l as θ l . the activation at layer l for the k th head or neuron is denoted as a l,k . the gradients for these activations are denoted as ∂l ∂a l,k . the gradient for the weight parameters of the k th neuron in the up-projection ffn at layer l is given as ∂l ∂θ l,k . current predictor-based sparsity research investigates the impact of magnitude-based criteria, such as the l2norm of the head and neuron activation on a subset of data d. the intuition is that the heads that are more activated should be more important. there is significant research on other criteria for pruning weights and activations @xcite . beyond activation magnitude being a criterion for importance, the process of pruning can also be framed as an optimization problem, with the goal of approximating the change in loss from removing parameters. methods such as optimal brain damage (obd) @xcite rely on the gradient of the loss with respect to the feature maps. while obd evaluates the second-order term (hessian), works such as @xcite come up with similar metrics based on the taylor expansion of the change in loss. in this paper, we evaluate pruning criteria of varying complexity, that use (1) activation methods, (2) first-order gradient (jacobian) methods, (3) activation + jacobian methods, (4) obd-style hessian methods and (5) sensitivity-based methods for pruning llms. among these methods, we find that a gradientbased sensitivity method we call plainact outperforms activation-based magnitude methods adapted in prior dynamic pruning research @xcite . the l2norm activation-magnitude based criterion assesses the importance of neurons by simply taking the l2 norm of the head and neuron activation as ||a l,k || 2 . the plainact criterion measures the expected sensitivity of the model on the loss if a head or neuron is removed. for the head and neuron, this can be described simply as ## predictors for neuron ranking when deploying a large language model, for a given input, we will not have access to the activations or the gradients. thus, calculating the l2norm or plainact criterion is not possible. however, it is possible to create a calibration dataset of inputs and their corresponding l2norm or plainact for each head and neuron. such a dataset can be used to train a predictor, which can take the input and predict the sparsity pattern of the model at deployment. sparsity prediction can reduce the end-toend latency of transformer inference by predicting which operations can be skipped. we propose a method called shadowllm that uses the first layer's attention output to predict the sparsity pattern for the entire model. this reduces the overhead of repeatedly calling the predictor at each layer and cuts the total flops of the predictor by 20%, as shown in we also explore a full sequence shadowllm, which uses a small transformer to take in the entire input token embedding and predict the sparsity pattern, allowing pruning of the first transformer layer as well. however, the full sequence shadowllm requires an additional 2(2e 2 + el 2 ) flops, making it as costly as running an entire dense attention layer and impractical due to the high computational cost. ## evaluation we find that the activation-gradient based pruning criteria that we use in shadowllm are effective for downstream evaluation tasks as well as perplexity. further, we demonstrated in section 4 that shadowllm can predict the sparsity pattern for the entire llm given just the input to the first layer. we find that the accuracy and predictability tradeoff is excellent for the plainact criterion, whereas other criteria were harder to learn due to outliers and high variance. in this section, we evaluate the effectiveness of combining shadowllm predictor design with the plainact criterion, compared to our implementation of dejavu-style foot_0 predictors, trained on a magnitude based pruning criteria.