MoEとは？複数の専門家を使うLLM構造

Routerがtokenごとに一部のExpertを選ぶMoEの全体像

Routerがtokenごとに一部のExpertを選ぶMoEの全体像

LLMのparameterを増やすと、一般には一回の推論で実行する計算も増えます。MoE（Mixture of Experts：複数のExpertから一部を選ぶ構造）は、この結び付きを弱めるための条件付き計算です。

3行で分かるMoE

  • MoEでは、Transformer内部のFFNを複数のExpertへ分け、Routerがtokenごとに上位\(k\)個を選びます。
  • 全Expertのparameterをmodel capacityとして持ちながら、一つのtokenでは選ばれたExpertだけを計算できます。
  • ただし、全weightを保持するmemory、Routerの偏り、Expert間の負荷分散、device間通信は残ります。総parameterが大きいのに、あらゆるcostが小さくなるわけではありません。

この記事ではSparsely-Gated MoEの原論文から、Routerの数式、Top-2の数値例、負荷分散、expert capacity、分散実装までを順に説明します。

Dense LLMで容量を増やすと計算も増える

通常のTransformer blockでは、Self-Attentionの後にFFN（Feed-Forward Network：各tokenのhidden stateを個別に変換する層）があります。単純化したDense FFNの出力は次の形です。

\[\operatorname{FFN}(h)=W_2\,\phi(W_1h)\]

\(h\)は一つのtoken位置のhidden state、\(W_1\)と\(W_2\)は学習するweight、\(\phi\)はactivation functionです。Dense（すべてのparameterを毎回使う構造）では、各tokenが同じFFNの全weightを通ります。

全FFNを使うDense構造と一部Expertだけを使うMoE構造の比較

FFNの中間次元を大きくすればparameter容量は増えますが、matrix multiplicationも大きくなります。layer数を増やした場合も、原則として一回のforward passで通る計算が増えます。

MoEが狙うのは、保持するparameterの総量を増やしながら、一つのtokenで動かすparameterを限定することです。Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layerは、conditional computation（入力に応じてnetworkの一部だけを動かす計算）によって、計算量に比例せずmodel capacityを増やす問題へ取り組みました。

MoE layerを構成する4つの役割

LLMで使われる代表的なSparse MoE layerは、次の4つへ分けると追いやすくなります。

部品 役割 注意点
Router / Gate tokenのhidden stateから各Expertのscoreを出す 予測するのは次tokenではなく、処理を任せるExpert
Dispatcher 選択結果に従ってtokenをExpertへ並べ替える 分散環境ではdevice間通信が発生する
Expert tokenを変換するFFN subnetwork 多くの場合、同じ形で別parameterを持つ
Combiner 選択Expertの出力をweight付きで合成する 元のtoken順へ戻す処理も必要

ここでいうExpertは、「数学専門LLM」「法律専門LLM」のように別々に完成したchatbotを並べたものとは限りません。Sparsely-Gated MoE原論文では、同じ入出力shapeを持ち、parameterだけが異なるfeed-forward networkとして定義されています。

ExpertとRouterは通常、model全体のlossからend-to-endで一緒に学習されます。人間が各Expertへ分野名を付けてから振り分ける仕組みではありません。学習後にsyntaxやsemanticsに関する分化が観察される場合はありますが、すべてのExpertが人間に説明しやすい専門分野を持つとは限りません。

RouterはtokenごとにTop-k Expertを選ぶ

一つのMoE layerに\(E\)個のExpertがあるとします。Routerはtokenのhidden state \(h_t\)から、各Expertに対応するlogitを計算します。

\[z_t=W_rh_t\]

\(W_r\)は学習可能なRouter weight、\(z_t\in\mathbb{R}^{E}\)は\(E\)個のscoreです。このうち大きい\(k\)個だけを残し、残りを\(-\infty\)にmaskしてsoftmaxを取れば、選択されなかったExpertのweightは0になります。

\[G(h_t)=\operatorname{softmax}(\operatorname{KeepTopK}(z_t,k))\]

MoE layerの出力は、選ばれたExpertの出力をRouter weightで合成したものです。

\[y_t=\sum_{i=1}^{E}G_i(h_t)\,E_i(h_t)\]

\(G_i(h_t)=0\)なら、\(E_i(h_t)\)を計算する必要はありません。\(E=64\)でも\(k=2\)なら、Expert部分で実行するのは原則2個です。

一つのpromptを丸ごと一人へ渡すわけではない

