1
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

MoE(Mixture of Experts)概要入門 🧠🔀

Posted at

Introduction

MoEとは?

Mixture of Experts(MoE)は、巨大なニューラルネットワークを「たくさんの小さな専門家(Expert)の集合」として捉え、入力ごとに必要な専門家だけを動かす考え方です。従来の“密”なモデルは、どんな入力でも全パラメータを一斉に計算します。これだとモデルを大きくするたびに計算コストとメモリが直線的に、時には指数的に膨らみます。MoEはここに「条件付き計算(Conditional Computation)」を持ち込みます。すべて計算するのではなく、入力の内容に合った一部の経路だけを通す。たとえば文章の中に数学的な表現が出てきたら“数学が得意な専門家”を、会話文なら“対話に強い専門家”を選んで処理するイメージです。こうすることで、総パラメータは大きいまま表現力を保ちつつ、1回の計算で実際に使うパラメータはごく一部に抑えられます。

この仕組みは、内部で小さな“振り分け役”が働くことで成立します。入力ベクトルを受け取ったルーター(ゲーティングネットワーク)が「どの専門家に任せるか」をスコア化し、上位の専門家(Top-1やTop-2など)だけを呼び出します。選ばれた専門家たちがそれぞれ出力を計算し、その結果を適切な重みで合成して最終出力にします。ポイントは、モデル全体は巨大でも“同時に動くのは一部だけ”という設計です。これがMoEの「混合(Mixture)」と「専門家(Experts)」という名前の由来であり、同時にコスト削減の源泉でもあります。


なぜ現代のLLMで重要か?

LLMは性能向上のためにパラメータを増やし続けてきましたが、密な計算のままでは学習・推論のコストが限界に達します。MoEはこのボトルネックに対する現実的な打ち手です。まず、計算量が入力に応じて“必要な分だけ”になるため、同じ推論時間でより大きなモデル(=より豊かな表現力)を扱えます。次に、専門家が役割分担することで、モデルが得意領域を分散的に学びやすくなり、複雑なタスクでも適切な専門家が呼び出されるほど出力の質が上がる傾向があります。結果として、MoEは「スケールさせたいがコストは抑えたい」という相反する要求を両立し、研究用途だけでなく実運用でも意味のあるトレードオフを提供します。

もちろん課題もあります。全専門家をメモリに載せておく必要があること、ルーティングの偏りによる負荷不均衡を防ぐ工夫が要ること、学習の安定化や微調整の設計がやや難しいことなどです。それでも、密モデルだけでは届かなかった“より大きく・より賢い”モデルを、実用的な計算コストで回す道を開いたのがMoEです。だからこそ、近年の代表的な大規模モデル群でMoE的な設計が注目され、広く採用されつつあります。

Historical Background 🕰️

Mixture of Experts(MoE)の発想は、実は非常に古くまでさかのぼります。そのルーツは1991年にJacobsらが提案した Adaptive Mixture of Local Experts という研究にあります。当時のニューラルネットワークはまだ規模が小さく、すべての入力に対してひとつのモデルで処理するのが一般的でした。しかし彼らは「問題を分割し、それぞれの部分を異なる専門家(Expert)に任せる」ことで、より柔軟で効率的な学習ができるのではないかと考えました。そして入力ごとにどの専門家を使うかを決定する Gating Network を導入し、条件に応じて異なる計算経路を通すという革新的な仕組みを生み出したのです。この考え方は、今日でいう「条件付き計算 (Conditional Computation)」の最初の萌芽でした🌱。

その後しばらくは、計算資源の制約もあってMoEは広く普及することはありませんでした。しかし2010年代にディープラーニングが台頭し、モデル規模が指数関数的に拡大していく中で、この古いアイデアは再び脚光を浴びることになります。特に2017年にShazeerらが発表した Sparsely-Gated MoE は大きな転機でした。この研究では、1370億ものパラメータを持つ超巨大モデルを提案しつつ、実際には入力ごとに一部のExpertしか計算しないという仕組みを導入しました。これにより計算効率が飛躍的に向上し、当時の常識を覆すような「巨大でありながら軽量なモデル」を実現したのです🚀。

さらに2021年にはGoogleが Switch Transformer を発表し、MoEの実用化が一気に加速しました。このモデルではRouterをシンプルにして「Top-1」のExpertだけを選択する方式を採用し、学習の安定性を高めつつ大規模データセットで優れた性能を示しました。こうしてMoEは、もはや研究上のアイデアにとどまらず、実際の大規模言語モデルに組み込まれる現実的な技術へと進化していったのです。

現在では、Mixtral 8×7B のようなオープンソースモデルや、GPT-4 をはじめとする商用の先端モデルでもMoEが使われていると広く信じられています。つまり、30年以上前に生まれた小さなアイデアが、計算効率とスケーラビリティという課題に対する解決策として、現代のAIの根幹を支える技術へと成長したわけです✨。


Core Components of MoE :robot:

MoEレイヤは、Transformer の FFN(Feed-Forward Network)層を置き換える形で挿入されるのが一般的です。ひとつの MoE 層の中には、複数の Expert(通常は独立な小さな FFN)が並び、各トークンに対して Router / Gate が「どの Expert に送るか」を決めます。計算はで、選ばれたごく一部の Expert だけが実行されます(Top-1/Top-2 routing など)。この「選択計算=条件付き計算」が、同じ FLOPs でモデル“総容量”を拡張できる理由です。概念整理としては、専門家(Expert)=FFN群Router=確率分布を出す小ネットSparse Activation=選択実行、の三点で捉えると分かりやすいです。


Router(Gating)の数学

入力トークン埋め込み $x \in \mathbb{R}^{d}$ に対し、最も基本的なゲーティングは線形写像→Softmaxです。すなわち

$$
z = x W_g \in \mathbb{R}^{E},\quad
p = \mathrm{softmax}(z),
$$

