[논문 리뷰] ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

Haoran You, Huihong Shi외 2명|arXiv (Cornell University)|2023. 06. 10.

Advanced Neural Network Applications인용 수 6

한 줄 요약

이 논문은 시각 트랜스포머의 주의 메커니즘과 MLP 레이어를 재구성하여 비트단위 시프트와 덧셈 연산의 혼합을 사용해 가속화하는 ShiftAddViT를 제안한다. 이는 재학습 없이도 GPU에서 엔드 투 엔드 추론 속도 향상을 가능하게 한다. 혼합 전문가(MoE) 프레임워크 내에서 지연 시간을 고려한 로드 밸런싱 손실을 도입하여 토큰을 효율적인 시프트 전문가 또는 정확도가 높은 곱셈 전문가로 동적으로 라우팅함으로써, 최대 5.18배의 지연 시간 감소와 42.9%의 에너지 절감을 달성하면서도 모델 정확도를 유지한다.

ABSTRACT

Vision Transformers (ViTs) have shown impressive performance and have become a unified backbone for multiple vision tasks. However, both the attention mechanism and multi-layer perceptrons (MLPs) in ViTs are not sufficiently efficient due to dense multiplications, leading to costly training and inference. To this end, we propose to reparameterize pre-trained ViTs with a mixture of multiplication primitives, e.g., bitwise shifts and additions, towards a new type of multiplication-reduced model, dubbed $ extbf{ShiftAddViT}$, which aims to achieve end-to-end inference speedups on GPUs without requiring training from scratch. Specifically, all $ exttt{MatMuls}$ among queries, keys, and values are reparameterized using additive kernels, after mapping queries and keys to binary codes in Hamming space. The remaining MLPs or linear layers are then reparameterized with shift kernels. We utilize TVM to implement and optimize those customized kernels for practical hardware deployment on GPUs. We find that such a reparameterization on attention maintains model accuracy, while inevitably leading to accuracy drops when being applied to MLPs. To marry the best of both worlds, we further propose a new mixture of experts (MoE) framework to reparameterize MLPs by taking multiplication or its primitives as experts, e.g., multiplication and shift, and designing a new latency-aware load-balancing loss. Such a loss helps to train a generic router for assigning a dynamic amount of input tokens to different experts according to their latency. Extensive experiments on various 2D/3D Transformer-based vision tasks consistently validate the effectiveness of our proposed ShiftAddViT, achieving up to $ extbf{5.18$ imes$}$ latency reductions on GPUs and $ extbf{42.9}$% energy savings, while maintaining a comparable accuracy as original or efficient ViTs.

연구 동기 및 목표

• 자신의 주의 메커니즘과 MLP에서 발생하는 조밀한 곱셈으로 인한 높은 하드웨어 지연 시간과 에너지 소비를 줄이기 위해.
• 사전 훈련된 ViT를 재구성함으로써 재학습 없이도 GPU에서 엔드 투 엔드 추론 속도 향상을 달성하기 위해.
• 곱셈 연산을 더 저렴한 시프트와 덧셈 연산으로 대체할 때 모델 정확도를 유지하기 위해 토큰을 적응적으로 라우팅하기 위해.
• 하드웨어 우수한 원리와 새로운 로드 밸런싱 메커니즘을 활용해 주의 메커니즘과 MLP를 동시에 최적화하여 효율성을 높이기 위해.
• 계산 비용과 정확도가 다른 전문가가 존재하는 혼합 전문가 프레임워크를 설계하여 지연 시간에 따라 토큰을 동적으로 할당할 수 있도록 하기 위해.

제안 방법

• 해밍 공간에서 쿼리와 키를 이진 코드로 매핑한 후, 자기 주의 메커니즘의 모든 행렬 곱셈을 덧셈 커널을 사용해 재구성한다.
• 남은 MLP와 선형 레이어를 재활용하여 곱셈을 비트단위 시프트와 덧셈으로 대체하는 시프트 커널을 사용한다.
• TVM를 통해 커스텀 CUDA 커널을 구현하여 GPU 배포에 최적화된 시프트 및 덧셈 연산을 구현한다.
• 곱셈 연산과 시프트 연산이 서로 다른 계산 비용을 가지는 비균형 전문가로 구성된 혼합 전문가(MoE) 프레임워크를 도입한다.
• 지연 시간을 고려한 로드 밸런싱 손실을 설계하여 중요한 토큰은 정확도가 높은 전문가(예: 곱셈 전문가)에 할당하고, 덜 중요한 토큰은 빠른 전문가(예: 시프트 전문가)에 할당함으로써 정확도를 유지한다.
• 라우터를 사용해 지연 시간과 중요도에 따라 중요한 민감한 토큰을 곱셈 전문가로, 덜 중요한 토큰을 시프트 전문가로 동적으로 라우팅한다.

실험 결과

연구 질문

• RQ1 사전 훈련된 시각 트랜스포머를 재학습 없이도 곱셈 연산을 시프트와 덧셈 원리로 대체하여 GPU에서 가속화할 수 있는가?
• RQ2 ViT에서 조밀한 곱셈 연산을 근사적인 시프트-덧셈 연산으로 대체할 때 정확도를 어떻게 유지할 수 있는가?
• RQ3 비균형 전문가(곱셈 대비 시프트)를 가진 혼합 전문가 프레임워크가 고정된 대체 전략보다 더 나은 정확도-효율성 트레이드오���을 달성할 수 있는가?
• RQ4 지연 시간을 고려한 로드 밸런싱 손실이 동기화 오버헤드를 최소화하고 최대 속도 향상을 달성하는 데 효과적으로 기여하는가?
• RQ5 시프트와 덧셈과 같은 하드웨어 기반 원리가 2D 및 3D 트랜스포머 기반 시각 모델에서 지연 시간과 에너지 소비를 얼마나 줄일 수 있는가?

주요 결과

• ShiftAddViT는 원래 ViT나 효율적인 ViT와 유사한 정확도를 유지하면서 GPU에서 최대 5.18배의 지연 시간 감소를 달성한다.
• 기본 ViT 대비 에너지 소비를 42.9% 절감하여 뛰어난 에너지 효율성을 입증한다.
• 지연 시간 인식 로드 밸런싱 손실 덕분에 PVTv2-B0에서 지연 시간을 100%에서 85.4%로 14.6% 감소시켰고, 정확도 손실는 단 0.01%에 그쳤다.
• 시각화 결과 라우터가 객체를 포함한 토큰을 곱셈 전문가로, 배경 토큰을 시프트 전문가로 성공적으로 라우팅하는 것으로 확인되었다.
• 덧셈 커널을 통한 주의 레이어 재구성은 모델 정확도를 유지하지만, MLP 레이어의 재구성만으로는 정확도 저하가 발생하므로 MoE 프레임워크가 필수적이다.
• 비균형 전문가를 가진 MoE 프레임워크는 부드러운 대체 전략을 제공하여 중요한 토큰의 정확도를 유지하면서 덜 중요한 토큰은 가속화한다.