はじめに
新緑の音色がそよぎに運ばれる今日この頃、みなさまいかがお過ごしでしょうか？
六花 牡丹（りっか ぼたん）と申します。
本記事は
で公開した、IMoE: Indefinacy Mixture of Expertsに関する簡単な技術メモになります。
IMoEは日本語で「不定形性専門家混合」という意味であり、通常のMoE: Mixture of Experts *1が決まった数のExpertsを毎回使うのに対して、IMoEは入力とモデル自身の自信度から用いるExpertsの数を動的に変えるというモジュールになっています。
これは後程説明しますが、確率グラフ理論と脳神経の成長過程から着想を得たものであり、かつ効率的な推論を行うための機構でもあります。
学習の過程で、IMoEは脳神経と同じようなネットワーク形成メカニズムとなり、かつ情報分布のエントロピーという視点からもよい性質の機構となっています。
本記事は技術メモであるため、分量は少ないです。
行間をいくつか飛ばしている部分があります。
拙筆ではございますが、皆様のお役に立つことを心から願っております。
未熟者故、記事中にて誤記・欠落などが見られることがございます。
もし発見しました場合には、コメント等にてご指摘いただきますようお願い申し上げます。
六花牡丹のX(Twitter)アカウント
ここで最新の進捗や技術に関する情報（時々近況）を共有していきますので、もしよろしければフォローなどよろしくお願いいたします。（Also, question, issue, correction request, joint research, recruit, etc...）
目次
１．背景：確率グラフ理論
２．IMoEの構築動機
３．IMoEの原理
４．GPU上での高速化のための工夫
５．最後に
参考論文・記事等
１．背景：確率グラフ理論
IMoEは確率グラフ理論の考えを背景に持ちます。
本章ではその理論について簡単に説明します。
１ー１．Erdős-Rényi-Gilbert model *2 *3
Erdős-Rényi-Gilbert modelはランダムグラフを生成するための基本的な数学的モデルであり、グラフ理論やネットワーク科学の分野で広く利用されています。
生成原理はn個の頂点を持つグラフの中から、等確率かつ独立にM本の辺を生成するというものです。
\displaylines{
G(n): \; Undirected \; graph \\
(i, j): \; specific \; node\\
s(n): \; expected \; number \; of \; connected \; node \\
\mathbb{P}[(i, j) \in G(n)] = p \quad \Rightarrow \quad s(n) = (n-1)p
}
Erdős-Rényi-Gilbert modelはランダムグラフの基礎的な生成原理であり、生物における胎児期（幼少期）の脳神経回路生成の仕組みとよく一致します。
１－２．fitness model (Bianconi-Barabási model) *4
Bianconi-Barabási modelもランダムグラフを生成するための基本的な数学的モデルですが、Erdős-Rényi-Gilbert modelとは生成原理の異なるモデルです。ノードは「次数が高いほど新たなリンクを得やすい」（優先的選択）というルールに従って成長します。
各ノードには「適応度（fitness）」
\eta_j \in (0, 1]
が与えられ、この$ \eta_j $がノードごとの違いを表すパラメータとなります。
Bianconi-Barabási modelは等確率ではない不均衡なネットワークとなることで、活動依存的なシナプス形成につながります。
また、Bianconi-Barabási modelは「一緒に発火するニューロンはつながりやすくなる」という 使用依存的可塑性（Hebbian learning rule） とも一致します。
さらに、異なる機能を持つ領域に対応した異質なノードが形成されること（機能分化・局在性）はこのモデルで自然に表現することができます。
LLMでも最近、「ハイパーパラメータ」と呼ばれる、機能の集中するパラメータが存在し、神経回路に不均衡性があることが示唆されています。
なお、機能分化および局在性は成熟期の脳神経回路をよく説明できます。
Bianconi-Barabási modelにおける新規ノードがノード$ j $にリンクされる確率$ \Pi_j $は
\Pi_j = \frac{\eta_j k_j}{\sum_i \eta_i k_i}
と表されます。
２．IMoEの構築動機
ここまで、確率グラフ理論および脳神経科学から神経回路の生成についてお話してきました。
ここで神経回路の不均衡性が存在するならば、それを考慮したモデルをあらかじめ作成すればよいのではないかと考えるかもしれません。
しかし、神経回路の不均衡性は外部からは観察できず、かつ数百億にもおよぶノードの不均衡性を設計するのは現実的でもありません。
また、すべての不均衡性はあくまで「確率的」に存在するのであって、定まった形があるものではありません。
したがって、学習の過程で確率的な不均衡性を実現する手法が必要になります。
この手法こそが IMoE: Indefinacy Mixture of Experts となります。
IMoEはErdős-Rényi-Gilbert modelからBianconi-Barabási modelへの遷移を非明示的に行えるモデルであり、かつ通常のMoEと比較して計算効率の良い手法となります。
さらに、人の脳は自信がない場合はより多くの領域から記憶を探索し、逆に自信がある場合は特定領域のみから記憶を引く出すことで、効率的に推論を行います。
このモジュールは人の脳のこの仕組みを再現するためのものでもあります。
そして、自信とは情報分布のエントロピーと同値であり、簡単なアルゴリズムで達成することが可能です。
（エントロピー（乱雑度）が大きいとは、確率が一様的であり、判断ができていないことと同値となります）
「自信」を定式化すると下記になります。
\displaylines{
H(p(y \; | \; x)) = - \sum_{i=0}^K p_i(y \; | \; x) \ln p_i(y \; | \; x) \\
Confidence(x) := 1- \frac{H(p(y \; | \; x))}{log|y|}
}
IMoEは情報分布のエントロピーを神経回路生成のトリガーにしています。
３．IMoEの原理
IMoEの数学的な定式化は下記のように簡単に行えます。
\displaylines{
W_g \in \mathbb{R^{M \times N}}, \; W_o \in \mathbb{R^{N \times L}}, \; x \in \mathbb{R^N} \\
E = \left( f_1, f_2, \cdots f_G \right) \\
\forall f_i (f_i \in \mathbb{R^{L \times N}}) \\
s: \; gate \; thres \\
Softmax(x) = \frac{e^{x_k}}{\sum_{k=0}^K e^{x_k}} \\
G(x) = Softmax(W_g x) \\
\xi_k(x), \; \Gamma_k(x) = TopP(G(x), s) \\
y = W_o \sum_{i \in \Gamma_k(x)}\xi_i(x)f_i(x)
}
IMoEは通常のMoEがTop-kでゲーティングを行う部分をTop-pで行う形で実装されます。
学習初期は、ランダムなゲーティングとなるため、すべてのノードはより等確率で現れます。学習に従って、特定の情報では特定のノードが使用されるようになり、機能分化・局在性が増大していきます。
ここで、Top-pの前にSoftmax関数を用いている点が重要となります。
Softmax関数はスコアの合計を１に保つため、ある閾値が定められているとき、自信のあるノードがあれば、他のノードは選ばれにくくなり、逆に自信のあるノードがない場合はより多くのノードから情報を取得するという機構が自然と成立します。
これにより自信がない場合はより多くの領域から記憶を探索し、逆に自信がある場合は特定領域のみから記憶を引く出すことで、効率的に推論を行うという人の脳の仕組みを再現することができます。
前章の定式化から、
\displaylines{
1- \frac{H(p(y \; | \; x))}{log|y|} \approx \xi_k (x) \\
\Rightarrow Confidence(x) \approx \xi_k (x)
}
であることは明らかでしょう。
Top-pの値を$ \frac{1}{k} $に設定すれば、Nノードあるうちの最大kノードがアクティベートされることとなり、$ \frac{N}{k} $を任意に設定することで通常のMoEより効率的に推論することが可能になります。
４．GPU上での高速化のための工夫
ここからは、実際の実装に関する事柄について説明していきます。
MoEにおいて、下記のような単純な実装では速度が著しく下がることが知られています。
for expert_idx in range(self.gate_num):
mask = (top_k_indice == expert_idx)
outputs[mask] = self.experts[expert_idx](x[mask])
速度が下がるのは計算量が大きいからではありません。
GPUには同一メモリへの頻繁なアクセスが速度を低下させるという特徴があります。
速度低下の要因は下記です。
- DRAMのにおける同一メモリアクセスによる帯域幅の競合
- SIMT（Single Instruction, Multiple Thread）での分散的なスレッド処理ができなくなる
したがって、同一メモリへのアクセスを最小限にとどめるという工夫がSparse MoEでは重要になります。
IMoEでは下記のような実装となります。
final_hidden_states = torch.zeros(
(bsz*seql, self.inter_dim), dtype=x.dtype, device=x.device
)
for expert_idx in range(self.gate_num):
mask = (top_p_indices == expert_idx).nonzero(as_tuple=True)[0]
if mask.numel() == 0:
continue
current_state = x.reshape(-1, embs).index_select(0, mask)
current_hidden_states = (
self.experts[expert_idx](current_state)
* weight_probs[:, expert_idx].index_select(0, mask).unsqueeze(-1))
final_hidden_states.index_add_(
0,
mask,
current_hidden_states.to(x.dtype))
特に index_add_ プロセスによる実装は広く採用されており、Sparse MoEを実装する際は使用が推奨されます。
なお、Dense MoEの場合はアインシュタイン縮約記法 einsum により高速に実装できます。
（余談ですが、Mambaでも高速化としてこの手法が採用されています）
また、ゲーティングが偏りすぎることによる学習速度低下に対処するため、IMoEではスコアに外乱を加えています。
z-lossを用いていない理由は、不均衡性を実現したいという意図があるからです。
５．最後に
計算機構について研究を行うにつれて、単に機械学習や統計学だけではない、広範な知識・他分野に対する鋭い洞察力の必要性を感じています。
広い視野を持てるように研鑽を積む重要性をひしひしと感じています。
参考論文・記事等
*1 'FastMoE: A Fast Mixture-of-Expert Training System', https://arxiv.org/abs/2103.13262
FastMoEの原論文です。アルゴリズムも記載されているため、MoEの学習に適しています。
*2 'On random graphsⅠ.', https://www.renyi.hu/~p_erdos/1959-11.pdf
ErdősとRényiによる論文です。（Erdős-Rényi-Gilbert modelに関する）
*3 'Random Graphs', https://projecteuclid.org/journals/annals-of-mathematical-statistics/volume-30/issue-4/Random-Graphs/10.1214/aoms/1177706098.full
Gilbertによる確率グラフ理論の論文です。
*4 'Competition and multiscaling in evolving networks', https://arxiv.org/abs/cond-mat/0011029
Bianconi-Barabási modelの原論文です。
*5 'The organization of behavior; a neuropsycholocigal theory', https://www.taylorfrancis.com/books/mono/10.4324/9781410612403/organization-behavior-hebb
使用依存的可塑性（ヘブ則 Hebbian learning rule）の本です。
*6 'Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity', https://arxiv.org/abs/2101.03961
MoE + Transformerによる精度検証を行った論文です。
執筆者：六花 牡丹（りっか ぼたん）
おさげとハーフツイン・可愛いお洋服が好きで、基本的にふわふわしている変わり者。
結構ドジで何もないところで転ぶタイプ。
人工知能に関しては独学のみ。
アラビアンナイトという期間限定の一番好きな紅茶が販売される季節になって、手に入るのがとっても楽しみになっている。
