[論文レビュー] ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

Haoran You, Huihong Shi+2 more|arXiv (Cornell University)|Jun 10, 2023

Advanced Neural Network Applications被引用数 6

ひとこと要約

本稿では、ビット単位のシフトと加算の組み合わせである乗算プリミティブの混合を用いて、自己注意機構とMLP層を再パラメータ化することで、微調整なしにGPU上でエンドツーエンドの推論高速化を実現するShiftAddViTを提案する。混合専門家（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 シフトや加算といったハードウェアにインspiredなプリミティブを用いることで、2次元および3次元のトランスフォーマーに基づくビジョンモデルにおける遅延とエネルギー消費はどの程度削減できるか？

主な発見

• ShiftAddViTは、オリジナルまたは効率的なViTと同等の精度を維持しながら、GPU上で最大5.18倍の遅延削減を達成した。
• ベースラインのViTと比較して、エネルギー消費を42.9％削減し、優れたエネルギー効率を示した。
• 遅延に配慮した負荷バランス損失により、PVTv2-B0では遅延が100％から85.4％に14.6％削減され、精度低下はたった0.01％にとどまった。
• 可視化の結果、ルーターが物体を含むトークンを乗算専門家に、背景のトークンをシフト専門家に適切に割り当てていることが確認された。
• 加法的カーネルによる注意機構の再パラメータ化はモデル精度を維持したが、MLPの再パラメータ化のみでは精度低下が生じたため、MoEフレームワークの導入が不可欠であった。
• 非対称な専門家を備えたMoEフレームワークはソフトな置き換え戦略を提供し、重要度の高いトークンの精度を保持しながら、重要度が低いトークンの処理を高速化した。