编辑推荐:
迭代剪枝因梯度失衡导致权重重要性评估不稳定,我们提出LAGA通过局部动态调整与全局学习率缩放机制,有效平衡训练稳定性与模型压缩效果,实验在多个模型和数据集上验证其有效性。
段连|张佳文|李崇新|王汉章
上海大学计算机工程与科学学院,上海上大路99号,200444,中国
摘要
迭代剪枝虽然在降低模型复杂性方面非常有效,但往往会引入权重重要性估计的不稳定性。这种不稳定性源于剪枝过程中产生的结构失衡:被剪掉的权重通常会被置零或屏蔽,而未被剪掉的权重则保持其之前优化后的状态。这种失衡导致重要性得分在迭代过程中波动,使得剪枝决策变得不可靠,并影响重新训练的效率。为了解决这个问题,我们提出了LAGA(局部调整和全局缩放),这是一种动态学习率调整策略,旨在稳定迭代剪枝中的重要性估计。LAGA根据每个权重的剪枝状态在两个维度上进行学习率调整:在局部层面,它平滑时间敏感性信号并补偿重新激活的动态;在全球层面,它根据模型的稀疏性来调整学习率,以适应结构压缩。这些机制共同作用,重新平衡训练动态,并在剪枝过程中恢复稳定的重要性轨迹。广泛的实验表明,LAGA显著提高了重要性评估的稳定性和剪枝模型的性能。在CIFAR-100数据集上,将ViT-B/16模型剪枝60%后,其Top-1准确率比使用AdamW算法时提高了4.71%;在CIFAR-100数据集上,将DeiT-S模型剪枝40%后,Top-1准确率提高了0.68%,超过了未剪枝的基线模型。在CIFAR-100和ImageNet-1K数据集上,LAGA在ViT-B/16、DeiT-S和Swin-B模型上的表现一致,证明了其跨不同剪枝场景的鲁棒性和泛化能力。
引言
深度神经网络(DNNs)在各种任务中表现出色,然而它们庞大的参数空间需要大量的计算和存储资源(Bohnstingl等人,2021年;Devlin等人,2018年;Dosovitskiy等人,2020年;He等人,2020年;Liu等人,2019年;Wu等人,2023年;Zhang和Freris,2024年)。这激发了人们对高效模型压缩技术的兴趣(Cheng等人,2024年;Frantar和Alistarh,2023年;Han等人,2015a年),其中神经网络剪枝作为一种广泛采用的方法,可以在保留性能的同时消除冗余权重(Cheng等人,2024年;Han等人,2015b年;Ma等人,2023年;Yang等人,2023年)。现有的剪枝方法大致可以分为一次性剪枝(Chen等人,2020年;Frantar和Alistarh,2023年;Lee等人,2018年;Liang等人,2021年)和迭代剪枝(Han等人,2015b年;Sanh等人,2020年;Zhang等人,2022年;Zhu和Gupta,2017年)。迭代剪枝通过在评分、剪枝和重新训练之间交替进行,能够根据网络结构的变化不断调整,从而产生更有效的稀疏模型(Molchanov等人,2019年;Renda等人,2020年)。然而,这一过程也引入了与训练动态相关的新挑战。
迭代剪枝中的一个关键挑战是剪掉的和未剪掉的权重之间的梯度幅度存在显著不平衡。如图1所示,被剪掉的权重梯度通常很小,限制了它们的更新能力,并阻碍了重新训练过程中的恢复。从梯度传播的角度来看,一旦权重被剪掉并置零,它就不再参与前向激活,导致其在反向传播中的梯度显著减小。这种减小有两个原因:(1)被剪掉的权重产生的激活值接近零,通过链式法则导致梯度衰减;(2)剪枝破坏了表示路径的平衡,使得学习动态集中在剩余的权重上,进一步减少了流向被剪掉连接的梯度。随着迭代的进行,未剪掉的权重继续接收到较强的梯度,而被剪掉的权重则训练不足,导致梯度不平衡加剧。这种梯度不平衡直接削弱了重要性估计的稳定性。许多敏感性指标(如梯度幅度、梯度-权重乘积或基于运动的信号)依赖于可靠的梯度信息。当被剪掉的权重持续接收到极小的梯度时,它们的重要性得分变得不可靠,本应被认为是重要的权重可能由于梯度信号不足而被错误地分类为不重要。
为了验证这个问题,我们在图2中展示了使用标准优化器(AdamW,Loshchilov和Hutter,2017年)时重要性得分的演变。观察到的波动幅度很大,表明敏感性评估和排名存在显著的不稳定性。这种不一致性会损害剪枝过程,因为重要的权重可能会被错误地判断并移除,最终降低模型性能。
这些观察结果突显了现有剪枝框架的一个根本局限性:它们通常依赖于静态的或每次迭代的敏感性评估,这些评估容易受到剪枝引入的随机波动和结构扰动的影响。然而,权重敏感性本质上是动态的——它在训练过程中不断演变,特别是在多次剪枝循环中。因此,我们认为剪枝策略不仅应该评估参数在某一时刻的重要性,还应该考虑其重要性随时间的变化。因此,提高敏感性评估的时间稳定性对于提高重要性估计的可靠性至关重要。
为了解决这个问题,我们提出了一种基于轨迹的学习率调节策略,该策略明确地模拟了剪枝迭代过程中权重敏感性的演变。我们引入了LAGA(局部调整和全局缩放),这是一个统一的框架,它结合了时间趋势平滑和波动性估计来指导自适应学习率调度。LAGA包括两个关键组成部分:(1)基于平滑敏感性趋势及其偏差的局部调整机制,旨在捕捉动态稳定性并抑制敏感性估计中的噪声;(2)全局学习率缩放策略,根据剪枝的稀疏性比例增加学习率,以适应模型容量的逐渐减少。我们的贡献总结如下:
1.我们确定了重要性得分估计的不稳定性和收敛性能下降是迭代剪枝中的关键挑战,并将这些问题归因于训练过程中观察到的梯度不平衡和敏感性动态的变化。
2.我们提出了LAGA,这是一种考虑剪枝的动态学习率调整策略。它根据平滑的敏感性和敏感性变化为每个参数分配学习率,并结合稀疏性进行全局缩放,以改善重新训练并保持重要性估计的一致性。
3.我们在CIFAR-100和ImageNet-1K数据集上的ViT-B/16、DeiT-S和Swin-B模型上进行了广泛的实验。我们的方法在各种稀疏性和模型架构下都一致地提高了剪枝性能和训练稳定性。
本文的其余部分安排如下:第2节回顾了与剪枝和自适应优化相关的工作。第3节详细介绍了我们提出的方法。第4节展示了实验结果和消融研究。第5节总结了本文并讨论了未来的研究方向。
章节片段
迭代剪枝
网络剪枝会显著降低模型性能,尤其是在高剪枝比例下,因此标准的剪枝流程通常包括三个步骤:训练原始神经网络,对其进行剪枝以达到目标稀疏度,然后重新训练剪枝后的模型以提高性能。与一次性剪枝(Frantar和Alistarh,2023年;Han等人,2015b)不同,后者只进行一次剪枝和重新训练,迭代剪枝(Han等人,2015b)会重复剪枝和重新训练的过程
方法
在本节中,我们介绍了LAGA(局部调整和全局缩放),这是一个为迭代剪枝设计的统一动态学习率调整框架。该框架通过解决重要性估计中的不稳定性来提高剪枝性能,这种不稳定性主要是由于剪掉权重在重新训练过程中的恢复能力有限。与直接修改剪枝标准的方法不同,LAGA从学习动态的角度进行干预,明确调节权重
实验
我们在CIFAR-100(Krizhevsky,2009年)和ImageNet-1K(Deng等人,2009年)数据集上,使用ViT-B/16(Dosovitskiy等人,2020年)、DeiT-S(Touvron等人,2021年)和Swin-B(Liu等人,2021年)模型评估了所提出的LAGA学习率策略的有效性。这些实验旨在展示我们的方法在多种模型和数据集上的有效性。
结论
LAGA证明了动态学习率调整对于稳定迭代剪枝至关重要,这种调整能够响应即时的敏感性变化和长期的结构稀疏性。为此,我们提出了一种权重自适应学习率策略,该策略结合了局部和时间敏感性信号以及每个参数的剪枝状态。
CRediT作者贡献声明
段连:撰写——审阅与编辑,撰写——初稿,可视化,方法论,概念化。张佳文:验证,软件,形式分析,数据管理,概念化。李崇新:可视化,验证,软件,方法论,概念化。王汉章:监督,项目管理,方法论,资金获取,概念化。
利益冲突声明
作者声明他们没有已知的竞争性财务利益或个人关系可能影响本文所述的工作。
致谢
本工作得到了中国国家自然科学基金(编号62206167)的支持。