UniF-MoE、共有ブロックを先に抽出して残余計算をルーティングし、top-2 GMoE比で推論時間を45.2%短縮

従来の疎なMoEでは、top-kルーターによって各tokenを固定数の完全なFFNエキスパートへ送る。共有エキスパート、エキスパート内部の枝刈り、動的kも、通常はそれぞれ異なるコントローラーで処理される。新たに提案されたUniF-MoEは、これら3要素には実際には順序上の依存関係があると考える。再利用可能な計算を取り除くと、残りの内容、最適なエキスパート、必要なエキスパート数のすべてが変わり得るため、「まず共有し、その後で残余部分をルーティングする」必要がある。

著者らはまず、同一の密なFFNを複製した後、疎構造へアップグレードしたエキスパートを調査し、各FFNを整列したkey-valueチャネルブロックに分割した。同時に活性化されることが多いエキスパートほど、value位置の重複率が高い。これらの共有位置を取り除くと、元のtop-2エキスパートの組み合わせを維持したtokenはわずか5.7%だった。共有率の高いエキスパートペアでは、2つの残余エキスパートだけで元の出力を再構成できる割合が56.3%だったのに対し、共有率の低いグループでは11.8%だった。この結果は、元の判断をそのまま利用するのではなく、再ルーティングすべきことを裏付けている。

UniF-MoEはtokenごとに共有需要αを計算し、それを使って共有ブロック数と共有パスの重みを決定する。各ブロックのup-projection keyの平均をプロトタイプとして、実際に共有する内容を選択する。続いて残余需要β=1−αを累積ルーティング確率のしきい値とし、確率の合計がβを満たすまで、最小限の残余エキスパートを活性化する。Gram正則化によってルーティングベクトルを直交に近づけ、複数のエキスパートが同じ処理を繰り返し担当することを防ぐ。

5つのDomainBedデータセットでの平均精度は69.5%で、完全な比較が可能な掲載ベースラインを上回った。GLUEの5タスク平均も82.76%で、タスクごとに固定kを選んだ場合の81.95%を上回った。VLCSでの測定では、top-2 GMoEと比べて活性化パラメータ数とFLOPsがそれぞれ9.1%と16.1%減少し、1推論ステップ当たりの時間は0.31秒から0.17秒へ、メモリ使用量は0.55 GiBから0.26 GiBへ低下した。ただし、密なDeiTは依然として0.04秒、0.29 GiBで済む。この結果はMoE内部での改善を示すものであり、密モデルを全面的に上回ったことを意味するわけではない。

Apache 2.0ライセンスのコードには、ビジョン、言語、テスト、再現用のスクリプトが含まれているが、学習済みcheckpointや結果配列は提供されていない。今後は、このルーティングをデコーダー型LLM、長いシーケンスのbatching、GPU間のエキスパート通信、さらに同一FFNから初期化されない異種エキスパートへ拡張できるかが焦点となる。