[論文レビュー] 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 は soft thresholding によって per-layer pruning thresholds を学習し、非均一な sparsity を誘導します。CNNs（ResNet50、ImageNet 上の MobileNetV1）で最先端の非構造的 sparsity 精度を達成しつつ FLOPs を削減し、RNN における構造化 sparsity へも拡張します。

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.

研究の動機と目的

• パラメータ予算の下で精度を最大化するために層間で非均一な sparsity budgets を動機づける。
• Backpropagation によって各層の prune 阈値を学習する Soft Threshold Reparameterization (STR) を導入する。
• STR が CNN における非構造的 sparsity の最先端の性能を達成し FLOPs を削減する。
• STR が構造化 sparsity を誘導し、学習した sparsity budgets をタスク間で移転できることを示す。

提案手法

• STR を soft-thresholded projection S_g(W_l, s_l) = sign(W_l) * ReLU(|W_l| - g(s_l)) with learnable per-layer s_l.
• 最適化を L(W, D) ではなく L(S_g(W, s), D) を最小化するよう再パラメータ化し、 prune threshold への backpropagation を可能にする。
• 連続関数 g (非構造稀疎化にはシグモイド、構造化 sparsity には指数関数) を用いて s_l を層ごとの閾値 alpha_l = g(s_l) に写像する。
• W_l を SGD と weight decay で訓練しつつ、STR projection に対する L の勾配を介して s_l を更新する。
• STR が層間で sparse な勾配と非均一な budgets を誘導することを示し、後向 pruning への結びつきから収束の直感を論じる。

実験結果

リサーチクエスチョン

• RQ1 全体のパラメータ制限の下で精度を最適化するように層ごとの sparsity budgets を学習できるか。
• RQ2 STR は安定した学習と既存の均一またはヒューリスティックな非均一 sparsity 法より精度を向上させるか。
• RQ3 STR は構造化 sparsity や他のアーキテクチャ（例: RNNs）へ拡張し、利得を維持できるか。
• RQ4 学習された sparsity budgets はタスクやデータセット間で転移可能か。
• RQ5 STR は Baselines と比較して層間の 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 における非構造的 sparsity の最先端精度を複数の sparsity レベルで達成した。
• STR は層ごとの非均一な budgets を学習し、場合によっては FLOPs を最大 about 50% 減らす。
• 超 sparse レジーム（99%）では STR がベースラインを最大約 10% 精度で上回る。
• STR は最終的な層ごとの閾値を明確に分割させ、 Earlier 層での計算削減を集中的にし、他の層で密度を保持する非均一な sparsity budgets を生み出す。
• 学習された sparsity budgets はタスク間で転送可能で、構造化 sparsity（RNN の低ランク化）へ拡張可能である。