同じ文のtokenがMoE layerごとに異なるExpertへ送られる様子

LLMのMoEを理解するときに大切なのは、routingの粒度です。代表的なtoken-level MoEでは、一つのpromptにつき一つのExpertを選ぶのではありません。

たとえば「東京でPython勉強会を探す」というtoken列があっても、同じMoE layerで各tokenのhidden stateは異なります。Routerは位置ごとに異なるExpertの組を選べます。次のMoE layerへ進むとhidden stateが変わるため、同じtoken位置でも別のExpertが選ばれる場合があります。

つまりroutingは、概念的には次の二重の単位で起こります。

  1. token位置ごと
  2. MoE layerごと

Routerが見るのはtoken IDだけではなく、前段までの文脈を含んだhidden stateです。このため、同じ表記のtokenでも周囲の文脈が違えばscoreが変わり得ます。

Top-2 routingを具体的な数値で追う

8個のExpertを持つlayerで、あるtokenに対するRouter logitが次の値だったとします。これは計算手順を説明するための独自例です。

Expert logit
E1 0.2
E2 1.1
E3 -0.4
E4 2.0
E5 0.7
E6 1.6
E7 -0.1
E8 0.5

8 ExpertのlogitからTop-2を選び重み付き出力を計算する過程

Top-2ではE4の2.0とE6の1.6を残します。選択後の2値だけでsoftmaxを計算すると、weightは次のとおりです。

\[g_4=\frac{e^{2.0}}{e^{2.0}+e^{1.6}}\approx0.599\]
\[g_6=\frac{e^{1.6}}{e^{2.0}+e^{1.6}}\approx0.401\]

Expert出力を\(E_4(h_t)\)、\(E_6(h_t)\)とすると、最終出力は次の重み付き和です。

\[y_t\approx0.599E_4(h_t)+0.401E_6(h_t)\]

実装によっては全Expertへsoftmaxを適用してからTop-kを取るか、Top-kの後に選択値を再正規化するかが異なります。score functionもsoftmaxだけとは限りません。重要なのは、Routerが疎な割り当てと合成weightを作り、選択外Expertの計算を省くという共通構造です。

Sparsely-Gated MoE原論文のNoisy Top-k gatingでは、学習中のlogitへ学習可能なscaleのGaussian noiseを加えます。これは選択の探索と負荷分散を助けるための設計です。すべての現代MoEが同じnoise式を使うわけではありません。

TransformerのどこがMoEになるのか

Transformer blockでDense FFNがMoE FFNへ置き換わる位置

2017年のSparsely-Gated MoE論文は、stacked LSTMの間にMoE layerを置きました。現在のLLMでよく見る構成は、Transformer blockのFFNをMoEへ置き換える形です。GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingは、Transformerのposition-wise MoEを複数のFFN Expertとして説明しています。

単純化すると、Dense blockは次の流れです。

hidden state
  → Self-Attention（共有）
  → Dense FFN（毎回1個を全tokenが実行）
  → next block

MoE blockではFFN部分が置き換わります。

hidden state
  → Self-Attention（共有）
  → Router
  → Top-k Expert FFN
  → weighted combine
  → next block

したがって、MoE modelの全parameterが条件付きになるわけではありません。embedding、Attention、normalization、出力headなどのshared parametersは、通常どおり各tokenで使われます。また、すべてのTransformer blockをMoEにするとは限らず、一定間隔でMoE layerを配置する設計もあります。

Transformer全体の復習が必要な場合は、Transformerとは？Attention Is All You Need論文からLLMの基本構造をやさしく解説を先に読むと、MoEが置き換える位置を確認できます。

総parameterとactive parameterを分けて考える

MoEの総parameter・active parameter・weight memory・通信costを分けた計器盤

MoEの「大きいのに計算が軽い」という説明では、少なくとも4つの指標を分ける必要があります。

指標 何を数えるか Expertを増やしたとき
Total parameters modelが保持する全parameter ほぼExpert数に応じて増える
Active parameters 一つのtokenで実際に通るparameter Top-kが固定ならExpert総数ほど増えない
FLOPs forwardで行う演算 Expert部分は主に\(k\)へ依存する
Weight memory weightを保存・配置する容量 選ばれなくても全Expert分が必要

active parametersの集計範囲は資料によって異なります。選択Expertだけを指す場合もあれば、Attentionなど毎回使うshared parametersを加えた値を指す場合もあるため、model同士を比べるときは定義を確認してください。

