跳到正文
原文
Hacker News· alphax314·· 4 小时前精选AI 评分58

WorldTUI:作者微调 360M 参数模型模拟可交互终端

WorldTUI: I trained a 360M model to pretend to be a terminal

AI 导读

WorldTUI 将 Qwen2.5-Coder-0.5B 微调为约 360M 参数的终端世界模型,接收原始按键并预测屏幕行与光标变化,背后没有真实 shell 或文件系统。

推荐理由

文章拆解了真实终端数据构造、模型错误修正和屏幕行并行解码的实现路径,并给出 ShellBench、记忆测试与浏览器延迟数据。

正文 · AI 翻译

TL;DR: 我微调了一个 360M 参数语言模型,让它充当终端。它背后没有 shell 或文件系统:每次按键都会被输入,模型则预测屏幕将如何变化。它可以创建文件、运行 vim 和 tmux、玩贪吃蛇,而且通常仍记得你在几条命令前写下的内容。最大的提升来自让它学会修复自身错误,以及并行解码屏幕的各行。在浏览器中试用 (208 MB),然后输入 pythoon,就会发现它并不是真正的 shell。

这个演示专为支持 WebGPU 且配有物理键盘的桌面浏览器设计。Chrome 或 Edge 是最佳尝试环境。它也应该能在现代手机上运行,但速度会慢得令人痛苦。

去年,我第一次看到 Genie 3 时,便开始想象这样一个未来:应用程序不再由代码编写,而是通过学习得到。场景不再由游戏引擎渲染,而是由模型根据用户输入预测下一帧最可能的内容。

Naturally,我自然想训练自己的世界模型。但我没有大型 GPU 集群,也没有大量视频数据,所以只能降低目标:不再打造游戏世界,而是打造一个终端。

事实证明,终端非常合适。它只是一个由字符而非像素组成的网格;通过运行真实的 shell,我想生成多少完美标注的训练数据就能生成多少。预训练语言模型也已经大致知道 ls 或 cat 应该输出什么。

我真正想知道的是,它能否维持一个前后一致的世界,还是只会生成看似可信的终端输出。

最终模型基于 Qwen2.5-Coder-0.5B,使用从 Docker 内运行的真实程序中记录下来的终端会话进行了微调。 跳转到示例视频

它如何理解终端

标准 LLM 会接收一段提示词,然后逐个将词元追加到末尾。这与终端的工作方式并不匹配。单次按键本身就可能产生作用,例如 vim 中的 esc;输出也不一定出现在底部:clear 会清空屏幕,并将光标送回顶部。因此,我需要预测的不是文本流,而是屏幕的变化。

我把终端视为一个字符网格,每个单元格放一个字符。1 模型看不到该网格本身,而是会看到一串描述其内容及针对它所做更改的词元。

一次会话

每个会话都从一个快照开始:屏幕尺寸、屏幕内容以及光标位置。

屏幕尺寸是输入的一部分,而不是固定为 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 token,而不是按键,这就是贪吃蛇自己保持移动的方式。

复制行

当你到达终端窗口底部并按下回车时,每一行都需要向上移动一行。我本可以让模型在新位置重新输出每一行,但单次按键就要约 3,000 个 tokens。

于是我改用了一个 COPY 操作:

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

每个操作都会从上一屏复制一行。在这个例子中,旧的第 1 行变成第 0 行,第 2 行变成第 1 行,依此类推。

现在每行滚动只需三个 tokens,而不是约 96 个。

同样的技巧也适用于清屏、在 vim 中插入和删除行,以及浏览 man 页面。你还可以复制一行然后覆盖其部分内容。

总体而言,模型使用了一个仅 333 个 token 的小型词表:256 个字节 token、64 个行 token,以及少量特殊操作。解码过程受语法约束,因此每个预测都是有效的屏幕更新。

这个小型词表之后还会派上用场,当我将模型压缩到浏览器中时。

生成训练数据

我最初寻找了现有的终端录制数据集,包括 asciinema,但没能找到满足我所需的精确按键-屏幕对齐的数据集。

幸运的是,自己创建并不太难。