ここで $E$ は Expert 数、$p_i$ は「Expert $i$ に送る重み(確率)」です。MoE 出力

$$
y = \sum_{i\in\mathcal{K}(x)} p_i , E_i(x)
$$

の形をとります。疎化(Top-$k$)では、$\mathcal{K}(x)$ は上位 $k$ 個の Expert のみを含み、他は計算しません(対応成分の $p_i$ を 0 に)。これにより、パラメータは巨大でも、実計算は小さいままにできます。

Shazeer ら(2017)はNoisy Top-$k$ Gatingを提案し、学習初期に特定 Expert に偏ってしまう問題を緩和しました。具体的には、各ロジットに学習可能なスケールのガウス雑音を加えた $H(x)$ を作り、その上で Top-$k$ を取り、それ以外を $-\infty$ にして Softmax(=確率 0)に落とします:

$$
H_i(x) = (x W_g)i + \mathcal{N}(0, \sigma_i^2),\quad
\sigma_i=\mathrm{softplus}!\big((xW
{\text{noise}})_i\big),\quad
p=\mathrm{softmax}!\big(\mathrm{KeepTopK}(H(x),k)\big).
$$

この設計はどの Expert にも一定のチャンスを与えるため、ルーティングの偏り(負荷不均衡)を抑えます。


負荷分散(Load Balancing)と補助損失

実運用では、Router が「学習しやすい一部の Expert だけを選び続ける」傾向が生まれがちです。これを防ぐために導入されるのが補助損失(Auxiliary / Load-balancing loss)です。Shazeer らの原型では「重要度(ゲート確率の総和)」や「負荷(割当てトークン数)」のばらつきを抑える形で定式化され、Switch Transformer(2021)ではより簡潔な形が広く使われます。Switch の代表的な定義は、バッチ内トークン数を $T$、Expert 数を $N$ とし、各 Expert の割当て比率 $f_i$ と 確率比率 $P_i$ を用いて

$$
\mathcal{L}{\text{aux}}
= \alpha \cdot N \cdot \sum
{i=1}^{N} f_i , P_i
$$

とするものです。理想は $f_i=P_i=\tfrac{1}{N}$(一様)で、この損失を小さくするほど均等なルーティングに近づきます($\alpha$ は重み)。


Expert Capacity(容量制約)とオーバーフロー

トークンが一箇所に偏ると、通信/メモリの観点でも非効率です。そこで各 Expert に「同時に処理できる最大トークン数」=容量 $C$ を設けます。バッチ当たりトークン数を $T$、Expert 数 $N$、Top-$k$ を $k$、容量係数を $\text{cf}$ とすると、実装上よく用いられる経験式は

$$
C ;\approx; \left\lceil \frac{T \cdot k}{N} \cdot \text{cf}\right\rceil
$$

です(Top-$k$ を含めない単純形 $\lceil \tfrac{T}{N}\cdot \text{cf}\rceil$ も文献に見られます)。容量を超えたトークンはオーバーフローとして次層へ残差経路で渡す/ドロップする、といった扱いになります。容量係数を上げればドロップは減りますが通信とメモリが増え、下げれば軽くなる一方で性能低下リスクが出る——このトレードオフを管理するのが運用上の肝です。


実務で効く工夫(Router安定化・選択戦略)

Switch Transformer は Top-1 の単純化により学習を安定させ、容量係数補助損失でバランスを取る設計を確立しました。また、GShard は Top-2 で「2つ目の Expert を確率的に選ぶ(Random routing)」や容量制約の導入により大規模分散で扱いやすくしました。近年の実装ガイド(HF のまとめ等)でも、Z-loss(Router のログスペース正則化)などの安定化テクが紹介され、前述の補助損失と合わせて使われます。


参考になる直感(専門化は“意味領域”より“トークン文脈”)

「数学専門家」「歴史専門家」という人間の科目感ではなく、特定のトークンや構文パターンに対する応答としての“専門性”が観察されることが多いです。Mixtral 8×7B の可視化では、トークンごとに安定して選ばれやすい Expert が存在し、構文的役割にひもづく傾向が示されています。これは「MoEは巨大な辞書の“索引”を、Router が文脈に応じて素早く引く」ような感覚で捉えると腹落ちします。


ちょい進んだ PyTorch 風サンプル(Top-k+容量+補助損失の計算)

学習用の概念実装です。分散や通信は省き、1 GPU・ミニバッチ内でのルーティング/容量割当/合成・補助損失計算の流れだけを示します。

import torch
import torch.nn as nn
import torch.nn.functional as F

class ExpertFFN(nn.Module):
    def __init__(self, d_model, d_hidden):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_hidden)
        self.w2 = nn.Linear(d_hidden, d_model)
    def forward(self, x):
        return self.w2(F.relu(self.w1(x)))

