[論文レビュー] Sparse-MLP: A Fully-MLP Architecture with Conditional Computation.

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

Domain Adaptation and Few-Shot Learning参考文献 41被引用数 5

ひとこと要約

Sparse-MLP は、MLP-Mixer にスパースな混合エキスパート（MoE）機構を統合した完全なMLPベースアーキテクチャを提案する。特定の密なMLPブロックを2段階のMoEレイヤーに置き換えることで、次元を跨ぐ特徴の混合を実現する。再表現（Re-represent）レイヤーを導入することでエキスパートの能力とルーティング効率を向上させ、パラメータ数が少なく、計算コストも低いにもかかわらず、密なMLPベースラインより2.5%高いImageNet Top-1精度を達成する。

ABSTRACT

Mixture-of-Experts (MoE) with sparse conditional computation has been proved an effective architecture for scaling attention-based models to more parameters with comparable computation cost. In this paper, we propose Sparse-MLP, scaling the recent MLP-Mixer model with sparse MoE layers, to achieve a more computation-efficient architecture. We replace a subset of dense MLP blocks in the MLP-Mixer model with Sparse blocks. In each Sparse block, we apply two stages of MoE layers: one with MLP experts mixing information within channels along image patch dimension, one with MLP experts mixing information within patches along the channel dimension. Besides, to reduce computational cost in routing and improve expert capacity, we design Re-represent layers in each Sparse block. These layers are to re-scale image representations by two simple but effective linear transformations. When pre-training on ImageNet-1k with MoCo v3 algorithm, our models can outperform dense MLP models by 2.5\% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On small-scale downstream image classification tasks, i.e. Cifar10 and Cifar100, our Sparse-MLP can still achieve better performance than baselines.

研究の動機と目的

• 混合エキスパート（MoE）を用いたスパースな条件付き計算を導入することで、MLP-Mixerの効率性とスケーラビリティを向上させること。
• 視覚変換器やMLP-Mixerアーキテクチャにおける密なMLPブロックの計算非効率性と能力限界を解決すること。
• スパースブロックにおける2段階のMoEレイヤーを用いて、チャネル次元とパッチ次元の両方を跨ぐより効果的な特徴混合を実現すること。
• 表現再スケーリングのためのRe-representレイヤーを導入することでルーティングのオーバーヘッドを低減し、エキスパートの利用効率を向上させること。
• パラメータ数とFLOPsを密なMLPモデルより低く抑えながら、画像分類ベンチマークで最先端の性能を達成すること。

提案手法

• MLP-Mixer内の一部の密なMLPブロックを、2段階のMoEを実装するスパースブロックに置き換える。1つはパッチを跨ぐチャネルワイドな混合を、もう1つはチャネルを跨ぐパッチワイドな混合を担当する。
• MoEルーティングの前後で単純な線形変換を適用するRe-representレイヤーを導入し、特徴表現の再スケーリングを実現することで、エキスパートの能力とルーティング効率を向上させる。
• スパースなMoE設定下でも堅牢な自己教師付き表現学習を保証するため、ImageNet-1kでMoCo v3を用いて事前学習を実施する。
• 各MoEレイヤーでトークンごとに上位-kのエキスパートを動的に選択するルーター機構を適用し、条件付き計算を実現し、FLOPsを削減する。
• エキスパートのスパースな活性化により高容量の特徴学習を可能にしつつ、計算効率を維持するためのスパースブロックの設計を行う。
• 入力表現の安定化と精錬を図るために、Re-representレイヤーを適用することでルーティングプロセスを最適化する。

実験結果

リサーチクエスチョン

• RQ1 MLP-Mixerのような完全なMLPベースアーキテクチャにスパースなMoE機構を効果的に統合できるか？ その結果、効率性と性能が向上するか？
• RQ2 チャネル次元とパッチ次元を別々に処理する2段階のMoE設計は、特徴表現とモデル精度にどのように影響を与えるか？
• RQ3 Re-representレイヤーは、スパースなMLPアーキテクチャにおけるエキスパートの利用効率をどの程度向上させ、ルーティング関連のオーバーヘッドを低減できるか？
• RQ4 スパースなMLPアーキテクチャは、ImageNet-1kおよび小規模ベンチマークにおいて、パラメータ数とFLOPsの点で密なMLPベースラインを上回ることができるか？
• RQ5 提案されたSparse-MLPは、CIFAR-10やCIFAR-100といった下流タスクにおいて、最小限のアーキテクチャ変更で強力な一般化性能を維持できるか？

主な発見

• MoCo v3を用いてImageNet-1kで事前学習した結果、Sparse-MLPは密なMLP-Mixerベースラインよりも2.5%高いImageNet Top-1精度を達成した。
• この性能向上は、パラメータ数が少なく、計算コストも低いという点で、密なMLPモデルよりも優れている。
• CIFAR-10およびCIFAR-100においても、これらのデータセットのスケールが小さいにもかかわらず、既存のベースラインを上回る性能を示した。
• Re-representレイヤーは、MoE機構におけるエキスパートの能力向上とルーティング関連の計算オーバーヘッド低減に効果的に寄与した。
• 2段階のMoE設計により、FLOPsを著しく増加させることなく、効率的な次元跨ぎ特徴混合が可能になり、表現学習が向上した。
• スパースアーキテクチャは、複数のベンチマークで強力な一般化性能を維持しており、スケーラビリティと効率性の優位性を示した。