带世界状态寄存器的流式多智能体自回归扩散模型

用持久的世界状态寄存器让多智能体视频世界模型保持跨视角一致性

前置知识

世界状态寄存器 (World State Registers, WSR)

一组 K 个可学习 token（本文 K=256），在生成过程中与帧 token 交织排列为 $[x_0, r_0, x_1, r_1, \ldots, x_{F-1}, r_{F-1}]$。每个寄存器 $r_i$ 在第 $i$ 个上下文帧生成后被'提交'，承担两个职责：跨智能体、跨 rollout 步骤持久化地保存共享世界信息（全局布局、物体配置、智能体状态），并在新观测到来时被增量更新。形式化为 $r_i = G_\theta(r_{i-1}, x_{i-W+1}, \ldots, x_i, a_i)$，并条件化下一帧 $p_\theta(x_{i+1} \mid x_{i-W+1}, \ldots, x_i, a_{i+1}, r_i)$。它们在因果 Transformer 中与帧 token 共同参与注意力，因此一次前向传播既去噪下一帧又刷新世界状态。推理时寄存器直接进 KV 缓存，训练时由 agent status / bird's-eye view / scene text 三个辅助预测头监督，推理时这些头被丢弃。

它是整篇论文的核心创新——把'世界状态'从隐含在帧历史里的副产品，提升为显式、可持久化、可监督、可验证的一等表示。理解它才能理解本文如何让多智能体的观察在共享世界状态下保持一致。

自回归/流式视频扩散 (Autoregressive / Streaming Video Diffusion)

基于 chunk 的因果视频扩散生成范式。把过去若干帧的潜变量作为条件上下文放入滑动窗口 KV 缓存，模型以联合时空注意力去噪下一个视频 chunk，每个潜在帧独立采样噪声水平，配合 block 级因果掩码。代表方法有 Diffusion Forcing、Self-Forcing、Causal Forcing。它能做流式长时段生成但只携带有限的观测历史，不显式建模世界状态，需要每步重新从像素推断世界信息。

本文的基线和改进对象就是这类 pipeline，WSR 正是嫁接在其上；理解滑动窗口、因果掩码、KV 缓存才能理解寄存器如何插入因果序列并被增量更新。

Self-Forcing / DMD 自蒸馏

解决 teacher-forcing 训练与自回归推理之间的 train-test mismatch。训练时让 student 在自己生成的帧和提交的寄存器上 rollout，再用 Distribution Matching Distillation (DMD) 损失对齐：$\mathcal{L}_{sf}=\tfrac{1}{2}\mathbb{E}\|\hat{x}_0 - \mathrm{sg}(\hat{x}_0-(s_{fake}(\hat{x}_t,c,t)-s_{real}(\hat{x}_t,c,t)))\|_2^2$，其中 $s_{real}$ 是冻结的 Stage-1 双向 teacher 评分网络，$s_{fake}$ 是在 student rollout 上训练的可学习 critic，$\mathrm{sg}$ 为 stop-gradient。

Stage-3 是本文让长程 rollout 稳定的关键，因为帧漂移和寄存器漂移会被同时暴露并校正。

Mixture-of-Transformers (MoT)

一种稀疏多模态架构。不同角色/模态的 token 使用各自的投影、FFN、归一化参数（分支），但通过共享的自注意力在交织序列上联合交互。本文中世界状态寄存器走'state 分支'，内容帧走'visual 分支'。由于分支形状匹配，每个 token 仍只激活一份参数，单 token FLOPs 几乎不变（除路由开销），但总参数量增加。本文用预训练的视觉分支权重来初始化状态分支。

它是解决'帧生成目标 vs 寄存器状态目标竞争'的关键架构改进，理解它才能看懂第 4.4 节的架构消融——MoT 在寄存器承载丰富语义时才变得必要。

Flow Matching 条件流匹配