基本思路是:在 Docker 容器中运行真实命令来生成逼真的终端会话,并记录每个按键以及之后的屏幕状态。

合成数据的配方是:

  1. 随机化用户名、主机名、提示符样式、终端尺寸。
  2. 根据虚构的角色(开发者、系统管理员等)随机化目录和文件。
  3. 选取 6–24 个带权重分布的活动(修改文件、vim、python、tmux 等)。
  4. 跟踪状态,使会话大多数时候打开的都是存在的文件。
  5. 像真人一样打字:每次一个按键,并使用 Tab 补全等。我将打错字率设为 1.2%,这对我来说已经很宽松了。
  6. 35% 的会话从中间状态开始,迫使模型学习从任意状态继续工作。

终端尺寸主要为 32×96,也有一些 24×80 和 40×120,避免模型过拟合到单一尺寸。

随着模型能力提升,我持续扩充数据集,最终达到 336k 个 episodes,约合 2B tokens。作为对比,基座模型在训练中见过约 18T tokens。

训练

我以 Qwen2.5-Coder-0.5B 作为起点。我需要一个小型模型以接近实时性能。作为一个编码模型,它也已经见过大量的 CLI 使用场景。

我一开始进行了直接的监督微调:给定屏幕历史和一次按键,预测真实终端产生的屏幕变化量。

普通打字很快就学会了:字母出现,退格删除,光标移动。

然而在更长的会话中,错误开始累积:文件名中的一个小错会贯穿会话的其余部分。

从自身错误中学习

为了解决这个问题,我借鉴了 DAgger 的思路:模型从它所访问到的状态中学习,而不仅仅是从理想演示中学习。2

对每个训练样本:

  1. 在录制的会话中选取一个点,通常接近产生输出的命令。
  2. 让模型预测接下来的 1-4 个事件。
  3. 将得到的屏幕与真实录制的屏幕进行比较。
  4. 如果差异足够小,就构造一个修正,在下一次按键时修复模型的屏幕。
  5. 在修正后的画面以及随后几个真实事件上继续训练。

差分会把模型错误的屏幕带到正确的下一屏。本质上,它学会“修复”自己的错误。

有一项重要限制:模型必须掌握足够的信息,才能推断出修正方案。修正所引入的任何词语都必须曾出现在更早的真实屏幕上或用户的输入中,而不能只出现在模型自己的预测中。

例如,如果模型把一个此前出现过的文件名拼错了,我就能教它修正。但如果它列出了一个从未见过的文件名,就不能指望它能写对。

用这类样本进行训练实际上会奖励幻觉,而我实在不需要更多幻觉了。

采集到的会话中约有三分之一产生了有用的修正样本。我以三轮重复了这一过程,每一轮都从最新模型中收集错误。

自我纠正是效果较为显著的训练改动之一。再配合一些额外的训练数据,它将 PresetBench(定义见下文)在 int4 下的得分从 24/47 提升到了 38/47。

总的来说,训练模型及其后续改进版本,在 RTX 4090 上大约耗费了 75 个 GPU 小时。

测试

我使用三项测试。ShellBench包含 24 段手写 bash 会话中的 106 条命令,按输出完全一致来评分。PresetBench包含 47 个脚本化测试用例,取自我在演示中遇到的故障。两项测试都没有出现在训练数据中。

记忆测试包含 32 个世界。每个世界会将两个随机值分别写入两个文件,然后在不同会话中读回。只有两次回读都完全准确,该世界才算通过,共计进行 64 次回读。名称和数值从未出现在训练中。3

微调能带来多大帮助?

模型4 ShellBench /106 文件回读 /13
Qwen2.5-Coder-0.5B,提示调用 52 2
Qwen3-0.6B,提示调用 60 3
WorldTUI 80 9

差异最大的是记忆能力:WorldTUI 正确完成了 9/13 次回读,而两个提示调用的模型分别只有 2/13 和 3/13。

它距离成为可靠的文件系统仍然很远。它能较好地处理熟悉的命令,但将这些命令以新方式组合起来,很快就会出问题。例如,在 echo green >> colors 之后运行 cat colors 通常不会显示追加的那一行。

让它变快

