WorldTUI: я обучил модель с 360M параметров притворяться терминалом

Как модель видит терминал

Стандартные LLM берут подсказку и приписывают токены в конец по одному за раз. Для терминала это не очень подходит. Одно нажатие клавиши может самостоятельно выполнить действие, как esc в vim, а вывод не всегда идёт в конец экрана: clear стирает всё и отправляет курсор в верх. Вместо предсказания потока текста нужно предсказывать изменения экрана.

Терминал рассматривается как сетка символов, по одному в каждой ячейке.¹ Модель не видит сетку напрямую, но последовательность токенов, которые описывают её содержимое и произошедшие изменения.

Сессия

Каждая сессия начинается со снимка: размер экрана, его содержимое и позиция курсора.

Размер — часть входных данных, а не фиксированные 32×96. В демо можно установить количество строк и столбцов при создании собственного стартового экрана.

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

Это позволяет начать работу с середины приложения. Дай ей bash-приглашение, буфер vim или игру в змейку — модель попробует продолжить. (С 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 занимает один токен, а не два.

Время — это тоже входные данные. Если ничего не нажато, модель получает токен WAIT с продолжительностью вместо клавиши — так змейка сама продолжает двигаться.

Копирование строк

Когда достигаешь дна окна терминала и нажимаешь enter, каждая строка должна сдвинуться вверх на одну позицию. Можно просто попросить модель вывести все строки снова в новых позициях, но это около 3000 токенов за одно нажатие.

Вместо этого добавлена операция COPY:

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

Каждая операция копирует строку с предыдущего экрана. В этом примере старая строка 1 становится строкой 0, строка 2 становится строкой 1 и так далее.

Скролл теперь стоит три токена на строку вместо примерно 96.

Этот же трюк подходит для очистки экрана, вставки и удаления строк в vim и навигации по справочным страницам. Можно также скопировать строку и затем перезаписать часть её.

Всего модель использует крошечный словарь из 333 токенов: 256 токенов байтов, 64 токена строк и несколько специальных операций. Декодирование имеет ограничения грамматикой, поэтому каждое предсказание — валидное обновление экрана.

Этот крошечный словарь пригодится позже, когда потребуется втиснуть модель в браузер.

Создание тренировочных данных

Изначально искали готовые наборы записей терминала, включая asciinema, но не нашли с нужным мне точным выравниванием от нажатия клавиши к изменению экрана.

К счастью, создать собственный набор оказалось несложно.

Базовая идея — генерировать реалистичные сессии терминала, запуская настоящие команды внутри Docker-контейнеров и записывая каждое нажатие клавиши и экран после каждого из них.

Синтетический рецепт был таким:

1. Рандомизировать имена пользователей, имена хостов, стили подсказок, размеры терминала.
2. Рандомизировать директории и файлы на основе вымышленной персоны (разработчик, системный администратор и т.д.).
3. Выбрать 6–24 взвешенных активностей (модификация файлов, vim, python, tmux и т.д.).
4. Отслеживать состояние, чтобы сессии в основном открывали существующие файлы.
5. Печатать как человек: по одному нажатию за раз, с автодополнением и т.д. Установил частоту опечаток на 1,2%, что щедро по отношению ко мне.
6. 35% сессий начинаются с середины, заставляя модель учиться работать с любого состояния.

Размер терминала в основном 32×96, с некоторыми 24×80 и 40×120, так что модель не переобучается на одном размере.

Набор постоянно расширялся по мере того, как модель становилась способнее, и в итоге остановился на 336k эпизодов, что соответствует примерно 2B токенов. Для сравнения: базовая модель видела около 18T токенов во время обучения.

Обучение

Использовал Qwen2.5-Coder-0.5B в качестве отправной точки. Нужна была небольшая модель, чтобы получить близкую к реальному времени производительность. Будучи моделью кодирования, она уже видела много использования CLI.

Начал с прямолинейного supervised fine-tuning: даю историю экрана и нажатие клавиши, предсказываю изменение экрана, произведённое реальным терминалом.

Обычный ввод текста усваивается быстро: буква появляется, backspace её удаляет, курсор движется.

В более длинных сессиях ошибки начали накапливаться: ошибка в имени файла распространялась через остаток сессии.

Обучение на собственных ошибках

Чтобы решить это, взял идею из DAgger, где модель учится на состояниях, которые она посещает, а не только на идеальных демонстрациях.²

Для каждого тренировочного примера:

1. Выбрать точку в записанной сессии, обычно близко к команде, которая произвела вывод.
2. Позволить модели предсказать следующие 1–4 события.
3. Сравнить полученный экран с реальным записанным экраном.
4. Если разница достаточно мала, создать коррекцию, которая восстанавливает экран модели при следующем нажатии клавиши.
5. Продолжить обучение на исправленном кадре и нескольких последующих реальных событиях.

Дельта переводит неправильный экран модели в правильный следующий экран. По сути, модель учится «исправлять» собственные ошибки.

Есть одно важное ограничение: модель должна иметь достаточно информации, чтобы вывести коррекцию. Любые слова, введённые коррекцией, должны были появиться на более раннем реальном экране или при вводе пользователя, а не просто в собственных предсказаниях модели.

Например, если модель неправильно написала имя файла, которое появилось ранее, можно научить её это исправить. Но если она перечислит имя файла, которое никогда не видела, нельзя ожидать, что она его правильно напишет.

Обучение на таких примерах на самом деле награждало бы галлюцинации, которые мне совсем не нужны.

Примерно треть отобранных сессий дали полезные примеры коррекции. Повторил это в три раунда, каждый раз собирая ошибки из последней модели.

Самокоррекция была одним из более эффективных изменений обучения. Вместе с дополнительными тренировочными данными это улучшило PresetBench (определён ниже) с 24/47 до 38/47 на int4.

Всего обучение модели и её последующие улучшения заняло примерно 75 GPU-часов на RTX 4090.

Тестирование

Использую три теста. ShellBench — 106 команд по 24 написанным вручную bash-сессиям, оценивается по точному выводу. PresetBench — 47 скриптовых случаев из сбоев, с которыми столкнулся в демо. Ни один тест не появляется в тренировочных данных.

Тест памяти — 32 мира. Каждый записывает два случайных значения в два файла, затем читает их назад в отдельных сессиях. Мир проходит только если оба прочтения точны, что дает 64 прочтения в сумме. Имена и значения никогда не появляются в обучении.³

Насколько fine-tuning помогает?

Самая большая разница в памяти: WorldTUI правильно получила 9/13 прочтений против 2/13 и 3/13 для prompted-моделей.

Это всё ещё далеко от надёжной файловой системы. Справляется хорошо со знакомыми командами, но их комбинация новым способом может быстро что-то сломать. Например, после echo green >> colors запуск cat colors часто не показывает добавленную строку.

Оптимизация для скорости

Какое-то время избегал авторегрессивных моделей вообще. Полная перерисовка — это около 3000 токенов, и генерировать их по одному казалось безнадёжным для чего-то, что должно ощущаться как терминал.

Но они показывали так много лучшие результаты, чем альтернативы, которые пробовал, что сдался. При генерации по одному токену за раз печать ощущалась нормально при ~38 мс на клавишу, но что-то, что перерисовывает весь экран (vim test.txt, man, clear), занимало до 12 сек на RTX 4090 — ровно как опасался. Было время пересмотреть, нужна ли справка. И хотелось, чтобы это работало на Mac в близком к реальному времени.

Строки параллельно

Что упустил — большинство строк в кадре не зависят друг от друга, поэтому нет причины генерировать их одну за другой. Запускаю один batch-элемент на строку. Каждый получает ту же предыдущую историю (включая KV cache), но генерирует только своё собственное обновление.

Это меняет стоимость декодирования с суммы всех строк, до 3000 шагов, на длину самой длинной строки, примерно 100 шагов.

Курсор затем становится собственным шагом в конце, который может видеть весь набор обновлений строк.

Это потребовало соответствующей маски внимания при обучении, но улучшение скорости того стоило. p95 latency больших кадров упала с 6.7 с до 0.87 с.

Строки, которые не могут разговаривать друг с другом

Звучит слишком хорошо, чтобы быть правдой — и так оно и было.

Проблема в том, что так как строки независимы, они не могут координироваться, что привело к открытию vim на строке 3 вместо строки 0 или к нескольким строкам одновременно добавляющим еду для змейки.

Решение — спекулятивное декодирование.

Обычно спекулятивное декодирование использует маленькую модель для набросания токенов и большую модель для их проверки. Благодаря causal attention несколько токенов можно проверить в одном forward pass.

Здесь использую ту же модель для обоих: параллельные строки создают черновик, который модель затем проверяет последовательно.

Verifier оставляет токены до первого несогласия. Оттуда генерирует свою собственную коррекцию и переделывает оставшийся черновик.

В браузере на M4 Mac печать одного символа занимает около 60ms. Большее обновление вроде man test занимает 2.5s последовательно или 0.9s с параллельной черновой и проверкой.⁵

Сжатие модели

Начал с Qwen2.5-Coder-0.5B, у которого было примерно 494M параметров. Очень хотел веб-демо, поэтому нужно было её сжать. Таблица вложений для её 151,936-токенного словаря — 136M параметров. Нужна была только 333 моих токена,⁶ поэтому мог исключить почти всё остальное, что уменьшило модель с 494M до примерно 360M параметров (990MB → 716MB).

Затем квантовал до int4, что привело к 208 МБ.

Веб-демо использует пользовательскую JavaScript-реализацию Qwen-декодера с compute-шейдерами, написанными на WGSL и запущенными через WebGPU. Шейдеры распаковывают 4-битные веса во время инференса. Активации и cache внимания используют 32-битные флоаты. Весь инференс работает локально в закладке.

Дистилляция

Дистилляция в меньше слоёв кое-что научила о том, на каких слоях зависит память. Учитель был более ранней версией checkpoint, набравшей 78/106 на ShellBench. Удаление только слоя 7 без переобучения едва изменило next-token accuracy: 99.3% упала до 98.9%. Но точные прочтения памяти упали с 64/64 на 0/64. Эти тесты подавали в истину, поэтому отказ не был вызван ошибками, которые накапливались на выводе модели.

Первый дистиллированный ученик сохранил каждый чётный слой, но использовал финальный слой, 23, вместо 22. Всё ещё писал правдоподобный вывод терминала, но вернул правильное значение только в 31/64 прочтениях памяти. Всегда печатал один из двух значений, которые записал, но выбирал неправильное примерно половину времени. Это шанс, и не восстановился после 480M токенов обучения.

Этот ученик выпустил слои 1, 7, 15 и 22. 16-слойный ученик, который держал эти четыре, прошёл все 32 мира теста памяти с 64/64 точными прочтениями. Всё ещё прошёл все 32 мира при 4-битном, примерно 138MB. Была цена в другом месте, хотя: получил 65 ShellBench команд против 78 у учителя. Таким образом, ученик может выглядеть хорошо и тихо потерять свою память.

В конце концов отказался от дистилляции.

Что умеет делать

Память

Это была часть, в которой был больше всего заинтересован. Создать убедительный листинг директории просто. Хотелось создать файл, отредактировать его, поискать что-нибудь и вернуться, найдя его неизменным.

Первый раз, когда запустил tmux, открыл файл в vim, изменил его, сохранил и затем в другом tmux-окне прочитал обновленное содержимое с cat, почувствовал себя как доктор Франкенштейн: он жив!

Другие приложения

Было очень важно, чтобы это работало за пределами bash-приглашения, поэтому обучал на vim, python, справочных страницах, snake и даже tmux.

В vim обрабатывает переключение режима на esc и позволяет сохранять, выходить или возвращаться в режим вставки.

Что действительно поразило — tmux: можно создавать несколько панелей, делать разные вещи в каждой и прыгать между ними. Не идеально, но это маленькая модель.

Даже управляет переменными окружения и расширением переменных.

Python

Так как базовая модель — модель кодирования, интересовало, как она справится с Python REPL. Пытается прилично: небольшие дополнения обычно правильны, и list comprehensions возвращают список, хотя часто с неправильным содержимым. Также немного прощает опечатки: pythoon запускает Python, но ython не запускает.

Что не работает

Не давал Docker-контейнерам доступ в интернет, поэтому инструменты вроде lynx довольно скучны. Python может обрабатывать очень простые операции, но что-нибудь посложнее выдаёт ответы, которые напоминают мне галлюцинации эпохи GPT-3.

Иногда при создании вложенных директорий теряется и внезапно оказываюсь в другом пути.

Очевидно, есть много места для улучшения консистентности.

При тестировании демо работала примерно в 20 раз медленнее в Firefox, чем в Chrome. Использует WebGPU в обоих. Не смог установить причину.

Связанные работы

Как всегда с такими проектами, сначала попытался решить это самостоятельно, а потом только посмотрел, как другие к этому подходили.

Ещё в 2022 году люди уже подсказывали ChatGPT притворяться Linux-терминалом. Недавно GameNGen симулировал Doom, используя диффузию, и NeuralOS взялся за GUI рабочего стола.

WorldTUI ближе к оригинальному трюку с ChatGPT-терминалом, кроме того, что обучена конкретно на нажатиях клавиш и обновлениях экрана, включая полноэкранные приложения вроде vim и tmux.

Резюме

Огромное удовольствие это создавать. Полезна ли она? Не особо. Но для 360M-параметровой модели, работающей в закладке браузера, намного лучше, чем ожидал.

Больше всего интересует, что она управляет поддерживать относительно консистентный мир, не реализуя никаких правил внизу.

Очевидно, это довольно слабая модель мира, и хотелось бы увидеть, как далеко этот подход мог бы зайти с лучшей консистентностью и другим способом генерации обновлений экрана, возможно, используя диффузию вместо авторегрессивного декодирования.

Один эксперимент, который хочу попробовать дальше — обучение на тысячах случайно сгенерированных приложений, каждое со своими собственными элементами управления и правилами. Идея в том, чтобы увидеть, может ли модель научиться предсказывать поведение приложения, которое никогда не встречала, при достаточно контекстных подсказках.

Попробуй здесь. Это загрузка 208 МБ, работает полностью локально и лучше всего в Chrome. Попробуй сломать, и дай знать, что найдёшь.