8ドル（約1300円）ほどのマイコンで、どこまで大きな言語モデルを動かせるか。GitHubでslvDevと名乗る開発者が、2890万パラメーターのモデルをESP32-S3の上で動かし、毎秒約9トークンで文章を生成させることに成功した。ネットワーク接続は使わず、すべてチップ内部で完結する。生成した文字は、配線でつないだ小型OLEDに1語ずつ表示される。
これまでマイコン上で動かされてきた言語モデルは26万パラメーター程度で、今回はその約100倍にあたる。壁になるのはメモリーだ。ESP32-S3が持つ高速なSRAMは512KBしかない。一方、2890万パラメーターのモデルは4ビットに量子化しても約15MBある。モデル全体を高速メモリーに置く必要がある限り、小さなモデルしか選べない。
開発者が使ったのは、モデル全体を高速メモリーに置くのをやめるという発想だ。言語モデルのパラメーターの大半は「埋め込みテーブル」と呼ばれる巨大な表で、計算に使うのではなく参照するだけの部分になる。そこで2500万行のテーブルは低速なフラッシュメモリーに置いたまま、メモリーマップして必要な行だけを読み出す。1トークンあたり読むのはおよそ6行、約450バイトで済む。実際に計算する中核部分は約56万パラメーターしかなく、これだけを高速メモリーに置けばよい。
この手法はGoogleがGemma 3nとGemma 4で採用した「Per-Layer Embeddings」で、スマートフォンやGPUではなくマイコンのメモリー構成に持ち込んだ形になる。開発者によれば、これほど小さなチップで試した例はこれまで見当たらないという。
メモリーの使い分けは3層だ。SRAMには毎トークン触れる中核部分、PSRAMには起動時に展開する出力層とKVキャッシュ、フラッシュには2500万パラメーターのテーブルを置く。推論処理はC言語でゼロから書き、PyTorchの参照実装と演算単位で突き合わせ、書き込み前にホスト側で誤差1e-5まで一致することを確認した。
速度は最初0.57トークン毎秒だった。出力層を起動時に4ビットから8ビットへ展開してPSRAMに置き、トークンごとの展開処理をなくす。活性化を8ビットに量子化し、精度への影響がほぼないことをホスト側で検証する。出力層を2つのCPUコアに分散させる。こうした積み重ねで9.7トークン毎秒まで引き上げた。ただしコア分散の効果は1.36倍にとどまった。出力層は1トークンあたり2.4MBを読み出すため、PSRAMの帯域が上限になって約40ミリ秒の下限が生じていたからだ。
モデルの学習にはTinyStoriesという短い物語のデータセットを使っており、簡単な物語なら破綻せずに書けるが、質問に答えたり指示に従ったりはできない。開発者は、重要なのはこのモデルが何を言えるかではなく、大きなモデルを小さなチップに載せる構造の方だと述べている。学習と量子化のコード、C言語の推論処理、実機での測定結果はすべて公開した。自身のパラメーター数の計算に誤りがあって初期の数値が過大になっていた経緯も、修正の履歴ごとリポジトリに残したと明かしている。