9月11日、Giles Thomasが「Extending Raschka's GPT-2: an MoE trained from scratch on an RTX 3090」と題した記事を公開した。この記事では、Sebastian RaschkaのGPT-2実装をベースにMixture-of-Experts（MoE）アーキテクチャを一から実装し、RTX 3090でトレーニングした実践的な記録が詳しく紹介されている。コンシューマーGPUでどこまでやれるかを地道に検証した記録として、実装面でも示唆に富む内容だ。
RTX 3090でMoEをスクラッチ実装する
LLMの学習コストが高騰するなか、コンシューマーGPUで大規模モデルに近い性能を出そうとする試みへの関心が高まっている。Giles Thomasは、Sebastian Raschkaの著書「Build a Large Language Model (from Scratch)」のGPT-2実装を出発点に、MoEアーキテクチャを自力で組み込み、446Mパラメータのモデルをゼロから訓練した。エキスパート数6、各トークンにつき2エキスパートを使用する構成で、アクティブパラメータは220Mに抑えられている。
トレーニング期間は約8日間。結果として、テストセットの損失はこれまでに著者が訓練した他のどのモデルよりも良好だった。比較対象のGPT-2 small（124Mパラメータ）を上回り、GPT-2 medium（355Mパラメータ）には及ばないが近い水準に達した。著者のモデルはGPT-2 mediumより総パラメータ数が多いにもかかわらず、アクティブパラメータ数は少ない点が興味深い。
MoEの仕組み：FFNを複数に分割する
MoEの核心はシンプルだ。GPT-2スタイルのTransformerブロックには、Attentionの後にFeed-Forward Network（FFN）が置かれている。MoEではこのFFNを複数の独立したFFN（エキスパート）に置き換え、各コンテキストベクトルを一部のエキスパートにだけ流す。
GPT-2のFFNは「線形層→次元を4倍に拡張→GELU活性化→線形層で元の次元に戻す」という2層構造で、一見シンプルに見える。しかしFFNはAttentionの2倍のパラメータ数を占め、モデルの「思考」を担う部分だと著者は説明する。MoEはこの部分を複数に増やすことで、メモリ使用量をほぼ変えずに計算量を削減しながらキャパシティを拡大する。
全エキスパートをRAMに保持する必要はあるが、各入力に対して全エキスパートを通す必要はない。これがMoEの速度上の利点の源泉だ。オープンウェイトモデルではDeepSeekがMoEを採用していることで知られており、本記事の「妥当性確認」セクションに登場するKimi K3も同様だ。Kimi K3はMoonshotAIが開発したMoE構成のオープンウェイトモデルで、本記事ではコードレビュー役として使用されている（詳細は後述）。
最大の実装上の難所：ルーターを訓練可能にする
概念は単純だが、実装で最も引っかかりやすいのがルーター（ゲーティングネットワーク）の訓練だ。
ルーターは入力コンテキストベクトルを受け取り、どのエキスパートに振り向けるかを決める線形層だ。例えば6エキスパートに対してルーターが以下のような出力を返し、上位2つを選ぶ：
tensor([ 0.0418, -0.1140, 0.4254, 0.1342, 0.5106, -0.1385])
→ インデックス4と2を選択
問題は、「どのエキスパートを選ぶか」という選択操作が計算グラフに接続されない点にある。バックプロパゲーションはエキスパートの出力から入力へとさかのぼるが、ルーターが行った選択はその経路に含まれない。結果として、ランダム初期化されたルーターはトレーニングを通じてもランダムなままになってしまう。
この問題を解決するのが、ルーターの出力をソフトマックスした確率値をエキスパート出力の重みとして使う手法だ。選ばれたエキスパートの出力に、ルーターが出力したスコアを掛け合わせて合算する。こうすることで、ルーターのスコアが計算グラフの中に明示的に組み込まれ、バックプロパゲーションがルーターを通過できるようになる。
Auxiliary Loss：エキスパートの偏りを防ぐ
もう一つの重要な実装ポイントが補助損失（auxiliary loss）だ。ルーターを普通に訓練するだけでは、特定のエキスパートばかりが選ばれる「崩壊」が起きる。これを防ぐため、全エキスパートが均等に使われるように促す損失項を通常の損失に加算する。
著者はSwitch Transformers論文（2021年）のアプローチを参考にしており、実装後にMixtralのソースコードと比較したところ、auxiliary lossの計算方法に若干の違いがあった以外はほぼ同一の実装になっていたという。
実装の参照として著者が挙げている論文は以下の4本だ：
- Adaptive Mixtures of Local Experts（1991年） — MoEの名称を導入した原論文
- Sparsely-Gated MoE（2017年） — スパースゲーティングによる計算削減を提案
- GShard（2020年） — Transformerへの適用と簡略化
- Switch Transformers（2021年） — さらなる簡略化、著者が最も参考にした論文
実装の妥当性確認：ChatGPT・Claude・Kimi K3に聞く
著者はコードの妥当性を確認するため、ChatGPT、Claude、Kimi K3の3つのLLMにコードをレビューさせた。プライベートセッションで過去の会話履歴を排除した上で「このコードを見て、特にMoEの部分について教えてほしい」と質問したところ、3モデルともほぼ同じ回答を返したという：
「GPT-2の上に実装された標準的なMoEで、ロードバランシングはSwitch Transformersから採用している」
この確認作業自体、実装のセルフレビュー手法として参考になる。
結果まとめ
| モデル | パラメータ数 | アクティブパラメータ |
|---|---|---|
| 著者のMoEモデル | 446M | 220M |
| GPT-2 small（OpenAI） | 124M | 124M |
| GPT-2 medium（OpenAI） | 355M | 355M |
※編集部の考察：表中のGPT-2 mediumのパラメータ数は元記事では345Mと記載されているが、OpenAI公式およびHugging Faceモデルカードでは一般に355Mとされている。元記事の記載に沿って読む場合はその点に留意されたい。
テストセット損失はGPT-2 smallを上回り、GPT-2 mediumに肉薄。Instruction fine-tuningでも著者の過去モデル中で最良の結果を出した。ただしOpenAIのGPT-2モデルには及ばず、その理由は著者が引き続き調査中とのことだ。
詳細はExtending Raschka's GPT-2: an MoE trained from scratch on an RTX 3090を参照していただきたい。