class MoELayer(nn.Module):
    def __init__(self, d_model, d_hidden, num_experts=8, k=2, capacity_factor=1.25, aux_alpha=1e-2):
        super().__init__()
        self.experts = nn.ModuleList([ExpertFFN(d_model, d_hidden) for _ in range(num_experts)])
        self.router_w = nn.Linear(d_model, num_experts, bias=False)
        self.noise_w  = nn.Linear(d_model, num_experts, bias=False)  # noisy top-k 用
        self.num_experts = num_experts
        self.k = k
        self.capacity_factor = capacity_factor
        self.aux_alpha = aux_alpha

    def forward(self, x):
        """
        x: (B, T, D)
        return: (B, T, D), aux_loss
        """
        B, T, D = x.shape
        x_flat = x.reshape(B*T, D)

        # --- Router logits & noisy top-k gating ---
        logits = self.router_w(x_flat)  # (BT, E)
        noise  = torch.randn_like(logits) * F.softplus(self.noise_w(x_flat))
        noisy_logits = logits + noise

        topk_vals, topk_idx = torch.topk(noisy_logits, k=self.k, dim=-1)  # (BT, k)
        mask = torch.full_like(noisy_logits, float("-inf"))
        mask.scatter_(1, topk_idx, topk_vals)
        gate = F.softmax(mask, dim=-1)  # (BT, E)  上位kのみ確率>0

        # --- Expert capacity 計算 ---
        # C ≈ ceil((T_total * k / E) * cf)
        total_tokens = B*T
        C = int((total_tokens * self.k / self.num_experts) * self.capacity_factor + 0.9999)

        # --- Token を Expert に割当(単純版:各expertごとに上位Cを残し、それ以外はdrop)---
        # expert別に gate[:, i] の高い順にC個だけ選び、それ以外は確率0に
        gate_np = gate.detach().clone()
        keep_mask = torch.zeros_like(gate_np, dtype=torch.bool)
        for i in range(self.num_experts):
            vals = gate_np[:, i]
            if vals.numel() == 0: 
                continue
            topC = torch.topk(vals, k=min(C, vals.numel()), dim=0).indices
            keep_mask[topC, i] = True
        gate = gate * keep_mask.float()  # drop overflow
        gate = gate / (gate.sum(dim=-1, keepdim=True) + 1e-9)  # 再正規化

        # --- Expert ごとに前向き計算(バッチ分割してまとめ実行)---
        outputs = torch.zeros(B*T, D, device=x.device, dtype=x.dtype)
        for i, expert in enumerate(self.experts):
            sel = keep_mask[:, i]
            if sel.any():
                y_i = expert(x_flat[sel])                  # (n_i, D)
                outputs[sel] = outputs[sel] + y_i * gate[sel, i:i+1]

        y = outputs.reshape(B, T, D)

        # --- 補助損失(Switch系の簡潔形)---
        # f_i: 割当トークン比率, P_i: ゲート確率の総和比率
        f = keep_mask.float().mean(dim=0)               # (E,)
        P = gate.mean(dim=0)                            # (E,)
        aux_loss = self.aux_alpha * self.num_experts * torch.sum(f * P)

        return y, aux_loss
  • Noisy Top-k(上の noise_wtopk 部分)で偏りを緩和。
  • Capacity($C$ の計算)で各 Expert の同時処理トークン数を制限。オーバーフローはドロップ(実務では残差で“素通り”させる実装もあります)。
  • 補助損失は Switch で一般的な $\alpha N \sum f_i P_i$ 形。

注:これは学習挙動を掴むための簡易版です。実務では「Dispatch/Combine」の効率的な疎演算カーネル、通信分散(Expert Parallelism)、Bucket 化、Z-loss などが入ります。


Benefits(MoEの利点)

MoEの最大の魅力は、計算量をほぼ増やさずに表現容量(パラメータ総量)だけを拡張できる点にあります。通常のTransformerでは、1トークンあたりのFFN計算量はおおむね $\mathrm{FLOPs}{\text{dense}} \approx 2,d,d{\mathrm{ff}}$ です(線形→活性化→線形の近似)。MoEでFFNを置き換えると、1層あたりに $E$ 個の専門家が並び、実際に動かすのは Top-$k$ のみです。このときの計算量は概ね

$\mathrm{FLOPs}{\text{MoE}} \approx 2d\left(kd{\mathrm{ff}}^{(\mathrm{expert})}\right) + O(dE)$

となり、各エキスパートの幅を $d_{\mathrm{ff}}^{(\mathrm{expert})} \approx d_{\mathrm{ff}}/k$ に調整すれば、FLOPsは密モデルと同程度に抑えつつ、パラメータ総量は $E$ 倍(厳密にはバイアス等を含めて係数は変わる)まで増やせます。つまり 「計算量≒据え置き、容量だけ拡大」 ができる。これは “大きくしたいが計算は増やしたくない” というLLMの根源的ジレンマに対する、実務的で強力な打ち手です。こうしたスパース活性化によるスケールの利点は、Switch/ ST-MoE などの系譜で再現性高く示され、近年のMoE実装知見として定着しています。

さらに、専門家の“ゆるい専門化”が学習効率と性能の底上げに寄与します。誤解されがちですが、エキスパートは「数学担当」「生物担当」のような領域というより、トークンや構文的文脈に同調する傾向が強い(句読点、固有名、接続語など)。この“粒度の細かい分担”は辞書の索引を引くように働き、入力に応じた最適経路を通せるため、同じ計算量でも情報の通り道がリッチになります。Mixtralの可視化を含む解説でも、同じ種類のトークンに似たエキスパートが繰り返し選ばれる様子が示されており、MoEが“計算の選択性”を通じて表現力を高めるメカニズムの実体を裏づけています。

学習・推論の実務面でも恩恵は明確です。まずプリトレでは、固定計算予算でのサンプル効率が上がりやすく、専門家数を増やすほど前半の伸びが良い(ただし256~512付近で逓減)といった傾向が整理されています。推論では、同等品質に対して密70Bより高速なMixtralのような実例も報告され、さらに蒸留でMoEの知識を密モデルへ移して一部の疎性利益を保持する、ブロック疎行列カーネル(Megablocks)で動的割当のオーバーヘッドを抑える、といった効率化スタックが整備されています。これらにより、MoEは「学術的には強いが重い」という段階を越えて、「運用で回る」技術へと成熟しつつあります。

最後に、タスク特性との相性という観点もあります。密モデルに比べ、MoEは知識依存型の課題で優位に出やすく、推論系の高度課題では密が強いという観察が報告されています。この“得意不得意”を踏まえると、企業現場ではプリトレをMoEで加速→蒸留して提供、あるいはサービングは用途別にMoE/密を使い分け、といった現実的な設計が見えてきます。


Challenges(MoEの課題)

