WorldTUI: ターミナルを模倣するように3億6000万パラメータのモデルをトレーニングしました

WorldTUI: ターミナルを模倣するように3億6000万パラメータのモデルをトレーニングしました

TL;DR: 3億6000万パラメータの言語モデルをファインチューニングして、ターミナルとして動作するようにしました。背後にはシェルやファイルシステムはありません。すべてのキー入力がモデルに入力され、モデルは画面がどのように変化するかを予測します。ファイルを作成したり、vimやtmuxを実行したり、snakeをプレイしたりでき、通常は数コマンド後の入力内容を記憶しています。最大の成果は、自身の誤りを修正する方法を学習させたことと、画面行を並列でデコードできるようにしたことです。ブラウザで試す (208 MB) で pythoon と入力して、これが実際のシェルではないことを確認してください。

このデモは、WebGPUと物理キーボードを備えたデスクトップブラウザ向けに設計されています。ChromeまたはEdgeで試すのが最適です。最新のスマートフォンでも動作するはずですが、非常に遅くなるでしょう。

［動画］

昨年初めてGenie 3を見たとき、アプリケーションがコードではなく学習によって作成される未来を想像し始めました。ゲームエンジンがシーンをレンダリングする代わりに、モデルがユーザーの入力から次に最も可能性の高いフレームを予測します。

当然、自分自身のワールドモデルをトレーニングしたくなりました。しかし、私には大規模なGPUクラスターや大量のビデオデータがありません。そのため、野心をスケールダウンしました。ゲームの世界ではなく、ターミナルを対象にしました。

ターミナルは良い選択肢であることがわかりました。ピクセルではなく文字のグリッドであるため、実際のシェルを実行することで、ラベル付けされたトレーニングデータを好きなだけ生成できます。事前学習済みの言語モデルは、ls や cat が何を出力すべきかについて、すでにかなり良いアイデアを持っています。

私が本当に知りたかったのは、一貫した世界を維持できるのか、それとも単に説得力のあるターミナル出力を生成できるだけなのか、ということでした。

最終的なモデルはQwen2.5-Coder-0.5Bをベースにしており、Docker内で実行される実際のプログラムから記録されたターミナルセッションでファインチューニングされています。例の動画へジャンプ

ターミナルをどのように「見る」か

標準的なLLMはプロンプトを受け取り、トークンを1つずつ末尾に追加します。これはターミナルにはあまり適していません。単一のキー入力がそれ自体で何かを行うことがあります（例: vimでのesc）、また出力が常に下部に行われるわけではありません（clearは画面を消去してカーソルを一番上に送ります）。そのため、テキストストリームを予測する代わりに、画面の変更を予測する必要があります。

私はターミナルを文字のグリッドとして扱います。1。モデルはグリッドを直接見るのではなく、その内容とそれに加えられた変更を記述するトークンのシーケンスを見ます。

セッション

すべてのセッションはスナップショットから始まります。画面サイズ、その内容、カーソルの位置です。

サイズは入力の一部であり、32×96に固定されていません。デモでは、カスタム開始画面を作成する際に、行数と列数を設定できます。

BOS SCREEN_START 32 96
ROW_0 0 'simedw@laptop:~$ '
CURSOR 0 17 1
SCREEN_END

これにより、アプリケーションの途中で開始することもできます。bashプロンプト、vimバッファ、またはsnakeゲームを与えると、モデルはそこから続行しようとします。（bashからはよりうまく機能します。）

