FlexiGPT: Pruning and Extending Large Language Models with Low-Rank Weight Sharing

📄 arXiv: 2501.14713v2 📥 PDF

作者: James Seale Smith, Chi-Heng Lin, Shikhar Tuli, Haris Jeelani, Shangqian Gao, Yilin Shen, Hongxia Jin, Yen-Chang Hsu

分类: cs.CL, cs.LG

发布日期: 2025-01-24 (更新: 2025-01-31)

备注: Accepted to NAACL 2025 - Main Conference

💡 一句话要点

FlexiGPT：通过低秩权重共享剪枝和扩展大型语言模型

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

关键词: 大型语言模型 模型剪枝 权重共享 低秩适配器 模型压缩 模型扩展 资源受限设备

📋 核心要点

  1. 现有大型语言模型部署在资源受限设备上时，面临着模型体积大、计算成本高的挑战。
  2. FlexiGPT通过选择性剪枝模型块，并利用权重共享和低秩适配器进行高效替换，实现了模型压缩。
  3. 实验结果表明，该方法在多种基准测试中显著优于现有方法，并在模型压缩和扩展方面均有提升。

📝 摘要（中文）

大型语言模型（LLM）在自然语言处理（NLP）领域的快速普及，对在内存受限设备上高效部署提出了迫切需求，同时还要保证性能。本文提出了一种剪枝LLM的方法，该方法基于重要性评分选择性地剪枝模型块，并用低参数替换策略替换它们。具体来说，我们提出了一种原则性的度量标准，使用权重共享机制（利用模型中未剪枝的对应部分）和块特定的低秩适配器来替换每个剪枝块。此外，我们通过输出特征归一化和基于低秩SVD重建的适配器初始化方案来促进这些替换块的学习。经验评估表明，与现有方法相比，该方法取得了显著的性能提升，在30%压缩率下，在6个基准测试中的5个上实现了最先进的性能，在40%压缩率下，在所有6个基准测试上实现了最先进的性能。我们还证明了我们的方法可以扩展较小的模型，仅使用约0.3%的扩展训练token和最小的额外参数成本，即可在所有6个基准测试中提高性能。

🔬 方法详解

问题定义：现有的大型语言模型参数量巨大，难以在内存受限的设备上部署。传统的模型压缩方法，如剪枝和量化，可能会导致显著的性能下降。因此，如何在保证性能的前提下，有效地压缩和扩展大型语言模型是一个关键问题。

核心思路：FlexiGPT的核心思路是选择性地剪枝模型中不重要的块，并使用低参数的替换策略来弥补剪枝带来的性能损失。通过权重共享机制，利用模型中未剪枝的块的信息，并结合低秩适配器，可以有效地学习替换块，从而在压缩模型的同时保持甚至提升性能。

技术框架：FlexiGPT的整体框架包括以下几个主要阶段：1) 重要性评估：对模型中的每个块进行重要性评估，确定需要剪枝的块。2) 块剪枝：根据重要性评估结果，选择性地剪枝模型块。3) 替换块学习：使用权重共享机制和低秩适配器来学习替换被剪枝的块。4) 微调：对整个模型进行微调，以进一步提升性能。

关键创新：FlexiGPT的关键创新在于：1) 提出了一种原则性的重要性评估指标，用于选择性地剪枝模型块。2) 引入了权重共享机制，利用模型中未剪枝的块的信息来学习替换块。3) 结合了低秩适配器，进一步降低了替换块的参数量。4) 提出了输出特征归一化和基于低秩SVD重建的适配器初始化方案，加速了替换块的学习。

关键设计：在重要性评估方面，论文提出了一种基于梯度信息的指标。在权重共享方面，论文选择与被剪枝块具有相似功能的未剪枝块进行权重共享。在低秩适配器方面，论文使用了 bottleneck 结构，降低了适配器的参数量。在损失函数方面，论文使用了交叉熵损失函数和正则化项，以防止过拟合。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

FlexiGPT在多个基准测试中取得了显著的性能提升。在30%压缩率下，在6个基准测试中的5个上实现了最先进的性能，在40%压缩率下，在所有6个基准测试上实现了最先进的性能。此外，该方法还可以通过少量额外训练数据提升小型模型的性能，例如，仅使用约0.3%的扩展训练token，即可在所有6个基准测试中提高性能。

🎯 应用场景

FlexiGPT具有广泛的应用前景，包括在移动设备、嵌入式系统等资源受限的平台上部署大型语言模型，以及加速模型推理和降低计算成本。该方法还可以用于模型蒸馏，将大型模型的知识迁移到小型模型中。此外，FlexiGPT还可以用于模型扩展，通过少量额外训练数据提升小型模型的性能。

📄 摘要（原文）

The rapid proliferation of large language models (LLMs) in natural language processing (NLP) has created a critical need for techniques that enable efficient deployment on memory-constrained devices without compromising performance. We present a method to prune LLMs that selectively prunes model blocks based on an importance score and replaces them with a low-parameter replacement strategy. Specifically, we propose a principled metric to replace each pruned block using a weight-sharing mechanism that leverages unpruned counterparts from the model and block-specific low-rank adapters. Furthermore, we facilitate the learning of these replacement blocks with output feature normalization and an adapter initialization scheme built on low-rank SVD reconstructions. Empirical evaluations demonstrate substantial performance gains over existing methods, achieving state-of-the-art performance on 5/6 benchmarks for a compression rate of 30% and 6/6 benchmarks for a compression rate of 40%. We also demonstrate that our approach can extend smaller models, boosting performance on 6/6 benchmarks using only ~0.3% tokens of extended training with minimal additional parameter costs.