Fast and Effective Weight Update for Pruned Large Language Models

📄 arXiv: 2401.02938v2 📥 PDF

作者: Vladimír Boža

分类: cs.CL, cs.LG

发布日期: 2024-01-01 (更新: 2024-07-22)

🔗 代码/项目: GITHUB

💡 一句话要点

提出快速有效的权重更新算法以解决大语言模型剪枝问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 模型剪枝 权重更新 交替方向乘子法 微调技术

📋 核心要点

  1. 现有剪枝方法在微调阶段面临高昂的计算成本，难以恢复模型性能。
  2. 本文提出了一种基于交替方向乘子法的快速权重更新算法，结合逐步剪枝掩码选择。
  3. 实验结果表明，所提方法在多种大型语言模型上实现了最先进的剪枝性能。

📝 摘要（中文）

剪枝大型语言模型（LLMs）是一项具有挑战性的任务，主要难点在于剪枝后模型的微调，以恢复因权重丢失而导致的性能下降。现有方法要么完全忽视微调，专注于高效的剪枝标准，要么尝试逐层权重更新，保持每层的行为。然而，即使是逐层权重更新，对于LLMs来说也可能代价高昂，之前的研究已采用各种近似方法。本文提出了一种基于交替方向乘子法（ADMM）的快速有效的剪枝层权重更新算法，并通过简单的逐步剪枝掩码选择进行扩展，在广泛的LLMs上实现了最先进的剪枝性能。代码可在https://github.com/fmfi-compbio/admm-pruning获取。

🔬 方法详解

问题定义：本文旨在解决大型语言模型剪枝后微调的高成本问题。现有方法在微调阶段往往忽视性能恢复，或采用逐层更新但仍面临计算开销。

核心思路：提出了一种基于交替方向乘子法（ADMM）的权重更新算法，旨在通过高效的权重调整来恢复剪枝后的模型性能，同时引入逐步剪枝掩码选择以优化剪枝过程。

技术框架：整体框架包括权重剪枝、ADMM权重更新和逐步掩码选择三个主要模块。首先进行剪枝，然后应用ADMM算法进行权重更新，最后通过逐步选择掩码来优化剪枝效果。

关键创新：最重要的创新在于结合ADMM算法与逐步剪枝掩码选择，显著提高了剪枝后的模型性能恢复效率，与传统方法相比，降低了计算成本。

关键设计：在算法设计中，设置了适当的超参数以平衡剪枝率与性能恢复，损失函数则考虑了剪枝前后的性能差异，确保更新过程的有效性。具体的网络结构和参数设置在实验部分进行了详细描述。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果显示，所提算法在多个大型语言模型上实现了最先进的剪枝性能，相较于基线方法，性能恢复效果提升了显著的百分比，验证了其有效性和高效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过有效的剪枝技术，可以在保持模型性能的同时，显著减少计算资源消耗，提升模型在实际应用中的可用性和效率。未来，该方法有望推动更大规模语言模型的应用与发展。

📄 摘要（原文）

Pruning large language models (LLMs) is a challenging task due to their enormous size. The primary difficulty is fine-tuning the model after pruning, which is needed to recover the lost performance caused by dropping weights. Recent approaches have either ignored fine-tuning entirely, focusing on efficient pruning criteria, or attempted layer-wise weight updates, preserving the behavior of each layer. However, even layer-wise weight updates can be costly for LLMs, and previous works have resorted to various approximations. In our paper, we propose a fast and effective weight update algorithm for pruned layers based on the Alternating Direction Method of Multipliers (ADMM). We further extend it with a simple gradual pruning mask selection and achieve state-of-the-art pruning performance across a wide range of LLMs. Code is available at https://github.com/fmfi-compbio/admm-pruning.