利点の裏側には、ルーティング安定性・負荷不均衡・通信コスト・メモリ常駐 という実務上の難所が並びます。第一に、ルーターの“偏り” はMoEの構造的リスクです。初期に強いエキスパートへ確率が集中し、他が“餓死”する(十分に訓練されない)と、モデル全体の汎化が崩れます。これに対し、2017年のSparsely-Gated MoEはNoisy Top-$k$2種の補助損失(重要度の変動係数CVに基づく $L_{\mathrm{importance}}$ と、割当数の平準化を促す $L_{\mathrm{load}}$)を導入しました。

という形で、専門家ごとの確率総和のばらつきを直接抑えます。一方でSwitch TransformerはTop-1に単純化し、割当比率 $f_i$確率比率 $P_i$ を近づける簡潔な補助損失で安定化させました。さらに Z-loss(ルーターロジットの絶対値を抑える正則化)を併用することで、指数関数(softmax)周りの数値不安定を弱められることが示されています。これらの安定化テクは、MoE運用の実質的必須装備だと考えてよいでしょう。

第二に、Expert Capacity(容量)とトークン・ドロップ/オーバーフローの扱いです。各エキスパートが同時処理できるトークン数 $C$ を制限しないと、人気エキスパートに入力が偏って通信滞留やVRAM上限を招きます。そこで、バッチ内トークン総数 $T$、エキスパート数 $N$、Top-$k$ を用いて

$$
C ;\approx; \Big\lceil \frac{T\cdot k}{N}\cdot \mathrm{cf} \Big\rceil
$$

のように容量係数 $\mathrm{cf}$ で上限を張り、超過分はドロップ(あるいは残差で素通し)します。実務では $\mathrm{cf}$ を上げると品質は上がりやすいが通信量とメモリが増え、下げると軽くなる代わりにドロップ増→品質悪化のリスクが出る——という三つ巴のトレードオフを調律します。さらに、MoEはトークンが専門家へ分散されるため**1エキスパートあたりの実効バッチが小さくなりがち(Shrinking Batch Problem)で、大きめのグローバルバッチやブロック疎行列カーネル(Megablocks)**での対処が定石になっています。

第三に、分散・通信が密モデルより構造的に難しいこと。Expert Parallelism では、MoE層ごとにトークンをエキスパートの居るワーカーへ All-to-All で輸送します。容量係数を上げると品質は伸びるが通信・アクティベーションメモリが膨らむため、ネットワーク帯域が律速の環境ではCFを下げるほうが良いTop-2で1.25を初期値にする、評価時はCFを下げて計算を節約するなど、運用ガイドラインが共有されています。FasterMoE はトポロジーに配慮したゲートでレイテンシの低いエキスパートを選ぶなどして高速化(理論・実測の両面)、Megablocks は不均衡割当を落とさず処理するGPUカーネルで実効スループットを大きく引き上げました。通信がボトルネックになりやすいMoEで、これらの工夫は直結の効果があります。

第四に、メモリ常駐の高さです。推論時に計算するのは一部でも、重みは基本“全エキスパートぶん”載るため、VRAM所要は密より重くなりがちです(とくにEが大きい構成)。ここは実務で悩ましい点で、蒸留で密へ戻す量子化(QMoE)エキスパートの統合(Aggregation of Experts)センテンス/タスク単位でルーティングして“切り出し可能なサブネット”を抽出するなどのサービング最適化が提案されています。プリトレはMoEで速く、提供は密or縮約という“二段構え”は、計算資源が限られる現場で現実的な選択です。

最後に、汎化とファインチューニングの難しさ。ST-MoE系の知見では、小規模タスクでは過学習しやすく、知識依存の大規模タスクで強みが出やすい微調整では密よりハイレート・小バッチが効く補助損失ON/OFFの挙動はタスク依存MoE層を凍結しても品質劣化は限定的など、密と異なる“流儀”が観察されています。命題は一つに収束せず、タスク・データスケール・専門家配置で最適地が動くため、実プロジェクトでは少数の実験で“どちらが有利か”の当たりを素早く引く設計(探索計画)が鍵になります。


Engineering Innovations(実装・運用の革新点)

MoE を“プロダクションで回す”ための革新は大きく、①ルーティングの安定化(Noisy Top-k、Switch、Z-loss、負荷分散損失)、②容量と通信の制御(Capacity Factor、Overflow 処理、Expert Parallelism)、③分散最適化(All-to-All の削減、トポロジー意識のゲート、専用カーネル)、④サービング最適化(蒸留・量子化・サブネット抽出)という四本柱に集約されます。以下、それぞれの要所を“何が効くのか/なぜ効くのか”の順に掘り下げます

1) ルーティングの安定化:Noisy Top-k → Switch(Top-1)→ Z-loss

初期の Sparsely-Gated MoE では、ゲートが学習初期に「強い」専門家へ偏り、他が十分に訓練されない“餓死”が発生しがちでした。Noisy Top-k はロジットにノイズ $H_i(x)=(xW_g)_i+\mathcal{N}(0,\sigma_i^2)$ を加え、$\mathrm{KeepTopK}$ 後に softmax を当てることで、上位 $k$ 以外の確率を 0 にしつつ探索性を確保します。これが負荷不均衡の初期収束を抑え、学習過程の多様性を担保します。これをさらに単純化したのが Switch TransformerTop-1 ルーティングで、通信・計算の両面でコストを下げつつ精度を維持しました(Top-1 により各エキスパートのミニバッチが大きくなり、安定しやすい)。加えて Router Z-loss は、ゲートのロジットの絶対値を罰して指数関数(softmax)周りの数値不安定を抑え、精度劣化を伴わずに安定性を上げます。実装では ゲートは FP32、エキスパートは bfloat16 のような選択精度も効きます(ゲートを低精度にすると不安定化しやすい)。