有一段时间,我完全避开了自回归模型。完整重绘一次大约需要 3,000 个 token,对于一个本应让人感觉像终端的程序来说,逐个生成这些 token 似乎毫无希望。

但它们的表现实在远好于我尝试过的其他方案,所以我还是妥协了。每次生成一个 token 时,~38 ms 按键一次的输入体验还不错,但任何需要重绘整个屏幕的操作(vim test.txt、man、clear)在我的 RTX 4090 上最多要花 12 s,完全和我担心的一样。这留给我充足的时间重新考虑是否真的需要那本手册。而且,我希望它能在我的 Mac 上以接近实时的速度运行。

并行处理各行

我忽略了一点:一帧中的大多数行互不依赖,因此没有理由逐行生成。我让每一行对应一个批处理元素。每个元素都会获得相同的先前历史(包括 KV 缓存),但只生成自己的更新。

这样一来,解码成本就从所有行的步数之和——最高 3,000 步——降到了最长行的长度,大约为 100 步。

随后,光标会在末尾成为一个单独的步骤,并且能够看到全部行更新。

这要求训练时采用相匹配的注意力掩码,但速度提升值得。大帧的 p95 延迟从 6.7 s 降到了 0.87 s。

大帧 p95 延迟(RTX 4090)↓

顺序处理

6.7 s

并行处理各行

0.87 s

彼此无法通信的各行

这听起来也许好得不太真实,而事实正是如此。 问题在于,由于各行彼此独立,无法协调,这导致 vim 从第 3 行而不是第 0 行开始,或多个行同时为蛇添加食物。

解决方案是推测解码。

通常,推测解码会使用一个小模型生成词元草稿,再由更大的模型进行验证。得益于因果注意力,单次前向传播即可验证多个词元。

这里我对两者使用同一个模型:并行各行生成一份草稿,再由模型按顺序进行验证。

验证器会保留到首次出现分歧为止的词元。随后,它生成自己的修正内容,并重新提交剩余的草稿。

在我的 M4 Mac 上,于浏览器中输入一个字符大约需要 60ms。较大的更新(如 man test)顺序处理需要 2.5s,而采用并行草拟与验证只需 0.9s。5

将其缩小

我最初使用 Qwen2.5-Coder-0.5B,它大约有 494M 个参数。我非常希望做一个网页演示,因此需要缩小模型。其 151,936 个词元的词表所对应的嵌入表包含 136M 个参数。我只需要自己的 333 个词元,6所以几乎可以将嵌入表全部删去,使模型从 494M 个参数降至约 360M 个参数(990MB → 716MB)。

随后我将其量化至 int4,体积降至 208 MB。

浏览器演示采用自定义 JavaScript 实现的 Qwen 解码器,计算着色器使用 WGSL 编写,并通过 WebGPU 运行。这些着色器会在推理期间解包 4-bit 权重。激活值和注意力缓存使用 32-bit 浮点数。所有推理均在浏览器标签页内本地运行。

蒸馏

将其蒸馏为更少层数,让我对记忆能力依赖于哪些层有了一些认识。教师模型是一个较早的检查点,在 ShellBench 上的得分为 78/106。仅移除第 7 层而不重新训练,几乎没有改变下一词元准确率:99.3% 降至 98.9%。但精确记忆回读从 64/64 降至 0/64。这些测试输入的是真实历史,因此失败并非由模型基于自身输出不断累积错误所致。

我的第一个蒸馏学生模型保留了每个偶数编号的层,但使用最后一层 23,而非 22。它仍能写出看似合理的终端输出,但在 64 次记忆回读中只有 31 次返回了正确值。它总会打印自己写下的两个值中的一个,但约有一半时候会选错。这属于碰运气,而且经过 480M 个训练词元后仍未恢复。

那个学生模型去掉了第 1、7、15 和 22 层。一个保留了这四层的 16 层学生模型通过了全部 32 个记忆测试场景,精确回读达到 64/64。在 4-bit 下,它仍通过了全部 32 个场景,模型大小约为 138MB。不过在其他方面也付出了代价:它在 ShellBench 上答对了 65 个命令,而教师模型答对了 78 个。因此,学生模型可能表面表现正常,实际上却已悄然丧失记忆能力。

