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

📄 arXiv: 2407.05690v1 📥 PDF

作者: Bowen Shen, Zheng Lin, Daren Zha, Wei Liu, Jian Luan, Bin Wang, Weiping Wang

分类: cs.CL, cs.AI

发布日期: 2024-07-08

备注: Findings of ACL 2024

💡 一句话要点

TransAct：通过模块内低秩架构剪枝LLM，降低激活值冗余

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

关键词: 大型语言模型 结构化剪枝 低秩分解 激活值分析 模型压缩

📋 核心要点

  1. 现有结构化剪枝方法对高度互连的LLM造成较大损害，难以实现高压缩比。
  2. TransAct通过降低MHA和MLP模块内部的过渡激活，将LLM剪枝为模块内低秩架构。
  3. 实验表明，TransAct在LLaMA模型上实现了高压缩比，并在下游任务中保持了良好的性能。

📝 摘要（中文）

本文提出了一种与任务无关的结构化剪枝方法TransAct，并结合了紧凑的Transformer架构设计，旨在降低大型语言模型（LLM）的计算和内存开销。TransAct减少了多头注意力（MHA）和多层感知机（MLP）模块内部的过渡激活，同时保留了对扰动敏感的模块间激活。通过这种方式，LLM被剪枝成模块内低秩架构，显著减少了权重、KV Cache和注意力计算。TransAct在LLaMA模型上实现，并在下游基准测试中进行了评估。结果验证了该方法在高压缩比下的效率和性能的优越性。消融研究揭示了激活引导迭代剪枝的优势，并对MHA和MLP模块的冗余性进行了实验分析。

🔬 方法详解

问题定义：大型语言模型（LLM）的结构化剪枝旨在降低计算和内存开销，使其能够部署在终端设备上。然而，由于LLM内部高度互连，粗粒度的结构化剪枝容易对模型造成较大损害，导致难以在保持性能的同时实现高压缩比。现有方法难以有效降低模型内部的冗余激活，从而限制了压缩效果。

核心思路：TransAct的核心思路是减少Transformer模块内部的过渡激活，同时保留模块间的激活。作者认为，模块间的激活对模型的性能至关重要，而模块内部的激活存在冗余。通过降低MHA和MLP模块内部的激活维度，可以将LLM转化为模块内低秩架构，从而减少权重数量、KV Cache大小和注意力计算量。这种方法旨在在不显著影响模型性能的前提下，实现更高的压缩比。

技术框架：TransAct的整体框架包括以下几个主要步骤：首先，对LLM进行激活值分析，确定MHA和MLP模块内部需要剪枝的激活维度。然后，采用激活引导的迭代剪枝策略，逐步降低这些激活维度的重要性。在剪枝过程中，模型会进行微调，以恢复因剪枝造成的性能损失。最终，得到一个模块内低秩的LLM，其权重数量、KV Cache大小和注意力计算量都显著降低。

关键创新：TransAct的关键创新在于其激活引导的结构化剪枝方法。与传统的基于权重或梯度的剪枝方法不同，TransAct直接针对激活值进行剪枝，能够更有效地降低模型内部的冗余。此外，TransAct还提出了一种紧凑的Transformer架构设计，进一步提高了模型的压缩率。

关键设计：TransAct的关键设计包括：1) 激活值分析方法，用于确定需要剪枝的激活维度；2) 激活引导的迭代剪枝策略，通过逐步降低激活维度的重要性来实现剪枝；3) 损失函数的设计，用于在剪枝过程中保持模型的性能；4) 模块内低秩架构的设计，通过降低MHA和MLP模块内部的激活维度来实现压缩。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TransAct在LLaMA模型上进行了实验，并在下游基准测试中取得了显著的性能提升。实验结果表明，TransAct能够在保持模型性能的同时，显著降低模型的权重数量、KV Cache大小和注意力计算量。例如，在某些任务上，TransAct可以将模型的参数量减少50%以上，而性能损失仅为1%左右。此外，消融研究还验证了激活引导迭代剪枝的有效性。

🎯 应用场景

TransAct具有广泛的应用前景，尤其是在资源受限的设备上部署大型语言模型。例如，可以将TransAct应用于移动设备、嵌入式系统和边缘计算设备，从而实现本地化的LLM推理。此外，TransAct还可以用于降低LLM的训练和部署成本，使其能够更广泛地应用于各种自然语言处理任务。

📄 摘要（原文）

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.