Numerical Pruning for Efficient Autoregressive Models

TL;DR

This work addresses the high computational and memory burden of decoder-only autoregressive transformers by introducing a training-free, structural pruning framework that computes a globally optimal pruning mask via Newton's method for Attention and MLP modules. It further includes a lightweight compensation step to recover performance using a small calibration set, and demonstrates state-of-the-art results on both language and image-generation tasks (LLaMA family and LlamaGen) with notable speedups and memory savings on GPUs. The approach is theoretically grounded and empirically validated across diverse models and tasks, highlighting robust cross-domain applicability from language to image generation. Overall, the method provides a practical, hardware-friendly means to compress autoregressive models without major retraining requirements.

Abstract

Transformers have emerged as the leading architecture in deep learning, proving to be versatile and highly effective across diverse domains beyond language and image processing. However, their impressive performance often incurs high computational costs due to their substantial model size. This paper focuses on compressing decoder-only transformer-based autoregressive models through structural weight pruning to improve the model efficiency while preserving performance for both language and image generation tasks. Specifically, we propose a training-free pruning method that calculates a numerical score with Newton's method for the Attention and MLP modules, respectively. Besides, we further propose another compensation algorithm to recover the pruned model for better performance. To verify the effectiveness of our method, we provide both theoretical support and extensive experiments. Our experiments show that our method achieves state-of-the-art performance with reduced memory usage and faster generation speeds on GPUs.

Figures (6)

• Figure 1: Pruning overview. Blue modules denote column pruning and green modules denote row pruning.
• Figure 2: Perplexity ($\downarrow$) results for LLaMA-2 and LLaMA-3 models on WikiText2 dataset with 2048 sequence length. Comprehensive detailed results are included in Table 6 and Table 7 at Appendix A.2 and A.3.
• Figure 3: Visualization of generated images through LlamaGen-3B in 384$\times$384 resolution (cfg=1.65) with 10% sparsity.
• Figure 4: Ablation for number of samples for compensation.
• Figure 5: GPU memory v.s. generation speed.

Theorems & Definitions (18)

• Lemma 1: informal version of Lemma 9 at Appendix D.2
• Theorem 2: Mask optimization, informal version of Theorem 10 at Appendix D.3
• Theorem 3
• Remark 4
• Remark 5
• Lemma 9: formal version of Lemma \ref{['lem:XW_upper_bound:informal']}
• proof
• Theorem 10: Mask optimization, formal version of Theorem \ref{['thm:mask_optimize:informal']}
• proof
• Definition 11