💬 NLP

OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

OmniMoEは、Cartesian Product RouterやExpert-Centric Schedulingといった斬新なコンポーネントを通じてMixture-of-Expertsの粒度をベクトルレベルまで押し上げることで、高い精度と大幅な推論速度向上を両立させたシステム・アルゴリズム協調設計フレームワークである。

原著者： Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo

公開日 2026-07-02

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、数百万冊の本（データ）を瞬時に処理しなければならない、巨大で高速な図書館を運営していると想像してください。人工知能の世界では、この図書館は「モデル」であり、本はそれが理解すべき情報です。

長い間、これらの図書館には主に2つの働き方があり、どちらにも大きな問題がありました。

「大型チーム」方式（粗粒度/Coarse-Grained）: 質問をするたびに、256人の司書からなるチーム全体を呼び出すようなものです。たとえ特定の事実を一つだけ必要としている場合でも、256人全員が作業を開始します。これはグループとして一斉に動くため速いのですが、質問の内容に無関係な司書がほとんどであるため、膨大なエネルギーの無駄遣いになります。
「ソロ・スペシャリスト」方式（細粒度/Fine-Grained）: エネルギーを節約するために、何百万人もの極めて専門化された小さな司書を雇います。質問をするたびに、その事実を正確に知っている「たった一人の完璧な司書」だけを呼び出します。これはリソースの面では非常に効率的ですが、物流上の悪夢です。何百万もの散らばった専門家を探し回ることは非常に時間がかかります。司書たちは、本を読むことよりも、自分のデスクまで歩くことに多くの時間を費やしてしまいます。

OmniMoE：究極の司書システム

この論文は、両方の良いところを組み合わせた新しいシステム、OmniMoEを紹介しています。これは、驚異的な精密さと、電光石火の速さを兼ね備えた図書館のようなものです。その仕組みを、3つのシンプルなパートに分けて解説します。

1. 「アトミック（原子レベル）」な司書（小さなスペシャリストたち）

チーム全体を雇う代わりに、OmniMoEは「アトミック・エキスパート」を雇います。これらは知識の最小単位であり、章全体ではなく、単一の完璧な文章や事実のようなものです。

革新性: 質問を受けたとき、システムは単に一人の司書を選ぶのではありません。その質問のためだけに、これらの小さなスペシャリストを動的に組み合わせて、カスタムチームを編成します。それは、あなたが入力する一文字一文字に対して、カスタムのパズルピースを組み立てるようなものです。

2. 「デカルト積」マップ（スマートな住所システム）

何百万もの小さなスペシャリストがいる場合の問題は、「どうやって素早く見つけるか？」です。もし1,000万人の名前のリストをチェックしなければならないとしたら、時間がかかりすぎます。

解決策: OmniMoeは「デカルト積ルーター（Cartesian Product Router）」を使用します。図書館が巨大な名前のリストではなく、巨大なグリッド（スプレッドシートの行と列のようなもの）であると想像してください。
なぜ役立つのか: 何百万もの名前の中から特定の名前を探す代わりに、システムは司書が座っている「行」と「列」を見つけるだけです。これは、「司書ボブ」を探すのではなく、「5行目の3列目へ行け」と言うようなものです。これにより、大規模で低速な検索が、極めて小さく即時的な計算へと変わります。

3. 「エキスパート中心」のバス時刻表（交通整理）

スマートなマップがあったとしても、もし1,000人の散ら散らばったスペシャリストを一人ずつ迎えに行くバスを走らせれば、バスは渋滞に巻き込まれるでしょう（これは「メモリ・ボトルネック」と呼ばれます）。バスは、停止と出発を繰り返すことに時間を費やしてしまいます。

解決策: OmniMoEは、処理の順序を変更します。質問ごとに一人ずつスペシャリストを迎えに行くのではなく、まずスペシャリストをグループ化します。
比喩: バス運転手が、「よし、科学セクションへ行く人は全員先にバスに乗ってください。それから歴史セクションの人たちを乗せます」と言う場面を想像してください。バスは一度にグループとしてのスペシャリストを積み込み、作業エリアへと運び、彼らは一つの大きな効率的なブロックとして共に働きます。これにより、混沌としたストップ＆ゴーの交通渋滞が、スムーズで高速なハイウェイへと変わります。

結果：速く、安く、賢い

この論文は、このシステムを現在の最高の手法と比較検証しました。

スピード: 従来の最高の「小さなスペシャリスト」システムよりも10.9倍速いです。73ミリ秒かかっていたものが、わずか6.7ミリ秒になりました。
賢さ: また、より正確です。「共有高密度MLP（一般的な常識や文法を扱う汎用的な脳）」を、小さなスペシャリスト（稀で特定の事実を扱うもの）と併用することで、推論能力を失うことがありません。
効率性: コンピュータのメモリをより効果的に使用し、システムを遅らせる原因となる「交通渋滞」を回避しています。

