PAT: Pruning-Aware Tuning for Large Language Models

2024-08-27

arXiv.org

AI Summary

To address the severe performance degradation and poor recoverability of large language models (LLMs) after structural pruning, this work proposes a pruning-aware optimization paradigm that jointly optimizes structured pruning and fine-tuning. Our method introduces: (1) a Hybrid Sparsification Module (HSM), integrating lightweight operators with globally shared, learnable channel masks to enable efficient channel-level pruning; and (2) an identity loss function that decouples linear transformations from scaling operations, enhancing training robustness. Evaluated on Llama2-7B, our approach achieves 1.33× inference speedup at 25% pruning ratio while outperforming LoRA baselines by 1.26% in accuracy—without increasing training overhead. To the best of our knowledge, this is the first work to unify structured pruning, plug-in sparse modules, and identity loss for LLMs, establishing a novel paradigm for efficient LLM deployment.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Large language models (LLMs) excel in language tasks, especially with supervised fine-tuning after pre-training. However, their substantial memory and computational requirements hinder practical applications. Structural pruning, which reduces less significant weight dimensions, is one solution. Yet, traditional post-hoc pruning often leads to significant performance loss, with limited recovery from further fine-tuning due to reduced capacity. Since the model fine-tuning refines the general and chaotic knowledge in pre-trained models, we aim to incorporate structural pruning with the fine-tuning, and propose the Pruning-Aware Tuning (PAT) paradigm to eliminate model redundancy while preserving the model performance to the maximum extend. Specifically, we insert the innovative Hybrid Sparsification Modules (HSMs) between the Attention and FFN components to accordingly sparsify the upstream and downstream linear modules. The HSM comprises a lightweight operator and a globally shared trainable mask. The lightweight operator maintains a training overhead comparable to that of LoRA, while the trainable mask unifies the channels to be sparsified, ensuring structural pruning. Additionally, we propose the Identity Loss which decouples the transformation and scaling properties of the HSMs to enhance training robustness. Extensive experiments demonstrate that PAT excels in both performance and efficiency. For example, our Llama2-7b model with a 25% pruning ratio achieves 1.33$ imes$ speedup while outperforming the LoRA-finetuned model by up to 1.26% in accuracy with a similar training cost. Code: https://github.com/kriskrisliu/PAT_Pruning-Aware-Tuning

Problem

• Large Language Models
• Memory and Computational Demand
• Performance Degradation after Pruning

Innovation

• Pruning-Aware Fine-tuning
• Mixed Sparsification Modules
• Identity Loss