[论文解读] A Unified Framework for Soft Threshold Pruning

Yanqi Chen, Zhengyu Ma+4 more|arXiv (Cornell University)|Feb 25, 2023

Sparse and Compressive Sensing Techniques被引用 6

一句话总结

本文提出了一套统一的理论框架，将软阈值剪枝与迭代收缩-阈值算法（ISTA）联系起来，将其视为L1正则化优化问题的隐式解。通过推导出一种最优、时不变的阈值调度策略，稳定L1正则化系数，该方法在ImageNet上的ResNet-50、MobileNet-V1和SEW ResNet-18模型上实现了最先进的剪枝性能，仅需极少的超参数调优，并可广泛适用于任何SGD训练的模型。

ABSTRACT

Soft threshold pruning is among the cutting-edge pruning methods with state-of-the-art performance. However, previous methods either perform aimless searching on the threshold scheduler or simply set the threshold trainable, lacking theoretical explanation from a unified perspective. In this work, we reformulate soft threshold pruning as an implicit optimization problem solved using the Iterative Shrinkage-Thresholding Algorithm (ISTA), a classic method from the fields of sparse recovery and compressed sensing. Under this theoretical framework, all threshold tuning strategies proposed in previous studies of soft threshold pruning are concluded as different styles of tuning $L_1$-regularization term. We further derive an optimal threshold scheduler through an in-depth study of threshold scheduling based on our framework. This scheduler keeps $L_1$-regularization coefficient stable, implying a time-invariant objective function from the perspective of optimization. In principle, the derived pruning algorithm could sparsify any mathematical model trained via SGD. We conduct extensive experiments and verify its state-of-the-art performance on both Artificial Neural Networks (ResNet-50 and MobileNet-V1) and Spiking Neural Networks (SEW ResNet-18) on ImageNet datasets. On the basis of this framework, we derive a family of pruning methods, including sparsify-during-training, early pruning, and pruning at initialization. The code is available at https://github.com/Yanqi-Chen/LATS.

研究动机与目标

• 为解决现有软阈值剪枝方法缺乏理论基础的问题，这些方法依赖于临时或可学习的阈值调度策略。
• 将软阈值剪枝中多样的阈值调优策略统一于单一优化框架之下。
• 推导出一种最优的阈值调度器，保持L1正则化系数在时间上恒定，确保目标函数的时间不变性。
• 在无需模型特定调优的前提下，实现对包括人工神经网络与脉冲神经网络在内的多种模型的系统性剪枝。

提出的方法

• 利用迭代收缩-阈值算法（ISTA）将软阈值剪枝重新表述为L1正则化优化问题的隐式解。
• 识别出所有先前的阈值调度策略均对应于不同方式的L1正则化项调优。
• 推导出一种最优的阈值调度器，使L1正则化系数在整个训练过程中保持恒定，从而确保目标函数的稳定性。
• 在统一框架下构建一系列剪枝方法，包括训练中剪枝、早期剪枝和参数初始化时剪枝。
• 在ResNet-50、MobileNet-V1和SEW ResNet-18上验证该框架，证明在所有设置下均实现一致的性能提升。
• 将该方法实现为LATS（可学习自适应阈值调度器），代码已公开。

实验结果

研究问题

• RQ1 软阈值剪枝如何与ISTA等成熟优化算法建立正式关联？
• RQ2 先前软阈值剪枝方法中采用的不同阈值调度策略的理论基础是什么？
• RQ3 能否推导出一个统一框架，以解释并改进现有的软阈值剪枝技术？
• RQ4 是否可以推导出一种最优的阈值调度器，使L1正则化系数在整个训练过程中保持稳定？
• RQ5 所提出的框架是否能在不同网络架构和稀疏度水平（包括脉冲神经网络）上实现泛化？

主要发现

• 所提出的框架表明，软阈值剪枝等价于L1正则化优化问题的隐式ISTA求解器，为现有方法提供了理论基础。
• 所有先前的阈值调度策略均被证明是L1正则化项调优的不同形式，从而在单一优化视角下统一了其行为。
• 推导出的最优阈值调度器保持L1正则化系数恒定，实现时间不变的目标函数，从而提升泛化性能。
• 在ImageNet上，该方法在ResNet-50上实现了99.04%稀疏度下的88.88% top-1准确率，达到SOTA性能。
• 对于SEW ResNet-18，该方法在71.18%稀疏度下达到60.11% top-1准确率，在92.57%稀疏度下达到53.74%准确率，优于先前方法。
• 该框架支持一系列剪枝策略——训练中剪枝、早期剪枝和初始化时剪枝——而无需微调或超参数搜索。