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

TL;DR

TransAct introduces activation-guided structured pruning to convert large LLMs into an intra-module low-rank architecture by shrinking transitional dimensions within MHA and MLP while preserving inter-module activations. The method relies on salience metrics derived from transitional activations to guide iterative, non-gradient pruning and achieves substantial reductions in FLOPs and KV cache with competitive downstream performance, demonstrated on LLaMA2-7B-base. Ablation studies show the importance of iterative pruning, calibration sample size, and joint MHA/MLP pruning for robust compression, including strong results on TriviaQA and TruthfulQA at high compression. The work advances edge-deployable LLMs by enabling faster inference and lower memory footprints, with practical implications for on-device AI and efficient server deployment.

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.

Figures (8)

• Figure 1: An illustration of TransAct model architecture. The model weights and activations are colored green and blue, respectively. Dashed hollow blocks represent the weights and activations that are pruned out.
• Figure 2: Detailed TransAct workflow on a Transformer layer with MHA and gated MLP. Bar charts indicate the activation-based pruning metric.
• Figure 3: Inference FLOPs of the original and pruned models with variable context length. LLM-Pruner is omitted as the implemented architecture is the same as TransAct counterparts. LLaMA2-7B with 4K-token context is omitted.
• Figure 4: Edge device E2E generation latency (seconds) tested on a Xiaomi 14 mobile phone. The models are quantized to W4A16 by MLC-LLM and tested with variable context length (L, prefix+decode). LLM-Pruner is omitted as the implemented architecture is the same as TransAct counterparts.
• Figure 5: LAMBADA perplexity and accuracy on every checkpoint of TransAct-2.6B, LLM-Pruner-2.6B and Sheared-LLaMA-2.7B post-training.