まとめ
OmniMoEは、AIモデルが速度を落とすことなく、何百万もの小さく超専門的なエキスパートを使用できるようにする巧妙なトリックです。これは、エキスパートをグリッドのように整理し（素早く見つけるため）、バスの時刻表のように仕事をグループ化する（素早く移動するため）ことで実現しています。その結果、特定の詳細について驚くほど知識が豊富でありながら、実用的なほど高速なAIが誕生しました。

技術要約: OmniMoE

1. 問題提起

Mixture-of-Experts (MoE) アーキテクチャは、トークンごとの計算量とモデル容量を切り離すことで、言語モデルをスケールアップするための標準となっています。しかし、既存の設計には、エキスパートの粒度とハードウェア実行効率の間の根本的なトレードオフが存在します。

粗粒度MoE: (例: DeepSeek-V3, KIMI-K2) 大規模で高密度なフィードフォワードネットワーク (FFN) ブロックをエキスパートとして利用します。これらはハードウェアに適した高密度行列演算と連続的なメモリ・アクセスの恩恵を受けますが、不正確な活性化という課題を抱えています。大きなエキスパート・ブロックを活性化させることは、特定のトークンに無関係なパラメータの計算を伴うことが多く、計算の無駄や硬直的なスケーリング制約につながります。
細粒度MoE: (例: PEER, PKM) 数百万の軽量なエキスパート（例: 静的なベクトル）を利用して、パラメータ効率とルーティングの精度を最大化します。しかし、これらを大規模にスケールさせると、深刻なシステム・ボトルネックが発生します：
1. 表現力の制限: エキスパートを静的なベクトルにまで縮小すると、複雑な言語的依存関係に不可欠な、トークン依存の非線形変換（MLP投影など）が失われます。
2. ルーティングのオーバーヘッド: 数百万のエキスパートへのルーティングは、法外な計算コストとメモリ・コスト ( $O(N)$ ) を引き起こします。
3. ハードウェアの非効率性: 細粒度のルーティングは、散在的で非連続的なメモリ・アクセスを引き起こし、実行のボトルネックを計算からメモリ帯域幅へとシフトさせます。その結果、GPUの利用率が低下し、レイテンシが増大します（メモリ・バウンド）。

中心となる課題は、細粒度モデルのパラメータ効率と、粗粒度アーキテクチャのハードウェア効率をいかにして両立させるかという点です。

2. メソドロジー: OmniMoE

OmniMoEは、ハードウェア効率を維持しながら、エキスパートの粒度を論理的な極限まで押し上げるシステム・アルゴリズム協調設計フレームワークです。これは、共有された高密度MLPと、大規模なルーティング対象の「アトミック・エキスパート (Atomic Experts)」のプールを組み合わせたハイブリッド・アーキテクチャを採用しています。

2.1 コア・アーキテクチャ構成要素

アトミック・エキスパート (Atomic Experts): 計算の基本単位は、「アトミック・エキスパート」と定義されます。これは、一対のベクトル ( $w_{in}, w_{out}$ ) によってパラメータ化された最小限のルーティング可能ユニットです。単一のアトミック・エキスパートは、軽量な非線形変換 $E_i(x) = \sigma(x w_{in}^T) w_{out}$ を実行します。
動的エキスパート・アセンブリ (Dynamic Expert Assembly: DEA): 数百万の独立したエキスパート・ブロックを個別に保存する代わりに、OmniMoEはすべての $N$ 個のアトミック・エキスパートを2つのグローバルなパラメータ行列 ( $W, V$ ) に集約します。各トークンに対して、システムはこれらの行列から疎なサブセットの行を動的に抽出し、それらをトークン固有の融合された計算ブロックへと組み立てます。これにより、動的な構成を通じて高い表現力を維持しながら、極限のパラメータ効率を実現します。
共有高密度MLP (Shared Dense MLP): 標準的な高密度MLPブランチは、一般的な意味論的推論を処理し、安定したバックボーンを提供するために、すべての入力に対して普遍的に活性化されます。これは、ロングテールな知識の検索に特化したルーティング・ブランチを補完します。

2.2 アルゴリズムの革新

数百万のアトミック・エキスパートの規模を管理するため、OmniMoEは2つの主要なアルゴリズム・メカニズムを導入しています。

