[论文解读] Prospect Pruning: Finding Trainable Weights at Initialization using Meta-Gradients

Milad Alizadeh, Shyam A. Tailor+5 more|arXiv (Cornell University)|Feb 16, 2022

Advanced Neural Network Applications被引用 11

一句话总结

Prospect Pruning (ProsPr) 是一种新颖的结构化剪枝方法，通过在训练初期的前几步中利用元梯度来识别可训练子网络，明确考虑了剪枝后的可训练性。该方法在单次剪枝设置下实现了视觉分类任务的最先进性能，且数据需求更低，通过将优化动态纳入显著性标准，显著优于现有方法。

ABSTRACT

Pruning neural networks at initialization would enable us to find sparse models that retain the accuracy of the original network while consuming fewer computational resources for training and inference. However, current methods are insufficient to enable this optimization and lead to a large degradation in model performance. In this paper, we identify a fundamental limitation in the formulation of current methods, namely that their saliency criteria look at a single step at the start of training without taking into account the trainability of the network. While pruning iteratively and gradually has been shown to improve pruning performance, explicit consideration of the training stage that will immediately follow pruning has so far been absent from the computation of the saliency criterion. To overcome the short-sightedness of existing methods, we propose Prospect Pruning (ProsPr), which uses meta-gradients through the first few steps of optimization to determine which weights to prune. ProsPr combines an estimate of the higher-order effects of pruning on the loss and the optimization trajectory to identify the trainable sub-network. Our method achieves state-of-the-art pruning performance on a variety of vision classification tasks, with less data and in a single shot compared to existing pruning-at-initialization methods.

研究动机与目标

• 为解决现有剪枝初始化方法性能不佳的问题，这些方法未能考虑剪枝后的可训练性。
• 克服当前显著性标准的短视性，这些标准仅在初始化时评估单一步骤，而未考虑后续训练过程。
• 开发一种方法，显式建模剪枝对前几步训练过程中优化轨迹和损失的影响。
• 实现在初始化阶段的高精度单次剪枝，无需迭代微调或渐进式剪枝。
• 通过识别从一开始就既稀疏又可训练的子网络，提升神经网络剪枝的实用性和效率。

提出的方法

• ProsPr 基于元梯度计算显著性分数，元梯度是损失对初始权重在经过若干优化步骤后的梯度。
• 通过反向传播前几步随机梯度下降过程，估计剪枝某一权重如何影响早期训练中的损失。
• 该方法采用元梯度的一阶近似，以提升计算效率和梯度稳定性。
• 为所有权重计算显著性分数，并根据这些分数剪枝重要性最低的权重。
• 剪枝掩码直接应用于初始权重，剩余子网络从头开始训练。
• 该方法设计为单次剪枝，避免了先前方法中常见的迭代剪枝与微调循环。

实验结果

研究问题

• RQ1 一种考虑剪枝后优化轨迹的剪枝初始化方法，是否能比现有方法取得更好的性能？
• RQ2 与静态显著性度量相比，通过在早期训练步骤中引入元梯度，如何提升对可训练子网络的识别能力？
• RQ3 ProsPr 在单次剪枝步骤和减少数据使用的情况下，能在多大程度上实现最先进精度？
• RQ4 通过元梯度建模可训练性，是否能减少对迭代或渐进式剪枝策略的依赖？
• RQ5 元梯度能否有效捕捉剪枝对损失和优化动态的高阶影响？

主要发现

• ProsPr 在多个视觉分类基准上实现了最先进精度，优于现有的剪枝初始化方法。
• 该方法通过单次剪枝过程实现高精度，消除了对迭代微调或渐进式剪枝的需求。
• 即使在高稀疏度下，ProsPr 的性能退化也低于先前方法。
• 元梯度的使用通过建模剪枝对早期训练动态的影响，提升了对可训练子网络的识别能力。
• 该方法在比以往方法更少的数据下依然有效，表明其在子网络发现中具有更高的样本效率。
• 当在超过三步训练步骤上计算元梯度时，会出现梯度稳定性问题，因此该方法仅限于少量初始训练步骤。