この記事は中国のNLP研究者Jianlin Su氏が運営するブログ「科学空間」で掲載された解説記事の日本語訳です。
苏剑林. (Feb. 08, 2025). 《MoE环游记：1、从几何意义出发 》[Blog post]. Retrieved from https://kexue.fm/archives/10699
過去二年間、ひょんなことから筆者は「Transformerレベルアップの道」というシリーズの連載を始め、主流的なTransformer構造の改善や個人的な考察を多く紹介した。一連の記事は、一部の読者には喜んで貰えたと思っている。そこで本記事から、同じスタイルで目下の新たな主流構造であるMoE（Mixture of Experts）を紹介したいと思う。
MoEの流行は改めて言うまでもないだろう。最近大バズりしたDeepSeek-V3はまさにMoEを採用しているし、GPT-4もMoEであると噂されている。国内（注：中国）で最近発表された多くのモデルもMoEを使い始めている。しかし、MoE自体は長く研究されてきたにもかかわらず、長い間特に注目を集められずにいた。去年年初に「Mixtral of Experts」が発表されたあたりから、MoEはようやく徐々に注目を集め始めたのである。
MoEの主な長所は、より大きなパラメーター量を持ちながら、学習と推論コストが顕著に低い点である。一方、学習が不安定・負担のバランスが悪い・あまり性能向上しないなどの課題もあり、それが故に当初はあまり流行らなかったのだ。しかし、ここ数年で注目が集まるにつれて、課題もだいぶ解決されている。この改善策を、これから逐一紹介していこうと思う。
問題の定義
始めに断っておくが、ここでは筆者自らの解釈に基づいてMoEを紹介しようと思っている。必要な所では参考文献を添付するが、MoE構造の根源まで深く掘り下げるつもりはないので、ご理解願いたい。
ご存知の通り、TransformerモデルはAttention層とMLP層からなる。MoEはMLP層に替わるものだ。MLP層はさらにFFN（FeedForward Network）とGLU（Gated Linear Unit）の二種類があり、GLUが主流だが、ここでは分かりやすいFFNを例に挙げよう。
$$
\boldsymbol{y} = f(\boldsymbol{x}\boldsymbol{W}^{(A)})\boldsymbol{W}^{(B)}
$$
$\boldsymbol{x}\in\mathbb{R}^d$は入力ベクトル、$\boldsymbol{W}^{(A)}\in\mathbb{R}^{d\times D},\boldsymbol{W}^{(B)}\in\mathbb{R}^{D\times d}$はパラメーター行列、$f$はelement-wiseな活性化関数である。さて、$D$を割り切れる何らかの整数$n$を使い、上の式をブロック行列でこう書き換えてみよう。
\boldsymbol{y} = f\big(\boldsymbol{x}\begin{bmatrix}\boldsymbol{W}^{(A)}_1 & \boldsymbol{W}^{(A)}_2 & \cdots & \boldsymbol{W}^{(A)}_n\end{bmatrix}\big)
\begin{bmatrix}\boldsymbol{W}^{(B)}_1 \\ \boldsymbol{W}^{(B)}_2 \\ \vdots \\ \boldsymbol{W}^{(B)}_n\end{bmatrix}=\sum_{i=1}^n \underbrace{f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i}_{\boldsymbol{v}_i}
\boldsymbol{W}^{(A)}_i = \boldsymbol{W}^{(A)}_{[:,(i-1)c:ic]}, \boldsymbol{W}^{(B)}_i = \boldsymbol{W}^{(B)}_{[(i-1)c:ic,:]},c= D/n
添え字はPythonの規則に従って書いている。つまり、FFNは$n$個のベクトル$\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$の和であり、各$\boldsymbol{v}_i$は小さなモデル$f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i$の出力である。この小さなモデルが、MoEの"Expert"である。
MoEの思想はこうだ。「$n$個のうち$k$個の$\boldsymbol{v}_i$だけで、$n$個の$\boldsymbol{v}_i$の和を近似できないだろうか？それができれば、計算量は$k/n$に減らせる！」
ノルム長で並べる
この問題を数式に書き換えると、
\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \lambda_i \boldsymbol{v}_i - \sum_{i=1}^n\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \lambda_i = k\
$\gamma_i = 1 - \lambda_i$とすると、
\mathop{\text{argmin}}_{\gamma_1,\gamma_2,\cdots,\gamma_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \gamma_i = n - k
この式を厳密に解くのは困難だが、簡単な近似解がある。仮に$\boldsymbol{v}_i$が互いに直交である場合、以下が成立する。
\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2 = \sum_{i=1}^n \gamma_i^2 \Vert\boldsymbol{v}_i\Vert^2 = \sum_{i=1}^n \gamma_i \Vert\boldsymbol{v}_i\Vert^2
明らかに、この場合最適解はノルム$\Vert\boldsymbol{v}_i\Vert$の小さい順から$n-k$個の$\gamma_i$を1とすることである。それはつまり、ノルムが大きい順の$k$個のベクトルの和で、$n$個のベクトルの和を近似するということだ。$\boldsymbol{v}_i$が直交条件を満たさなくても、これを近似解と見なしてもいいだろう。幾何学的に考えても分かりやすい。ベクトルのノルムが大きいほど、和を求める際に打ち消されにくくなるので、より大きな影響力を持つということだ。
高次元空間内では、ランダムに選んだベクトル同士の内積がゼロに近付く傾向があり、この性質からみても直交性を仮定した近似は合理的といえます。
MoEの出現
「ノルム長が長い順に$k$個のベクトルを選ぶ」という方針は決まった。だがよく考えてみると、これは実用的じゃない。ノルム長が長いベクトルを選ぶためには、全てのベクトルのノルム長を求めなければならない。それはつまり、まずは全ての$\boldsymbol{v}_i$を求めることになる。でも、本来の目的は$\boldsymbol{v}_i$の計算量を減らすことであるはずだ。
この矛盾を解決するために、ノルム長を低コストで計算できるよう、各Expertのモデルを見直す必要がある。どういうことかというと、まず、$\boldsymbol{v}_i$を正規化し、$\boldsymbol{e}_i = \boldsymbol{v}_i/\Vert\boldsymbol{v}_i\Vert$を求める。各$\boldsymbol{e}_i$のノルム長は同じである。続いて、以下の記号を定義する。
\underbrace{[\rho_1,\rho_2,\cdots,\rho_n]}_{\boldsymbol{\rho}} = h(\boldsymbol{x}\cdot\boldsymbol{W}^{(R)})\quad\in\mathbb{R}_{\geq 0}^n
$\boldsymbol{W}^{(R)}\in\mathbb{R}^{d\times n}$はパラメーター行列、$h(\cdot)$は$\mathbb{R}\to \mathbb{R}_{\ge0}$の活性化関数である。平たく言えば、$d$次元から$n$次元への線形変換に活性化関数を加えただけのものだ。このモデルを、MoEでは「Router」と呼んでいる。
$\boldsymbol{\rho}$の役割は、各Expertのノルムを予測することである。つまり、$\rho_i$を$i$番目のExpertのノルムと見なし、$\rho_i\boldsymbol{e}_i$をExpertの出力として扱うのだ。この出力は、計算量が小さいノルム$\rho_i$と、計算量が比較的大きい方向ベクトル$\boldsymbol{e}_i$に分けられる。計算量を削減するには、まず$\boldsymbol{\rho}$を計算し、大きい順に$k$個選んでから$\boldsymbol{e}_i$を計算し、最後に$\rho_i$を掛けて和を求めればいい。
\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i
これがMoEモデルの基本的な計算式である。Top-k部分だけ保留したため、このモデルは本質的にはSparseなモデルである。対して従来のFFN、つまり$k=n$の時のモデルは、Denseモデルと呼ばれている。
一旦まとめ
MoEをよく知る読者も知らない読者も、上記の説明は初めて聞いたかもしれない。これは筆者が独自に考えたMoEの解釈だからだ。ただ、明確な幾何学的な意味があるので、本質的には分かりやすい説明になっていると思う。
改めて、これまでの推理をまとめてみよう。
- 一般的なDenseなFFNは、$n$個のExpertベクトル$\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$の和に書き換えることができる
- 計算量を節約するために、$k$個のベクトルだけを選んで、元の$n$個のベクトルの和を近似したい
- 数学的に分析すると、ノルムが長い順に$k$個のベクトルを選べばいいことに気付く
- 全てのベクトルのノルムを計算すると計算量の節約にならないので、Expert自体を見直す必要がある
- $\boldsymbol{v}_i$を正規化して$\boldsymbol{e}_i$にしてから、別の軽量モデル（Router）でノルム$\rho_i$を予測し、Expertを$\rho_i\boldsymbol{e}_i$とする
- まず$\rho_i$を計算してから、ノルムが長い順に$k$個の$\boldsymbol{e}_i$を計算すれば、計算量を節約できる
動機
まだ疑問に思う読者がいるだろう。なぜこのような複雑なことをするのか？本来のMoEでも十分わかりやすいのではないか？
一般的なMoEは以下の形式である。
\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{v}_i
和を求める前に、特に$\boldsymbol{v}_i$を正規化したりせず、$\rho_i$もノルム長という意味を持っていない。Routerは単純にExpertの順番を決めるために用意されたモデルである。では、どうして$\rho_i$をExpertに乗算するだけで、RouterがExpertsの正しい順位を学習できるのか？論文「Sparse Backpropagation for MoE Training」で一つの解釈が提示されているが、まだ少し分かりづらい。
しかし本記事で示したような幾何学的な視点から見れば、一気に分かりやすくなる。Expertを$\rho_i\boldsymbol{e}_i$に再パラメーター化(Reparameterize)すると、Denseモデルは$\rho_i\boldsymbol{e}_i$全体の和、MoEは$\rho_i$のTop-kに対する和であり、これはDenseモデルに対する理論的に保証された近似である。「RouterはどのExpertを選べばいいのか」ということは特に考慮せず、ただDenseモデルになるべく近付けることだけを目指した結果なのだ。「大きいパラメーター数」と「小さい計算量」を両立させたいなら、これが最良の選択であることは明白だろう。
$\rho_i$は確率値ではないので、活性化関数$h(\cdot)$は正規化の要件を満たす必要はない。Softmaxのほかに、SigmoidやReLUも考えられるだろう。Routerで非正規化な活性化関数を使うと、$k>1$の時にExpert間で「過当競争」が起こる現象を避けられるので、より良い効果が得られる場合もある。
ちなみに、先ほどの説明ではノルムを揃えるために$\boldsymbol{e}_i = \boldsymbol{v}_i/ \Vert\boldsymbol{v}_i\Vert$と定義したが、実際の実装はL2正規化でなくてもいい。たとえばgamma=1のRMS Normなど、普段使い慣れた正規化手段を選んでもいいだろう。
全体まとめ
本記事は「Denseモデルの最良な近似」という視点から、少し特殊な形のMoEを導出した。一般的なMoEと比べて正規化処理が増えているが、MoEが幾何学的な意味を持つようになった。もちろん、正規化があっても無くても、ここはまだMoEの序の口である。まだ多くの困難が我々を待ち受けている。