**補助損失(Aux/Load-balancing)**は二系統あります。Shazeer 系は専門家ごとの「重要度(ゲート確率総和)」の 変動係数 $CV$ を下げる $L_{\text{importance}}$ と、割当て数の平準化を促す $L_{\text{load}}$。Switch 系は 割当比率 $f_i$確率比率 $P_i$ の整合を罰する簡潔形($\mathcal{L}_{\text{aux}}=\alpha N \sum_i f_i P_i$)で、学習を安定させつつ実装を軽くします。どちらも“餓死”や極端な偏りを抑え、**全エキスパートが“育つ”**環境を作るのが狙いです。

2) 容量と通信の制御:Capacity Factor と Overflow の設計

Expert Capacity は、各エキスパートが 1 フォワードで処理可能なトークン上限 $C$ を決める仕組みです。バッチ内トークン総数を $T$、エキスパート数を $N$、Top-$k$ を $k$、容量係数を $\mathrm{cf}$ とすると、

$$
C ;\approx; \Big\lceil \frac{T \cdot k}{N} \cdot \mathrm{cf} \Big\rceil
$$

が実装上の目安です。cf を上げれば品質↑だが通信・活性メモリ↑cf を下げれば軽量化できるが Overflow(ドロップ or 残差素通し)↑ という綱引きになります。実務では Top-2 かつ cf≈1.25、1 デバイス 1 エキスパートから入り、評価時は cf を下げて計算を絞る、といった運用指針が共有されています。GShard は Top-2+**確率的セカンド選択(Random routing)**+容量制御により、大規模分散での効率を実証しました。

この容量設計は通信コストとも不可分です。MoE は「各トークンを担当エキスパートへ送る」ため、All-to-All が律速になりやすい。cf を上げると“受け皿”が広がる一方で 送受信バイト数と活性の滞留が増えるため、ネットワーク帯域が弱い環境では cf を抑えるのが定石です。さらに、トークンがエキスパートへ分散されるため Per-Expert の実効バッチが小さくなりやすい(Shrinking Batch)。ここは 大きめのグローバルバッチトークン整列の工夫でカバーします。

3) 分散最適化:Expert Parallelism とトポロジー意識の高速化

MoE の“第四の並列化”が Expert Parallelism(EP) です。データ・モデル・パイプラインに 専門家次元を加え、エキスパートをワーカー間に配置して、トークンを該当ワーカーへ配送します。EP は非 MoE 層ではデータ並列のように振る舞い、MoE 層でだけトークンをルーティングするのが肝です。EP とテンソル並列を併用する場合は Sequence Parallelism を有効化する、といった実装要件もベンダー実装では明記されています。
All-to-All の本質的な遅さに対しては、二系統のアプローチが有効です。ひとつは トポロジー意識のゲートで、「同等スコアなら近い(レイテンシの低い)エキスパートを選ぶ」ようにルーティングを微調整し、理論限界に迫るスケジューリングで通信を“隙間時間”に追いやる手法(FasterMoE)。もうひとつは GPU カーネルで、MoE を「バッチ化 GEMM の集合」と見なす従来前提を外し、割当て不均衡をそのまま飲み込むブロック疎行列演算に落とす(Megablocks)。前者は最大 17×の高速化報告、後者はトークンドロップなしでダイナミズムを処理でき、実効スループットを大きく引き上げます。ポイントは「通信を減らす」か「通信・計算をハードに合わせて再表現するか、の二刀流」であること。

運用では Straggler(最遅エキスパート)がレイテンシを支配する問題も無視できません。最近の研究は、推論時の容量感知ドロップ/再ルートで過負荷エキスパートを抑え、平均 1.9× の高速化を示しています。これは数式的には、層レイテンシを $L=\max_i L_i$ とみなし、最大値を直接下げに行く制御(容量で上限拘束+再配分)だと捉えると腹落ちします。

4) サービング最適化:蒸留・量子化・サブネット抽出

MoE の推論は「計算は疎でも全エキスパート重みを常駐」というメモリ問題がつきまといます。ここで効くのが 蒸留です。プリトレは MoE で速くやり、本番は密へ縮約する設計は、疎性の恩恵の 30–40% を保持しつつメモリ常駐を下げる現実的な落としどころです。さらに、文(タスク)単位でルーティングして「抽出可能なサブネット」を得る発想は、**“MoE を学習器、提供は抽出ネット”**という分業に道を開きます。これらは Serving techniques の章でも強調され、実運用の勘所になっています。量子化(QMoE)については、ゲートまわりは高精度を維持しつつ、エキスパート本体を低精度に寄せるのが安定的です。


MoE in Practice(実践)

1) まず決めること:アーキ設計とスケールの落としどころ

実運用で最初に決めるのは、どの層をMoE化するか(FFN置換の位置)Expertsの数 $E$Top-$k$(1 or 2)Capacity Factor(cf)補助損失の係数ルーターの数値安定化(Z-loss など)といった“骨格”です。既存の実務ガイドでは、FFN層をMoE化し、Top-2+cf≈1.25から入り、補助損失を有効ゲートはFP32(本体はbf16)で安定化、という出発点が扱いやすいと整理されています。最初は1GPU=1Expert(Expert Parallelismの単位と一致)でトポロジーを単純に保ち、通信とメモリの挙動を見ます。これらは Hugging Face の総括と実装解説でも強調される、いわば“定石”です。

**Top-1(Switch)**はさらに実装が軽く学習も安定しやすい一方、Top-2は表現力の“余裕”が出て性能が伸びやすい——このあたりのトレードオフは Switch Transformer 論文の設計思想とも対応します(“同程度FLOPsで疎にしてサイズは大きく”)。探索初期は Top-1で当たりを取り、Top-2で詰めるのが現実的です。


2) 学習レシピ:データ、最適化、ハイパラの勘どころ

データ面では、密モデルと同じく高品質のコーパスが基本。ただし MoE はトークン分布の偏りに敏感なので、ドメイン偏重のデータでは特定Expertへの集中が起きやすく、Aux Lossを強めても解けない場合があります。ここはデータの混合比とルーター安定化(Noisy Top-$k$、Z-loss)の両輪で見るのがコツです。ビジュアルガイドの直観とも合致します(専門家は“意味領域”よりトークン/構文的文脈に結びつきやすい)。(newsletter.maartengrootendorst.com)