一つのMoE layerで、各Expertが\(P_e\) parameters、Expert数が\(E\)、Top-kが\(k\)だとします。Expert部分だけを単純化すると、総parameterとtokenごとのactive parameterは次のように分かれます。

\[P_{\mathrm{total,expert}}=E\,P_e\]
\[P_{\mathrm{active,expert}}\approx k\,P_e\]

\(E\)を8から64へ増やしても、\(k=2\)を維持すれば一つのtokenが通るExpert部分は2個です。これがparameter capacityと計算を部分的に切り離す中心です。

ただし、次のcostは式から消えていません。

  • Routerが全Expertのscoreを計算するcost
  • Attentionなどshared layerの計算
  • 全Expert weightをGPU memory、CPU memory、または複数deviceへ置くcost
  • tokenを選択Expertへ移動し、元の順へ戻すcost
  • Expertごとのtoken数が小さいときのGPU利用率低下

このため、total parametersとactive parametersだけを比べて「同じ比率で速い」と結論づけることはできません。Switch Transformersも、algorithm上の計算削減と最終的なspeed差を分け、実測速度がlow-level optimizationに依存すると説明しています。

RouterとExpertはどう学習されるのか

MoEは、完成済みExpertをRouterが後から選ぶだけの仕組みではありません。基本的にはRouter、Expert、shared layerを同じ事前学習lossからbackpropagationで更新します。事前学習そのものは、LLMの事前学習とは？大量テキストから言語能力を獲得する仕組みで解説しています。

あるtokenがE4とE6へ送られた場合、そのstepでは主にE4、E6とRouterの選択weightへ勾配が流れます。選ばれなかったExpertは、そのtokenから直接の更新を受けません。これによりExpertごとに異なる役割が生まれる余地があります。

一方、hardなTop-k index選択は不連続です。Sparsely-Gated MoE原論文では\(k>1\)のとき、選択されたExpertの非zero gate値を通じてRouterへ勾配が流れると説明しています。後続研究では、routing、探索、補助loss、精度の扱いに複数の設計があります。

「Expertが自然に専門化する」という表現にも条件が必要です。原論文ではsyntaxやsemanticsに基づくspecializationを観察していますが、専門化の形はdata、objective、Router、Expert granularityに依存します。Expert番号だけを見て、人間が意味を決めつけるべきではありません。

一部のExpertへ集中すると何が起こるか

tokenがExpertへ分散する正常系と一つへ集中する負荷偏りの比較

Routerが純粋にscore上位だけを選ぶと、初期に少し有利だったExpertへtokenが集中することがあります。選ばれたExpertは多くのdataで更新され、さらに選ばれやすくなります。反対に、選ばれないExpertは学習機会を失います。

この状態では、次の問題が同時に起こります。

  • busy Expertの処理待ちが長くなる
  • idle Expertの計算資源が使われない
  • 一部Expertのbufferだけがあふれる
  • 選ばれないExpertが十分に学習されない
  • device間でstep完了を待つ時間が増える

Sparsely-Gated MoE原論文はExpertのimportanceとloadを均す補助lossを使いました。GShardも、実際のdispatch比率と平均gate値を使うauxiliary loss（主目的を助ける追加損失）をmodel lossへ加えています。

補助lossの目的は、「すべてのtokenを均等にランダム配布すること」ではありません。Routerがtask lossを改善する選択を学びつつ、一部Expertへの極端な集中を抑えることです。係数が強すぎれば、内容に応じたroutingより均等化を優先する可能性があります。弱すぎれば偏りを抑えられません。最適な係数はmodel、batch、Expert数などに依存し、原論文の値を別構成へそのまま移せるとは限りません。

ST-MoE: Designing Stable and Transferable Sparse Expert Modelsでは、load balancingとは別に、Router logitが大きくなりすぎることを抑えるrouter z-lossも検討されています。これはnumerical stability（丸め誤差などで学習が不安定になる問題）への対策であり、負荷分散と同じ役割ではありません。

Expert capacityとoverflowは別の安全装置

負荷分散lossがあっても、各batchでtoken数が完全に均等になる保証はありません。そこで分散実装では、一つのExpertが一度に受け取れるtoken数へcapacity（bufferの上限）を設けます。

一般化した目安は次の形です。

\[C=\left\lceil c\frac{kT}{E}\right\rceil\]

Expert capacityを超えたtokenがresidual経路へ迂回する様子