训练扩散模型去预测速度场 $v_\theta(x_t, c, t)$，损失为 $\mathcal{L}_{flow}=\mathbb{E}\,\|v_\theta(x_t,c,t)-(\epsilon-x_0)\|_2^2$。它是本文 Stage-1/2 帧生成的主目标，与寄存器损失 $\mathcal{L}_{reg}$ 组合成总损失 $\mathcal{L}_{S2}=\mathcal{L}_{flow}+\lambda_{state}\mathcal{L}_{reg}$。

论文所有阶段的核心帧级监督信号，理解它才能理解各阶段损失如何组合、以及它与寄存器监督的权重配比。

研究动机

多智能体交互世界模型有本质要求：每个智能体只能看到同一底层 3D 世界的一个 2D 投影，但跨智能体的观察必须与共享世界状态相容；且世界持续演化，无论某观察者是否看到。然而现有 chunk-based 自回归视频扩散（Diffusion Forcing、Self-Forcing、Solaris）只能把过去若干帧潜变量搬进大小为 W 的滑动窗口 KV 缓存，每个生成步都先把 chunk 解码到像素、再从像素"重新推断"世界信息。这带来两个问题：它携带的是有限观测历史，而非能解释"已观察+未观察"变化的显式世界状态；多视角下每个智能体观察与共享状态间存在 mismatch，状态被混入视觉 token 且偏向近期。Solaris 双人 Minecraft 数据量化了痛点：Building 任务最强基线 VLM 仅 9.4%、Consistency 仅 57.8%、Building 类 FID 高达 83.4——缺乏显式世界状态时，模型很难维持"两个玩家看到同一个世界"的逻辑。

本文的目标是提出一个流式多智能体视频扩散模型 WorldWeaver (W2)，在生成过程内部显式建模世界状态，而非把它当作每帧视频生成的副产品。具体目标有三层：第一，设计一组既持久（跨智能体、跨 rollout 步骤）又可动态更新（新观测到来时刷新）的世界状态寄存器，承载共享世界信息；第二，用一组互补的监督信号来"落地"寄存器语义，分别覆盖个体动力学（agent status 的位置/速度/朝向）、全局几何（bird's-eye view 俯视图）、跨模态语义（scene text 文本描述），使寄存器编码的世界状态可验证；第三，用 Mixture-of-Transformers 架构把状态建模与帧生成分到不同权重分支，缓解二者目标竞争。最终目标是在 Solaris 双人 Minecraft 测试集上，既提升视觉保真度（FID）又提升逻辑一致性（VLM accuracy），尤其是状态敏感类别（Building/Consistency/Grounding）。

与已有工作不同的是，与已有工作相比，本文的独特切入在于"把记忆放在生成器内部、但用显式监督把它训练成隐式记忆"。显式记忆方法（Worldmem、MosaicMem、3D/4D 场景表示）依赖可查询的外部缓冲或 3D 表征，受可扩展性和动态物体状态更新限制；隐式记忆方法（past-frame conditioning、KV 缓存、recurrent rollout）把历史塞进生成模型，但记忆与视觉 token 纠缠且偏向近期观测，约束很弱。Solaris 虽也重视多智能体一致性，但只是把多玩家帧拼成更长序列并用双向注意力，没有把"世界状态"作为独立的持久通道。本文在两者之间开辟新路径：寄存器是生成器内部的隐式记忆（推理时直接进 KV 缓存，不引入外部可查询结构），却由 agent 坐标、俯视图、文本三股显式监督锚定，让世界状态获得超越帧历史的直接 grounding。此外，半监督设定（1K 标注 + 大量无标注）验证了该路径能向真实世界扩展——少量标注锚定语义，大量无标注视频继续改善生成。

核心方法

