本地 AI 世界模型 第二部分 - 将图像变成可操作角色的深度神经网络，支持生成过程中切换文本提示

上次我在这个 subReddit 上分享我的工作时，反响几乎让我落泪，因为一个小小的 demo 竟然让这么多人开始讨论这个话题。过去两个月里，我一直在训练一个新模型，但这次加入了真正的文本引导。先聊聊旧模型。普通视频模型太大了，不适合在消费级硬件上实时运行。你可以生成一段静态视频片段，甚至速度很快，但在本地实时生成视频仍然尚未解决。虽然出现了很多世界模型 demo，但它们要么需要在巨大的数据中心 GPU 上运行，要么只是像 WAN 或 LTX 这样的热门模型，被蒸馏成自回归方式工作（顺便说一下，这在本地也不是实时的）。上面那段视频是在 RTX 5090 上运行的，利用率仅为 30%。这个 1B 模型在 RTX 5090 上的最高帧率为 50-60，但我用软件强行将其限流到 12fps。根据一些测试，这意味着该模型也可以在其他 RTX 40、30 系列显卡上运行（我很快会尝试）。我有一台 MacBook，还没有把模型移植到 MLX 上，但我做了一个基准测试，模型在我的 M5 MacBook 上运行时可达 30 fps。 关于架构 该模型是一个纯 transformer，使用块因果掩码（block causal mask），这意味着在训练过程中过去的帧看不到未来的帧，就像 LLM 一样进行学习。我训练它时使用的另一个重要方法是"diffusion forcing"，这意味着与普通的视频模型训练不同，在训练时我们对每一帧独立加噪，因此模型学会了适应嘈杂的历史帧和各种情况。该模型有 28 个 block、20 个 head，约 960M 参数。在推理时，我们每帧运行 2-5 步扩散，一旦一帧完成去噪，就将其加入 KV 缓存。这类似于 LLM 的解码步骤。与 LLM 最大的区别在于，我们不会保留所有的 KV 上下文（即所有历史上下文），而是使用滑动窗口，只保留过去 80 帧的上下文。 上一个模型是 MMDiT，这意味着文本没有交叉注意力（cross attention）。这在世界模型中是个大问题，因为历史帧的 KV 和文本的 KV 在 softmax 中直接竞争，你永远无法可靠地实现文本引导。上一个模型也没有在文本-视频数据上训练过，所以这个问题本身也不成立。当前的模型恢复了文本交叉注意力，并且我进行了大量的文本-视频预训练。它能实时接受我给出的键盘操作（一个 adaLN 额外项有助于用 WASD 操作引导生成），最有趣的部分是文本提示切换。"给沙漠加一个池塘""穿红色连帽衫""把环境改成冰雪世界"。由于模型经过了大量的文本-视频对齐训练，它现在能够真正跟随提示词了。我知道仍有很多局限性，比如一致性和质量提升，但我真诚地希望在今年年底前，我可以发布一个任何拥有 RTX GPU 或新款 MacBook 的人都能尝试的东西。我特意选择了这个初始图像，因为在我上次在这个 subReddit 的帖子中也用了同一张。 P.S. 在我上次的帖子中，很多人问起我的情况和资金来源。我住在 Bangalore，大学最后一年（处于部分退学状态），由一个学生孵化器资助。我完全独立工作，没有团队，也没有真实公司之类的。上面的模型是在 8x H100 SXM 上训练了大约 3-4 周。我制作的每一个模型都将明确面向本地推理，永远不会面向数据中心。