[論文レビュー] Cross-token Modeling with Conditional Computation

Yuxuan Lou, Fuzhao Xue+2 more|arXiv (Cornell University)|Sep 5, 2021

Advanced Neural Network Applications被引用数 4

ひとこと要約

本稿では、空間的文脈をモデル化するためのトークン混合MoEとチャネル別情報フローをモデル化するためのチャネル混合MoEを用いた、すべてのMLPに基づくビジョンモデルSparse-MLPを提案する。ViT-L/16と比較してパラメータが62.8%少なく、事前学習コストが85.5%低く、ImageNet Top-1精度79.2%を達成し、MLP-Mixerを2.5%上回りながらも計算コストは低くなる。

ABSTRACT

Mixture-of-Experts (MoE), a conditional computation architecture, achieved promising performance by scaling local module (i.e. feed-forward network) of transformer. However, scaling the cross-token module (i.e. self-attention) is challenging due to the unstable training. This work proposes Sparse-MLP, an all-MLP model which applies sparsely-activated MLPs to cross-token modeling. Specifically, in each Sparse block of our all-MLP model, we apply two stages of MoE layers: one with MLP experts mixing information within channels along image patch dimension, the other with MLP experts mixing information within patches along the channel dimension. In addition, by proposing importance-score routing strategy for MoE and redesigning the image representation shape, we further improve our model's computational efficiency. Experimentally, we are more computation-efficient than Vision Transformers with comparable accuracy. Also, our models can outperform MLP-Mixer by 2.5\% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On downstream tasks, i.e. Cifar10 and Cifar100, our models can still achieve better performance than baselines.

研究の動機と目的

• 自己注意機構に基づくMoEの訓練の不安定性を解消するため、クロストークン注意機構を条件付き計算に置き換えること。
• 空間的（トークン）およびチャネル次元の両方でMoEを適用するスケーラブルで効率的なすべてのMLPアーキテクチャを設計すること。
• 重要度スコアルーティングと表現形状の再設計により、MoEの効率を向上させること。
• ViT や MLP-Mixer よりも低いパラメータ数と計算コストで、ImageNetおよび下流のビジョンベンチマークで最先端の性能を達成すること。

提案手法

• 各チャネル内でパッチ次元にわたってスパースなエキスパートを適用するトークン混合MoE（MoE_S）を導入し、空間的文脈をモデル化する。
• 各パッチ内でチャネルにわたってスパースなエキスパートを適用するチャネル混合MoE（MoE_C）を導入し、チャネル別相互作用をモデル化する。
• ルーティング計算を削減するため、重要度スコアでトークンまたはチャネルを順位付けする重要度スコアルーティング戦略を提案する。
• トークン混合MoEゲーティングネットワークの効率を向上させるために、Sparseブロック内の画像表現形状を再設計する。
• 計算コストを低減し、訓練の安定性を向上させるために、Sparseブロックにリスケールサブレイヤーを採用する。
• 下流の性能を評価するため、ImageNet-1kでMoCo v3を用いて自己教師付き事前学習を実施する。

実験結果

リサーチクエスチョン

• RQ1 自己注意機構に代えて、ビジョントランスフォーマーにおけるクロストークンモデリングにMixture-of-Expertsを効果的かつ安定的に適用できるか？
• RQ2 空間的およびチャネル次元の両方でMoEを適用することで、すべてのMLPアーキテクチャにおけるモデル容量と効率が向上するか？
• RQ3 重要度スコアルーティングにより、モデル性能を損なわずMoEのルーティング計算を削減できるか？
• RQ4 画像表現形状の再設計により、トークン混合MoEの効率が向上するか？
• RQ5 2段階のMoEを備えたすべてのMLPモデルは、精度と効率の両面でViT や MLP-Mixer を上回ることができるか？

主な発見

• Sparse-BはImageNet Top-1精度77.9%を達成し、計算コストは同等のMLP-Mixer-B/16より2.0%高く、ViT-B/16より1.2%高く、コストは低く抑えられている。
• Sparse-LはTop-1精度79.2%を達成し、MLP-Mixer-L/16を2.5%上回り、パラメータが62.8%少なく、事前学習コストが85.5%低い。
• Sparse-LはViT-L/16を1.8%上回り、事前学習コストは半分未塔。
• アブレーションスタディの結果、リスケールサブレイヤーは事前学習コストを30.5%削減し、精度を76.9%から77.9%に向上させた。
• 重要度スコアルーティングは、モデル容量を維持したままMoEのルーティング計算を削減し、効率性を向上させた。
• 空間的およびチャネル別条件付き計算を備えた2段階MoE設計は有効であり、すべてのMLPモデルがViTおよびMLP-Mixerを性能と効率の両面で上回ることを可能にした。