[논문 리뷰] Dynamic Sparse Training: Find Efficient Sparse Network From Scratch With Trainable Masked Layers

Junjie Liu, Zhe Xu외 3명|arXiv (Cornell University)|2020. 05. 14.

Advanced Neural Network Applications참고 문헌 31인용 수 34

한 줄 요약

논문은 trainable pruning thresholds를 통해 희소 네트워크 구조와 가중치를 함께 학습하는 엔드-투-엔드 방법인 Dynamic Sparse Training(DST)을 도입하며, 훈련 중 단계별로 미세한 가지치기와 복구를 가능하게 한다.

ABSTRACT

We present a novel network pruning algorithm called Dynamic Sparse Training that can jointly find the optimal network parameters and sparse network structure in a unified optimization process with trainable pruning thresholds. These thresholds can have fine-grained layer-wise adjustments dynamically via backpropagation. We demonstrate that our dynamic sparse training algorithm can easily train very sparse neural network models with little performance loss using the same number of training epochs as dense models. Dynamic Sparse Training achieves the state of the art performance compared with other sparse training algorithms on various network architectures. Additionally, we have several surprising observations that provide strong evidence for the effectiveness and efficiency of our algorithm. These observations reveal the underlying problems of traditional three-stage pruning algorithms and present the potential guidance provided by our algorithm to the design of more compact network architectures.

연구 동기 및 목표

• 추론 시 메모리와 계산을 줄이기 위해 효율적인 희소 네트워크의 필요성을 동기화한다.
• 가중치와 계층별 가지치기 마스크를 모두 학습하는 엔드-투-엔드 희소 학습 프레임워크를 제안한다.
• 각 학습 단계 내에서 epochs 간이 아니라 미세한 단계별 가지치기 및 복구를 가능하게 한다.
• 역전파와 trainable threshold 메커니즘을 통해 계층별 가지치기 비율을 자동으로 조정한다.
• 다양한 아키텍처에서 MNIST, CIFAR-10, ImageNet에 대해 최첨단 성능을 입증한다.

제안 방법

• 각 계층마다 뉴런/필터에 대한 trainable한 계층별 임계값으로 가지치를 표현한다.
• |W| - t에 단위 계단 함수 S를 적용해 얻은 이진 마스크 M으로 희소한 W ∘ M를 얻는다.
• 임계값 벡터 t의 학습을 위한 Straight-Through Estimator(스루-스루 추정) 기반 도함수를 도입한다.
• 더 높은 희소성을 촉진하기 위해 Ls = sum exp(-ti)의 희소 정규화 항을 도입한다.
• Dense 계층을 trainable masked layer로 대체하여 성능 및 구조 그래디언트를 역전파로 전달할 수 있게 한다.
• 각 학습 단계에서 임계값과 마스크를 업데이트할 수 있음을 보여주어 미세한 가지치기 및 복구를 가능하게 한다.

실험 결과

연구 질문

• RQ1 가중치와 희소 구조를 trainable pruning thresholds를 갖춘 엔드-투-엔드 프레임워크에서 함께 학습할 수 있는가?
• RQ2 학습 중 단계별 가지치기와 복구가 희소 학습에서 미리 정의된 가지치기 스케줄보다 더 우수한가?
• RQ3 계층별 학습 가능한 임계값이 다양한 아키텍처에서 최종 희소성 패턴과 모델 성능에 어떤 영향을 미치는가?
• RQ4 DST가 관찰된 희소 패턴을 통해 컴팩트한 아키텍처 설계에 어떤 지침을 제공하는가?

주요 결과

• DST는 MNIST Lenet-300-100에서 파라미터의 거의 98%를 가지치기하고도 성능 손실이 거의 없으며(희소 결과: 남은 비율 2.48%).
• MNIST Lenet-5-Caffe에서 희소 학습은 남은 비율 1.64%를 기록하고 거의 Dense 수준의 정확도를 달성한다(희소: 99.11%).
• 시퀀스 MNIST용 LSTM 모델은 99% 이상의 파라미터를 가지치기하더라도 동등하거나 더 좋은 희소 정확도를 달성한다.
• CIFAR-10의 VGG-16 및 WideResNet에서 DST는 높은 희소도에서 Sparse Momentum 및 Dynamic Sparse Reparameterization을 능가한다(예: VGG-16: 남은 비율 8.82%로 93.93% 희소 정확도, 다른 방법은 10% 남김).
• ImageNet(ResNet-50): DST는 베이스라인보다 약간 더 높은 희소도에서 더 높은 top-1/top-5 정확도를 달성한다(설정에 따라 남은 비율 약 9.87–19.24% 등).
• DST는 서로 다른 α 값에서도 일관된 희소 패턴을 보이며 계층별 중복성을 시사하고 아키텍처 설계에 지침을 제공한다.