整体直觉：与其每步从像素反复推断世界，不如在模型内部维护一组"状态 token"，让所有智能体共同读写、随时间增量更新，并用监督告诉它该存什么。技术分三层。第一层是世界状态寄存器：K=256 个可学习 token 与帧交织成 $[x_0, r_0, x_1, r_1, \ldots]$，使 rollout 在状态层面也因果——提交的 $r_i$ 先于并条件化下一帧 $x_{i+1}$；结合窗口 W=6，每个寄存器查询只关注 $[x_{i-W+1},\ldots,x_i]$ 和上一寄存器 $r_{i-1}$。第二层是监督：训练时三个临时预测头从寄存器分别解码 agent status、bird's-eye view、scene text，推理时丢弃这些头、不增加 rollout 成本。第三层是 MoT 架构与三阶段训练：状态 token 走独立参数分支、帧走视觉分支；先用双向 teacher 学跨玩家结构，再用因果 student 加寄存器与状态监督，最后 Self-Forcing 让 student 在自己的 rollout 上训练，把帧漂移与寄存器漂移一起暴露并校正。

核心创新点与已有方法的本质区别在于：把"世界状态"做成一条独立、持久、可监督、跨智能体共享的 token 通道，与帧生成并列存在于同一次前向传播中。已有隐式记忆（KV 缓存、past-frame conditioning）的记忆与视觉 token 纠缠且偏向近期；已有显式记忆（3D/4D 场景、检索缓冲）是模型外部的可查询结构，受扩展性和动态更新限制。本文的寄存器两者都不是——它在生成器内部（推理时直接进 KV 缓存），却由 agent 坐标、俯视图、文本三股显式监督"锚定"成可验证的世界状态。这一设计带来三个本质变化：其一，状态层因果性——交织排列使状态更新与帧预测交替进行，把 Diffusion Forcing 的 block 因果思想从"帧"推广到"状态"；其二，监督可落地——寄存器编码什么不再完全交给像素损失决定；其三，架构解耦——MoT 给寄存器单独的参数分支，承载更丰富语义时不再与帧生成抢参数。

方法步骤详情

三阶段训练。Stage-1 双向训练：从 Solaris 双向 checkpoint 初始化并加玩家轴，用双向+跨玩家注意力联合去噪玩家潜序列作 teacher。Stage-2 因果训练：先无寄存器训练，再加入 WSR+MoT+三个状态头；每潜帧独立采样噪声、block 因果掩码、W=6 滑窗；寄存器损失 $\mathcal{L}_{reg}=\sum_{i,m}\lambda_m d_m$ 中 agent stats 用 MSE、BEV 用与冻结 DINOv2 特征的余弦相似、scene text 用交叉熵，总损失 $\mathcal{L}_{S2}=\mathcal{L}_{flow}+\mathcal{L}_{reg}$。Stage-3 Self-Forcing：冻结的 $s_{real}$ 与可学习 $s_{fake}$ 都从 Stage-1 初始化；先无梯度 rollout（early-exit），再对生成上下文做梯度 pass 算 DMD 损失与寄存器损失，后期冻结 visual 分支、只更新 register。推理每步从 KV 缓存+寄存器+动作去噪下一帧并刷新寄存器。

技术新颖性

技术新颖性体现在四个交叉点。其一，状态因果性：把寄存器与帧交织成 $[x_0, r_0, x_1, r_1, \ldots]$，使状态更新与帧预测在同一因果序列交替——把 Diffusion Forcing 的 block 因果思想从"帧"推广到"状态"；$r_i$ 只看局部窗口加 $r_{i-1}$，既控计算量又保留持久性。其二，监督空间设计：首次系统探索"用什么落寄存器语义"，三股信号分别覆盖个体动力学（agent 坐标）、全局几何（俯视图）、跨模态语义（文本），消融证明互补、组合最优。其三，架构与目标解耦：MoT 让寄存器走独立分支，关键发现是 MoT 无监督时不一定更好（93.8 vs Dense 95.5），但加 scene text 后 Dense 掉到 91.3、MoT 升到 103.2——承载语义时必须解耦。其四，把 Self-Forcing 同时用于帧和状态：Stage-3 让两种漂移一起被暴露并校正，配合分阶段学习率（1500 步后冻结视觉分支，mean VLM 63.8→68.1）。整套设计把"隐式记忆+显式监督"这条中间路线具象化为可工程化的 pipeline。

**Figure 2**: Pipeline overview. (A) A standard streaming autoregressive diffusion model... (B) Our model streams with world state registers... (C) Causal attention mask during Stage-2 training. (D) Auxiliary decoders ground the registers.

