ブラウザで学習するMoE言語モデルを作った

言語モデルを「使う」記事は世に溢れているが、「その場で学習させる」となると急に少なくなる。当然で、学習には普通GPUサーバとPythonの環境構築が要る。でもtf.jsを使えば、HTMLファイル一枚で、ブラウザの中で、スマホでも、小さなTransformerを本当に学習できる。作ったのがこれ。

TinyLM Lab を開く → /works/tinylm-lab/

テキストを貼るかファイルを読み込ませて「初期化」→「学習開始」を押すと、lossが下がっていくのが見え、途中で止めて文章を生成させられる。全部端末内で完結し、通信は一切しない。

発端: 動いていなかった初版

もともとはGPT-2構造そのままの素朴な文字レベルモデルをHTML一枚に詰めたものがあった。これをClaudeに「適当な長文でテストして」と投げたところ、開口一番「このファイル、ブラウザで開いてもボタンが一切効かないはずです」と言われた。

原因はジェネレータ関数の中でforEachのアロー関数越しにyieldしていたこと。アロー関数はジェネレータではないのでyieldは不正で、"use strict"下では構文エラー、つまりスクリプト全体が読み込み時点で死ぬ。書いた本人（AI）も、持ち込んだ本人（私）も気づいていなかった。人間もAIも「読めば正しそうに見えるコード」には弱い。テストは実行してこそ。

全部載せ: 現代アーキテクチャの縮小模型

直すだけでは面白くないので、アーキテクチャを現代化した。LayerNorm→RMSNorm、学習位置埋め込み→RoPE、GELU→SwiGLU。ここまでやるとLlama系ミニチュアと言える構成になる。

そしてMoE（Mixture of Experts）。各層のFFNを「ルーター＋N個の専門家FFN」に置き換え、トークンごとにルーターが選んだtop-k個の専門家だけを使う。総パラメタ数は専門家の数だけ膨らむのに、1トークンが実際に通る「有効パラメタ」は一定——現代の大規模モデルの核心的な仕組みが、画面上の二つの数字（計/有効）としてヘッダーに常時表示される。学習中は各専門家の稼働率がバーで動き、Switch Transformer式の負荷分散補助損失がルーターの偏りを矯正していく様子が目で見える。ここが一番作りたかった部分。

規模は5段階のプリセットにした。ナノ（約18万パラメタ・スマホで10分）からラージ（数億パラメタ・サーバ級GPUで数日、ブラウザでは動作保証外の実験領域）まで。正直に書いておくと、ブラウザで現実的に学習できるのはスモール（千数百万）あたりが上限で、それ以上は「構造は本物と同じだが、賢さは伴わない」領域になる。Chinchilla則的に言えば、パラメタに見合うデータを食わせる速度がブラウザにはない。これは賢さの縮小模型ではなく、仕組みの縮小模型である。

技術メモ2題

MoEの勾配とtop-kマスク。 tf.jsでtopkのインデックスからマスクを作ると、逆伝播が勾配未定義の演算に突っ込んで死ぬ。通常モードではインデックスを一度CPUに読み出して定数マスクを作り、勾配経路をルーター確率側だけに限定した。ただこのCPU往復は毎ステップのGPU同期点になる。そこで高速モードではtf.customGradでマスク構築全体を包み、「K番目に大きい確率を閾値にgreaterEqual」でGPU上のままマスク化して勾配はゼロを返す。あわせてlossの読み出し自体を10ステップに1回へ間引き、GPUコマンドキューを切らさない。両モードのlogitsは完全一致することを確認済み。

改行が出続ける事件。 学習後の生成で改行ばかり出る、という症状に遭遇した。診断してみると原因候補が二つあり、どちらも「バグ」ではなかった。一つは学習発散でNaN化した重み——サンプラーが同一トークンに固着し、語彙表は辞書順ソートなので先頭は大抵改行になる。もう一つはコーパス自体に空行が多い場合——モデルは段落区切りの連続改行を忠実に学習して忠実に再現する。対策として、サンプラーをCPU実装に置き換えてNaNを明示エラー化し、top-pサンプリングと空白正規化オプションを足した。「おかしな出力」の原因が壊れた重みなのか正しい学習なのかを切り分ける、というのは大きなモデルの評価でも本質的に同じ話だと思う。

開発方法について

今回の開発で一番効いたのは、HTMLの中のモデル数学部分をマーカーで区切ってDOM非依存にし、出荷するコードそのものを抽出してNode.jsでテストする方式だった。パラメタ数の計算式が実変数の総要素数と一致するか、保存→復元でlogitsが完全一致するか、学習ループでテンソルがリークしないか。この過程でテンソルリーク（120文字の生成で2万6千個残留）を含む3件のバグが出荷前に捕まった。テスト用に書き直したコードをテストしても、出荷物の保証にはならない。

道具は嘘をつかない。手を抜けば、必ずどこかで音が変わる。