Tinyworldformer: Raspberry Pi Pico 2 Wでループ型Transformerを動かす

概要

Tinyworldformerは、ループ型Transformerとint8量子化を用いることで、リソースが制限されたハードウェア上で複雑な世界モデルや強化学習エージェントを動作させることが可能であることを実証しています。このシステムにより、Raspberry Pi Pico 2 Wはゲームのフレームを生成し、それをリアルタイムでプレイすることを学習できます。つまり、マイクロコントローラーが独自のシミュレーション環境内で「夢を見」て行動することが実質的に可能になります。

アーキテクチャ: ループ型Transformerとアクター

このシステムは、RP2350チップ上で動作する2つの異なるニューラルネットワークに依存しています。

• 世界モデル: 現在のフレームとアクションに基づいて、次のフレームと報酬を予測するループ型Transformerです。単一のTransformerブロックを同じ重みで複数回再利用することで、モデルは大幅なメモリ節約を実現し、マイクロコントローラーでの実行を可能にしています。ゲームの複雑さに応じて26k〜131kのパラメータを使用します。
• アクター: 生成されたフレームを観察し、アクションを選択する多層パーセプトロン（MLP）です。9k〜16kのパラメータしか使用しない軽量な設計です。

両方のネットワークはRP2350のコア1で動作し、コア0はUSB通信とディスプレイ出力を管理します。このシステムは推論中に「実際の」ゲームロジックを使用せず、完全に世界モデルの予測に依存しています。

ゲームの実装とパフォーマンス

このプロジェクトは、Breakout、Doom-lite、Flappy Bird、Pong、Pac-Manの5つのゲームをサポートしています。各ゲームは数百行のNumPyコードで実装されており、小さなトークンIDのグリッドにレンダリングされます。世界モデルの精度はゲームによって異なり、Flappy Birdのような単純なタイトルでは高い忠実度（フレーム予測精度100%）を達成し、Pac-Manのような複雑なタイトルでは（2ループで86.4%）の精度を達成しています。

モデルの安定性を維持するため、このプロジェクトでは「シングルトン」オブジェクト（ボールやプレイヤーなど）を配置し、モデルが同じオブジェクトを複数生成する幻覚を防ぐといった特定のデコードルールを採用しています。世界モデルは、エージェントのトレーニング中に実際のゲームルールと予測を比較することで検証されます。

量子化とデプロイ

Pico 2 Wの512 KBのSRAMと4 MBのフラッシュメモリに収めるため、このプロジェクトでは量子化対応トレーニング（QAT）を使用しています。世界モデル内のすべての行列はint8形式に変換されます。ファームウェアはno_std Rustで記述されており、動的なメモリ割り当てなしで決定論的なパフォーマンスを保証します。

最大のゲームであるPac-Manのパフォーマンス指標を見ると、システムは約600 KBのフラッシュと460 KBのSRAMを消費しており、ハードウェアの制限内に十分に収まっています。RP2350のFPUを使用することでノルムや残差の効率的な処理が可能になり、整数ベースの内積は特定の命令セット（SMLAD）を使用して最適化されています。

マイクロコントローラーにおけるAIへの洞察

コミュニティの議論では、将来のAIアプリケーションにおけるこのアプローチの可能性が強調されています。コメント投稿者は、最小限のローカル実行モデルに焦点を当てたこのようなプロジェクトが、組み込みAIの将来にとって重要な方向性を示していると指摘しました。

"I honestly believe that these neat projects (smallest possible, locally run LLMs) are the real future of AI on the long run." — skrellm

また、ハードウェアの制約は印象的であるものの、その根底にあるロジックは従来のゲーム開発にも統合できるだろうという意見もありました。これらのモデルをプロシージャル生成や適応型NPCの動作に使用することで、開発者は標準的なプロシージャル生成技術に見られる一般的な制限を解決できる可能性があります。