**Figure 5**: Visualization of the data collection and annotation streams used in our multi-agent Minecraft setting.

实验结果

核心发现在 Table 1：WorldWeaver 把聚合 WorldScore 从 Solaris 的 81.0 提到 105.1。对状态最敏感类别提升最大——Grounding VLM 81.3→93.8、Building VLM 9.4→28.1、Consistency VLM 57.8→76.6；FID 也全面下降，说明持久寄存器不仅提升视觉保真，更维持跨玩家逻辑一致世界。Table 2 监督消融：仅加寄存器就把 WorldScore 从 81.0 拉到 93.8，证明寄存器本身有用；单信号里 scene text 最强(103.2)、BEV 次之(102.4)、agent stats 单用略降到 88.1；三者组合最优(105.1)。Table 3 架构消融：MoT 在 scene text 监督下把 WorldScore 从 Dense 的 91.3 拉到 103.2，证明语义越丰富越需解耦。Table 4 半监督：固定 1K BEV 标注，无标注 0→5K→10K 时 WorldScore 63.2→82.3→90.3，少量标注锚定语义后无标注视频持续改善生成。

**Table 1**: Comparison with multi-agent world model baselines. Category-level VLM accuracy, FID, and the aggregate world score.

**Table 2**: Supervision ablation for world state registers. We compare registers without explicit targets and registers with different signals.

**Table 3**: Dense and MoT backbone ablation. We compare a dense transformer with a Mixture of Transformers under no register supervision and scene text supervision.

**Table 4**: Semi-supervised training. Each row varies the number of data clips without bird's-eye view annotation (unlabeled data).

**Table 5**: Configuration of the three world state register supervision signals. Each signal has its own training-only prediction head, target, distance metric, and loss weight.

**Table 6**: Hyperparameter setup for the three stage training curriculum used by the primary combined world state model.

**Figure 3**: Qualitative visualization. We show generated multi-agent rollouts with grounded world states: agent coordinates, bird's-eye-view maps, and scene text.

**Figure 4**: VLM accuracy over Stage 3 self-forcing.

查看结构化数据

任务	指标	本文	基线	提升
双人 Minecraft 多智能体世界模型（聚合指标）	WorldScore↑（VLM/FID 之比×100，K=5 类均值）	105.1	Solaris*: 81.0（Frame concat: 49.1）	+24.1（较 Solaris），相对提升约 30%
Building（双人对方块结构的协同建造，状态敏感）	VLM accuracy↑	28.1%	Solaris*: 9.4%	+18.7 个百分点，相对翻 3 倍
Consistency（跨玩家观察一致性）	VLM accuracy↑	76.6%	Solaris*: 57.8%	+18.8 个百分点
Grounding（世界关系/物体相对位置）	VLM accuracy↑	93.8%	Solaris*: 81.3%	+12.5 个百分点
Movement（个体运动一致性）	FID↓	34.0	Solaris*: 43.3	-9.3（视觉保真度提升）
寄存器监督消融（Table 2）	WorldScore↑	+All（三监督组合）105.1	Registers only 93.8 / +scene text 103.2 / +BEV 102.4 / +agent stats 88.1	组合最优，证明三类监督互补
MoT vs Dense 架构消融（scene text 监督下，Table 3）	WorldScore↑	MoT 103.2	Dense 91.3	+11.9，证明语义越丰富越需解耦
半监督训练（固定 1K BEV 标注，Table 4）	WorldScore↑	1K+10K 无标注: 90.3	1K+0K 无标注: 63.2	+27.1，无标注视频持续改善生成

局限与改进

