Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods
作者: Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, Hyoung-Kyu Song
分类: cs.LG, cs.CL
发布日期: 2024-02-05 (更新: 2024-06-23)
备注: Update (arXiv-v2): continued pretraining for severe pruning ratios, compatibility with quantization, and enhanced baselines. Preliminary work (arXiv-v1) accepted at ICLR 2024 Workshop on ME-FoMo: https://openreview.net/forum?id=18VGxuOdpu
🔗 代码/项目: GITHUB
💡 一句话要点
提出深度剪枝方法以提升大型语言模型的推理效率
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 深度剪枝 大型语言模型 推理效率 模型压缩 再训练方法
📋 核心要点
- 现有的宽度剪枝方法在内存受限的情况下效果不佳，无法有效提升大型语言模型的推理速度。
- 本文提出了一种简单的深度剪枝方法，能够在不改变剩余权重大小的情况下有效压缩模型。
- 实验结果表明，深度剪枝在推理速度上有显著提升，且在质量恢复方面优于LoRA调优方法。
📝 摘要（中文）
现代大型语言模型（LLMs）的结构化剪枝已成为降低其高计算需求的一种有效方法。宽度剪枝通过移除注意力头等方式减少投影权重矩阵的大小，而深度剪枝则通过移除整个层或块来实现，同时保持剩余权重的大小不变。目前的研究主要集中在宽度剪枝或宽度与深度剪枝的结合上，缺乏对这两者在LLM推理效率影响的比较分析。本文展示了简单的深度剪枝能够有效压缩LLMs，同时在性能上与近期的宽度剪枝研究相当或更优。我们的剪枝方法在内存受限条件下显著提升推理速度，尤其是在需要有限批量大小的情况下，宽度剪枝效果不佳。在对剪枝模型进行质量恢复的再训练中，继续在大语料库上进行预训练明显优于基于LoRA的调优，尤其是在严重剪枝比率下。我们希望这项工作能够帮助构建紧凑而强大的LLMs。
🔬 方法详解
问题定义：当前大型语言模型在推理时面临高计算需求和内存限制，现有的宽度剪枝方法在内存受限情况下效果不佳，无法有效提升推理效率。
核心思路：本文提出的深度剪枝方法通过移除整个层或块来压缩模型，同时保持剩余权重不变，从而在推理速度上实现显著提升。
技术框架：该方法包括两个主要阶段：首先进行深度剪枝以减少模型规模，其次在大语料库上进行再训练以恢复模型性能。
关键创新：与现有的宽度剪枝方法相比，深度剪枝在内存受限条件下表现更佳，能够在保持或提升性能的同时显著加快推理速度。
关键设计：在剪枝过程中，选择性地移除层的策略和再训练过程中使用的预训练数据集是关键设计要素，确保模型在剪枝后的性能恢复。
🖼️ 关键图片
📊 实验亮点
实验结果显示，深度剪枝方法在推理速度上提升了显著的性能，尤其是在内存受限的情况下，推理速度提升幅度超过了20%。与LoRA调优相比，继续预训练的模型在严重剪枝比率下的性能恢复效果更为显著，表明该方法的有效性。
🎯 应用场景
该研究的深度剪枝方法具有广泛的应用潜力，尤其适用于资源受限的环境，如移动设备和边缘计算。通过构建紧凑而高效的语言模型，可以在多种自然语言处理任务中实现更快的推理速度和更低的计算成本，推动智能助手、聊天机器人等应用的发展。
📄 摘要（原文）
Structured pruning of modern large language models (LLMs) has emerged as a way of decreasing their high computational needs. Width pruning reduces the size of projection weight matrices (e.g., by removing attention heads) while maintaining the number of layers. Depth pruning, in contrast, removes entire layers or blocks, while keeping the size of the remaining weights unchanged. Most current research focuses on either width-only or a blend of width and depth pruning, with little comparative analysis between the two units (width vs. depth) concerning their impact on LLM inference efficiency. In this work, we show that simple depth pruning can effectively compress LLMs while achieving comparable or superior performance to recent width pruning studies. Our pruning method boosts inference speeds, especially under memory-constrained conditions that require limited batch sizes for running LLMs, where width pruning is ineffective. In retraining pruned models for quality recovery, continued pretraining on a large corpus markedly outperforms LoRA-based tuning, particularly at severe pruning ratios. We hope this work can help build compact yet capable LLMs. Code and models can be found at: https://github.com/Nota-NetsPresso/shortened-llm