ShiftAddViT: Mixture of Multiplication Primitives Towards Efficient Vision Transformer

要約

ビジョン トランスフォーマー (ViT) は優れたパフォーマンスを示し、複数のビジョン タスクの統合されたバックボーンとなっています。
ただし、ViT のアテンション メカニズムと多層パーセプトロン (MLP) は両方とも、密な乗算のために十分な効率が得られず、コストのかかるトレーニングと推論につながります。
この目的を達成するために、$\textbf{ShiftAddViT}$ と呼ばれる新しいタイプの乗算削減モデルに向けて、ビット単位のシフトや加算などの乗算プリミティブを組み合わせて事前トレーニングされた ViT を再パラメータ化することを提案します。
-ゼロからトレーニングする必要がなく、GPU でのエンドツーエンドの推論が高速化されます。
具体的には、クエリ、キー、および値のすべての $\texttt{MatMuls}$ は、クエリとキーをハミング空間のバイナリ コードにマッピングした後、加法カーネルを使用して再パラメータ化されます。
残りの MLP または線形層は、シフト カーネルで再パラメータ化されます。
TVM を利用して、GPU 上で実際のハードウェアを展開するためにカスタマイズされたカーネルを実装および最適化します。
このような注意に関する再パラメータ化はモデルの精度を維持しますが、MLP に適用すると必然的に精度の低下につながることがわかりました。
両方の長所を組み合わせるために、乗算またはそのプリミティブ (乗算とシフトなど) をエキスパートとして取り、新しいレイテンシを考慮したロード バランシング損失を設計することにより、MLP を再パラメータ化するための新しいエキスパート混合 (MoE) フレームワークをさらに提案します。
このような損失は、遅延に応じてさまざまなエキスパートに動的な量の入力トークンを割り当てるように汎用ルーターをトレーニングするのに役立ちます。
さまざまな 2D/3D Transformer ベースのビジョン タスクに関する広範な実験により、当社が提案する ShiftAddViT の有効性が一貫して検証され、GPU で最大 $\textbf{5.18$\times$}$ のレイテンシ削減と $\textbf{42.9}$% のエネルギー節約が達成されます。
オリジナルまたは効率的な ViT と同等の精度を維持しながら。

要約(オリジナル)

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 $\textbf{ShiftAddViT}$, which aims to achieve end-to-end inference speedups on GPUs without requiring training from scratch. Specifically, all $\texttt{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 $\textbf{5.18$\times$}$ latency reductions on GPUs and $\textbf{42.9}$% energy savings, while maintaining a comparable accuracy as original or efficient ViTs.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google