作者承认：主要增益来自额外状态监督，但真实世界状态远比 Minecraft 复杂、且并不直接可得；本研究只验证双人 Minecraft，未推广到更多玩家、真实世界视频、或更细粒度状态（低层 3D 细节、高层跨智能体语义关系）。我自己还有四点观察：第一，WorldScore 作为无界指标对 FID 尺度敏感，作者自己也只当辅助排序用，Building/Consistency 类样本量小、VLM 评估有噪声，9.4→28.1 这类小样本变化的统计显著性不明；第二，scene text 由 Qwen2.5-VL-72B 离线生成且 caption 时条件化了真实控制器输入，本质是"伪标注"，可能引入标注模型偏差；第三，256 个寄存器用 64/64/128 切分给三个头，切分是否最优未消融，固定 K 可能在复杂场景饱和；第四，计算成本高——三阶段+双评分网络($s_{real},s_{fake}$)+MoT 翻倍参数，论文未报告训练 GPU 时长与推理延迟，复现门槛高。

独立分析的弱点

弱点一，只覆盖 2 个智能体：Minecraft 双人场景相对简单（低多边形、语义有限），多智能体协同未验证。改进方向：扩展到 >2 玩家与开放世界，寄存器按智能体数量结构化（每智能体一组 slot 加一组全局 slot），并引入图注意力建模智能体间关系。弱点二，监督依赖仿真器：agent status 和 BEV 来自 SolarisEngine，scene text 依赖大 VLM 且条件化了真实输入，难迁移真实世界。改进方向：用 SLAM/3D 重建、可微渲染或自监督几何（DUSt3R 风格）从真实视频反推俯视图与位姿，把 BEV/状态监督从仿真器解放。弱点三，寄存器容量固定 K=256：复杂场景信息超载会饱和，64/64/128 切分未消融。改进方向：动态寄存器数量，或与检索式外部记忆结合（Worldmem 路线），用寄存器存"当前激活状态"、外部库存"长期记忆"。弱点四，训练成本高：三阶段+双评分网络+MoT 翻倍参数。改进方向：合并评分网络、用 LoRA/适配器做 MoT 分支、把 Stage-1 teacher 复用作 $s_{real}$ 减少显存。

未来方向

作者明确提出的方向：让世界模型感知更复杂的状态，从低层 3D 一致视觉细节到高层跨智能体/玩家语义关系。基于该成果可延伸五方向：其一，把 WSR 思想迁移到具身机器人多机器人协同——每台机器人一个 agent，BEV 由自车感知重建，寄存器维护共享地图与任务状态。其二，把寄存器从"被动状态槽"升级为"可被动作直接读写"的接口，支持更细粒度交互控制（类似 Worldplay 的实时交互建模）。其三，探索半监督之外——用世界状态寄存器做主动标注（不确定性高的状态主动请求 BEV/文本标注），形成标注高效闭环。其四，与长上下文状态空间模型（如 Long-context SSM world model）结合，处理分钟级以上超长 rollout，并用 attention sink 稳定流式生成。其五，把寄存器解码的俯视图/文本作为 RL 奖励或可验证信号，训练可规划的世界模型而非纯生成模型，使寄存器同时服务生成与决策。

复现评估

复现评估：项目页 https://vail-ucla.github.io/worldweaver/ 已开放，但正文未明确代码/权重是否完全开源。数据基于 SolarisEngine 收集约 126 小时 832×480@20fps 双人视频，并用 Qwen2.5-VL-72B 对每 4 帧 chunk 标注场景文本——需 Solaris 访问权限和可观算力。模型基于 Solaris 官方双向 checkpoint 初始化，依赖冻结的 Llama-3.2-1B（文本评分）和 DINOv2 ViT-B/14（BEV 特征）。训练三阶段分别 120K/80K/3-4K 步、batch 32、bf16，Stage-3 还需冻结 teacher $s_{real}$ 和训练 critic $s_{fake}$，算力门槛高（多卡 A100/H100 级）。优点是超参表完整（学习率、损失权重、窗口、寄存器切分 64/64/128）、伪代码清晰。综合判断：若 Solaris 与基础 checkpoint 可获取，方法层面可复现；但 126 小时双人数据采集和三阶段算力是主要障碍，难度偏高。

带世界状态寄存器的流式多智能体自回归扩散模型 Streaming Multi-Agent Autoregressive Diffusion Model with World State Registers