その後、各キー入力は、実際のターミナルが送信する生のバイトとして入力されます。文字はそれ自体、Enterは\r、Escapeは0x1b、←はESC [ Dなどです。

INPUT_START 'l' INPUT_END

モデルは、どの行を書き換えるか、カーソルがどこに配置されるかを出力します。

OUTPUT_START
ROW_0 17 'l'
CURSOR 0 18 1
OUTPUT_END

言葉で言うと、行0の列17から始めてlを書き込み、その後カーソルを行0、列18に移動させます。最後の1はカーソルが表示されていることを意味します。フルスクリーンのアプリケーションではカーソルが非表示になることがあるためです。

何も最初にクリアされません。新しい文字は単に古い文字の上に描画されます。

文字のために新しいトークンは必要ありませんでした。Qwenはすでに256バイト値すべてのトークンを持っているため、それらを文字、座標、その他の小さな数値に再利用します。したがって、17のような座標は、2つではなく1つのトークンで表現されます。

時間も入力です。何も押されない場合、モデルはキーの代わりに期間を指定したWAITトークンを受け取ります。これにより、snakeは自分で動き続けることができます。

行のコピー

ターミナルウィンドウの最下部に達してEnterキーを押すと、すべての行が1つ上に移動する必要があります。モデルに新しい位置ですべての行を再度出力させることもできますが、これは1回のキー入力で約3,000トークンになります。

そこで、代わりにCOPY操作を追加しました。

ROW_0 COPY 1
ROW_1 COPY 2  
ROW_2 COPY 3
...

各操作は前の画面から行をコピーします。この例では、古い行1が新しい行0になり、行2が新しい行1になります。

スクロールは、約96トークンではなく、行あたり3トークンで済みます。

同じトリックで、画面のクリア、vimでの行の挿入と削除、manページのナビゲーションもカバーできます。行をコピーして、その一部を上書きすることもできます。

合計すると、モデルは333トークンの非常に小さな語彙を使用します。256バイトトークン、64行トークン、およびいくつかの特殊操作です。デコーディングは文法的に制約されているため、すべての予測は有効な画面更新です。

その小さな語彙は、モデルをブラウザに収める際に再び役立ちます。

トレーニングデータの生成

当初、asciinemaを含む既存のターミナル録画データセットを探しましたが、必要な正確なキーストロークから画面へのマッピングを持つものを見つけることができませんでした。

幸いなことに、自分で作成するのはそれほど困難ではありませんでした。

基本的なアイデアは、Dockerコンテナ内で実際のコマンドを実行し、各キー入力とそれに続く画面を記録することで、説得力のあるターミナルセッションを生成することでした。

合成レシピは以下の通りです。

1. ユーザー名、ホスト名、プロンプトスタイル、ターミナルサイズをランダム化する。
2. 架空のペルソナ（開発者、システム管理者など）に基づいてディレクトリとファイルをランダム化する。
3. 6〜24個の重み付けされたアクティビティ（ファイルの変更、vim、python、tmuxなど）を選択する。
4. 状態を追跡し、セッションがほとんどの場合存在するファイルを開くようにする。
5. 人間のようにタイプする。タブ補完など、1回のキー入力ごとに行う。タイプミス率は1.2%に設定した。これは私には寛大すぎる設定だ。
6. セッションの35%は途中から開始し、モデルがあらゆる状態から作業できるように強制する。

ターミナルサイズは主に32×96で、一部24×80と40×120もあり、モデルが単一サイズに過学習しないようにしました。

モデルの能力が向上するにつれてデータセットを拡張し続け、最終的に336kエピソードに到達しました。これは約20億トークンに相当します。比較として、ベースモデルはトレーニング中に約18兆トークンを学習しました。

トレーニング

開始点としてQwen2.5-Coder-0.5Bを使用しました。リアルタイムに近いパフォーマンスを得るために、小さなモデルが必要でした。コーディングモデルであるため、すでに多くのCLI使用例を見ていました。

最初は単純な教師ありファインチューニングを行いました。画面履歴とキー入力を与え、実際のターミナルが生成する画面差分を予測させます。

通常のタイピングはすぐに学習されました。文字が表示され、バックスペースで削除され、カーソルが移動します。

しかし、長いセッションでは、エラーが蓄積し始めました。ファイル名の間違いがセッションの残りの部分に波及しました。

自身の誤りから学習する

これに対処するために、モデルが理想的なデモンストレーションだけでなく、訪問した状態からも学習するDAggerから着想を得ました。2。

各トレーニング例について：

1. 記録されたセッション内のポイントを選択します。通常は、出力を生成するコマンドの近くです。
2. モデルに次の1〜4イベントを予測させます。
3. 結果の画面と実際に記録された画面を比較します。
4. 差が十分に小さい場合、次のキー入力時にモデルの画面を修復する修正を作成します。
5. 修正されたフレームとそれに続くいくつかの実際のイベントでトレーニングを続行します。

差分は、モデルの誤った画面を正しい次の画面に変換します。本質的に、モデルは自身の誤りを「修復」することを学習します。

1つの重要な制限があります。モデルは修正を推論するのに十分な情報を持っている必要があります。修正によって導入された単語は、モデル自身の予測だけでなく、以前の実際の画面に表示されていたか、ユーザーの入力に含まれている必要があります。

例えば、モデルが以前に表示されたファイル名を誤って入力した場合、それを修正するように教えることができます。しかし、モデルが見たことのないファイル名を出力した場合、それを正しく予測することは期待できません。

そのような例でトレーニングすると、実際には幻覚を助長することになります。それは本当にこれ以上必要ありません。

サンプリングされたセッションの約3分の1が有用な修正例を生成しました。これを3ラウンド繰り返し、各ラウンドで最新モデルからの誤りを収集しました。

自己修正は、より効果的なトレーニング変更の1つでした。追加のトレーニングデータとともに、int4でPresetBench（以下で定義）を24/47から38/47に改善しました。

合計で、モデルとその後の改良のトレーニングには、RTX 4090で約75 GPU時間かかりました。

テスト

3つのテストを使用します。ShellBenchは、24個の手書きbashセッションにわたる106個のコマンドを含み、正確な出力で採点されます。PresetBenchは、デモで遭遇した失敗から抽出された47個のスクリプト化されたケースを含みます。どちらのテストもトレーニングデータには含まれていません。

メモリテストには32個の世界があります。それぞれが2つのランダムな値を2つのファイルに書き込み、その後、別々のセッションでそれらを読み戻します。世界は、両方の読み戻しが正確である場合にのみ合格し、合計64回の読み戻しが行われます。名前と値はトレーニングに現れません。3

ファインチューニングはどの程度役立つか？

最大の差はメモリでした。WorldTUIは13回の読み戻し中9回を正しく行いましたが、プロンプト付きモデルでは2/13および3/13でした。

それでも信頼性の高いファイルシステムにはほど遠いです。馴染みのあるコマンドはかなりうまく処理しますが、それらを新しい方法で組み合わせるとすぐに問題が発生する可能性があります。例えば、echo green >> colorsの後にcat colorsを実行すると、追加された行が表示されないことがよくあります。

高速化

しばらくの間、自己回帰モデルを完全に避けていました。完全な再描画は約3,000トークンであり、それらを1つずつ生成するのは、ターミナルのように感じられるべきものとしては絶望的に思えました。

しかし、それらは私が試した他の選択肢よりもはるかに優れたパフォーマンスを発揮したため、私はそれに屈しました。1トークンずつでは、タイピングは〜38 ms/キーで良好でしたが、画面全体を再描画するもの（vim test.txt、man、clear）は、私のRTX 4090で最大12秒かかりました。これはまさに恐れていた通りです。マニュアルが必要かどうかを再考するのに十分な時間でした。そして、これを私のMacでほぼリアルタイムで実行したかったのです。

並列処理される行

私が見落としていたのは、フレームのほとんどの行は互いに依存しておらず、逐次生成する必要がないということです。各行に対して1つのバッチ要素を実行します。各要素は同じ以前の履歴（KVキャッシュを含む）を受け取りますが、自身の更新のみを生成します。

これにより、デコードコストがすべての行の合計（最大3,000ステップ）から、最も長い行の長さ（約100ステップ）に変わります。

その後、カーソルは、行の更新セット全体を見ることができる最後のステップになります。

これにはトレーニング中に対応するアテンションマスクが必要でしたが、速度の向上はそれだけの価値がありました。大規模フレームのp95レイテンシは6.7秒から0.87秒に低下しました。

大規模フレームのp95レイテンシ（RTX 4090）↓

逐次

6.7 s

並列行

0.87 s

通信できない行

これは本当すぎて信じられないように聞こえるかもしれませんが、そうでした。問題は、行が独立しているため、調整できず、vimが行0ではなく行3で開いたり、複数の行が同時にsnakeの餌を追加したりすることにつながったことです。

解決策は投機的デコーディングでした。

通常、投機的デコーディングでは、小さなモデルがトークンをドラフトし、大きなモデルがそれらを検証します。因果的アテンションのおかげで、複数のトークンを単一のフォワードパスで検証できます。

ここでは両方に同じモデルを使用します。並列行がドラフトを生成し、モデルがそれを逐次検証します。

検証者は、最初の不一致までのトークンを保持します。そこから、独自の修正を生成し、残りのドラフトを再提案します。

私のM4 Mac上のブラウザでは、1文字の入力は約60msかかります。man testのような大きな更新は、逐次処理では2.5秒、並列ドラフトと検証では0.9秒かかります。5

小型化

約4億9400万パラメータを持つQwen2.5-Coder-0.5Bから始めました。Webデモを強く望んでいたため、縮小する必要がありました。151,936トークンの語彙の埋め込みテーブルは1億3600万パラメータです。私の333トークン6しか必要ないため、そのほとんどを削除でき、モデルは4億9400万から約3億6000万パラメータ（990MB → 716MB）になります。

次にint4に量子化し、208 MBになりました。

ブラウザデモは、QwenデコーダーのカスタムJavaScript実装を使用しており、WGSLで書かれたコンピュートシェーダーがWebGPU経由で実行されます。シェーダーは推論中に4ビット重みを展開します。アクティベーションとアテンションキャッシュは32ビット浮動小数点数を使用します。すべての推論はタブ内でローカルに実行されます。

分割蒸留

より少ないレイヤーに分割蒸留することで、どのレイヤーがメモリに依存しているかについて学びました。教師は以前のチェックポイントで、ShellBenchで78/106を記録しました。レイヤー7を再トレーニングせずに削除しても、次のトークン精度はほとんど変わりませんでした（99.3%から98.9%に低下）。しかし、正確なメモリ読み戻しは64/64から0/64に低下しました。これらのテストは実際の履歴を入力したため、失敗はモデル自身の出力でエラーが蓄積したことによるものではありませんでした。

私の最初の分割蒸留された学生は、すべての偶数番号のレイヤーを保持しましたが、レイヤー22の代わりに最終レイヤー23を使用しました。まだもっともらしいターミナル出力を生成しましたが、メモリ読み戻しでは31/64の正しい値を返しただけでした。常に書き込んだ2つの値のうちの1つを出力しましたが、約半分の確率で間違った方を選択しました。これは偶然であり、4億8000万トレーニングトークン後も回復しませんでした。

その学生はレイヤー1、7、15、22を削除しました。それらの4つのレイヤーを保持した16レイヤーの学生は、すべての32個のメモリテスト世界に合格し、64/64の正確な読み戻しを行いました。約138MBで、4ビットでもすべての32個の世界に合格しました。しかし、他の部分でコストがかかりました。教師の78に対して、65個のShellBenchコマンドを正しく実行しました。したがって、学生は問題なく見えても、静かに記憶を失っている可能性があります。

結局、分割蒸留は諦めました。

何ができるか

メモリ

これが最も興味があった部分です。説得力のあるディレクトリリストを生成するのは簡単です。私が望んでいたのは、ファイルを作成し、編集し、離れて、後で戻ってきたときにファイルが変更されていないことを見つけることでした。

tmuxを起動し、vimでファイルを開き、変更して保存し、別のtmuxウィンドウでcatを使って更新された内容を読み込んだ最初の時、私はフランケンシュタイン博士のように感じました。それは生きている！

［動画］

その他のアプリケーション

これがbashプロンプト以外でも機能することは私にとって非常に重要でした。そのため、vim、python、manページ、snake、さらにはtmuxでもトレーニングしました。

vimでは、escでのモード切り替えを処理し、保存、終了、または挿入モードに戻ることができます。

［動画］

おそらく世界で最も計算負荷の高いsnakeです。

tmuxが私を本当に驚かせたのは、複数のペインを作成し、それぞれでさまざまなことを行い、それらを切り替えることができることです。完璧ではありませんが、これは小さなモデルです。

［動画］

環境変数や変数展開さえも管理します。

［動画］

Python

ベースモデルはコーディングモデルなので、Python REPLをどのように処理するか興味がありました。まともな試みです。小さな追加は通常正しく、リスト内包表記はリストを返しますが、多くの場合内容が間違っています。タイポにも多少寛容です。pythoonはPythonを起動しますが、ythonは起動しません。

［動画］

何がうまくいかないか

Dockerコンテナにインターネットアクセスを許可しなかったため、lynxのようなツールはかなり退屈です。Pythonは非常に単純な操作を処理できますが、それ以上のものはGPT-3時代の幻覚を連想させるような回答を生成します。

時々、ネストされたディレクトリを作成すると、迷子になり、突然別のパスにいることがあります。一貫性を向上させる余地は明らかにたくさんあります。

私のテストでは、デモはChromeよりもFirefoxで約20倍遅く動作しました。どちらもWebGPUを使用しています。原因を特定できませんでした。

このようなプロジェクトでは、常にまず自分で解決策を探し、その後で他者がどのようにアプローチしたかを見てきました。

2022年には、人々はすでにChatGPTにLinuxターミナルとして動作するようにプロンプトを与えていました。より最近では、GameNGenが拡散モデルを使用してDoomをシミュレートし、NeuralOSがデスクトップGUIに取り組みました。

WorldTUIは、元のChatGPTターミナルトリックに最も近いですが、キー入力と画面更新、vimやtmuxのようなフルスクリーンアプリケーションを含む、それらに特化してトレーニングされている点が異なります。

まとめ

これを構築するのはとても楽しかったです。役に立つか？あまり。しかし、ブラウザタブで実行される3億6000万パラメータのモデルとしては、予想をはるかに超える性能でした。

最も興味深いのは、基盤となるルールを実装せずに、ある程度一貫した世界を維持できていることです。

明らかにかなり弱いワールドモデルであり、このアプローチが、より良い一貫性と、画面更新の異なる生成方法（おそらく自己回帰デコーディングの代わりに拡散モデルを使用するなど）でどこまで進むかを見てみたいです。

次に試したい実験は、それぞれ独自の制御とルールを持つ、数千のランダムに生成されたアプリケーションでトレーニングすることです。アイデアは、十分なコンテキストの手がかりがあれば、モデルがこれまで遭遇したことのないアプリケーションの動作を予測できるかどうかを確認することです。

ここで試してみてください。208 MBのダウンロードで、すべてローカルで実行され、Chromeで最もよく機能します。壊してみてください、そして見つけたものを教えてください。

まもなくウェイトを公開し、ここにダウンロードリンクを追加します。

---

1. モノクロです。色やテキストスタイルはなく、文字のみです。tmuxの罫線文字を含むUnicode文字はエンコード前にASCIIにマッピングされるため、これらの罫線はマルチバイト文字トークンを必要としません。↩
2. 私のMIDI投稿のスケジュールサンプリングとそれほど似ていません。↩
3. これらは短いギャップのテストで、0〜1個のコマンドが介在します。ShellBenchは18個のコマンドを横断する永続性もテストします。↩
4. プロンプト付きモデルはそれぞれ3つの作業例を得ました。個々の画面更新ではなく、コマンド全体の出力を予測するため、タスクはまったく同じではありません。↩
5. 投機的デコーディングは、M4よりも4090で大幅な高速化をもたらします。ボトルネックが計算とメモリ帯域幅の間でシフトしているのではないかと疑っていますが、これを適切に調査していません。↩
6. GPUの要求を満たすために384にパディングされています。↩

---