デカルト積ルーター (Cartesian Product Router):
- 問題: $N$ 個のエキスパートへの標準的なルーティングには、 $d \times N$ サイズの投影行列が必要であり、$O(Nd)$ の複雑さが生じます。
- 解決策: ルーターは、1次元のエキスパート・インデックス空間を2次元グリッド ( $N_r \times N_c$ ) に分解します。2つの独立した低次元分布（行および列のロジット）を予測し、それらを組み合わせてエキスパートのスコアを算出します。
- 影響: この因数分解により、ルーティングの複雑さとパラメータ・ストレージが $O(Nd) $から$ O(\sqrt{N}d)$ に削減され、大規模なルーティングが実用的になります。
エキスパート中心型スケジューリング (Expert-Centric Scheduling):
- 問題: 標準的な「トークン中心型」の実行では、各トークンに対して個別にパラメータをフェッチするため、散在的なメモリI/Oが発生し、ベクトル化が妨げられます。
- 解決策: 実行パラダイムを「エキスパート中心型」へと反転させます。
  - タスク収集: バッチ内のすべてのルーティング決定をフラット化し、タスク・リストにします。
  - グルーピングとソート: タスクをエキスパートIDの連続した範囲ごとにグループ化し、各グループ内でトークンID順にソートします。
  - グループ化GEMM (Grouped GEMM): この再編成により、システムはエキスパートの重みブロックを一度ロードしてバッチ内の複数のトークンで再利用できるようになり、散在的なメモリ・ルックアップを、連続的で高スループットな Grouped GEMM 操作へと変換します。
- 影響: これにより、実行をメモリ・バウンドから計算量バウンドへと変容させ、Tensor Core の利用率を最大化し、メモリ帯域幅を最適化します。

3. 主な貢献

アトミック・エキスパートの定式化: 動的なアセンブリを通じて、非線形な表現力を維持しながら数百万のエキスパートへのスケーリングを可能にする、新しい最小限のルーティング可能ユニット。
デカルト積ルーター: 複雑さをエキスパート数に対して線形から平方根へと削減し、大規模なエキスパート・プールのインデックス作成におけるボトルネックを解決するルーティング・メカニズム。
エキスパート中心型スケジューリング: メモリ帯域幅のボトルネックを排除するために計算を再配置するシステムレベルのスケジューリング戦略であり、細粒度MoEが粗粒度設計と同等のハードウェア効率を実現することを可能にする。
システム・アルゴリズム協調設計: 高いモデル表現力とハードウェア効率の両立を実現するために、これらのコンポーネントを統合した包括的なフレームワーク。

4. 実験結果

OmniMoEは、7つの常識ベンチマークと5つの拡張推論/コーディングベンチマーク、および広範な効率性テストによって評価されました。

精度: 総パラメータ数6.4B / アクティブパラメータ数1.7BのOmniMoEモデルは、7つの常識ベンチマークにおいて**平均ゼロショット精度50.9%**を達成しました。これは、粗粒度ベースライン（例: DeepSeekMoE）を+0.7、細粒度ベースライン（例: PEER）を+2.0上回っています。
推論とスループット: 拡張ベンチマーク（GSM8K, MATHなど）において、OmniMoEは平均スコア42.5を記録し、PEER（+15.5）およびDeepSeekMoE（+2.0）を大幅に上回りました。スループットに関しては、OmniMoEは約14.2k tokens/secondに達し、高密度ベースラインと同等であり、細粒度のPEER（約11.6k）やPKM（約7.9k）よりも大幅に高速でした。
効率性: OmniMoEは、PEERと比較して推論レイテンシにおいて10.9倍の高速化（レイテンシを73msから6.7msへ短縮）を実証しつつ、粗粒度MoEと同等のメモリ・フットプリントを維持しました。
スケーリング則: 同等のトレーニングFLOPsおよび活性化パラメータ・予算の下で、OmniMoEは一貫してすべてのベースラインの中で最も低いパープレキシティを達成しており、優れた計算およびパラメータ効率を示しています。

5. 意義と主張

本論文は、包括的なシステム・アルゴリズム協調設計によってサポートされるならば、大規模な細粒度MoEは高精度かつ非常に効率的になり得ることをOmniMoEが示していると主張しています。

トレードオフの解消: 細粒度モデルのパラメータ効率と、粗粒度モデルのハードウェア効率の間の溝を埋めることに成功しました。
機能的分業: アブレーション研究とブランチ優位性分析により、明確な役割分担が明らかになりました。共有高密度MLPは一般的な言語パターンと推論の接続詞を扱い、ルーティングされたアトミック・エキスパートは事実的なエンティティやロングテールな知識に特化しています。
システムへの影響: ルーティングの複雑さとメモリ・アクセス・パターンが細粒度MoEの主要なボトルネックであり、大規模なエキスパート・プール（百万規模）の可能性を解き放つためには、アルゴリズムの革新（デカルト積ルーター）とシステムレベルの最適化（エキスパート中心型スケジューリング）の両方が必要であることを強調しています。

著者らは、現在の実装がNVIDIA GPU向けのカスタムTritonカーネルに依存しており、6.4Bを超える総パラメータへのスケーリングについては今後の検証領域であると述べています。