最適化は AdamW+Cosine(またはLinear)で問題ありませんが、ミクロバッチ設計が肝です。MoEではトークンが専門家に分散されるためPer-Expert の実効バッチが小さくなりがち(Shrinking Batch Problem)。大きめのグローバルバッチ勾配蓄積(gradient accumulation)で統計を安定化するとともに、Aux Loss係数 $\alpha$ は 1e-2 前後からスイープしてオーバーヘッド最小で偏り緩和できる点を探ります。cfは 1.0–1.25 を中心に調律、ゲートFP32+本体bf16は多くの実装で推奨。DeepSpeek-MoE のチュートリアル群や HF の実装ノートでも同様の指針が述べられています。


3) 実装の選択肢:Transformers+DeepSeek か Megatron-Core か

Hugging Face Transformers + DeepSeekは、既存のTrainerに DeepSeek 設定を差し込む形でZeRO最適化や混合精度と併用しやすく、研究~小~中規模の実験に向きます。DeepSpeed 側には MoE のチュートリアル(学習・推論)と実装パターンがまとまっており、MoE特有の**補助損失、容量制御、EP(Expert Parallelism)**を含む“動かせる”ガイドが揃っています。

Megatron-Core / Megatron-LMは、EP×TP×DP×SP×PPの並列を高スループットに張り付ける前提のスケール実装です。重要なのは EP(エキスパート並列)とTP(テンソル並列)を併用するならSP(Sequence Parallelism)必須という要件(メモリと勾配の分配が崩れないようにするため)で、公式ドキュメント/READMEにも明記されています。クラスタ全体を視野に入れたプロファイル前提の設計ができるなら Megatron のほうが“伸び代”は大きいです。


4) 実用的な設定テンプレ(Transformers + DeepSpeek)

以下は “モデル側でMoE層を定義済み”を前提に、Trainer で DeepSpeed を使うときの最小設定の例です(ZeROとbf16・offloadは好みで調整)。MoE固有のパラメータ(num_experts, router_topk, capacity_factor, aux_loss_coef など)はモデルconfig側で持つ想定にしています。

// deepspeed_moe_zero.json(一例)
{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },
  "bf16": { "enabled": true },
  "flops_profiler": { "enabled": false },
  "optimizer": { "type": "AdamW", "params": { "lr": 2e-4, "betas": [0.9, 0.95], "eps": 1e-8, "weight_decay": 0.1 } },
  "scheduler": { "type": "cosine", "params": { "warmup_min_lr": 0.0, "warmup_num_steps": 2000 } }
}
# train.py(抜粋):Hugging Face Transformers の Trainer を使用
from transformers import AutoConfig, AutoModelForCausalLM, Trainer, TrainingArguments

cfg = AutoConfig.from_pretrained(
    "your-init",
    # ---- MoE固有(モデル側で解釈される想定)----
    num_experts=8,
    router_topk=2,              # Top-1なら1
    capacity_factor=1.25,
    aux_loss_coef=1e-2,
    router_z_loss_coef=1e-3,    # 実装側にあれば
    gate_dtype="float32",       # ルーターはFP32
)

model = AutoModelForCausalLM.from_config(cfg)

args = TrainingArguments(
    deepspeed="deepspeed_moe_zero.json",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    bf16=True,
    learning_rate=2e-4,
    warmup_steps=2000,
    logging_steps=50,
    save_steps=2000,
    evaluation_strategy="steps",
)

trainer = Trainer(model=model, args=args, train_dataset=ds_train, eval_dataset=ds_eval)
trainer.train()

注:MoEパラメータの正確なフィールド名はモデル実装に依存します(Mixtral系・自作MoE層など)。DeepSpeedのMoEチュートリアルとHFのDeepspeed統合ガイドを合わせて参照し、モデル実装側にパラメータを持たせるのが衝突しにくいです。


5) 実用的な設定テンプレ(Megatron-Core / Megatron-LM)

Megatron はCLIフラグで MoE を細かく制御します(名称はバージョンで差異あり)。概念的には以下を指定します:

  • エキスパート数/層router top-kcapacity factoraux loss係数
  • EPサイズ(1ノード内に複数Expertを置くか、ノードまたぎで分散か)、TP/PP/DP/SP
  • bf16+ルーターの高精度Grouped GEMM / Megablocks などの最適化

Megatron-Core の MoE API/ガイドには、EP×TP×SPの組み合わせ、Mixtral 8×7B の高いTFLOPs実績、MoE固有の最適化がまとまっています。EP+TP時はSP必須という要件は README にも繰り返し登場します。環境差分が大きいので、**手元クラスタのトポロジー(NVLink/IB)**に合わせて EP サイズを調整し、All-to-All の衝突を避ける配置(近接GPUにExpertを寄せる)から詰めるのが定石です。


6) Throughput と通信の見積り(感覚を掴む計算)

MoE層での All-to-All は、概ね**「バッチ内トークン×隣接層次元」相当のアクティベーションを E/GPU 間で送る」**イメージです。Top-$k$=2 なら送信側複製も増えます。cf を上げるほど溢れにくいが送受信バイト数が増えるcf を下げると軽いがドロップやストラグラーが増える、という現実的なバランスに落ちます。Switch/ GShard 系の論文や実装ノートも、疎化の計算量は一定でも通信が支配しうる点を繰り返し指摘しています。


7) 観測とデバッグ:何を“必ず”ログるか

