虽然基于幅度的剪枝通过移除绝对值最小的权重提供了一种直接方法,但它基于一个可能具有局限的假设:在剪枝过程早期被认为不重要的权重会一直保持不重要。这种静态观点可能不是最优的,因为训练大型模型的复杂动态可能允许最初较小的权重在后期变得重要。移动剪枝和相关的动态稀疏技术通过允许被剪枝的权重集合在训练或微调过程中演变来应对这一局限。
移动剪枝的理念
设想优化过程如同在一个复杂的损失曲面上航行。标准的基于幅度的剪枝在特定检查点识别接近零的权重并永久移除它们。然而,移动剪枝承认流经网络的梯度可能会将一些被剪枝的权重“推离”零点,暗示它们可能再次变得有用。反之,一些最初较大的权重可能会趋向零并成为剪枝的候选。
移动剪枝(Movement pruning)的核心思想,由Sanh等人(2020)在BERT背景下提出,在于训练期间允许权重 进入 和 离开 被剪枝的集合。这种方法通常通过为权重W维护一个二元掩码M来实现。不同于静态剪枝在剪枝步骤后掩码M固定不变,移动剪枝会周期性地更新M,从而实现动态调整。
算法机制
移动剪枝通常将剪枝直接集成到优化循环中。一种常见的方法包含以下步骤:
- 初始化: 从一个密集模型或使用静态方法剪枝的模型开始。定义目标稀疏度 s。
- 稀疏度应用: 根据权重幅度应用稀疏掩码 M。仅当 Mij=1 时,权重 Wij 才活跃。
- 梯度计算: 计算关于 带掩码 权重 W⊙M 的梯度 ∇L (其中 ⊙ 是元素乘法)。
- 梯度累积 (重要步骤): 不仅仅更新活跃权重,还要累积 所有 权重的梯度,包括当前被掩码的权重 (Mij=0)。这通常通过使用标准优化器更新规则来实现,但在重新应用掩码之前将其应用于底层权重张量。例如,对于带有动量的SGD等优化器,更新可能如下所示:
Vt+1=βVt+∇L(Wt⊙Mt)
Wt+1∗=Wt−αVt+1// 更新应用于所有权重
- 掩码更新: 周期性地(例如,每几百步)重新评估掩码 M。
- 根据幅度识别 Wt+1∗ 中排名前 (1−s)×N 的权重 (其中 N 是权重总数)。将它们在 Mt+1 中的相应条目设为 1。
- 将 Mt+1 中剩余的 s×N 个条目设为 0。这使得先前被掩码 (Mt,ij=0) 但已累积显著梯度更新(导致 Wt+1∗ 中幅度更大)的权重可以变为未掩码状态 (Mt+1,ij=1)。反之,活跃但幅度减小的权重可以变为被掩码状态。
- 权重更新: 将新掩码应用于更新后的权重:Wt+1=Wt+1∗⊙Mt+1。
- 重复: 从步骤3继续训练过程。
这个过程允许稀疏模式动态适应。权重不仅仅被剪枝;如果优化过程在后期认为它们重要,它们有机会被“恢复”。渐进式剪枝计划(其中目标稀疏度 s 随时间增加)常与移动剪枝结合使用以稳定训练。
移动剪枝在训练循环中的流程。梯度影响所有权重(包括被掩码的),可能会充分改变其幅度,从而在下一次再生步骤中改变掩码。
移动剪枝中的动态稀疏性
移动剪枝是动态稀疏性的一个具体实例。这种更广泛的理念涵盖了稀疏模式不是预先固定,而是在训练期间甚至在推理时可以改变的任何技术。
- RigL (Rigged Lottery): 由Evci等人(2020)提出,RigL在训练循环中结合了基于幅度的剪枝和随机剪枝的元素。它周期性地根据幅度移除权重,但也随机添加新的连接,利用梯度信息来引导哪些连接被激活。这在保持计算成本的同时,允许尝试不同的稀疏结构。
- 动态稀疏训练: 一些研究着眼于根据除幅度以外的标准动态调整稀疏掩码,可能包含敏感性分析或Hessian信息,尽管这些方法通常计算量更大。
- 输入依赖稀疏性: 一种更高级的理念涉及根据输入数据变化的稀疏模式。例如,在MoE(专家混合)模型中,只有某些“专家”(子网络)针对给定的输入令牌被激活,创建了一种动态的、输入依赖的结构化稀疏形式。虽然通常不称为“剪枝”,但它通过选择性地使用模型部分来达到减少计算量的目的。
优势与权衡
优势:
- 提高准确性: 动态方法,特别是移动剪枝,在相同稀疏度下通常比静态剪枝方法获得更高的准确性。它们更善于识别和保留那些可能被静态方法错误剪枝的真正重要权重。
- 更好的稀疏方案: 通过允许稀疏结构在训练期间与权重共同适应,这些方法与剪枝预训练的密集模型相比,可以发现更有效的稀疏子网络。
挑战:
- 训练复杂性: 实现和管理动态掩码会增加训练过程的开销。追踪被掩码权重的梯度并周期性地再生掩码需要仔细实现。
- 超参数调优: 动态稀疏性引入了新的超参数,例如掩码更新频率、稀疏度计划(稀疏度增加的速度),以及可能与梯度累积或随机连接增长相关的参数(如RigL中)。寻找最佳设置可能需要大量实验。
- 潜在的不稳定性: 变化的稀疏模式如果管理不当,有时可能导致训练不稳定,特别是与高学习率或激进的稀疏度计划结合时。
- 硬件效率: 尽管动态非结构化稀疏性可以产生非常稀疏的模型,但在硬件上实现理论加速可能具有挑战性。静态、结构化稀疏模式通常更容易被当前硬件和编译器高效加速。然而,对动态稀疏性的运行时支持是一个活跃的研究方向。
何时考虑动态稀疏性
移动剪枝和动态稀疏性是强大的技术,当主要目标是在最小准确性损失下实现极高稀疏度(>80%)时,它们显得尤为有用。它们在研究场景或在拥有训练计算资源时微调模型时尤其适用,并且重点是最大化最终稀疏模型的性能。虽然静态结构化剪枝可能因更好的硬件兼容性而更受即时部署青睐,但动态方法代表了一种更先进的途径,用于找到高度优化的稀疏LLM。理解这些技术有助于了解模型参数、优化动态和网络架构在效率背景下的演变关系。