Pruning Large Language Models to Intra-module Low-rank Architecture with Transitional Activations (2024)

Abstract: Structured pruning fundamentally reduces computational and memory overheads of large language models (LLMs) and offers a feasible solution for end-side LLM deployment. Structurally pruned models remain dense and high-precision, highly compatible with further tuning and compression. However, as the coarse-grained structured pruning poses large damage to the highly interconnected model, achieving a high compression ratio for scaled-up LLMs remains a challenge. In this paper, we introduce a task-agnostic structured pruning approach coupled with a compact Transformer architecture design. The proposed approach, named TransAct, reduces transitional activations inside multi-head attention (MHA) and multi-layer perceptron (MLP) modules, while preserving the inter-module activations that are sensitive to perturbations. Hence, the LLM is pruned into an intra-module low-rank architecture, significantly reducing weights, KV Cache and attention computation. TransAct is implemented on the LLaMA model and evaluated on downstream benchmarks. Results verify the optimality of our approach at high compression with respect to both efficiency and performance. Further, ablation studies reveal the strength of activation-guided iterative pruning and provide experimental analysis on the redundancy of MHA and MLP modules.

## introduction deploying large language models (llms) locally on edge devices instead of relying on remote apis has been a pressing initiative. local deployment of llms ensures independence from network conditions and enhances privacy at an advanced level @xcite . nevertheless, deploying a scaled-up llm onto a resource-constrained end sheared-llama @xcite activation projection device poses multifaceted challenges, encompassing inference speed, memory footprint, and power consumption. therefore, comprehensive optimizations on the efficiency of llms are imperative, including architecture design @xcite , model compression @xcite , inference schemes @xcite , compilation and runtime @xcite . model compression emerges as the silver-bullet solution for reducing deployment costs given an accessible llm. to essentially reduce model computation and memory overhead, pruning aims to discard weights with low salience to the llm. @xcite suggest that state-ofthe-art (sota) unstructured pruning approaches i.e., sparsegpt @xcite and wanda @xcite , along with their semistructured variations, often underperform in downstream benchmarks. @xcite emphasize the significance of post-training after pruning to restore the capabilities of the llm. however, the post-training and inference of a sparse model are notably inefficient. also, an unstructured pruning with arbitrary sparsity pattern has no speedup or memory saving on the llm, whereas a semistructured sparse model heavily relies on specific hardware @xcite ). an alternative pruning category, i.e., structured pruning, has shown feasibility for llms. llm-pruner @xcite , the pioneering structured pruning of llm, incorporates the approximated taylor series as the pruning metric. however, this approximation loses accuracy when pruning a large ratio of the model @xcite . while taylor expansion assumes small perturbations, it is not applicable when a large number of parameters are pruned (i.e., set to zero). the sota approach sheared-llama @xcite , on the other hand, completely transfers the evaluation of the pruning metric to supervised training with masks. however, training with masks poses much more computation and memory footprint at training time, as well as the training unstableness. also, the pruned architecture of sheared-llama, as illustrated in the upper part of figure 1 , involves the unified pruning of layer normalization (ln) weights, disregarding the varying sensitivity of ln parameters to perturbation across layers @xcite . to address the challenges of efficient and effective llm pruning, we propose transact, a transitional activation-based structured pruning approach. from the perspective of pruning architectural design, transact reduces intra-module activations, which prunes the mha and mlp in llm into low intrinsic dimension as depicted in figure 1 . transact pruning metric is inspired by the observation of @xcite that a small proportion of activations within the llm exhibit outlier magnitudes, rendering them particularly sensitive to perturbations and need to be preserved. this approach effectively reduces the memory footprint of both model weights and kv cache, alleviating the memory constraints inherent in autoregressive generation on edge devices @xcite . specifically, the contributions of this paper are outlined as follows. ## related work extensive works have been proposed to optimize the efficiency of transformer-based lms, covering pruning, quantization, dynamic acceleration, etc. however, to generalize these approaches to the continually scaling-up llms remains challenging. quantization, which reduces the bit representation of values, stands out due to its ease of implementation. post-training quantization (ptq) approaches, e.g., gptq @xcite and awq @xcite , are without any further tuning after the quantization. on the contrary, quantization-aware training (qat) approaches train the model along with the quantization parameters and is still challenging when the llm is scaled up @xcite . quantizing an llm from float16 to int3 with weight-only ptq approaches like gptq @xcite can reach roughly 80% compression of model weights. however, the kv cache which contributes to a large amount of memory overheads is still in float16 and uncompressed. also, obtaining an acceptable quantization precision with int3 weights remains a challenge. @xcite proposed a w8a8 ptq approach where both weights and activations are quantized to int8, saving 50% memory footprint. the lack of flexibility poses a significant limitation to quantization. most general computing platforms and libraries primarily support low-bit representations such as int8 and int4 @xcite . however, opting for representations lower than 4 bits necessitates dequantization back to the supported higherbit representations, thereby introducing additional computation and memory overheads. apart from quantization, unstructured pruning is also an efficient approach to obtain a sparse llm. @xcite and @xcite enabled fully unstructured and semistructured n:m sparsity (i.e., n zeros in m consecutive weights) of llm across different sizes. however, there are two major obstacles hinder- ## methodology in this section, we first recap the preliminaries of transformer-based llm architecture and introduce the transitional activations. then, we propose our approach transact with the pruning metric and architecture design of the pruned model. ## conclusion in this paper, we introduce transact, an effective and efficient pruning approach coupled with an architecture designed for pruned llms. trans-act compresses the original llm into a compact dense model with an intra-module low-rank architecture, achieving the fastest inference speed and lowest overheads compared to models of similar sizes. the compression is guided by the magnitudes of the transitional activations within the mha and mlp modules. specifically, intramodule dimensions with small activations are structurally pruned out, while inter-module dimensions are preserved. experiments on open-source llms and downstream benchmarks demonstrate the strength of our approach, particularly at high compression rates. also, we thoroughly evaluated the pruned llm with respect to calibration samples, pruning ratio, and pruning shots. the results provide insights and experimental results for further development of compact yet powerful lms.