[論文レビュー] M$^3$ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design

Hanxue Liang, Zhiwen Fan+7 more|arXiv (Cornell University)|Oct 26, 2022

Advanced Neural Network Applications被引用数 33

ひとこと要約

M3ViT は、効率的なマルチタスク学習のために視覚トランスフォーマーへ混成専門家層(MoE)を統合し、訓練時および推論時のタスク特異的スパースルーティングを可能にし、FPGA ベースのハードウェア共設計による memory- and energy-efficient on-device deployment.

ABSTRACT

Multi-task learning (MTL) encapsulates multiple learned tasks in a single model and often lets those tasks learn better jointly. However, when deploying MTL onto those real-world systems that are often resource-constrained or latency-sensitive, two prominent challenges arise: (i) during training, simultaneously optimizing all tasks is often difficult due to gradient conflicts across tasks; (ii) at inference, current MTL regimes have to activate nearly the entire model even to just execute a single task. Yet most real systems demand only one or two tasks at each moment, and switch between tasks as needed: therefore such all tasks activated inference is also highly inefficient and non-scalable. In this paper, we present a model-accelerator co-design framework to enable efficient on-device MTL. Our framework, dubbed M$^3$ViT, customizes mixture-of-experts (MoE) layers into a vision transformer (ViT) backbone for MTL, and sparsely activates task-specific experts during training. Then at inference with any task of interest, the same design allows for activating only the task-corresponding sparse expert pathway, instead of the full model. Our new model design is further enhanced by hardware-level innovations, in particular, a novel computation reordering scheme tailored for memory-constrained MTL that achieves zero-overhead switching between tasks and can scale to any number of experts. When executing single-task inference, M$^{3}$ViT achieves higher accuracies than encoder-focused MTL methods, while significantly reducing 88% inference FLOPs. When implemented on a hardware platform of one Xilinx ZCU104 FPGA, our co-design framework reduces the memory requirement by 2.4 times, while achieving energy efficiency up to 9.23 times higher than a comparable FPGA baseline. Code is available at: https://github.com/VITA-Group/M3ViT.

研究の動機と目的

• コンパクトなモデル内でのマルチタスク学習における訓練勾配の衝突に対処する。
• タスク特異的エキスパートをスパースに活性化して、1タスクのみを効率的に推論できるようにする。
• エッジデバイス上のメモリとエネルギーコストを最小化するようなハードウェア認識型 MoE-ViT 設計を開発する。

提案手法

• ViT の密な MLP を、スパースに活性化する Mixture-of-Experts 層に置換する。
• 入力トークンごとにエキスパートのサブセットを選択するタスク依存ルーターを使用する。
• MTL ルーティングのためのタスク条件付きおよび多ゲート MoE 変種を導入する。
• FPGA 上でメモリ効率の良い計算再配置を適用して、エキスパートごとの処理を可能にする。
• ハードウェア設計によるゼロオーバーヘッドのタスク切替を実証する。

実験結果

リサーチクエスチョン

• RQ1 密なビジョントランスフォーマーを用いた MTL において、MoE 層はタスク間の勾配衝突を緩和できるか？
• RQ2 タスク特異的なスパースルーティングは、精度を犠牲にせずに単一タスク推論の効率を向上させるか？
• RQ3 ハードウェア共設計（メモリ認識的計算）は、FPGA 上で複数エキスパートとタスクに対して MoE をどのように拡張できるか？
• RQ4 標準の密な予測ベンチマークにおける M3ViT とエンコーダ中心の MTL モデルの精度、FLOPs、メモリ、エネルギーの比較的な利得はどれくらいか？

主な発見

• 評価されたタスクで、M3ViT はエンコーダ中心の MTL メソッドより高い精度を達成する。
• M3ViT を用いた単一タスク推論は推論 FLOPs を88%削減する。
• FPGA 上では、メモリ使用量が2.40×低減し、エネルギー効率はFPGAベースラインに対して最大9.23×、GPUベースラインに対して最大10.79×改善する。
• MoE ベースの変種（Single, Multi, Task-conditioned）は、タスク依存の MoE 設計が MTL の性能を改善することを示し、特にタスク数が増えるほど顕著。
• ハードウェア認識型計算再配置によりゼロオーバーヘッドのタスク/フレーム切替を可能にし、エキスパートの数に関係なくスケールする。