[論文レビュー] Spatial Mixture-of-Experts

Nikoli Dryden, Torsten Hoefler|arXiv (Cornell University)|Nov 24, 2022

Human Mobility and Location-Based Analysis被引用数 5

ひとこと要約

本論文は、テンソルルーティング関数を用いてピクセル単位でエキスパートをルーティングすることで、データ内の細粒度の空間構造を学習するスパarsely-gatedニューラルネットワーク部品、空間的混合エキスパート（SMoE）レイヤーを導入する。空間的インダクティブバイアスを自己教師付きルーティング損失とエキスパート誤差ダミングを用いて組み込むことで、密度の高い回帰タスクにおいて畳み込み型および標準的なMoEモデルを上回る、中距離天気予報およびアンサンブル後処理において最先端の性能を達成した。

ABSTRACT

Many data have an underlying dependence on spatial location; it may be weather on the Earth, a simulation on a mesh, or a registered image. Yet this feature is rarely taken advantage of, and violates common assumptions made by many neural network layers, such as translation equivariance. Further, many works that do incorporate locality fail to capture fine-grained structure. To address this, we introduce the Spatial Mixture-of-Experts (SMoE) layer, a sparsely-gated layer that learns spatial structure in the input domain and routes experts at a fine-grained level to utilize it. We also develop new techniques to train SMoEs, including a self-supervised routing loss and damping expert errors. Finally, we show strong results for SMoEs on numerous tasks, and set new state-of-the-art results for medium-range weather prediction and post-processing ensemble weather forecasts.

研究の動機と目的

• 位置に基づく構造を持つデータにおいて、標準的なニューラルネットワークが細粒度の空間的依存関係を捉えることのできないという制限に対処すること。
• 密度の高い回帰タスクにおいて、畳み込み型および標準的な混合エキスパートモデルが位置依存パターンを学習できないという問題を克服すること。
• 各空間的位置に応じてエキスパートが特化できるが、計算効率を維持できるルーティング機構を開発すること。
• 空間的インダクティブバイアスを新たなトレーニング技術を用いて組み込むことで、天気予測タスクにおけるモデル性能を向上させること。

提案手法

• 入力をピクセル単位でスパースにルーティングする学習可能なゲーティング関数を用いる空間的混合エキスパート（SMoE）レイヤーを提案する。
• 入力ドメイン内の空間的構造を明示的にモデル化する、軽量で効果的なルーティング関数「テンソルルーティング」を導入する。
• ルーティング誤差を予測することで、ゲートの学習信号を強化するため、自己教師付きルーティング分類損失を採用する。
• 誤ってルーティングされたサンプルからの勾配更新を抑制するエキスパート誤差ダミングを適用し、ロバスト性と性能を向上させる。
• 同じ空間的次元を持つ層同士でゲートを共有することで、パラメータ数を削減し、一般化性能を向上させる。
• U-NetおよびResNetアーキテクチャ内にSMoEレイヤーを導入し、標準的な畳み込みを置き換えることで、空間的特化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 位置に基づく構造を持つデータにおいて、混合エキスパートアーキテクチャは細粒度の空間的依存関係を効果的に学習できるか？
• RQ2 MoEにおける細粒度のピクセル単位のルーティングは、粗粒度のルーティングや標準的な畳み込みと比較して、密度の高い回帰タスクでより優れた性能を発揮するか？
• RQ3 エンドツーエンドのルーティングが空間的構造を捉えられない場合、自己教師付きルーティング損失はゲートの学習を改善できるか？
• RQ4 ルーティング誤りが生じる状況下で、エキスパート誤差ダミングはモデルの収束性と性能にどのように影響するか？
• RQ5 SMoEレイヤーは、パラメータ数を前例の手法より少なく抑えながら、天気予測や画像分類を含む多様なタスクに一般化可能か？

主な発見

• SMoEは、中距離天気予報のためのWeatherBenchベンチマークで、先行SOTA手法を上回る最先端の結果を達成した。
• アンサンブル天気予報の後処理を目的としたENS-10データセットにおいて、すべてのリードタイムおよびアンサンブルサイズで、CRPSおよびEECRPSスコアが低く、新たなSOTAを樹立した。
• 位置依存の拡散係数を有する制御された熱拡散タスクにおいて、SMoEは空間的構造を正しく学習したが、標準的なCNNおよびMoEは平均拡散係数に収束した。
• ImageNet-1kでは、3×3畳み込みをすべて置き換えた場合、SMoEはLRLCNを上回り、パラメータ数を56%も削減しながら、競争力ある精度を維持した。
• テンソルルーティング、自己教師付きルーティング損失、エキスパート誤差ダミングの組み合わせにより、特に細粒度の空間モデリングを要するタスクで顕著な性能向上が得られた。
• SMoEは、グローバル天気データなどの大規模な空間ドメインに対しても、性能や効率の低下を伴わず、効果的にスケーリング可能である。