备注好友:方向-学校/公司-姓名/昵称
【AIGC 先锋科技】交流群
大语言模型(LLMs)在各项任务中取得了显著性能,但由于其深度、多层架构,计算成本高昂。层剪枝作为一种缓解这些低效的策略应运而生,但传统的静态剪枝方法忽略了LLM推理中两个关键的动态特性:
(1) 水平动态,即token Level 的异质性要求上下文感知的剪枝决策;
(2) 垂直动态,即MLP层和自注意力层的不同功能角色需要组件特定的剪枝策略。作者提出了SkipGPT,这是一个动态层剪枝框架,通过两个核心创新优化计算资源分配:
(1) 全局token感知路由,优先处理关键token;
(2) MLP和自注意力组件的解耦剪枝策略。为缓解训练不稳定性,作者提出了一种两阶段优化范式:首先,通过软参数化进行解耦训练阶段,学习路由策略以避免过早剪枝决策,然后进行参数高效的LoRA微调,以恢复因层移除而受影响的性能。
大量实验表明,SkipGPT减少了超过40%的模型参数,同时在各项基准测试中达到或超越了原始密集模型的性能。
通过协调动态效率与保留的表达能力,SkipGPT推动了可扩展、资源感知LLMs的实用部署。
作者的代码已公开发布在:https://github.com/EIT-NLP/SkipGPT。
- 引言
大语言模型(LLMs)基于逐层的Transformer架构,其中每一层由自注意力机制和多层感知器(MLP)组成(Vaswani等人,2017)。扩大模型规模已在广泛任务中推动了重大突破。然而,这一进展伴随着高昂的计算成本,需要大量资源进行推理。相比之下,人脑——尽管拥有100万亿个突触连接,远超最大的LLMs——仅需30瓦的功率即可高效运行。这种显著差异凸显了当前LLM架构中的根本性低效,揭示了人工智能与人类认知之间的差距。
由于其逐层顺序结构,即使拥有丰富的计算资源,大语言模型(LLMs)也难以充分利用并行性。这一局限性使得层剪枝成为加速推理和提高效率的关键策略。尽管现有的剪枝方法提供了一些改进,但它们通常忽略了剪枝动态的两个关键方面(见图1):
- 水平动态:输入序列中的不同token需要不同 Level 的计算量。当前方法要么为每层分配资源给最相关的top-
个token(Raposo等人,2024;Zeng等人,2023),要么在所有token上强制执行固定的计算比例(Jiang等人,2024)。这些僵化的方法无法适应token的复杂性,导致效率欠佳。为解决这一问题,作者引入全局Sparse机制,允许计算预算在整个前向传递过程中灵活分配,而非施加固定的逐层或逐token约束。 2. 垂直动态:每层中的MLP和自注意力组件具有不同的功能,但大多数剪枝方法将它们同等对待。研究表明,MLP类似于局部的神经过程,捕捉特定任务的交互,而注意力机制则类似于High-Level认知功能,管理上下文相关性和信息流。受人类大脑为不同任务激活不同区域的启发,作者提出将MLP和自注意力的剪枝解耦,以实现更具针对性和效率的计算减少。
为实现动态剪枝,近期一些研究探索了自适应计算方法,在每个Transformer层引入路由器。该路由器作为决策模块,用于在推理过程中决定是否执行特定网络单元。
然而,这些方法通常在联合训练范式下同时优化路由器和模型参数,类似于专家混合(MoE)。然而,该方法未能考虑到剪枝与预训练之间的根本差异——在剪枝场景中,路由器从随机初始化开始,而模型参数已通过大量预训练收敛到最优或局部最优分布。这种不匹配可能导致联合训练不稳定,并阻碍动态剪枝充分发挥其潜力。
在本工作中,作者首先提供了实证证据,证明了水平动态和垂直动态在大语言模型(LLMs)中的重要性。为了整合这两种动态,作者提出了SkipGPT,这是一种新颖的方法,能够在每个token的基础上动态地剪枝层,并根据每个token的复杂性调整剪枝过程。此外,SkipGPT将每一层中的MLP和自注意力组件解耦,从而能够更精细地控制哪些模型部分被剪枝,进而优化计算效率和模型性能。为了充分释放动态剪枝的潜力,作者引入了双阶段训练范式。
首先,在路由器微调阶段,作者冻结模型参数,仅优化路由器,使其能够在不破坏模型预训练知识的情况下识别最关键的计算。其次,在LoRA微调阶段,作者冻结路由器,并使用LoRA微调模型,以弥补剪枝可能导致的性能下降。作者的结果表明,SkipGPT是一种非常有效的剪枝策略——即使参数量减少了40%,剪枝后的模型也能完全恢复其性能,甚至超越原始模型。
此外,由于路由调整不会修改预训练模型的参数,因此它允许对原始模型中的模块重要性进行直接分析。通过详细的路由行为分析,作者揭示了两个关键见解:
(1) 注意力模块比MLP模块具有更高的冗余性。
(2) 随着上下文长度的增加,后续的token需要更多的注意力计算,但需要更少的MLP处理。这些发现突出了当前大型模型架构中的固有低效性,为未来的架构设计和推理优化提供了宝贵的见解。
- 动机
使用余弦相似度衡量模块重要性 近期研究表明,模块的输入与输出之间的余弦相似度可作为评估每个模块重要性的可靠指标(Men等人,2024;Gromov等人,2024)。值得注意的是,这些研究通常将“模块”定义为整个Transformer层。为了实现更细粒度的分析,作者将模块的定义细化为层内的MLP块或注意力块。
使用余弦相似度的基本假设是,冗余模块产生的输出与其输入高度相似,表明其转换程度较低。相反,重要模块会显著改变其输入,表明它们在模型中发挥着关键作用,应予以保留。先前方法通常对所有token的余弦相似度进行平均,以得出每个模块的总体重要性指标。然而,这种聚合可能会掩盖模块重要性在不同token和上下文中的变化。为了更好地理解这种变化,作者分析每个模块在token层面的余弦相似度。具体而言,第i个模块在token t处的余弦相似度
计算如下:
其中
表示L2范数,
表示在 Token
处模块
前的隐藏状态。为了说明模块重要性的变化,作者使用BookCorpus数据集(Zhu等人,2015年)中随机选择的一个句子进行案例研究。作者分析了LLaMA-2-7B(Touvron等人,2023a)中15个连续 Token 的余弦相似度分布,结果如图2所示。
垂直动态的必要性现有的剪枝方法,无论是动态还是静态,通常将整个Transformer层视为最小的剪枝单元。然而,如图2所示,作者发现同一层内,注意力模块和MLP模块的余弦相似度分布差异显著。
例如,在第二层中,注意力模块中大多数token的余弦相似度范围在0.1到0.4之间,而MLP模块主要分布在0.4到0.7的范围内。这表明对于该层,MLP模块几乎普遍比注意力模块冗余得多。但在最后一层,情况完全相反:所有注意力模块的余弦相似度值都在0.9到1之间,而MLP模块则完全处于不同的范围,从0到0.2,使得MLP模块在这个阶段远比注意力模块关键。这些发现表明,即使在同一层内,注意力模块和MLP模块之间的冗余程度也会随层的变化而变化和转移,这突显了将注意力模块和MLP模块解耦以实现更有效剪枝的必要性。
- SkipGPT:动态层剪枝
在确立了水平和垂直动态的必要性之后,作者现介绍SkipGPT,这是作者提出的用于动态层剪枝的框架。在本节中,作者首先概述理解SkipGPT优化过程所需的必要预备知识,随后解释其Sparse机制、路由实现、损失函数,最后介绍作者用于稳定有效学习的两阶段训练范式。
3.1. 预备知识:Gumbel-Softmax和STE
Gumbel-Softmax重参化
为优化SkipGPT的动态剪枝决策,作者将剪枝过程构建为一个离散优化问题,其中每个模块(MLP或注意力机制)根据其计算出的重要性被执行或跳过。然而,直接优化离散决策是不可微分的,这使得基于梯度的标准优化方法不可行。Gumbel-Softmax分布是对离散分布的连续近似,能够实现可微分的采样(Jang等人,2022)。这通过重参化实现,将离散样本转换为可微分的连续样本,以进行基于梯度的优化。令
表示
类离散分布的类别概率。为从这个分布中采样,Gumbel-Max技巧(Gumbel,1954;Maddison等人,2015)选择
值最大的类别,其中
是从Gumbel分布
中独立同分布抽取的样本。
π
由于最大值函数(arg max)不可微,Gumbel-Softmax重参化将其替换为softmax函数,从而生成连续样本,近似于分类分布:
其中
控制分布的尖锐程度。当
时,样本类似于 one-hot 向量,恢复原始的 arg max 操作。这种可微近似使得标准反向传播优化成为可能。
直通估计器直通(ST)估计器(Bengio等人,2013年)能够在保持可微性以支持反向传播的同时实现离散采样。在前向传播过程中,作者使用Gumbel-Softmax生成连续样本。为了将这些样本离散化,作者应用arg max,但在反向传播过程中,梯度计算方式如同使用连续近似。这是通过以下方式实现的:
其中
是连续的Gumbel-Softmax采样结果。这确保了输出为one-hot编码,同时梯度跟随
,使得尽管采样是离散的,优化过程依然高效。
3.2. Sparse性概念
为了控制总浮点运算次数(FLOPs),作者引入了Sparse性的概念,该概念定义为在正向传播中跳过的模块计算(注意力或MLP)的比例,相对于全连接Transformer中所有层和序列位置的总计算量。在llm-SkipGPT_2506中,Sparse性通过动态路由实现,其中每次正向传播仅选择一部分模块进行计算。具体而言,假设在正向传播中,一个包含
个模块的
层大语言模型(LLM)(每层包含一个注意力模块和一个MLP模块)处理一个长度为
的序列,全连接Transformer将计算
个模块,对应于Sparse度为0。当Sparse度大于0时,仅计算
个模块。
与以往工作不同,以往工作通过将计算限制在每个层的top-
token(Raposo等人,2024;Zeng等人,2023)或对每个token执行相同的Sparse性(Jiang等人,2024)来定义计算预算,作者允许计算负载在宽度(每个层参与计算的token数量)和高度(每个token参与计算的模块数量)之间动态分配。
3.3. 路由实现
为了实现动态分配,SkipGPT 在每个独立模块之前分配一个路由器。具体来说,作者将 token 路由到两条计算路径:(1)自注意力(用于注意力路由器)或 FFN 模块(用于 FFN 路由器),以及(2)一个残差连接。后者在计算上成本较高,其输出完全由输入决定,而前者则会产生较高的计算成本。
假设在
-th 转换模块
之前有一个 Token Embedding
,该模块可以是自注意力或 FFN 。在通过模块
之前,
首先经过一个路由函数,这是一个简单的线性投影,产生一个分类分布
,其中第一个元素表示跳过该模块的概率,第二个元素表示执行该模块的概率。
一旦获得这种分类分布,一种自然的路由策略是Top-1路由。具体来说,作者有:
以便梯度能够反向传播。然而,这种路由策略对精确Sparse性控制提出了挑战,因为
仅仅是二元选择的软近似。
通过利用Gumbel-Softmax和ST Gumbel估计器,作者有效地解决了这个问题。具体来说,在将Gumbel-Softmax和ST Gumbel估计器应用于
之后,作者得到一个二进制独热向量
。因此,输入到下一个模块的计算如下:
在正向传播过程中,采样离散的二进制值,以确保清晰的剪枝决策。在反向传播过程中,
软概率有助于梯度传播,使得路由权重能够得到有效更新。
基于此路由设计,原则上作者可以独立训练路由器而不改变模型参数,从而为给定的Sparse率获得最优路由解决方案。
3.4. 损失函数
根据第3.2节中的定义,Sparse度
为:
其中
表示token索引,
代表模块索引,
是LLM中总层数,
表示序列长度。为了满足不同的计算需求,作者引入了一个Sparse正则化项:
其中
表示绝对值,
是用户定义的目标Sparse度。整体损失函数则表示为:
其中
是标准语言建模损失,表示平均预测下一个token的负对数似然。超参数
控制整体损失函数中Sparse性惩罚的强度。2
3.5. 两阶段训练范式
路由器从随机初始化开始,而大语言模型(LLM)已经过预训练,直接训练可能导致不稳定且次优的剪枝决策。为解决此问题,作者提出了一种两阶段训练范式。
路由器调优在训练的初始阶段,作者专注于调整路由器,同时冻结所有其他模型参数。这一阶段非常高效,因为它只需要在每个模块之前添加一个轻量级的线性层,所有路由器参数仅占LLaMA2-7B总参数的0.007%。通过这一阶段,即使丢弃了25%的参数,作者仍然实现了模型原始性能的90%以上,如表1所示。路由器调优过程中的前向和后向传递如图3所示。
LoRA微调(可选)虽然路由调整阶段足以保持模型的大部分性能,但作者也为那些希望完全恢复至原始模型性能水平的用户提供了一个可选的LoRA微调阶段。低秩适配(LoRA)(Hu等人,2021)能够以最小的计算开销高效地优化大语言模型。先前工作,如Ma等人(2023);Kim等人(2024),已经展示了LoRA在增强静态剪枝模型方面的有效性。在本研究中,作者证明了LoRA同样可以有效地恢复动态剪枝模型的性能。为了更清晰地理解整个训练过程,请参考附录B中所示的算法。
- 实验配置
4.1. 模型与基准
模型作者使用LLaMA2-7B、LLaMA2-13B和LLaMA3.1-8B(Touvron等人,2023a;b; Dubey等人,2024)进行实验
数据 在路由和LoRA调优过程中,作者使用RedPajama-Data-1T-Sample数据集(计算机,2023) ^ { 3 }$ C,该数据集包含850万个样本(10亿个token),每个样本截断为4096个token。该数据集有两个作用:(1)作为校准集(100个随机样本)来计算块级显著性,用于剪枝冗余层(静态方法),以及(2)作为动态方法的训练集,以及恢复静态方法性能(静态和动态方法的具体细节将在第4.2节中介绍)。
训练过程每个模型使用next-token预测进行10,000步训练,在路由器和LoRA微调中使用16的批处理大小。在路由器微调阶段,作者使用恒定的学习率
。此外,Gumbel-Softmax的softmax温度
从5线性退火到1。在LoRA微调阶段,学习
学习率设置为2e-4,预热比例为0.1,并使用余弦学习率调度器。梯度反向传播采用AdamW优化器(Loshchilov & Hutter, 2019),参数
和
。
4.2. Baseline
静态剪枝 Baseline 静态剪枝涉及永久移除冗余模型组件。ShortGPT(Men等人,2024年)基于块影响(BI)分数移除冗余层。缩短的LLaMA(Kim等人,2024年)使用PPL和泰勒展开剪枝层,产生两种变体:Shortened-PPL和Shortened-Taylor。
LaCo(杨等人,2024)从深到浅逐步合并层,并采用阈值防止过度合并。联合层丢弃(He等人,2024)是ShortGPT的一种细粒度变体,根据相同的BI指标分别移除注意力层和MLP层。LLM-Pruner(马等人,2023)使用基于梯度的标准选择性地移除非关键结构。SliceGPT(Ashkboos等人,2024)通过用较小的密集矩阵替换较大的权重矩阵来减小嵌入维度。与静态剪枝方法不同,SkipGPT考虑了水平动态性,提供了一种更自适应的方法。
动态剪枝 Baseline 动态剪枝允许在推理过程中选择性地激活模型组件。MoD(Raposo等人,2024)采用了一种top-
路由机制,为每个token动态激活层。为确保公平比较,作者引入了一个新变体MoD-
,它与MoD的不同之处在于将注意力机制和MLP模块解耦,而MoD将它们视为一个单元。D-LLM(Jiang等人,2024)在每个Transformer层引入了一个动态决策模块,决定该层是执行还是跳过。此外,作者还包含了SkipGPT-Joint,它采用联合训练范式而非两阶段训练。从水平动态的角度来看,MoD-D按层分配其计算预算,而D-LLM按token分配。在垂直动态方面,D-LLM没有将注意力机制和MLP模块解耦。请参考附录C以获取这些 Baseline 的详细描述。
- 结果
本节对每个训练阶段的模型进行了详细分析。为清晰起见,作者将第一阶段(路由调整)后得到的模型定义为SkipGPT-RT,将第二阶段(LoRA微调)后的最终模型定义为SkipGPT-RT-L。
5.1. SkipGPT-RT与 Baseline 的比较
作者在表1中展示了 Baseline 剪枝方法和SkipGPT-RT的准确率和PPL。对于所有作者评估的模型,包括LLaMA2-7B、LLaMA2-13B和LLaMA3.1-8B,注意力模块的参数量大约是MLP模块的一半。因此,在联合层剪枝和SkipGPT-RT中,作者按比例剪枝更多模块,以保持不同方法间平均参数比例的一致性。对于长上下文任务,注意力模块的FLOPs超过了MLP模块,使得SkipGPT-RT相较于其他 Baseline 实现了更低的计算开销。对于LLaMA2-7B和LLaMA2-13B,作者报告了在25%参数削减设置下的评估结果,而对于LLaMA3.1-8B,作者提供了25%和40%参数削减的结果。
结果表明,仅通过路由器调优(不修改模型参数),SkipGPT-RT显著优于 Baseline 方法。具体而言,对于LLaMA2-7B和LLaMA2-13B,SkipGPT-RT在25%参数剪枝下仍保留超过90%的稠密模型的性能。对于LLaMA3.1-8B,它在25%剪枝下保留超过95%的性能,在40%剪枝下保留超过80%的性能——在此水平下,几乎所有 Baseline 方法都会灾难性地崩溃。值得注意的是,虽然路由器调优涉及训练路由器参数,与基于预定义指标的静态剪枝永久移除层不同,但其成本极小。这是因为路由器参数仅占模型总参数的不到0.01%,并且能够快速收敛(图4)。在实践中,调优过程仅需单个A800(80GB)GPU,并在四小时内完成。此外,表1中的结果突出了三个关键观察:
层级剪枝方法——如ShortGPT、LaCo、Shortened LLaMA和Joint Layer Drop——通常与嵌入维度缩减方法(如LLM-Pruner和SliceGPT)相比或超越它们。这表明LLM在深度上的冗余性大于宽度,与先前工作(Men等人,2024)一致。这强烈支持在设计SkipGPT时专注于深度剪枝。
联合层丢弃,作为ShortGPT的细粒度变体,独立移除注意力机制和前馈神经网络,在大多数情况下表现优于ShortGPT。这一结果验证了作者强调垂直动态对于实现更有效剪枝的必要性。
SkipGPT-RT在准确率(反映了大语言模型作为通用任务解决者的能力)和困惑度(涉及生成连贯流畅句子的能力)方面,显著优于所有模型中的联合层Drop。这有力地支持了作者强调水平动态必要性的动机。
5.2. SkipGPT-RT-L与 Baseline 的比较
SkipGPT-RT-L
原始模型性能图4展示了SkipGPT-RT、SkipGPTRT-L、使用LoRA的静态剪枝方法以及LLaMA2-7B和LLaMA2-13B的动态剪枝联合路由器+LoRA训练的培训曲线。结果总结在表2中。值得注意的是,SkipGPT-RT实现了仅使用路由器调整的最佳微调 Baseline 的性能相当。
在LoRA微调后,llm-SkipGPT_2506不仅完全恢复了模型的原始性能水平,甚至超越了它(未经微调),在LLaMA2-7B中仅次于直接将LoRA应用于原始模型,排名第二。对于那些希望快速构建高性能剪枝模型的人来说,路由器调整非常有效。同时,为了完全匹配或超过原始模型,LoRA微调是一个绝佳的选择。值得重申的是,所有模型的整体过程仅需一个A800 GPU。
两阶段训练范式有效性研究
作者将SkipGPT-Joint(指路由参数和模型参数同时训练的变体)作为两阶段训练范式的消融 Baseline 。如表2所示,其性能显著劣于SkipGPT。虽然直观上直接采用MoE联合训练范式可能优于两阶段训练——因为路由参数和LoRA参数可以逐渐适应彼此的表征——但实验结果却呈现相反趋势。不仅SkipGPT-Joint表现不佳,遵循联合训练范式的动态剪枝方法始终无法取得令人满意的结果。
这一发现印证了作者第3.5节的观点:在联合训练范式下,随机初始化的路由器迫使模型参数过早适应次优的随机路由策略。
这种错位会扰乱模型已建立的参数分布,导致路由器越来越难以识别关键模块。随着时间的推移,这形成了一个强化反馈循环:模型适应不良路由时,路由器难以优化决策,进一步降低其效能和整体模型性能。
最终,这种机制阻碍了联合训练范式取得优异结果。相比之下,作者的两阶段训练范式有效缓解了这些挑战,使剪枝模型相比现有动态剪枝方法实现了更优性能。
- 路由行为分析
除了效率之外,作者探索了在SkipGPTRT中调整路由器是否能为原始LLM提供更深入的洞察。通过分析路由器的行为,作者旨在揭示不同模块如何对推理过程做出贡献的规律性模式。
注意力模块与MLP模块冗余性比较 最近的研究表明,大语言模型比前馈层更能抵抗自注意力层的移除,这表明注意力模块相比MLP模块具有更高的冗余性。为了进一步探索这一点,作者分析了五个通过在不同目标Sparse度水平
下进行路由调整生成的LLaMA2-13B剪枝模型。对于每个剪枝模型,作者使用WT2数据集(Merity等人,2016)中随机选取的50个句子来测量注意力模块和MLP模块的平均Sparse度。在这里,Sparse度被定义为剪枝模块数与注意力或MLP中模块总数的比率。如图5所示,SkipGPT
RT始终比MLP模块剪枝更多的注意力机制,这在低和高剪枝率下均成立。这一观察结果与先前的研究发现一致,可能表明当前Transformer架构存在结构上的低效性,该架构强制要求每层固定搭配一个注意力机制和一个MLP模块。作者假设未来的LLM可以通过重新审视这一设计,并可能减少注意力模块的比例来获得更高的效率。
注意力机制和MLP模块中的冗余变化Del Corro等人(2023)认为随着上下文随序列增长,后置token变得更加可预测,因此需要更少的计算量。与此同时,He等人(2025)表明MLP模块在解码阶段相比于预填充阶段的冗余度更高,而注意力机制的冗余度在两个阶段几乎保持不变。作者通过分析通过LLaMA2-13B路由器调优获得的45%Sparse模型,研究了注意力机制和MLP模块中冗余度的变化。作者从RedPajama中随机选择句子,将其截断至模型的最大上下文长度,并在每个句子上应用100-token滑动窗口。
在每个步骤中,作者记录滑动窗口内注意力机制和MLP的激活比例,并对50个样本的结果进行平均。最终结果如图6所示。与先前的研究发现相反,作者观察到随着上下文增长,后置token需要更多的注意力计算量,但需要更少的MLP计算量。作者假设在早期阶段,模型尚未确定任务且上下文信息有限。因此,它更依赖复杂的MLP计算进行任务识别,而较少依赖注意力机制提取上下文。后来,随着任务变得清晰且上下文积累,模型转向增加注意力计算量,同时减少MLP计算量。开发能够动态调整计算量的技术为未来研究提供了令人兴奋的方向。
- 结论
在这项工作中,作者介绍了SkipGPT,一个通过引入水平和垂直动态机制来解决静态层剪枝效率问题的创新性动态剪枝框架。SkipGPT能够按每个token动态分配计算资源,并将注意力机制和MLP剪枝解耦,从而实现更细粒度的优化。
为了稳定训练,作者提出了一种两阶段训练范式,首先单独调整路由器,然后再使用LoRA对模型进行微调。
实验表明,即使参数减少了40%,剪枝后的模型仍能完全恢复其性能,甚至在某些情况下超越原始模型。
此外,作者的路由器分析揭示了模块冗余和token级计算资源分配的关键见解,为未来模型效率提升指明了潜在方向。
参考
[1]. SkipGPT: Dynamic Layer Pruning Reinvented with Token Awareness and Module Decoupling