[논문 리뷰] Soft Threshold Weight Reparameterization for Learnable Sparsity

Aditya Kusupati, Vivek Ramanujan외 5명|arXiv (Cornell University)|2020. 02. 08.

Advanced Neural Network Applications참고 문헌 58인용 수 82

한 줄 요약

STR은 소프트 임계값 재매개화를 통해 계층별 가지치기 임계값을 학습하여 비균일한 희소성을 유도한다. CNN(ResNet50, ImageNet의 MobileNetV1)에서 FLOPs를 줄이면서 비구조적 희소성 최첨단 정확도를 달성하고 RNN에서 구조적 희소성으로 확장한다.

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이 CNN에서 비구조적 희소성 성능에서 최첨단을 달성하고 FLOPs를 감소시킨다.
• STR이 구조적 희소성을 유도하고 학습된 희소성 예산을 과제 간에 이전할 수 있음을 보인다.

제안 방법

• STR을 가중치 W_l와 각 계층당 학습 가능한 s_l를 갖는 소프트 임계값 투영 S_g(W_l, s_l) = sign(W_l) * ReLU(|W_l| - g(s_l))로 정의한다.
• 최적화를 L(S_g(W, s), D)를 최소화하도록 재매개화하여 L(W, D) 대신 가지치기 임계값을 통해 역전파가 가능하게 한다.
• 연속 함수 g를 사용하여 s_l을 계층별 임계값 alpha_l = g(s_l)으로 매핑한다(비구조적 희소성은 시그모이드, 구조적 희소성은 지수형).
• W_l을 SGD와 가중치 감소로 학습하는 동안 STR 투영에 대한 L의 그래디언트를 통해 s_l를 업데이트한다.
• STR이 희소 그래디언트와 계층 간 비균일 예산을 유도함을 입증하고 역방향 가지치기와의 연결을 통해 수렴 직관을 논의한다.

실험 결과

연구 질문

• RQ1 매개변수 예산 하에서 정확도를 최대화하기 위해 계층 간 비균일 희소성 예산을 학습할 수 있는가?
• RQ2 STR이 기존의 균일 또는 휴리스틱 비균일 희소성 방법에 비해 안정적인 학습과 정확도 향상을 제공하는가?
• RQ3 STR이 구조적 희소성과 다른 아키텍처(RNN 등)로 확장되면서 이득을 유지하는가?
• RQ4 학습된 희소성 예산이 과제 간 또는 데이터 세트 간에 이전되는가?
• RQ5 _STR가 기본선 대비 계층 간 FLOPs 분포에 어떤 영향을 미치는가_?

주요 결과

방법	Top-1 Acc (%)	Params	Sparsity (%)	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의 저랭크 등)으로 확장될 수 있다.