\(T\)はbatch内token数、\(E\)はExpert数、\(k\)はtokenごとの選択数、\(c\)はcapacity factorです。\(c=1\)なら平均割り当て相当、1より大きければ偏りを吸収する余裕を持たせます。実際の式と割り当て単位はimplementationにより異なります。

たとえば\(T=1{,}024\)、\(E=8\)、\(k=2\)、\(c=1.25\)なら、概念上のcapacityは次の値です。

\[C=\left\lceil1.25\times\frac{2\times1{,}024}{8}\right\rceil=320\]

あるExpertへ320件を超えて割り当てようとするとoverflowです。GShardでは、選択先がcapacityを超えたtokenについてMoE出力をzeroにし、residual connectionで次layerへ渡す構成を説明しています。Switch Transformerでも、overflow tokenがExpert計算をskipする設計が使われています。

capacity factorにはtradeoffがあります。

設定 利点 弱点
小さい buffer、memory、padding計算を抑えやすい overflowとtoken dropが増えやすい
大きい routingの偏りを吸収しやすい 空きslot、memory、通信が増える

補助lossは学習によって割り当て分布を整える仕組み、capacityは実行時に一つのExpertへ入れられる上限です。両者を同じものとして扱わないことが大切です。

分散実装ではtokenがdevice間を移動する

4台のGPU間でtokenをAll-to-All交換するExpert Parallelism

Expert数が増えて一台のacceleratorへ収まらなければ、Expertを複数deviceへ分けます。これがExpert Parallelism（Expertをdevice間で分割する並列化）です。

たとえばGPU 0上のtokenが、GPU 3にあるE6を選んだとします。処理は概念的に次の順になります。

  1. 各deviceでRouter scoreとTop-kを計算する
  2. tokenを選択Expertがあるdeviceへdispatchする
  3. 各deviceが担当Expertをbatch処理する
  4. Expert出力を元のtoken所有deviceへ返す
  5. 元のtoken順へ戻し、weight付きで合成する

複数deviceが互いに異なるtokenを交換するため、All-to-All（全device間で異なるdataを交換する通信）が使われます。NVIDIA Megatron CoreのMoE documentationも、Expert Parallelismのtoken dispatcherとしてAll-to-Allを扱っています。

MoEのExpert部分でFLOPsを抑えても、通信待ちが長ければGPUは計算を続けられません。さらに、Expertごとのtoken batchが小さすぎるとmatrix multiplicationが細切れになり、hardwareを使い切れません。Grouped GEMM、通信と計算のoverlap、Routerやpermute処理のfusionなどが必要になるのはこのためです。

最小RouterをPythonで確かめる

次のcodeは、一つのtokenに対するRouter logitからTop-kを選び、選択値だけでstable softmaxを計算する教育用の最小例です。tensorのbatch処理、capacity、分散dispatch、backpropagationは含みません。

from __future__ import annotations

import logging
import math
from dataclasses import dataclass

LOGGER = logging.getLogger(__name__)


@dataclass(frozen=True)
class Route:
    """一つのExpertへのrouting結果。"""

    expert_id: int
    weight: float


def route_top_k(logits: list[float], top_k: int) -> list[Route]:
    """Router logitからTop-k Expertと正規化weightを返す。

    Args:
        logits: Expertごとの未正規化score。
        top_k: 一つのtokenで選択するExpert数。

    Returns:
        weightの大きい順に並んだrouting結果。

    Raises:
        ValueError: logitsが空、有限値でない、またはtop_kが範囲外の場合。

    Example:
        >>> routes = route_top_k([0.2, 1.1, -0.4, 2.0, 0.7, 1.6, -0.1, 0.5], 2)
        >>> [route.expert_id for route in routes]
        [3, 5]
        >>> round(sum(route.weight for route in routes), 6)
        1.0
    """
    if not logits:
        LOGGER.error("routing failed: logits are empty")
        raise ValueError("logits must not be empty")
    if not 1 <= top_k <= len(logits):
        LOGGER.error("routing failed: top_k=%d experts=%d", top_k, len(logits))
        raise ValueError("top_k must be between 1 and the number of experts")
    if any(not math.isfinite(value) for value in logits):
        LOGGER.error("routing failed: non-finite logit found")
        raise ValueError("all logits must be finite")

    selected = sorted(
        enumerate(logits),
        key=lambda item: item[1],
        reverse=True,
    )[:top_k]
    max_logit = selected[0][1]
    exponentials = [math.exp(value - max_logit) for _, value in selected]
    denominator = sum(exponentials)

    routes = [
        Route(expert_id=expert_id, weight=exp_value / denominator)
        for (expert_id, _), exp_value in zip(selected, exponentials, strict=True)
    ]

    if top_k > 1 and routes[0].weight > 0.95:
        LOGGER.warning("routing is highly concentrated: weight=%.3f", routes[0].weight)
    LOGGER.debug("router logits=%s", logits)
    LOGGER.info(
        "selected experts=%s",
        [(route.expert_id, round(route.weight, 3)) for route in routes],
    )
    return routes

