[论文解读] Learned Threshold Pruning

Kambiz Azarian, Yash Bhalgat+2 more|arXiv (Cornell University)|Feb 28, 2020

Advanced Image and Video Retrieval Techniques参考文献 45被引用 23

一句话总结

该论文提出了一种可微分的、基于梯度的端到端训练方法——学习阈值剪枝（LTP），通过软剪枝和可微分的L₀正则化，联合训练每层的剪枝阈值与网络权重。LTP在结构化剪枝方面达到当前最优性能，对ResNet50实现9.1倍压缩、AlexNet实现26.4倍压缩，且精度损失极小，同时计算效率高，并与批量归一化兼容。

ABSTRACT

This paper presents a novel differentiable method for unstructured weight pruning of deep neural networks. Our learned-threshold pruning (LTP) method learns per-layer thresholds via gradient descent, unlike conventional methods where they are set as input. Making thresholds trainable also makes LTP computationally efficient, hence scalable to deeper networks. For example, it takes $30$ epochs for LTP to prune ResNet50 on ImageNet by a factor of $9.1$. This is in contrast to other methods that search for per-layer thresholds via a computationally intensive iterative pruning and fine-tuning process. Additionally, with a novel differentiable $L_0$ regularization, LTP is able to operate effectively on architectures with batch-normalization. This is important since $L_1$ and $L_2$ penalties lose their regularizing effect in networks with batch-normalization. Finally, LTP generates a trail of progressively sparser networks from which the desired pruned network can be picked based on sparsity and performance requirements. These features allow LTP to achieve competitive compression rates on ImageNet networks such as AlexNet ($26.4 imes$ compression with $79.1\%$ Top-5 accuracy) and ResNet50 ($9.1 imes$ compression with $92.0\%$ Top-5 accuracy). We also show that LTP effectively prunes modern extit{compact} architectures, such as EfficientNet, MobileNetV2 and MixNet.

研究动机与目标

• 为解决在非结构化权重剪枝中选择最优逐层剪枝阈值的挑战，该问题通常依赖人工设定或昂贵的迭代搜索。
• 开发一种可微分、端到端可训练的剪枝方法，联合优化网络权重与层特定的剪枝阈值。
• 实现对深层网络（尤其是含批量归一化层的网络）的高效、可扩展剪枝，而无需依赖L₁或L₂正则化。
• 生成一系列逐步剪枝的连续模型，以灵活应对稀疏度与精度之间的权衡。
• 在现代高效架构（如MobileNetV2、EfficientNet和MixNet）上实现高倍率压缩，而这些架构在现有方法下难以压缩。

提出的方法

• LTP引入了一种可微分的软剪枝机制，使梯度能够通过剪枝操作流动，从而支持端到端训练。
• 其提出了一种可微分的L₀正则化，通过梯度下降学习应剪枝的权重，从而促进稀疏性。
• 每层均具有可训练的阈值参数，在微调过程中进行优化，使模型能自动学习每层的最优剪枝水平。
• 该方法采用软分配策略，低于学习到的阈值的权重被逐步剪枝，避免训练期间使用硬性二值决策。
• 通过依赖可微分的L₀惩罚而非传统L₁/L₂惩罚，该方法与批量归一化层兼容，后者在传统方法中常导致正则化失效。
• LTP在训练过程中生成一系列检查点，每个检查点代表一个不同的稀疏度水平，从而可根据期望的精度-稀疏度权衡灵活选择模型。

实验结果

研究问题

• RQ1 能否通过梯度下降有效学习逐层剪枝阈值，以提升非结构化剪枝的性能？
• RQ2 可微分的L₀正则化是否能在批量归一化网络中实现有效稀疏性学习，从而克服L₁/L₂惩罚失效的问题？
• RQ3 可训练的阈值机制是否能比迭代剪枝-微调方法更快实现高倍率压缩？
• RQ4 LTP是否能有效应用于现代高效架构（如MobileNetV2、EfficientNet和MixNet），而无需修改网络结构？
• RQ5 LTP是否能生成一条连续的剪枝模型路径，从而支持基于精度与稀疏度需求的灵活选择？

主要发现

• 在ResNet50上，LTP实现9.1倍压缩，Top-5精度仅下降0.3%，在18轮剪枝和12轮微调后达到92.0%的精度。
• 在AlexNet上，LTP实现26.4倍压缩，Top-5精度无损失（79.1%），在压缩率与精度方面均优于先前方法。
• 在MobileNetV2上，LTP实现1.33倍压缩，Top-1精度下降不足1%，在相同压缩率下比全局剪枝高出9%的精度。
• LTP对EfficientNet-B0实现3倍压缩，对MixNet-S实现2倍压缩，Top-1精度下降均不足1%，这是首次报告的该架构剪枝结果。
• 该方法收敛迅速，ResNet50仅需18轮剪枝，MixNet-S甚至无需额外微调，展现出极高的计算效率。
• LTP生成一条连续的剪枝模型序列，支持根据稀疏度与性能权衡进行部署级模型选择。