发表:EMNLP_FINDING_2024
机构:Shanghai Jiao Tong University
连接:LaCo: Large Language Model Pruning via Layer Collapse - ACL Anthology
代码:https://github.com/yangyifei729/LaCo
Abstract
基于 Transformer 的大语言模型(LLMs)正呈现出明显的规模扩张趋势,这导致了训练和推理阶段的成本大幅上升。然而,现有的方法如模型量化、知识蒸馏和模型剪枝各自存在限制,包括硬件支持受限、对大量训练的依赖以及对模型内部结构的修改等问题。
本文提出了一种简洁的逐层结构化剪枝方法,称为 层折叠(Layer Collapse, LaCo)。该方法通过将模型的后续层“折叠”进前面的某一层,实现了在保留模型结构的同时快速压缩模型体积。
全面的实验表明,在剪枝率为 25%–30% 的情况下,LaCo 能保持平均任务性能在 80% 以上,显著优于当前最先进的结构化剪枝方法。我们还进行了一系列剪枝后的微调实验,验证 LaCo 能有效继承原始模型的参数。此外,我们还在不同设置下进行了消融实验。
最后,我们从层间相似性的角度出发,讨论了提出该方法的动机,并评估了 LaCo 在不同剪枝比例下的表现。
Introduction
近年来,基于 Transformer 架构的大语言模型(LLMs)(Vaswani et al., 2017)在各类任务中展现出令人瞩目的能力。然而,当前模型发展的趋势是不断向更大规模演进,这对计算资源提出了极高的要求。
为缓解上述挑战,已有多种方法被提出,以降低模型在训练和推理阶段的成本,或从 LLM 中派生出一个更小的模型,包括模型量化(Dettmers et al., 2022;Yao et al., 2022;Xiao et al., 2023)、知识蒸馏(Liu et al., 2022;Hsieh et al., 2023;Shridhar et al., 2023)以及模型剪枝(Zhang et al., 2022;Frantar and Alistarh, 2023;Ma et al., 2023)。然而,现有解决方案都存在一些显著的缺陷。模型量化通常需要特定的硬件支持,并且往往会影响模型性能;知识蒸馏通常需要对小模型进行再训练,这不仅成本高昂,而且任务相关性较强;至于模型剪枝,无论是非结构化的还是结构化的,也都有各自的问题。非结构化剪枝通常会带来模型稀疏化,这通常导致一定程度的性能损失,同时仍依赖硬件支持;结构化剪枝则意味着移除模型中的特定模块,通常会改变模型结构,从而降低模型的可移植性。
考虑到上述问题,我们设想采用一种新的方式直接对模型进行剪枝:从一个已经训练好的大模型中直接剪去某些层,并用一个层的参数去替代多个层的参数,从而实现有效的模型压缩。
具体而言,我们观察到,将某些层的参数差异与其后续层合并,通常并不会显著影响模型性能,如图 1 所示。
我们将这一操作称为“保留差异并寻找共性”(Reserving-Differences-while-Seeking-Common,简称 RDSC)层合并,它融合了参数差分与合并的思想。在这一发现的基础上,我们提出了一种简洁而强大的逐层剪枝器,称为层折叠(Layer Collapse, LaCo),其核心思想是将后续层“折叠”进前面的某一层,以尽可能保留模型的输出表示。
本文的主要贡献如下:
层折叠方法可在不依赖训练的情况下直接移除 30%–50% 的模型层,同时保持模型性能。在多个基准任务上的实验结果表明,在相同剪枝率下,我们的方法优于现有最先进的结构化剪枝方法。
层折叠方法保持了 LLM 的内部结构,例如中间维度,因此剪枝后的模型可以无缝集成到现有应用中,无需对系统实现进行修改。
我们通过微调实验验证了层折叠能够有效继承原始模型的参数,仅需极少的训练即可将剪枝模型恢复至原始模型的损失收敛水平。此外,我们还从动机出发,讨论了层间相似性的作用,并在不同剪枝比例下评估了 LaCo 剪枝模型的性能,并进行了一系列消融实验。
2 Method
2.1 Reserving-Differences-while-Seeking-Common Layer Merge
2.2 Layer Collapse
2.3 Complexity Analysis
3 Experiments
3.1 Models
为了评估所提出的 LaCo 方法的有效性,我们在多个流行的英文大语言模型(LLMs)上进行了实验,包括 Llama2-7B 和 Llama2-13B(Touvron 等,2023)。此外,我们还验证了该方法在双语 LLM 上的有效性,具体为 Baichuan2-7B 和 Baichuan2-13B(Yang 等,2023),这些模型支持中文和英文。我们使用这些模型的 base(基础)版本。上述所有模型均基于 Transformer 架构,属于 decoder-only 模型。
3.2 Benchmarks
为全面评估剪枝模型的能力,我们采用 OpenCompass 评估框架(Contributors,2023)。具体而言,依据 OpenCompass 的任务分类,我们从以下五个方面开展评估:推理能力(Reasoning)、语言能力(Language)、知识能力(Knowledge)、考试能力(Examination)、理解能力(Understanding)
我们从每一类中选取若干代表性基准任务进行评测:
Reasoning:CMNLI(Xu 等,2020)、HellaSwag(HeSw)(Zellers 等,2019)、PIQA(Bisk 等,2019)
Language:CHID(Zheng 等,2019)、WSC(Levesque 等,2012)
Knowledge:CommonSenseQA(CSQA)(Talmor 等,2018)、BoolQ(Clark 等,2019)
Examination:MMLU(Hendrycks 等,2021)、CMMLU(Li 等,2023)
Understanding:Race-High / Middle(Lai 等,2017)、XSum(Narayan 等,2018)、C3(Sun 等,2020)
我们使用 OpenCompass 的官方脚本进行评测,所有实验均为 zero-shot 或 few-shot 设定,不使用额外训练数据。评估采用两种模式:
PPL 模式(困惑度 Perplexity)
GEN 模式(文本生成 Generation)1
具体设置如下:
CHID 与 XSum 使用 GEN 模式;
WSC 数据集分别使用 PPL(WSCP)和 GEN(WSCG)两种模式;
其他任务使用 PPL 模式评估。
每个基准任务的得分会通过 OpenCompass 转换为标准化分值,分数越高代表性能越好。虽然 OpenCompass 提供了 Baichuan2 和 Llama2 系列模型的官方评测结果,但为避免由于硬件和软件环境差异以及官方结果可能存在的误差,我们全部自行复现评测结果以确保公平性。
3.3 Baselines
由于 LaCo 属于结构化剪枝方法(即直接从 LLM 中移除组成部分),我们选择了两个当前最先进的结构化剪枝方法作为对比基线:
LLM-Pruner(LLMPru.)(Ma 等,2023)
SliceGPT(Ashkboos 等,2024)
这两种方法的剪枝效果已超越之前最先进的稀疏剪枝方法 SparseGPT(Frantar 和 Alistarh,2023)。
在实验中,为了确保公平性,我们将两种基线方法的剪枝比例设置为与 LaCo 相同或略低。
3.4 Settings
由于现有工作大多将剪枝比例设置在 30% 以下,我们通过启发式方式调整超参数,使模型的剪枝比例尽可能接近 30%。具体的超参数设定可见附录 A 表 8。
我们从英文和中文维基百科数据集中分别随机选取了 5 条句子作为 Baichuan2 系列模型的 few-shot 校准样本,对于 Llama2 系列模型则选取了 10 条英文句子作为校准样本。所有实验均在配备 8 张 Nvidia A100 80GB 显存的服务器上进行。
3.5 Main Results
在表 1 中,我们展示了在不同剪枝方法下,4 个 LLM 模型在多个基准任务上的评测结果:
“Dense”表示 OpenCompass 官方公布的未经剪枝的 LLM 分数;
“Dense∗”表示我们自行复现的 Dense 分数;
“LLMPru.”和 “SliceGPT”分别对应两个基线方法;
“Ratio”表示剪枝比例,即被剪掉的参数数量占总参数量的比例;
“Lay.”表示模型的总层数。
将 Dense 与 Dense∗ 对比可以发现,二者差异不大,大多数任务的偏差都在 5% 以内,说明我们的实验设置没有引入明显误差。因此,后续的性能分析均以 Dense∗ 作为参考基准,以确保公平性。
将 LaCo 与两种基线方法进行对比,从表 1 可以观察到,尽管 LaCo 的剪枝比例略高,但它在大多数任务上都取得了最好的结果。
为更直观地展示表 1 中的数据,我们在表 2 中给出了每种剪枝器在所有基准任务下的平均分(Avg.)、各任务类别下的平均分(Reas., Lan., Know., Exam., Unde.),以及相较于 Dense∗ 的平均性能占比(Per.)。
整体来看,LaCo 的平均得分显著高于两种基线方法,并且在五个评估类别中有四个类别的得分优于基线方法。尽管在 Reasoning 类别中稍有下降,但表现仍然具有可比性。
此外,从所有数据集的平均性能百分比来看,LaCo 的剪枝模型表现远优于基线方法。在四个模型中,有三个模型的平均性能超过 Dense∗ 的 80%,最低的也达到了 73%(Baichuan2-7B),而两个基线方法在所有模型上都未能超过 70%。
为了进一步证明 LaCo 剪枝模型的稳定性,我们计算了各任务上的性能百分比(相对 Dense∗),如附录 D.4 表 22 所示。LaCo 剪枝模型在大多数基准任务上都保持在 70% 以上的性能,并且没有出现模型性能“崩溃”的情况,即没有任何一个任务的得分低于 30%。
特别值得一提的是,对于通过 GEN 模式评估的三个任务(CHID、XSUM 和 WSCG),LaCo 剪枝后的模型仍能保持稳定性能,而基线方法剪枝后的模型在这些任务中表现很差,甚至多项结果为 0.00。GEN 模式依赖模型生成的文本进行评分,基线方法剪枝后的模型容易输出无意义的重复文本。
在附录 D.5 表 23 中,我们展示了一个来自 XSum 的示例:LaCo 剪枝的 Llama2-7B 能生成正常句子,而基线方法剪枝后的模型输出重复无效内容。
此外,我们还在更大规模的 Llama2-70B 模型上进行了部分基准测试,结果见附录 D.1 表 19。实验结果显示,LaCo 在大模型上仍然优于基线方法。
总结而言,LaCo 是一种性能优异、稳定性强的结构化剪枝方法。它完全基于参数差分与加和操作实现,不会改变模型内部结构,因而是一种简洁高效的剪枝解决方案。
3.6 Comparison of Perplexity
困惑度(PPL)也是衡量语言模型性能的常用指标之一,我们希望进一步比较 LaCo 与基线方法在保持模型困惑度方面的差异。
我们以 Llama2-7B 模型为例,在 27% 剪枝率下,分别使用不同剪枝器进行评估。评估样本为从 Wikipedia 中选取的 500 条句子,每条句子长度为 512 个 token。
表 3 显示了不同剪枝方法下的 PPL 结果,结果表明:
LaCo 剪枝后的模型的 PPL 低于其他基线方法,进一步突出显示了 LaCo 的优势。
3.7 Pruning Time
为了验证 LaCo 的时间复杂度较低且剪枝速度快,我们以 Llama2-7B 在 27% 剪枝率下为例,在 单张 A100 GPU 上对 LaCo 与两种基线方法的剪枝耗时进行比较。
为了保证公平性,我们仅统计剪枝的核心过程时间,不包含加载模型、加载数据或保存模型的耗时。
表 4 中的结果表明:
LaCo 的剪枝速度明显快于两个基线方法,尤其远快于 SliceGPT。
4 Further Analysis
4.1 Post-training and Re-pruning
4.1.1 Post-training
由于剪枝不可避免地带来性能损失,我们进一步探讨使用 LaCo 剪枝后的模型是否能够有效继承原始模型的参数,并通过对全部参数进行后训练迅速恢复性能。具体而言,我们选取了主实验中通过 LaCo 剪枝得到的 Llama2-7B 和 Baichuan2-7B 模型,并对它们进行后训练。在对剪枝后的 Llama2-7B 模型进行训练时,我们使用了大约 10 亿个英文数据集中的 token;而对剪枝后的 Baichuan2-7B 模型,我们使用了约 12.5 亿个 token,其中 50% 来自英文,50% 来自中文数据集。更详细的实现细节可见附录 C。
图 3 展示了训练过程中的损失曲线。可以观察到,这两个模型在训练过程中都迅速收敛,在大约 250 步之后损失迅速下降,随后趋于稳定。两个剪枝模型的参数量都约为 50 亿,最终收敛的损失值分别约为 1.6 和 2.0,与技术报告中 Llama2-7B 和 Baichuan2-7B 所报告的损失值(分别为 1.75 和 1.9)非常接近。在 4 张 Nvidia A100 80GB 显卡上,对剪枝后的 Llama2-7B 和 Baichuan2-7B 模型进行后训练的耗时分别约为 28 小时和 35 小时。相比之下,从零开始训练一个 50 亿参数的大语言模型通常需要至少 5000 亿个 token,使用上百张 A100 显卡并耗时数月。而我们只用了原始成本的千分之一,就得到了一个已收敛、参数规模相当的模型。这表明,被剪枝的模型确实有效地继承了原始模型的参数,仅通过极小成本的后训练就能迅速恢复性能并实现收敛。
我们还对后训练后的模型在多个基准任务上进行了评估,详细结果见附录 E 的表 24。各个类别的平均分以及整体平均得分汇总如表 6 所示。从表中可以明显看出,剪枝后的 Llama2-7B 经过后训练后在各个基准任务上的性能均有显著提升。然而,剪枝后的 Baichuan2-7B 在后训练后的表现则较为复杂,部分基准任务表现有所提升,而其他任务则有所下降,其整体得分也出现了轻微下降。我们推测,这可能是因为 Baichuan2-7B 的预训练数据来源更加多样化,导致其数据分布与我们用于后训练的数据之间存在差异,从而削弱了后训练的效果。然而,Llama2-7B 剪枝模型在后训练后的得分提升非常一致,进一步证明了使用 LaCo 剪枝的模型确实能够有效继承参数,并通过低成本的后训练恢复性能。
LaCo 通过后训练实现了优异的表现,这促使我们进一步在相同的训练数据上,将其效果与当前最先进的剪枝方法 LLM-Pruner 进行了对比。我们的对比结果见附录 D.3 表 21,显示 LaCo 剪枝后的模型在后训练后的性能优于 LLM-Pruner 剪枝模型,同时 LaCo 所消耗的训练资源也显著更少。
4.2 Layer-wise Similarity
本节讨论了我们进行相邻层合并的动机。我们最主要的动机来自于观察到:在大语言模型(LLMs)中,相邻层之间的参数变化和输出表示的差异并不显著。
在图 4 中,我们展示了 Llama2-7B 和 Baichuan2-7B 模型中,每一层与其下一层对应的自注意力(SAN)中 q、k、v 矩阵,以及多层感知机(MLP)中的升维和降维矩阵之间的 L2 相似度。结果表明,相邻层之间对应矩阵的最大 L2 值通常不超过 200。考虑到 MLP 升维矩阵(11008×4096)和 SAN 的 q、k、v 矩阵(4096×4096)本身规模非常大,这意味着相邻层之间每个元素的变化是极小的。
在图 5(a)中,我们从 Wikipedia 中随机选取了 20 条句子,并计算了相邻层输出的隐藏状态之间的余弦相似度。结果显示,无论是 Baichuan2-7B 还是 Llama2-7B,从第 3 层到第 28 层,相邻层之间的表示相似度通常都非常接近于 1。这种在参数和表示上的高度相似性促使我们思考:是否可以用一个层来替代多个后续层。
此外,参数之间的高相似性也提示我们应该关注层与层之间的差异。受先前模型融合相关工作的启发(Yu 等,2023;Matena 和 Raffel,2022),我们提出收集相邻层之间的参数差异,并将这些差异合并进单个层中。
为验证 RDSC 层合并方法是否真的可以用一个层替代多个层,我们设计了一个实验:在第 10 层到第 19 层之间,我们将每连续 4 层合并为一个层,并计算该合并层的输出与原始最后一层输出之间的余弦相似度。如图 5(b)所示,在 4096 维向量上,最小的余弦相似度也超过了 0.996,这验证了 RDSC 层合并在保留表示能力方面的有效性。
4.1.2 Re-pruning
既然在对剪去了约 25%–30% 参数的大语言模型(LLM)进行后训练后可以部分恢复其性能,那么我们进一步提出一个问题:是否可以在后训练完成后,继续对模型进行剪枝,从而得到一个仅保留约 50% 参数的模型,同时仍保持较好的性能。
为此,我们使用 LaCo 方法对先前已经经过后训练的、剪枝后的 Llama2-7B 模型再次进行剪枝,最终得到一个包含 17 层、仅保留原始 Llama2-7B 模型 55% 参数的模型。
我们对这一再剪枝后的模型进行了评估。详细结果见附录 E 的表 24,平均结果见表 6。
从表格中可以看出,即使只保留了 55% 的参数,该模型依然保留了大约 70% 的原始 7B 模型性能。不过,我们当前用于后训练的数据在质量和规模上仍然有限。若能使用更大规模、更高质量的训练数据,LaCo 在“剪枝 + 后训练 + 再剪枝”这一流程中,尤其是在大规模模型上,应能展现出更大的潜力与实用价值。
4.2 Layer-wise Similarity
本节讨论了我们进行相邻层合并的动机。我们最主要的动机来自于观察到:在大语言模型(LLMs)中,相邻层之间的参数变化和输出表示的差异并不显著。
在图 4 中,我们展示了 Llama2-7B 和 Baichuan2-7B 模型中,每一层与其下一层对应的自注意力(SAN)中 q、k、v 矩阵,以及多层感知机(MLP)中的升维和降维矩阵之间的 L2 相似度。结果表明,相邻层之间对应矩阵的最大 L2 值通常不超过 200。考虑到 MLP 升维矩阵(11008×4096)和 SAN 的 q、k、v 矩阵(4096×4096)本身规模非常大,这意味着相邻层之间每个元素的变化是极小的。
在图 5(a)中,我们从 Wikipedia 中随机选取了 20 条句子,并计算了相邻层输出的隐藏状态之间的余弦相似度。结果显示,无论是 Baichuan2-7B 还是 Llama2-7B,从第 3 层到第 28 层,相邻层之间的表示相似度通常都非常接近于 1。这种在参数和表示上的高度相似性促使我们思考:是否可以用一个层来替代多个后续层。
此外,参数之间的高相似性也提示我们应该关注层与层之间的差异。受先前模型融合相关工作的启发(Yu 等,2023;Matena 和 Raffel,2022),我们提出收集相邻层之间的参数差异,并将这些差异合并进单个层中。
为验证 RDSC 层合并方法是否真的可以用一个层替代多个层,我们设计了一个实验:在第 10 层到第 19 层之间,我们将每连续 4 层合并为一个层,并计算该合并层的输出与原始最后一层输出之间的余弦相似度。如图 5(b)所示,在 4096 维向量上,最小的余弦相似度也超过了 0.996,这验证了 RDSC 层合并在保留表示能力方面的有效性。
4.3 Varying Pruning Ratio
在本节中,我们探讨 LaCo 在不同剪枝比例下的性能表现。我们在 Llama2-7B 和 Llama2-13B 上进行了实验,通过设置不同的超参数(见附录 A 表 9),将剪枝比例控制在大约 10%、25%(即我们的主实验设置)以及约 50%。随后我们对这些剪枝后的模型进行了评估。
其平均评估结果展示于表 7,详细结果见附录 E 表 25。
随着剪枝比例的增加,模型整体性能呈下降趋势。然而,从剪枝比例约 10–15% 增加到约 25% 的过程中,性能并未出现明显下降,说明在此范围内我们的方法具有良好的稳定性。进一步地,即使在剪去了接近一半参数、剪枝比例接近 50% 的情况下,模型依然能够保持大约 70% 的性能,这表明我们的方法即使在极端压缩下也能有效避免模型崩溃(性能骤降至不可用)。
5 Related Work
**模型量化(Model Quantization)**通过将权重从高精度浮点数转换为低精度浮点数或整数,从而减小模型体积。SmoothQuant(Xiao 等,2023)对权重和激活值同时进行量化,并平滑激活值中的异常值。GPTQ(Frantar 等,2022)使用近似的二阶信息进行量化。QLoRA(Dettmers 等,2023a)通过一个冻结的 4-bit 量化模型反向传播梯度到低秩适配器(Low Rank Adapters)。OmniQuant(Shao 等,2023)则在多个量化参数上进行了优化。
**知识蒸馏(Knowledge Distillation)**指的是将一个大模型的知识迁移到一个小模型中。Distilling step-by-step(Hsieh 等,2023)训练出的较小模型甚至能超过原始 LLM 的表现。DISCO(Chen 等,2023)从大语言模型中蒸馏出反事实知识。SOCRATIC COT(Shridhar 等,2023)蒸馏了大语言模型中的“思维链(Chain-of-Thought)”能力。ZEPHYR(Tunstall 等,2023)则将蒸馏后的偏好优化(Direct Preference Optimization)应用于聊天模型的训练中。
**模型剪枝(Model Pruning)**是一种通过稀疏化或参数移除来提升模型效率的技术。**非结构化剪枝(Non-structured pruning)**通常涉及模型稀疏化。例如,SparseGPT(Frantar 和 Alistarh,2023)将剪枝问题转化为大规模稀疏回归问题;SpQR(Dettmers 等,2023b)在剪枝过程中识别并隔离异常权重。
**结构化剪枝(Structured pruning)**则主要是移除模型模块中的部分结构。LLM-Pruner(Ma 等,2023)基于梯度信息选择性地删除非关键结构。DirectShare(Cao 等,2024)通过激活共享机制提升推理效率。ShearedL-LaMA(Xia 等,2023)对 Llama2 进行有针对性的结构化剪枝并引入动态批加载策略。FLAP(An 等,2024)在移除权重后评估输出特征图的可恢复性,利用波动剪枝指标标准化重要性评分,进而自适应定义压缩模型的全局结构。Shortened LLaMA(Kim 等,2024)则证明:深度剪枝能有效压缩 LLM,且性能可匹敌甚至超越最新的宽度剪枝方法。
然而,模型量化和稀疏化通常需要专门的硬件支持,并且往往会影响性能。知识蒸馏代价高昂,且高度依赖具体任务。现有的结构化剪枝方法常常会破坏模型的原生结构。
相比之下,LaCo 保持了模型结构的完整性,方法更简洁,同时性能优异。尽管已有一些研究(如 Din 等,2023;Fan 等,2019;Belrose 等,2023)利用跳层(layer-skipping)或删层(layer-dropping)方式加速推理,但 LaCo 在本质上是不同的:它是第一个基于层折叠(layer collapse)的剪枝器,能够生成更小、更快、更节省内存、同时性能强劲的模型。而且,这些传统方法通常需要训练新的参数,以判断在推理时应跳过/移除哪些层;而 LaCo 完全不需要任何训练过程。
6 Conclusion
本文提出了一种简洁的逐层结构化剪枝方法,称为层折叠(Layer Collapse,LaCo),该方法通过将模型的后续层合并至前面的层,从而实现模型体积的快速压缩。LaCo 不依赖任何特殊硬件支持,且能够保留模型的原生结构。
实验结果表明,LaCo 在性能上显著优于当前最先进的结构化剪枝方法,并揭示了现有大语言模型中可能存在的参数冗余。
我们还对 LaCo 的多个设置进行了消融实验,并对剪枝后的模型进行了后训练,验证了 LaCo 能有效继承原始模型的参数。此外,我们从层间相似性的角度出发,阐述了提出该方法的动机,并探讨了在不同剪枝比例下,LaCo 剪枝模型的性能表现。
Limitations
由于 LaCo 的剪枝过程主要依赖于逐层迭代,它无法像以往的方法那样直接控制剪枝比例。相反,它需要通过调整一些超参数(如表示相似度阈值 TT)来间接控制剪枝程度。我们计划在未来的工作中,总结更多实验规律,用于指导如何设置这些超参数以实现特定的剪枝比例。
我们的方法灵感来自当前的模型融合技术,但与现有的一些基线方法(如 LLM-Pruner(Ma 等,2023)和 SliceGPT(Ashkboos 等,2024))一样,我们的方法目前尚缺乏完整的理论证明,这也是我们未来工作的一个方向。
此外,尽管实验结果已经表明 LaCo 当前的层合并方法是有效的,但仍有可能存在更优的合并方式。因此,我们未来还将继续探索更优的层合并策略。