[論文レビュー] Scaling Vision with Sparse Mixture of Experts

Carlos Riquelme, Joan Puigcerver+6 more|arXiv (Cornell University)|Jun 10, 2021

Domain Adaptation and Few-Shot Learning参考文献 60被引用数 29

ひとこと要約

Vision MoE (V-MoE) を導入した Vision Transformer のスパース版で、一部の MLP ブロックを Mixture-of-Experts 層に置換し、大規模なビジョンモデルを、推論コストを抑えつつ密なモデルに匹敵させ、最大で 15B パラメータまでスケール可能にします。

ABSTRACT

Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are "dense", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-of-the-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet.

研究の動機と目的

• スパースな Mixture-of-Experts がビジョンモデルを効果的にスケールさせられるかを調査する。
• V-MoE が推論コストを削減しつつ、密な ViT の性能と同等またはそれを上回ることを示す。
• トレーニングを安定化させ転移を改善するためのルーティングとキャパシティ戦略を開発する。
• Batch Prioritized Routing を導入し、画像ごとまたはバッチごとに計算を適応させる。
• 最大 15B パラメータのビジョンモデルが ImageNet で高い性能を達成することを示す。

提案手法

• 選択された ViT の MLP ブロックをスパース MoE 層に置換し、各トークンをエキスパートの小さなサブセットにルーティングします。
• トーク関数 g(x) を用い、softmax(Wx+ε) に TOP_k を適用してトークンを k 個のエキスパートに割り当てます（k は通常 1 または 2）。
• ノイズ ε を追加し、学習中にエキスパートの負荷を均等化するため容量認識バッファ B_e を用います。
• 容量比 C によってエキスパートのバッファ容量を固定し、負荷分散を促す補助損失を用います。
• 大規模でノイズの多いデータ（JFT-300M）で訓練し、ImageNet と VTAB で Linear Probing とフルファインチューニングによる転移を評価します。
• 推論時に低ユーティリティなトークンをスキップすることを可能にし、バッチ全体でトークンを優先付けする Batch Prioritized Routing を導入します。

実験結果

リサーチクエスチョン

• RQ1 ビジョン・トランスフォーマーにおけるスパース MoE 層は、計算量を削減しつつ競争力のある精度を達成できるか？
• RQ2 ルーティング、容量制御、ノイズが V-MoE のトレーニングの安定性と性能にどのように影響するか？
• RQ3 Batch Prioritized Routing と可変容量が推論時の計算量と性能にどのような利点をもたらすか？
• RQ4 V-MoE モデルは下流タスクや few-shot / ファインチューニングの状況にどれだけうまく転移するか？
• RQ5 ImageNet におけるパラメータ数と精度の観点で V-MoE のスケーリング潜力はどの程度か？

主な発見

モデル	パラメータ	JFT prec@1	IN/1shot	IN/5shot	IN/10shot	IN/Fine-t.	ExaFLOPs	TPUv3-days
VIT-H/14	656M	56.68	62.34	76.95	79.02	88.08	4.27k	2.38k
V-MoE-L/16, Every-2	3.4B	57.65	62.41	77.10	79.01	87.41	2.17k	1.20k
V-MoE-H/14, Last-5	2.7B	60.12	62.95	78.08	80.10	88.23	4.75k	2.73k
V-MoE-H/14, Every-2	7.2B	60.62	63.38	78.21	80.33	88.36	5.79k	3.47k
V-MoE-15B, Every-2	14.7B	—	68.66	82.78	84.29	90.35	33.9k	16.8k

• V-MoE の派生は、推論計算量を概ね半分程度で抑えつつ、密な ViT の性能に匹敵するかそれを上回る。
• 15B パラメータの V-MoE モデル（V-MoE-15B）は、完全微調整時に ImageNet で 90.35% を達成。
• Batch Prioritized Routing はトレーニング FLOP を約 20% 減らし、画像ごとの計算量のトレードオフを可能にする。
• MoE 層を用いた JFT-300M での事前訓練は、few-shot および全ファインチューニング設定で強い転移性能を示す。
• V-MoE モデルは推論時に k と容量 C を調整する柔軟性を示し、性能の最小限の損失で大幅な計算削減を可能にする。
• 最大の V-MoE モデル（15B）は ImageNet 相当のベンチマークで競争力を発揮し、拡張可能なビジョンモデル容量を示している。