マイコンで対話型のAIを動かそうとすると、たいていはクラウドのAPIを叩くことになる。手元のチップは音声を送って結果を受け取るだけの端末で、ネットが切れれば何も返らない。ESP32に外付けのPSRAMを足して言語モデルを載せる例もあるが、標準的な開発ボードが積むSRAMは512KBしかなく、そのままでは大きな重みを置く場所がない。
ahmedbarakat207氏が公開した「ESP-LLM」は、この512KBの内側だけで文章を生成する推論エンジンだ。C++で書いたベアメタルの実装で、外付けメモリーもネット接続も使わない。推論中に使うSRAMは約81KBにとどまる。ESP32 Dev Kit V1で毎秒5トークン前後を生成すると本人は説明する。
容量を切り詰めた最大の工夫は、モデルを丸ごとメモリーに置くことをやめた点にある。使ったのはMixture-of-Expertsという構造で、層ごとに16個の小さな専門家を用意し、入力に応じてそのうち1つだけを選ぶ。選ばれた専門家の重みだけをフラッシュから読み出すので、専門家を何個並べても1トークンあたりの処理時間は変わらない。重みは4ビットまで量子化し、2つを1バイトに詰め込んである。
注意の計算にはMulti-Query Attentionを充てた。過去のトークンを覚えておくキャッシュはSRAMを最も食う部分で、ここのヘッドを1つに絞ることで、6層48トークン分のキャッシュを36864バイトずつに抑えた。文脈が48トークンで埋まると、先頭の例示プロンプトだけを残して古い会話から押し出し、位置情報を計算し直す。
より非力なESP8266でも動く版を用意した。起動と同時にWi-Fiを切って内部SRAMを約15KB取り戻し、行列積の途中で細かく処理を中断してウォッチドッグに応答を返す。SDKが約1.5秒でリセットをかけてしまうため、長い計算を一息に回せない。
モデルの規模はおよそ520万パラメータで、語彙は1024トークンしかない。本人も会話が成立するぎりぎりの大きさだと認めており、簡単な計算は言語モデルに解かせず、入力を先に数式パーサーへ通して正確に答えさせている。書き込みはPythonのスクリプト1本で、重みの書き出しからビルドと転送までを済ませられる。手元のマイコンだけで完結する対話を試したい人には、出発点になる作りだろう。ライセンスは設定されていない。
