[论文解读] Soft Threshold Weight Reparameterization for Learnable Sparsity

Aditya Kusupati, Vivek Ramanujan+5 more|arXiv (Cornell University)|Feb 8, 2020

Advanced Neural Network Applications参考文献 58被引用 82

一句话总结

STR 通过软阈值化在每一层学习剪枝阈值，以实现非均匀稀疏。它在 CNNs（ResNet50、MobileNetV1 在 ImageNet 上）上实现了最先进的非结构化稀疏性精度，同时降低 FLOPs，并扩展到 RNNs 的结构稀疏性。

ABSTRACT

Sparsity in Deep Neural Networks (DNNs) is studied extensively with the focus of maximizing prediction accuracy given an overall parameter budget. Existing methods rely on uniform or heuristic non-uniform sparsity budgets which have sub-optimal layer-wise parameter allocation resulting in a) lower prediction accuracy or b) higher inference cost (FLOPs). This work proposes Soft Threshold Reparameterization (STR), a novel use of the soft-threshold operator on DNN weights. STR smoothly induces sparsity while learning pruning thresholds thereby obtaining a non-uniform sparsity budget. Our method achieves state-of-the-art accuracy for unstructured sparsity in CNNs (ResNet50 and MobileNetV1 on ImageNet-1K), and, additionally, learns non-uniform budgets that empirically reduce the FLOPs by up to 50%. Notably, STR boosts the accuracy over existing results by up to 10% in the ultra sparse (99%) regime and can also be used to induce low-rank (structured sparsity) in RNNs. In short, STR is a simple mechanism which learns effective sparsity budgets that contrast with popular heuristics. Code, pretrained models and sparsity budgets are at https://github.com/RAIVNLab/STR.

研究动机与目标

• 在参数预算下激励跨层的非均匀稀疏预算以最大化准确率。
• 引入 Soft Threshold Reparameterization (STR)，通过反向传播学习每层的剪枝阈值。
• 展示 STR 在 CNNs 上实现最先进的非结构化稀疏性表现并减少 FLOPs。
• 证明 STR 可以诱导结构化稀疏性并将学习到的稀疏预算迁移到不同任务。

提出的方法

• 将 STR 定义为 soft-thresholded 投影 S_g(W_l, s_l) = sign(W_l) * ReLU(|W_l| - g(s_l))，并具有可学习的每层 s_l。
• 将优化重新参数化为最小化 L(S_g(W, s), D) 而不是 L(W, D)，使剪枝阈值可以反向传播。
• 使用连续函数 g（未结构化使用 sigmoid，结构化稀疏使用指数函数）将 s_l 映射到每层阈值 alpha_l = g(s_l)。
• 在使用 SGD 和权重衰减训练 W_l 的同时，通过对 STR 投影的 L 关于 s_l 的梯度进行更新来更新 s_l。
• 证明 STR 会引入稀疏梯度和跨层的非均匀预算，并通过与向后剪枝的联系来讨论收敛直觉。

实验结果

研究问题

• RQ1 我们是否可以学习在全局参数限定下优化准确率的逐层稀疏预算？
• RQ2 STR 是否提供稳定的训练并在现有的均匀或启发式非均匀稀疏方法之上提升精度？
• RQ3 STR 是否能够扩展到结构化稀疏性和其他架构（如 RNN）同时保持收益？
• RQ4 学习到的稀疏预算是否能够跨任务或数据集迁移？
• RQ5 与基线相比，STR 如何影响各层的 FLOPs 分布？

主要发现

方法	Top-1 精度 (%)	参数	稀疏度 (%)	FLOPs
ResNet-50	77.01	25.6M	0.00	4.09G
GMP	75.60	5.12M	80.00	818M
DSR ∗#	71.60	5.12M	80.00	1.23G
DNW	76.00	5.12M	80.00	818M
SNFS	74.90	5.12M	80.00	-
SNFS + ERK	75.20	5.12M	80.00	1.68G
RigL ∗	74.60	5.12M	80.00	920M
RigL + ERK	75.10	5.12M	80.00	1.68G
DPF	75.13	5.12M	80.00	818M
STR	76.19	5.22M	79.55	766M
STR	76.12	4.47M	81.27	705M

• STR 在 ImageNet-1K 的 ResNet50 和 MobileNetV1 的多种稀疏水平下实现了未结构化稀疏性的最先进精度。
• STR 学习到跨层的非均匀预算，在某些情况下将 FLOPs 降低最多约 50%。
• 在极度稀疏区间（99%）时，STR 相较基线在精度上提升高达约 10%。
• STR 诱导最终的逐层阈值差异，导致非均匀稀疏预算，能在前几层集中计算资源的降低，而在其他层保持密度。
• 学习到的稀疏预算具备跨任务的可迁移性，并且可以扩展到结构化稀疏性（RNN 的低秩）。