LOG ENTRY / post
TinyLM LabからTensorFlow.jsを卒業した——手書き逆伝播とWebGPU、そして日本語対話コーパスの公開
TinyLM LabからTensorFlow.jsを卒業した
午前中に公開したブラウザで学習するMoE言語モデルの続きです。あのあと勢いで「tensorflow.js非依存にできます?」とAIに聞いたら、その日のうちに依存ゼロ版とWebGPU版ができてしまったので、経緯ごとまとめて公開します。おまけに、学習教材として作った日本語対話コーパスも配布します。
公開物
- TinyLM Lab Zero — 依存ゼロ・WebGPU対応の単一HTML(約76KB)。CDN読み込みなし、機内モードでも動きます。
- TinyLM JP Dialogue Corpus 100k — 約10万字の日本語対話コーパス一式(生成・検証スクリプト、評価文30問付き)。
- 学習済みチェックポイント(ナノ・1,000 step) — 上のコーパスで学習したモデル。Zeroに読み込めばすぐ生成を試せます。
tf.jsを審判にして、tf.jsを卒業する
午前の版はtensorflow.jsに行列演算と自動微分を任せていました。依存ゼロ版では、RMSNorm・RoPE・causal attention・SwiGLU・MoEルーター・Adamの逆伝播をすべて手書きし、TypedArrayだけで学習エンジンを組んでいます。メモリはgradient checkpointing方式(各層の入力だけ保存し、逆伝播時に層内を再計算)で、活性化メモリを1層分に抑えました。
面白いのは検証の構図です。手書き逆伝播は、符号一つ・転置一つの間違いで「静かに学習が劣化する」世界なので、三段構えで検証しました。
- 数値微分との突き合わせ — Float64モードで中心差分と比較。3構成・261点すべて絶対誤差1e-9未満。
- tf.js版を審判にした相互検証 — 同一重み・同一入力でlogits最大差2.7e-7、CE lossは小数6桁一致。二つの独立実装が同じ数学を計算している強い証拠です。
- 実学習 — 150 stepでloss 5.18→0.13、生成はコーパスの文体を再現。
卒業したいフレームワーク自身を審判に使って卒業する、という構図が気に入っています。
副産物として速度も上がりました。同一環境のCPU比較でtf.js-cpuの約4.5倍(1,040 vs 230 tok/s)。フレームワークのオーバーヘッドが消えたこと、そしてMoEが「選ばれた専門家だけを計算する」本物のスパース実行になったことが効いています(tf.js版は勾配の都合で全専門家を密計算していました)。保存形式は旧版と完全互換で、相互に読み込めます。
WebGPUを手書きする
「GPUも使えます?」の答えがこちらです。tf.jsには戻らず、WGSLコンピュートシェーダを21カーネル手書きして、学習ステップ全体——順伝播・逆伝播・Adam更新——をGPU常駐で回します。CPUへの読み出しはlossと専門家稼働率だけ。
開発環境にGPUがないという問題は、ソフトウェアVulkan(llvmpipe)+Node用WebGPUバインディングで突破しました。速度は測れませんが正しさは検証できる環境です。検証済みのCPUエンジンを審判にして、全36勾配テンソルの一致(最悪相対差1.4e-6)、GPU単独学習の収束、GPU→CPUの状態往復まで確認しています。
実機の速度保証ができない代わりに、初回学習時にその端末上でCPUエンジンと全勾配テンソルを突き合わせる自己診断を組み込みました。合格した場合だけGPUが有効になります。「開発側で保証できないことは、実行側で機械検査する」という設計です。
外部レビューで診断が本物になった
このプロジェクトは複数のAIを使い分けて進めています。実装と検証をClaudeが行い、別のAI(GPT-5.6)に成果物をレビューさせる、を2往復。指摘は的確でした。
- 自己診断が「loss比較だけ」で、逆伝播が壊れていても素通しになっていた → 全勾配テンソル比較へ昇格
- Adam再開の不整合(モーメント未保存なのにバイアス補正だけ過去から続く)→「重みは継続、optimizerは再始動」に統一
- ほぼゼロ勾配のテンソルで相対誤差判定が過敏 → 絶対誤差フロアを追加
- 例外時にボタンとWake Lockが取り残される → finallyで後始末
一度目の修正が半端で、二度目のレビューに取り残しを突かれる場面もありました。実装者と監査者を分ける普通のソフトウェア工学が、AI同士でも普通に機能します。
日本語対話コーパスも作った
10万字級の教材が欲しくなり、対話コーパスを一式作りました。中身は挨拶・聞き返し・計画・四則演算・感情サポート・安全な断り方など36カテゴリの型を意味対応固定のテンプレートで生成したオリジナル部分(CC0、100,006字・1,225会話)と、公開データLLM-jp OASST2 33k Japaneseから人手選別した補助28会話(Apache-2.0)です。ライセンスはファイル単位で分離し、SHA-256・選定マニフェスト・検証スクリプト・非リーク保証付きの評価文30問まで同梱しています。
ナノ構成(総226k・有効158kパラメタ)で1,000 step学習した結果が配布中のチェックポイントです。ここで面白いのは、二つのAIによる独立評価が異なる深さの結論を出したことです。
Claudeの評価(30問・ルーブリック8項目でおよそ9/16)は「応答の型は獲得、意味の対応付けは未達」でした。助手の文体は全問一貫し、曖昧な依頼への聞き返しや感情サポートの型は出る。しかし計画テンプレの中身が混線し、算数は形式だけ覚えて計算が壊れる。
一方、GPT-5.6 Solによる独立再現(HTMLと教材のSHA-256を固定し、同一seedで学習から再実行)は、もう一段深いところを掘りました。訓練データと完全一致するprefixでの生成16回のうち意味合格は2回、訓練回答の完全再現はゼロ、完全一致の算数でも「4+4→18」、安全拒否は全試行失敗。つまりloss 0.67は「会話の意味対応」どころか「暗記」すら反映していない。次文字予測lossの大部分は日本語断片の局所的な継続で下がるため、応答境界での正しい選択とlossが乖離する——という解釈まで含めて、「低い言語モデルlossだけで会話能力を判定してはいけない」という教訓が定量化されました。
このコーパスのREADMEは「158kパラメタで狙えるのは応答の型であり、知識ではない」と予告していましたが、独立診断はさらに「現構成では意図の分岐が広すぎる」ことを示しています。提案された次の一手(教材を6〜8意図へ絞る、学習窓を応答境界へ整列させる、利用者部分のlossをマスクする)は次回の実験テーマです。評価文とルーブリックは同梱してあるので、読者の追試も歓迎します。
独立テストがバグも見つけた
GPT-5.6 Solのレポートは性能評価に留まらず、実バグを一つ発見しました。モデル初期化後に教材を編集・差し替えても内部のトークン列が無効化されず、画面に表示された教材ではなく以前の教材で学習し続ける、という状態管理の穴です。公開版では修正済みで、教材の編集・読込・空白正規化の変更で必ず再トークン化し、モデルの語彙にない文字が混ざった場合は警告を出すようにしました。
実装者(Claude)が2往復のレビューを通過した後でも、視点の違う監査者が実行環境ごと変えて再現テストをすると、まだ出てくるものは出てくる。多重チェックの価値はこの「まだ出てくる」にあります。
締め
朝の「テストして」から始まった一日が、tf.js卒業・WebGPU手書き・コーパス公開・独立監査まで来ました。ブラウザとテキストエディタしかない環境でも、言語モデルの中身を全部さわれる時代です。手元の端末で自己診断が通ったかどうか、よければ教えてください。
なお、この小モデルは応答の型の教材です。医療・法律・金融など結果が重大な用途には使わないでください。