最后,我放弃了蒸馏它。

它能做什么

记忆

这是我最感兴趣的部分。 生成一份令人信服的目录列表很容易。我想要的是创建一个文件、编辑它、离开一段时间,再回来时发现文件仍未改变。

我第一次启动 tmux、在 vim 中打开文件、修改并保存,然后在另一个 tmux 窗口中用 cat 读取更新后的内容时,感觉自己就像弗兰肯斯坦博士:它活了!

其他应用

对它而言,能够在 bash 提示符之外正常工作非常重要,因此我在 vim、python、man 页面、snake 甚至 tmux 上进行了训练。

在 vim 中,它能处理 esc 对应的模式切换,并让我保存、退出或返回插入模式。

可能是世界上计算开销最大的贪吃蛇游戏。

真正让我惊呆的是 tmux:我可以创建多个窗格,在每个窗格做不同的事情,并在它们之间切换。它并不完美,但模型很小。

它甚至能管理环境变量并进行变量展开。

Python

由于基础模型是一个编程模型,我很好奇它会如何处理 Python REPL。它的尝试还算不错:小幅修改通常是正确的,列表推导式也会返回一个列表,尽管内容往往不对。它对拼写错误也有一点容错:pythoon 会启动 Python,但 ython 不会。

无法奏效的部分

我没有给 Docker 容器开放互联网访问权限,因此 lynx 之类的工具比较无聊。 Python 能处理非常简单的运算,但难度稍高一点,生成的结果就会让我想起 GPT-3 时代的幻觉。

有时,我创建嵌套目录时,它会迷失方向,突然间我就身处另一条路径。 显然,一致性仍有很大的改进空间。

在我的测试中,这个演示在 Firefox 中的运行速度比 Chrome 慢约 20×。两者都使用 WebGPU。我还没能确定原因。

像往常一样,面对这类项目,我先尝试自己解决,之后才去看其他人采用了什么方法。

早在 2022 年,人们就已经在提示 ChatGPT 充当 Linux 终端。最近,GameNGen 使用扩散模拟了 Doom,而 NeuralOS 则尝试处理桌面 GUI。

WorldTUI 更接近最初的 ChatGPT 终端把戏,只不过它是专门针对按键和屏幕更新进行训练的,其中包括 vim 和 tmux 这类全屏应用。

总结

构建这个项目时我玩得很开心。它有用吗?没什么用。但对于一个运行在浏览器标签页中的 360M 参数模型来说,它的表现远超我的预期。

最让我感兴趣的是,它竟然能在不实现任何底层规则的情况下,维持一个相对一致的世界。

它显然是一个相当薄弱的世界模型。我很希望看看,如果一致性得到改善,并采用不同的屏幕更新生成方式——或许使用扩散而非自回归解码——这种方法究竟能走多远。

接下来我想尝试的一个实验,是用数千个随机生成的应用进行训练,每个应用都有自己的控件和规则。我想看看,模型能否在获得足够上下文线索的情况下,预测一个此前从未接触过的应用的行为。

在这里试用。下载大小为 208 MB,完全在本地运行,并且在 Chrome 中效果最佳。尽管试着把它弄坏,并告诉我你的发现。

我很快会发布权重,并在这里添加下载链接。


  1. 它是单色的:只有字符,没有颜色或文本样式。Unicode 字符(包括 tmux 的方框绘制边框)会在编码前映射为 ASCII,因此这些边框不需要多字节字符词元。 ↩

  2. 与我那篇midi 文章中的计划采样并没有太大不同。 ↩

  3. 这些是短间隔测试,中间有 0–1 条命令。ShellBench 还会测试状态能否在 18 条命令之间保持。 ↩

  4. 两个基于提示的模型各自获得了三个完整解答示例。它们预测的是完整的命令输出,而非单次屏幕更新,因此任务并不完全相同。 ↩

  5. 推测解码在 4090 上带来的加速远大于在我的 M4 上。我怀疑瓶颈会在算力和内存带宽之间转移,但还没有对此进行认真调查。 ↩

  6. 填充至 384,以迎合 GPU。 ↩

来源:Hacker News · simedw.com