[논문 리뷰] Bi-directional Masks for Efficient N:M Sparse Training

Yuxin Zhang, Yiting Luo외 5명|arXiv (Cornell University)|2023. 02. 13.

Speech and Audio Processing인용 수 6

한 줄 요약

이 논문은 양방향 마스크(Bi-Mask)를 제안하며, 이는 전방 및 역방향 가중치 희소성에 별도의 마스크를 사용하여 전방 및 역방향 모두에서 N:M 희소 텐서 코어를 완전히 가속할 수 있도록 하는 효율적인 N:M 희소 학습을 위한 새로운 방법이다. 기존 방법에 비해 성능 저하 없이 전체 가속을 달성한다. 전방 마스크 적용 이전에 효율적인 행 순열을 적용함으로써 Bi-Mask는 고정밀도 모델 정확도를 유지하면서도 기존 방법보다 뛰어난 학습 효율성을 확보한다.

ABSTRACT

We focus on addressing the dense backward propagation issue for training efficiency of N:M fine-grained sparsity that preserves at most N out of M consecutive weights and achieves practical speedups supported by the N:M sparse tensor core. Therefore, we present a novel method of Bi-directional Masks (Bi-Mask) with its two central innovations in: 1) Separate sparse masks in the two directions of forward and backward propagation to obtain training acceleration. It disentangles the forward and backward weight sparsity and overcomes the very dense gradient computation. 2) An efficient weight row permutation method to maintain performance. It picks up the permutation candidate with the most eligible N:M weight blocks in the backward to minimize the gradient gap between traditional uni-directional masks and our bi-directional masks. Compared with existing uni-directional scenario that applies a transposable mask and enables backward acceleration, our Bi-Mask is experimentally demonstrated to be more superior in performance. Also, our Bi-Mask performs on par with or even better than methods that fail to achieve backward acceleration. Project of this paper is available at \url{https://github.com/zyxxmu/Bi-Mask}.

연구 동기 및 목표

• N:M 희소 텐서 코어 하드웨어 지원에도 불구하고 밀도 높은 역방향 전파로 인한 N:M 희소 학습의 비효율성 문제를 해결한다.
• 기존 방법에서 유연성 부족한 전치 가능한 마스크로 인해 발생하는 성능 저하 문제를 해결한다.
• 모델 정확도를 유지하면서도 N:M 희소 학습에서 전방 및 역방향 모두를 가속화할 수 있도록 한다.
• 희소 학습에서 전방 및 역방향 전달 간의 기울기 갭을 최소화할 수 있는 유연하면서도 효율적인 메커니즘을 개발한다.
• 양방향 마스크와 지능적인 순열 전략이 비가속화된 희소 학습 방법의 성능을 따라하거나 능가할 수 있음을 입증한다.

제안 방법

• 전방 전파(행)용과 역방향 전파(열)용으로 별도의 N:M 희소 마스크를 도입하여 희소성 패턴을 분리한다.
• 역방향 마스크를 계산하기 이전에 가중치 행렬에 효율적인 행 순열을 적용하여 유효한 N:M 블록 수를 최대화한다.
• 순열 처리된 가중치에 대해 크기 기반 프루닝을 적용하여 역방향 마스크를 생성함으로써 N:M 희소 텐서 코어 가속과의 호환성을 확보한다.
• 출력 기울기의 열 순서를 순열 처리된 가중치 행렬과 일치시켜 기울기 일致성을 유지함으로써 기울기 갭을 최소화한다.
• 비용이 많이 드는 탐색을 피하기 위해 작은 후보 집합(예: 12개)에서 N:M 블록 수를 기반으로 최적의 순열을 선택한다.
• 역방향 마스크의 이진화 기준으로 가중치 크기를 사용하여 전방 마스크와 일치시키고 기울기 불일치를 줄인다.

(a) Vanilla N:M Mask

실험 결과

연구 질문

• RQ1 전치 가능한 마스크 제약 조건을 부과하지 않고도 N:M 희소 학습에서 전방 및 역방향 모두를 완전히 가속화할 수 있는가?
• RQ2 전방 및 역방향 희소성 마스크를 분리함으로써 모델 성능과 학습 효율성에 어떤 영향을 미치는가?
• RQ3 효율적인 행 순열 전략은 역방향 전파에서 유효한 N:M 블록 수를 증가시키면서 기울기 갭을 최소화하는 데 얼마나 기여하는가?
• RQ4 Bi-Mask는 성능을 희생하여 가속을 달성하거나 역방향 전파를 가속화하지 못하는 기존 방법보다 뛰어나게 성능을 높이는가?
• RQ5 다양한 이진화 기준이 역방향 마스크의 효과성과 전체 모델 정확도에 어떤 영향을 미치는가?

주요 결과

• Bi-Mask는 전방 및 역방향 전달의 희소성 패턴을 분리함으로써 N:M 희소 텐서 코어에서 전방 및 역방향 모두를 완전히 가속화한다.
• 특히 1:8 및 1:16와 같은 높은 희소 수준에서 기존 전치 가능한 마스크 방법 대비 성능 저하를 줄였다.
• 제거 실험 결과, 역방향 마스크 도입으로 ResNet-18에서 Top-1 정확도가 0.3% 감소하는 데 그치며, 순열 업데이트는 기준선 대비 정확도를 0.3% 향상시켰다.
• 이진화 기준 중에서 가중치 크기를 사용할 경우 가장 우수한 성능(2:4 희소성에서 ResNet-18의 Top-1 정확도 70.73%)을 기록했으며, 기울기 크기(70.56%) 및 무작위 선택(67.76%)보다 뛰어났다.
• 한정된 후보 집합(예: 12개)을 사용한 순열 전략은 근사 최적의 N:M 블록 수를 확보하는 데에 충분하며, 계산적으로 효율적이다.
• 전치 가능한 마스크에 의존하는 최신 기술(SOTA) 방법을 뛰어넘으며, 비가속화 기준선조차도 정확도에서 앞서는 것으로 나타나 효율성과 효과성 양면에서 뛰어난 성능을 입증했다.

(b) Transposable N:M Mask