max_logitを引いてから指数を取るのは、softmaxのoverflowを避けるためです。sampleのExpert IDは0始まりなので、結果の3と5は表のE4とE6に対応します。

本番実装ではPython loopでtokenを一つずつ処理しません。Router projection、Top-k、token permutation、Expertごとのgrouped matrix multiplication、逆permutationをtensor operationとしてまとめます。また、Router logitや合成weightの精度をExpert計算と別に扱うimplementationもあります。

DenseとMoEは何を交換するのか

MoEはDense modelの一方向な上位互換ではありません。選択は、modelを置くhardware、batch、latency目標、training方法まで含めて行います。

比較軸 Dense FFN Sparse MoE FFN
tokenごとの経路 毎回同じFFN RouterがTop-k Expertを選ぶ
total parameter拡張 計算も増えやすい Expert数とactive数を分けられる
weight memory 総parameter分 総parameter分。非active Expertも保持が必要
演算の規則性 高い token割り当てで変動する
device間通信 model並列方式に依存 Expert Parallelではdispatch/returnが増える
load balance FFN間のrouting偏りはない 補助lossとcapacity設計が必要
training安定性 比較的単純 RouterとExpertの相互作用が増える
fine-tuning 全weightまたはPEFTを設計しやすい Router、Expert、shared layerの更新範囲を決める必要
小batch・低latency 規則的で最適化しやすい Expert batchと通信overheadが不利になり得る

一台に全weightが収まり、低latencyの小batch推論を優先するなら、Dense modelの規則性が有利な場合があります。大量のdataとdeviceを使い、計算量に対してmodel capacityを増やしたい場合はMoEが候補になります。

MoEのweight memoryが問題になる場合、量子化は別の軸で保存量を減らせます。ただしRouterやExpertごとの精度影響を含めた検証が必要です。量子化の基本はLLMの量子化とは？INT8・INT4でモデルを軽くする仕組みで確認できます。

Sparsely-Gated MoE論文が示したことと限界

2017年のSparsely-Gated MoE論文は、最大数千のfeed-forward Expertから少数を選び、model capacityを計算量に比例させず増やせることを大規模言語modelと機械翻訳で示しました。

1 Billion Word benchmarkでは、forward passを約8 million operations per timestepに揃えた系列でExpert数を増やし、最大の4096 Expert構成が4 Expert構成よりtest perplexityを24%下げたと報告しています。また、論文全体では最大137 billion parametersのMoE構成を扱いました。

この結果から直接言えるのは、同論文のLSTM architecture、dataset、training方法、K40 GPU clusterという条件で、疎なExpert追加がcapacity拡張に機能したことです。次の主張までは導けません。

  • 現在のすべてのTransformer MoEがDense modelより同じ割合で高品質になる
  • total parametersが同じなら品質も同じになる
  • active parametersが同じならlatencyも同じになる
  • Expert数を増やすほど単調に性能が上がる

現代のMoEを読むときにも残っている中心的な問いは、Routerがどの粒度で何個を選ぶか、負荷をどう均すか、Expertをどこへ配置し、通信をどう隠すかです。

まとめ

MoEは、Transformer内部のFFNを複数のExpertへ分け、Routerがtokenごと・MoE layerごとに一部だけを動かす条件付き計算です。Expert総数\(E\)とactive数\(k\)を分けることで、総parameter容量を増やしながらExpert部分の計算を主に\(k\)へ依存させられます。

一方で、全weightのmemory、Routerの学習、load balance、expert capacity、overflow、All-to-All通信は残ります。MoEを理解する際は「total parametersが大きい」というcapacityの話と、「一つのtokenで何を計算するか」というactive computeの話を分けてください。

次回は、Top-2などのroutingをTop-1へ単純化し、大規模なTransformerへ適用したSwitch Transformerを扱います。

参考資料

スポンサーリンク

コメント

タイトルとURLをコピーしました