[논문 리뷰] Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch

Aojun Zhou, Yukun Ma외 6명|arXiv (Cornell University)|2021. 02. 08.

Advanced Neural Network Applications참고 문헌 43인용 수 74

한 줄 요약

본 연구는 SR-STE를 사용하여 N:M 미세 구조적 희소 네트워크를 무에서 시작해 하드웨어 친화적 희소성을 가능하게 하고 Nvidia A100에서 최대 2배 속도 향상을 달성하는 동시에 정확도를 유지하며, 희소 토폴로지 변화를 분석하기 위한 SAD 지표를 도입한다.

ABSTRACT

Sparsity in Deep Neural Networks (DNNs) has been widely studied to compress and accelerate the models on resource-constrained environments. It can be generally categorized into unstructured fine-grained sparsity that zeroes out multiple individual weights distributed across the neural network, and structured coarse-grained sparsity which prunes blocks of sub-networks of a neural network. Fine-grained sparsity can achieve a high compression ratio but is not hardware friendly and hence receives limited speed gains. On the other hand, coarse-grained sparsity cannot concurrently achieve both apparent acceleration on modern GPUs and decent performance. In this paper, we are the first to study training from scratch an N:M fine-grained structured sparse network, which can maintain the advantages of both unstructured fine-grained sparsity and structured coarse-grained sparsity simultaneously on specifically designed GPUs. Specifically, a 2:4 sparse network could achieve 2x speed-up without performance drop on Nvidia A100 GPUs. Furthermore, we propose a novel and effective ingredient, sparse-refined straight-through estimator (SR-STE), to alleviate the negative influence of the approximated gradients computed by vanilla STE during optimization. We also define a metric, Sparse Architecture Divergence (SAD), to measure the sparse network's topology change during the training process. Finally, We justify SR-STE's advantages with SAD and demonstrate the effectiveness of SR-STE by performing comprehensive experiments on various tasks. Source codes and models are available at https://github.com/NM-sparsity/NM-sparsity.

연구 동기 및 목표

• GPU에서 DNN 속도를 높이기 위해 비구조적 희소성과 구조적 희소성을 결합하는 동기를 제시한다.
• 큰 성능 저손실 없이 무에서 시작하여 N:M 희소 네트워크를 학습하는 프레임워크를 제안한다.
• 훈련 중 그래디언트로 인한 아키텍처 교란을 완화하기 위해 SR-STE를 도입한다.
• 훈련 중 토폴로지 변화를 정량화하기 위해 Sparse Architecture Divergence(SAD)를 정의한다.
• 비전 태스크 및 기계 번역에 걸친 효과를 입증한다.

제안 방법

• 연속된 M개 가중치 그룹마다 최대 N개만 비제로인 N:M 희소성을 정의한다.
• 훈련 중 온라인 프루닝을 가능하게 하기 위해 Straight-through Estimator(STE)를 확장한다.
• 훈련 중 토폴로지 변화를 측정하기 위해 Sparse Architecture Divergence(SAD)를 도입한다.
• 훈련 중 아키텍처를 안정시키기 위해 프루닝된 가중치를 페널티하는 규제항을 가진 Sparse-refined STE (SR-STE)를 제안한다.
• 이미지 분류, 객체 탐지, 인스턴스 세분화, 광학 흐름, 기계 번역에서 평가하고 ASP, STE 및 다른 희소화 방법과 비교한다.

실험 결과

연구 질문

• RQ1 성능 저하 없이 무에서 시작해 N:M 희소 네트워크를 학습할 수 있는가?
• RQ2 SR-STE가 프루닝된 가중치의 그래디언트 불일치를 줄이고 훈련 중 희소 아키텍처를 안정화시키는가?
• RQ3 다른 N:M 패턴(예: 2:4, 4:8, 1:4, 2:8)이 태스크 전반에서 정확도와 속도 향상에 어떻게 영향을 미치는가?
• RQ4 제안된 방법이 희소 모델의 다운스트림 태스크로의 전이 가능성을 유지하는가?

주요 결과

모델	방법	희소 패턴	Top-1 Acc(%)	매개변수(M)	플롭스(G)
ResNet50	Dense	-	77.3	25.6	4.09
ResNet50	SR-STE	2:4	77.0	13.8	2.15
ResNet50	SR-STE	4:8	77.4	13.8	2.15
ResNet50	SR-STE	1:4	75.3	7.93	1.17
ResNet50	SR-STE	2:8	76.2	7.93	1.17
ResNet50 x1.25	SR-STE	2:8	77.5	11.8	1.79

• 2:4 희소 네트워크는 Dense 기준 대비 약 무손실에 가깝고 Nvidia A100 GPU에서 ResNet-50의 ImageNet에서 약 2배의 속도 향상을 달성할 수 있다.
• 4:8 희소성(同 50% 희소도)은 ImageNet의 ResNet-50에서 유사 FLOPs일 때 2:4보다 더 좋은 성능을 낸다.
• SR-STE는 다양한 패턴(예: 2:4, 4:8)에서 ImageNet에 대해 STE 및 ASP 기준선보다 일관되게 Top-1 정확도를 향상시킨다.
• COCO 객체 탐지에서 2:8 희소성은 dense 기준에 근접한 mAP를 제공하고, 4:8은 Faster R-CNN + ResNet-50의 경우 dense 성능을 능가하기도 한다.
• 광학 흐름(RAFT)과 신경 기계 번역(Transformer)에서도 SR-STE는 매개변수와 FLOP를 대폭 줄이면서 Dense 모델에 비해 비슷한 성능을 달성한다.
• SAD 지표는 성능과 상관관계가 있으며 SR-STE가 희소 아키텍처를 안정시키면 감소한다.