[논문 리뷰] DSA: More Efficient Budgeted Pruning via Differentiable Sparsity Allocation

Xuefei Ning, Tianchen Zhao외 4명|arXiv (Cornell University)|2020. 04. 05.

Advanced Image and Video Retrieval Techniques참고 문헌 25인용 수 9

한 줄 요약

DSA는 FLOPs 제약 조건 하에서 ADMM 기반 최적화를 통해 미분 가능한 스파arsity 할당 방법을 제안하여, 경량화된 프루닝을 위한 엔드 투 엔드, 프루닝-프롬-스크래치 최적화를 가능하게 한다. 이는 새로운 미분 가능한 프루닝 과정을 사용함으로써 CIFAR-10과 ImageNet에서 반복적 프루닝 방법 대비 최소 1.5배의 속도 향상을 달성하며, 최신 기술 수준(SOTA)의 성능을 기록한다.

ABSTRACT

Budgeted pruning is the problem of pruning under resource constraints. In budgeted pruning, how to distribute the resources across layers (i.e., sparsity allocation) is the key problem. Traditional methods solve it by discretely searching for the layer-wise pruning ratios, which lacks efficiency. In this paper, we propose Differentiable Sparsity Allocation (DSA), an efficient end-to-end budgeted pruning flow. Utilizing a novel differentiable pruning process, DSA finds the layer-wise pruning ratios with gradient-based optimization. It allocates sparsity in continuous space, which is more efficient than methods based on discrete evaluation and search. Furthermore, DSA could work in a pruning-from-scratch manner, whereas traditional budgeted pruning methods are applied to pre-trained models. Experimental results on CIFAR-10 and ImageNet show that DSA could achieve superior performance than current iterative budgeted pruning methods, and shorten the time cost of the overall pruning process by at least 1.5x in the meantime.

연구 동기 및 목표

• 이산적이고 검색 기반의 스파arsity 할당에 의존하는 전통적인 반복적 예산 기반 프루닝 방법의 비효율성을 해결한다.
• 프루닝-프롬-스크래치 최적화를 가능하게 하여 사전 훈련의 필요성을 제거한다.
• 연속 최적화를 위한 기울기 계산이 가능한 프루닝 비율에 대해 기울기 흐름을 허용하는 미분 가능한 프루닝 과정을 개발한다.
• FLOPs 또는 지연 시간 제약 조건 하에서 스파arsity 할당의 샘플 효율성과 계산 효율성을 향상시킨다.
• 스킵 연결로 인한 구조적 제약 조건을 처리하기 위해 군집화 절차를 통해 구조적 통합성을 확보한다.

제안 방법

• 채널 마스크에 대한 확률 분포를 사용하여 딱딱한 프루닝을 부드럽게 하는 미분 가능한 프루닝 과정을 도입하여 프루닝 비율을 통한 기울기 흐름을 가능하게 한다.
• 정확도의 유도 변수로 검증 손실을 사용하여 스파arsity 할당의 기울기 기반 최적화를 가능하게 한다.
• 예산 제약 조건 하에서 가중치와 프루닝 비율을 동시에 최적화하기 위해 다중수단의 교대 최적화(ADMM)를 적용한다.
• ADMM 업데이트를 주요 훈련 루프에 통합하여 가중치 최적화 단계 수십 번마다 스파arsity 할당을 업데이트함으로써 오버헤드를 줄인다.
• 스킵 연결 그룹을 단위로 간주하여 잔차 신경망에서 연결성을 유지하기 위한 위상적 군집 전략을 구현한다.
• 층별 민감도를 고려하여 프루닝 비율에 대한 기울기 크기를 정규화하고 정렬하여 합리적인 스파arsity 할당을 이끌어내는 데 기여한다.

실험 결과

연구 질문

• RQ1 연속적인 스파arsity 비율에 대한 기울기 기반 최적화가 이산적 검색에 비해 예산 기반 프루닝에서 더 나은 성능을 내는가?
• RQ2 사전 훈련 없이도 미분 가능한 프루닝 과정이 효과적이고 효율적인 스파arsity 할당을 가능하게 하는가?
• RQ3 작업 손실에 대한 프루닝 비율에 대한 기울기들이 실제 층별 민감도와 잘 일치하는가?
• RQ4 DSA는 CIFAR-10과 ImageNet에서 반복적 프루닝 기준선 대비 뛰어난 정확도와 더 빠른 수렴을 달성할 수 있는가?
• RQ5 DSA의 프루닝-프롬-스크래치 접근 방식은 고비용의 사전 훈련 단계를 제거하면서도 성능을 유지할 수 있는가?

주요 결과

• ResNet-18에서 FLOPs의 60%를 유지할 경우 DSA는 ImageNet에서 정확도가 1.11% 감소하는 것으로 나타났으며, 이는 모든 반복적 프루닝 기준선을 초월한다.
• DSA는 반복적 방법 대비 총 프루닝 시간을 최소 1.5배 이상 단축시켰으며, CIFAR-10에서는 5/300 GPU 시간, ImageNet에서는 150/300 GPU 시간을 기록했다.
• 프루닝 비율에 대한 기울기 크기가 실제 층별 민감도와 잘 일치함을 확인하여, 스파arsity 할당의 합리성을 검증했다.
• DSA는 잔차 블록과 같은 위상적으로 관련된 층들을 군집화하여 모델의 연결성을 성공적으로 유지함으로써 구조적 통합성을 확보했다.
• DSA의 프루닝-프롬-스크래치 접근 방식은 사전 훈련이 필요 없게 되어 총 계산 비용을 줄였으며, 300 에포크의 사전 훈련 단계를 회피함으로써 비용을 절감했다.
• CIFAR-10에서 DSA는 AMC 및 SSL과 같은 기준선들을 여러 FLOPs 예산에서 일관되게 뛰어넘어 강건성과 효율성을 입증했다.