MoE は“見えないところで壊れやすい”ので、以下の最低限メトリクスを継続ログに出すと事故が激減します。

  • **Aux Loss(負荷分散)**の推移:下がり切らない/発散は偏りの兆候。
  • Expert使用率の分散・エントロピー:あるExpertだけ突出していないか。
  • Capacity overflow率/ドロップ率:cf調整と直結。
  • Per-Expert 実効バッチStragglerの待ち時間:遅いExpertの特定。
  • ゲート数値安定(Z-loss、ゲートlogitの分布):極端な尖りは不安定のサイン。

これらは HFのMoE解説・ビジュアルガイド・システム系のまとめにも繰り返し登場する“現場の勘所”です。

簡易インスツルメント例(PyTorch風):学習の各イテでエキスパート選択分布とドロップ率を出す。

def moe_diagnostics(gate_probs, keep_mask):
    # gate_probs: (tokens, E) softmax後
    # keep_mask : (tokens, E) True=採用
    usage = keep_mask.float().mean(0)                  # 割当比率 f_i
    prob  = gate_probs.mean(0)                         # 確率比率 P_i
    entropy = -(prob * (prob.clamp_min(1e-9)).log()).sum().item()
    drop_rate = 1.0 - keep_mask.any(dim=1).float().mean().item()
    return {
        "expert_usage_mean": usage.mean().item(),
        "expert_usage_max" : usage.max().item(),
        "gate_entropy"     : entropy,
        "drop_rate"        : drop_rate,
        "aux_proxy"        : (len(prob) * (usage * prob).sum()).item()  # Switch系 aux の指標に近い
    }

8) 推論・提供(Serving):レイテンシ安定化の手筋

推論は**“計算は疎でも重みは全常駐”**が原則なので、VRAMとレイテンシがボトルネックになりがちです。現実解は三本柱:
(1) 推論時にcfを下げる(学習よりも厳しめ:1.0前後)
(2) トポロジー意識のルーティング(等価スコアなら“近い”Expertを選ぶ)
(3) 蒸留/量子化/抽出サブネット(MoEで学んだ表現を密に“落とす”)

DeepSpeed-MoEの推論ガイドは、密モデル向け最適化にMoEの工夫を追加する形で実務的テクニックをまとめています。Megablocks系の疎カーネルFasterMoE系のトポロジー最適化


9) 小規模で“まず動かす”:ミニMoEの練習帳

いきなり大規模は難しいので、文字レベル言語モデルで MoE を試すと勘所が掴めます。Hugging Face の “makeMoE from scratch” は、Karpathyのmakemore流儀で MoE を最小部品から実装する良い練習問題です。Top-kゲート、Aux損失、容量制御の“芯”が見えるので、あとから大規模実装(Transformers/Megatron)に乗り換えても迷いにくくなります。


10) ありがちな失敗と対処の地図

  • ゲート崩壊(特定Expertに極端集中):Noisy Top-$k$ を強める、Aux係数↑、学習率↓、データ混合の偏りを緩和。Switch系のTop-1+Z-lossで簡素化するのも手。
  • 通信律速/Straggler:cf↓、EP配置の見直し(近接配置)、トポロジー意識のルート、疎カーネル導入、評価時はcfさらに↓。
  • VRAM圧迫:蒸留で密へ縮約、QMoE、専門家の統合(Aggregation)、“文単位ルート→抽出サブネット”の提供戦略。
  • 微調整が不安定:バッチ拡大/勾配蓄積、学習率低め、Aux/Capacityは過度に触りすぎない、MoE層凍結の比較。

Future Directions(発展方向)

1) ルーティングの次世代化:Expert Choice と“距離感”の活用

従来のトークン選択型(token-choice)Top-$k$ は、各トークンが上位 $k$ 個の専門家を選ぶ設計でした。Expert Choice(EC)は逆に、各専門家が自分に適したトークン上位 $K_i$ を“指名”する方式で、固定バケット容量を保ちながら完全な負荷分散を達成します。スコア $s_i(x)=\langle x, w_i\rangle$ に基づき、専門家 $i$ はトークン集合 $S_i=\mathrm{TopK}_i{s_i(x)}$ を選択、各トークンは $\cup_i \mathbf{1}[x\in S_i]$ により可変個の専門家へ割当てられます。これにより、路由の不均衡→オーバーフロー/ドロップというMoE 固有の実務課題が構造的に緩和されます(NeurIPS 2022。Google の解説もわかりやすい)。EC の系統は最近、**トークン–ゲート重みのコサイン類似度(affinity)**を明示的に扱い、専門家の“分業境界”を理論的に整える議論へと広がっています。今後は EC × Capacity 動的制御や**EC × トポロジー意識(近い専門家を優先)**の統合が鍵になります。

数学の最小形:
EC は「専門家側の Top-$K_i$」でバケットを埋め、各トークンに割当てられる専門家数が可変になる。

Screenshot 2025-09-01 at 11.58.18 PM.png

—— 固定容量 $K_i$ と 可変エキスパート数の両立(=負荷分散と表現力の両取り)。

また、階層型MoE(ゲートの上にもう一段MoEを重ねる)も、分岐因子の抑制大規模分散の両面で引き続き有望です。2017年の古典(Shazeerら)は階層化で並列実行とデバイス配置を整理しており、EC/階層を融合した“階層EC”は次の自然な一手です。


2) 「ドロップしないMoE」:ブロック疎行列カーネルと動的容量

実運用では容量係数 $\mathrm{cf}$ を上げてもトークンドロップは残りがちで、品質と通信・メモリの綱引きになります。MegaBlocks は、MoE をブロック疎行列に再表現し、パディングやドロップなしで不均衡割当てをそのまま飲み込むカーネルを提示しました。報告ではTut.el 比で最大40%の学習高速化、Megatron-LM の強実装にも肉薄する性能を示しています。式の観点では、バッチ化 GEMM 前提を捨て、**可変ブロックサイズのブロック対角(+小ブロック分割)**で MoE のダイナミクスをそのまま計算可能にします。ドロップしない = 学習損失の低下が実測でも確認され、今後の“MoE既定路線”になり得ます。

