[论文解读] Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods

Bo-Kyeong Kim, Geonmin Kim+5 more|arXiv (Cornell University)|Feb 5, 2024

Natural Language Processing Techniques被引用 8

一句话总结

通过 LoRA 重训练并移除 Transformer 块的深度裁剪在零样本任务中可以与宽度裁剪相媲美，并在内存受限的小批量条件下实现更快的推理。

ABSTRACT

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

研究动机与目标

• 在内存受限、小批量设置下，提升大模型推理效率的动机。
• 提出一种简单的深度裁剪方法，移除整个 Transformer 块，同时保持其余权重形状不变。
• 在公开的大语言模型（LLaMA-7B 和 Vicuna-7B/13B）上，将深度裁剪与宽度裁剪基线（例如 Wanda-sp、FLAP、LLM-Pruner）进行对比评估。
• 证明深度裁剪结合 LoRA 重训练在提高生成速度的同时，能够实现具有竞争力的零样本任务性能。

提出的方法

• 将 Transformer 块视为可裁剪单元，以降低推理时延。
• 使用 Mag、Taylor 和基于 PPL 的标准评估块的重要性，其中 Taylor+ 与 PPL 被选定用于裁剪决策。
• 执行一次性裁剪，移除最不重要的块以达到目标参数数量（保持前 4 个和后 2 个块完好）。
• 在校准数据集上用 LoRA（低秩自适应）高效地对裁剪后的模型进行再训练，从而实现快速恢复性能。
• 在零样本任务中将深度裁剪与宽度裁剪基线进行比较，并在小批量条件下测量延迟、吞吐量和内存使用情况。

实验结果

研究问题

• RQ1 简单的 Transformer 块深度裁剪是否能够在大语言模型的零样本性能上媲美或超过宽度裁剪？
• RQ2 在内存受限的小批量场景下，深度裁剪模型是否在自回归生成中提供实际的加速？
• RQ3 哪些块级重要性标准和裁剪粒度能在准确性和效率之间给出最佳权衡？
• RQ4 一次性深度裁剪结合 LoRA 重训练在实践中是否具有与迭代裁剪方法竞争的表现？

主要发现

• 在小批量条件下，深度裁剪使生成速度快于原始模型。
• 在使用 LoRA 重训练时，深度裁剪在零样本任务性能上与宽度裁剪基线（Wanda-sp、FLAP、LLM-Pruner）相竞争。
• Taylor+ 提高了常识推理准确性，而 PPL 提高了生成质量。
• 深度裁剪模型降低了 GPU 内存需求，在硬件受限设置下可处理比未裁剪模型更大的批量或更长的输出。
• 一次性裁剪配合 LoRA 重训练可以接近迭代裁剪的性能，从而实现高效部署。
• 在更大规模时，裁剪整个 Transformer 块通常比裁剪单独的 MHA/FFN 模块获得更好的结果；对于较小的模型，块裁剪仍然具有优势。