Pruning Large Language Models with Semi-Structural Adaptive Sparse Training
作者: Weiyu Huang, Yuezhou Hu, Guohao Jian, Jun Zhu, Jianfei Chen
分类: cs.CL, cs.AI
发布日期: 2024-07-30 (更新: 2024-12-18)
备注: Accepted at AAAI25
💡 一句话要点
提出AST框架，通过半结构化自适应稀疏训练压缩大语言模型，显著降低性能损失。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 模型剪枝 稀疏训练 知识蒸馏 自适应学习 模型压缩 LLaMA2
📋 核心要点
- 现有一次性剪枝方法压缩LLM时，在复杂语言理解任务上性能下降明显，限制了LLM剪枝的实际应用。
- 提出自适应稀疏训练器（AST），在权重更新时学习最优掩码，无需额外计算开销，提升重训练效率。
- 实验表明，AST在LLaMA2-7B上实现了SOTA性能，显著缩小了稀疏模型与密集模型之间的性能差距，训练成本极低。
📝 摘要（中文）
大型语言模型（LLMs）的卓越性能依赖于其庞大的规模，这给模型部署带来了延迟和内存消耗方面的巨大挑战。最近，许多研究试图使用一次性剪枝方法来压缩LLMs。然而，这些方法在复杂的语言理解任务上常常遭受显著的性能下降，引发了对LLMs中剪枝可行性的担忧。为了解决这个问题，我们提出了一种新颖且高效的重训练框架——自适应稀疏训练器（AST），专为半结构化稀疏模型定制。AST使模型能够在权重更新过程中学习最佳掩码，而不会产生额外的计算开销。此外，我们证明了结合知识蒸馏可以显著提高重训练效率，并在固定的计算约束下增强模型性能。此外，还集成了一组经过良好初始化的补充参数，以进一步增强模型的有效性。AST以最小的训练成本实现了最先进的性能。当应用于LLaMA2-7B模型时，AST将密集模型和2:4半结构化稀疏模型之间的困惑度和零样本准确率差距分别降低到0.6和1.16%，使用的预训练tokens和GPU时间不到0.4%。我们的工作证明了部署半结构化稀疏LLMs的可行性，并为实现高度压缩的模型提供了一种有希望的替代方案，可以与现有的量化技术相结合。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLMs）压缩后性能显著下降的问题，尤其是在使用一次性剪枝方法进行压缩时。现有的剪枝方法虽然可以减少模型大小和计算量，但往往会导致模型在复杂语言理解任务上的性能大幅降低，限制了LLM在资源受限环境中的部署。
核心思路：论文的核心思路是通过自适应稀疏训练，使模型在训练过程中学习最优的稀疏结构，从而在压缩模型的同时尽可能地保留模型的性能。通过在权重更新过程中动态调整掩码，模型可以更好地适应稀疏结构，避免一次性剪枝带来的信息损失。同时，结合知识蒸馏和补充参数，进一步提升模型的性能和训练效率。
技术框架：AST框架主要包含以下几个关键模块：1) 自适应稀疏训练：在权重更新过程中，根据一定的策略动态调整权重掩码，使模型学习最优的稀疏结构。2) 知识蒸馏：利用预训练模型的知识来指导稀疏模型的训练，提高模型的泛化能力和性能。3) 补充参数：引入一组经过良好初始化的参数，作为原始参数的补充，进一步增强模型的表达能力。整个流程是在预训练模型的基础上进行微调，通过AST框架进行稀疏化训练，最终得到一个高性能的稀疏模型。
关键创新：论文最重要的技术创新点在于提出了自适应稀疏训练方法，该方法能够在权重更新过程中动态调整掩码，使模型学习最优的稀疏结构。与现有的一次性剪枝方法相比，AST能够更好地保留模型的性能，避免了因剪枝导致的信息损失。此外，结合知识蒸馏和补充参数，进一步提升了模型的性能和训练效率。
关键设计：AST的关键设计包括：1) 掩码更新策略：设计合理的掩码更新策略，例如基于梯度或重要性的策略，以确保模型能够学习到最优的稀疏结构。2) 知识蒸馏损失函数：选择合适的知识蒸馏损失函数，例如KL散度或MSE损失，以确保稀疏模型能够有效地学习预训练模型的知识。3) 补充参数初始化：采用合适的初始化方法，例如随机初始化或预训练初始化，以确保补充参数能够有效地增强模型的表达能力。4) 稀疏比例：选择合适的稀疏比例，以在模型大小和性能之间取得平衡。
🖼️ 关键图片
📊 实验亮点
AST在LLaMA2-7B模型上取得了显著的性能提升。与密集模型相比，AST将2:4半结构化稀疏模型的困惑度差距降低到0.6，零样本准确率差距降低到1.16%。更重要的是，AST仅使用了不到0.4%的预训练tokens和GPU时间，表明其具有极高的训练效率。这些结果证明了AST在压缩LLM方面的有效性和实用性。
🎯 应用场景
该研究成果可应用于各种需要部署大型语言模型的场景，例如移动设备、边缘计算设备等资源受限的环境。通过AST框架，可以将LLM压缩到更小的尺寸，降低计算和存储成本，同时保持较高的性能。这有助于推动LLM在更广泛的领域得到应用，例如智能助手、自然语言处理、机器翻译等。
📄 摘要（原文）
The remarkable success of Large Language Models (LLMs) relies heavily on their substantial scale, which poses significant challenges during model deployment in terms of latency and memory consumption. Recently, numerous studies have attempted to compress LLMs using one-shot pruning methods. However, these methods often suffer from considerable performance degradation on complex language understanding tasks, raising concerns about the feasibility of pruning in LLMs. To address this issue, we propose Adaptive Sparse Trainer (AST), a novel and efficient retraining framework tailored for semi-structured sparse models. AST enables models to learn optimal masks during the weight update process without incurring additional computational overhead. Furthermore, we demonstrate that incorporating knowledge distillation significantly improves retraining efficiency and enhances model performance under fixed computational constraints. Additionally, a supplementary set of well-initialized parameters is integrated to further augment the model's efficacy. AST achieves state-of-the-art performance with minimal training cost. When applied to the LLaMA2-7B model, AST reduces the perplexity and zero-shot accuracy gap between dense and 2:4 semi-structured sparse models to 0.6 and 1.16%, respectively, utilizing less than 0.4% of the pretraining tokens and GPU hours. Our work demonstrates the feasibility of deploying semi-structured sparse LLMs and offers a promising alternative for achieving highly compressed models when combined with existing quantization techniques.