通信そのものを削る方向では、FasterMoE のようにトポロジー意識のルーティング/スケジューリングAll-to-All の輻輳を避け、反復ごとの待ち時間を直接縮める系が成熟しています。MoE のスループットは本質的に「$\max_i$ の遅い専門家」に支配されるため、近接配置+輻輳回避の経路選択は引き続き中核テーマです。


3) 圧縮と提供(Serving):量子化・蒸留・サブネット抽出

推論では“計算は疎でも重みは全常駐”が悩みの種。QMoE1ビット未満/パラメータの極端圧縮と専用デコードカーネルで、Switch-c2048(1.6T)を <160GBに詰め、<5% のレイテンシ上乗せで動かす道を示しました。専門家側は量子化耐性が高いとの報告も出ており、MoQEや 2025年の微粒度MoE量子化は、専門家ごとの感度差を活かす“非一様ビット割当”でメモリと品質を最適化します。Hugging Face の総説が指摘するように、蒸留で密へ落とす実務路線は今後も主流で、量子化(QMoE/MoQE)×蒸留の二段圧縮はサービングの決定打になり得ます。

一方、Mixture Compressor(MC-MoE)のように量子化+動的プルーニングを組み合わせ、重要トークンを保持しつつ他は専門家選択を絞る系も登場。実行時の活性疎性をさらに押し広げる方向で、ECやドロップレス計算との組み合わせが見どころです。


4) 「密→疎」の架け橋:MoEfication と動的 $k$

MoEfication は、既存の密FFNを“機能分割”してMoE化するアプローチで、総パラメータの10–30%使用で95%+の性能維持を示しました。近年はDense→Dynamic-k MoE(D2D-MoE)のように、活性疎性を高めてから MoE 化し、ルータが $k$ を動的決定する設計も提案されています。既存モデルの推論コストを落とす実用ルートとして、微調整フェーズでの段階的 MoE 化やLoRA × MoEなどのハイブリッド手法が現実解になっていくはずです。


5) “MoEの外側”への拡張:ビジョン・注意・マルチモーダル

MoE は言語に限らずViTなどでも適用が進み、**パッチ(=トークン)単位の路由が自然にハマります。さらにMixture of Sparse Attention(MoSA)**のように、注意機構そのものをEC的に疎化する潮流も出てきました。MoE=FFNの置換という枠を越え、注意やモダリティ間融合に“混合と路由”を広げる研究線は今後も加速しそうです。


6) ルータの学習原理と安定化の理論化

実務で効く Z-loss選択精度(ゲートはFP32)は経験則として定着しましたが、なぜ効くかの理論は発展途上です。条件付き計算(Conditional Computation)の古典は強化学習的に“計算ポリシー”を学ぶ視点を与えており、EC/階層と結びつけた安定性の理論化が進むと、補助損失の自動調律学習前半の探索度管理など、現場寄りのノウハウがさらに体系化されるはずです。


Conclusion(まとめ)

MoE の本質は、“計算を増やさず容量を増やす”という矛盾の突破です。密FFNを専門家群に分解し、条件付きで一部だけ動かすことで、FLOPs を据え置いたまま表現力を拡張する。この発想は 1991 年の古典に始まり、2017 年の Sparsely-Gated MoE、2020–2022 年の GShard / Switch / ST-MoE を経て、Mixtral の隆盛とともに実務で回る技術として定着しました。現場から見れば、Router 安定化(Noisy Top-$k$, Aux, Z-loss)Capacity/通信のトレードオフEP(Expert Parallelism)×トポロジー最適化推論時のメモリ常駐という壁を、EC・Megablocks・FasterMoE・量子化/蒸留の“技術スタック”で一つずつ崩していくのが現在地です。

この先の実践指針を、**短く“要約の設計原則”**として置いておきます。

  1. まず当てにいく:Top-1/2 と $\mathrm{cf}\approx 1.0!-!1.25$、ゲート FP32+Z-loss、Aux は小さめに効かせる。
  2. 観測を欠かさない:使用率分散・ドロップ率・ゲートエントロピー・ストラグラー時間を全学習でログ
  3. 通信とメモリは“設計一次元”:EP 配置、近接優先、評価時は $\mathrm{cf}\downarrow$。必要なら EC / FasterMoE / Megablocks を導入。
  4. 提供は二段構え蒸留で密へ縮約 or QMoE/MoQE の非一様ビット。必要に応じて抽出サブネット
  5. 探索の最後に“格上げ”:EC や階層化で表現の余白を増やし、安定が保てるなら容量をもう一段積む。

最後にひとこと。MoE は“万能の速い密モデル”ではありません。強いのは知識依存・語彙豊富な領域で、推論一撃必殺系では密が良いことも多い。だからこそ、プリトレはMoEで攻め、提供は密/量子化/蒸留で守るという“二刀流”が現場の答えになりつつあります。あなたの記事では、ここまでの体系(歴史→仕組み→利点/課題→エンジニアリング→実践→発展→結論)がすでに整っているので、最後は各ハイパラの“初期値表”と計測スクリプト(使用率・ドロップ率・Z-loss など)を付録に置くと、読み手の“そのまま動かせる”価値がグッと上がります。

— 参考:
Hugging Face “Mixture of Experts Explained”(歴史・Z-loss・容量/通信・蒸留/量子化の見取り図).
A Visual Guide to MoE(専門家は“意味領域”ではなく構文/トークンに同期しやすい直観、ViTへの拡張)。
Bruno Magalhães のタイムライン(1991→深層MoE→Sparsely-Gated→階層MoE→重要度/負荷損失の式)。
Cameron R. Wolfe “Conditional Computation”(疎計算の系譜としてのMoEの位置づけ)。
MegaBlocks / FasterMoE / QMoE / Expert Choice(次世代の核になるドロップレス計算・トポロジー意識・極端圧縮・負荷分散ルート)。

1
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
1
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?