MemLearner: Learning to Query Context Memory for Video World ModelsMemLearner：为视频世界模型学习如何查询上下文记忆

视频世界模型要在用户控制下预测未来画面，但长时间生成时场景会逐渐失真，本质是缺一套记忆机制。已有做法靠规则从历史帧里挑上下文，比如按视场角重叠或点云估计，可一旦出现遮挡或运动物体就失效。MemLearner 把检索换成学习：插入一组查询 token，让它们从历史上下文 token 里自适应抽取信息，再交给预测 token 作为生成条件。关键设计是把查询交给视频生成模型本身完成，而非新建模块，从而复用预训练视觉先验。作者还收集了带遮挡、动态对象与相机位姿的长视频数据，并提出多数据集训练策略联合利用渲染与真实视频。

视频世界模型把历史状态与交互动作映射成未来画面，长视频生成中场景一致性下降，根因是上下文窗口有限、模型没有记忆。现有记忆手段分三类：3D 重建代价高且易出错、特征压缩对细粒度场景不够、上下文检索最直接但需先挑相关历史帧。检索类方法沿用 FOV 重叠或点云+面片匹配等硬编码规则，面对遮挡墙壁、运动物体与动态环境的状态变化都难以泛化。作者由此转向：与其手工设计检索规则，不如让网络通过端到端训练自己学会从历史帧中自适应查询信息。

关键贡献

• 提出 MemLearner，一种基于学习的自适应上下文查询方法，用查询 token 从上下文 token 中提取有效记忆，使视频世界模型具备记忆并提升一致性与泛化能力。
• 关键设计是让视频生成模型自身承担上下文查询，而非额外训练一个从零初始化的查询模块，从而复用预训练视觉先验、无需专门监督。
• 提出两项高效训练与推理策略：仅在浅层查询层做上下文查询，以及剔除冗余注意力计算，降低长序列上下文查询开销。
• 收集基于 Unreal Engine 的渲染长视频数据集，含定制遮挡、动态对象、相机位姿标注与充分的回访场景，满足现有数据集缺失的条件。
• 提出多数据集训练策略，为不同位姿标注质量的数据分配独立相机编码器，使渲染与真实视频的互补优势可以同时利用。
• 实验表明 MemLearner 在遮挡与动态场景下显著优于既有视频世界模型，并验证其架构无关性与真实世界迁移能力。

总体思路

MemLearner 把长视频生成建模成「给定历史上下文帧 C，预测未来帧 P」。它在 C 与 P 之间插入一组可学习查询 token Q，让 Q 先做「查询者」：Q 作为 query 去 attend 上下文 C（必要时也看 P），自适应抽出对当前生成有用的历史信息；再让 P 作为 query 去 attend Q，把抽出的信息当作生成条件。整个交互由视频生成模型自身的 3D 注意力完成，Q 不设专门监督，仅靠 P 上的扩散损失做端到端间接监督。

输入表示与 token 构成

基于带因果 3D VAE 与 DiT 的潜空间扩散模型，一个视频生成被切成块逐块自回归推进。第 t 步输入记为 Zt = {C, Q, Pt}：上下文 token 与查询 token 保持无扰动，只有预测 token 被随机高斯噪声扰动。每块 P 对应 77 帧（20 帧视频潜变量），Q 数量与 P 相同，且用随机采样噪声而非可学习参数初始化，以贴合扩散模型的输入分布；C 长度最多可到 P 的 9 倍，其最后一帧与 P 的第一帧对齐以保证连续性。

训练损失

扩散损失只施加在预测 token 上。记 ϵP ∼ N(0, I) 为施加在 P 上的噪声，模型预测 ϵθ，则： $$\mathcal{L}(\theta)=\mathbb{E}\big[\|\epsilon_\theta(Z_t,\mathrm{cam},p,t)-\epsilon_P\big]$$ 由于监督只落在 P，Q 的「如何查询」是由反向传播通过中间注意力隐式学出来的，不需要给查询模块单独配损失。

用生成模型自身做查询

一个直观但失败的设计是新增一个从零初始化的上下文查询模块（Fig. 2(b)）。实验发现它与预训练生成器联合训练时学不到东西：注意力相似度几乎为零，生成器随即忽略其输出，退化成纯文生视频模型。因此作者改用设计 (c)：所有 C、Q、P 一并送入视频生成模型，靠它已有的视觉先验完成查询。好处有二：一是复用先验知识、省数据和算力；二是端到端训练自然习得查询能力，不依赖为额外模块设计的输入输出与监督。这一取舍是本文区别于规则检索与旁路模块的核心。

标准 3D 注意力

把注意力块里的线性运算记为 q(·)、k(·)、v(·)、o(·)，对应 query/key/value/output，softmax 记为 sm(·)，用 g(·,·,·) 简记 query/key/value 三个输入，则一个标准 3D 注意力层为： $$F_{\mathrm{out}}=F_{\mathrm{in}}+o\big(\mathrm{sm}(q(F_{\mathrm{in}})k(F_{\mathrm{in}})^{\top})v(F_{\mathrm{in}})\big)=F_{\mathrm{in}}+g(F_{\mathrm{in}},F_{\mathrm{in}},F_{\mathrm{in}})$$ 当上下文很长时，每个位置都全局 attend 所有 token 开销很大，这正是两条效率策略要压缩的对象。

为什么只在浅层查询

信息查询在本质上接近「编码」：编码器通常远小于生成器，所需参数与计算也少得多。于是把 n+m 层的 DiT 切成 n 个浅层 Query Layers 和 m 个深层 Generative Layers，n≪m。交互只发生在 Query Layers；生成层只保留 Q 与 P。这样短上下文交互只占用模型的一小部分，既保留生成能力，又令成本近似与上下文长度解耦。

效率策略与汇总公式

策略一让 C、Q、P 只在浅层查询层交互： $$\text{Query: } C_{\mathrm{out}}=C,\quad Q_{\mathrm{out}}=Q+g(Q,\{C,P\},\{C,P\}),\quad P_{\mathrm{out}}=P+g(P,\{P,Q\},\{P,Q\})$$ $$\text{Generative: } Q_{\mathrm{out}}=Q+g(Q,P,P),\quad P_{\mathrm{out}}=P+g(P,\{P,Q\},\{P,Q\})$$ 策略二只保留三种必要注意力模式——Q 查 P（让 Q 知道该依据预测目标去提什么）、Q 查 C（真正提取上下文）、P 查 P 与 Q（P 从中提取信息），其余（尤其是 C 作为 query）全部剔除。消融显示去掉「Q 查 P」会显著恶化性能，说明 Q 必须从 P 获得「要提什么」的指引。

小结

MemLearner 的关键不是「加了 Q」本身，而是重新分配了注意力角色并借生成器自己完成查询，从而在遮挡与动态场景下实现真正可推广的记忆。

在自建遮挡动态数据集、无遮挡动态的 CaM 及真实世界 SpatialVID 上，与 DFoT、FramePack、VMem、Context-as-Memory、Separate Module 对比，评估 GT 与回访下的 PSNR/LPIPS 与 FID/FVD。MemLearner 在自建数据上全部指标最优。

关键数据

既有方法把「记忆」固化为规则检索，MemLearner 把这一步换成端到端可学习，代表记忆机制从硬编码检索转向学习式查询的范式转折。其「复用生成模型本身而非新增模块」避免了从零查询模块与预训练生成器联合训练难收敛的问题，为在生成骨架中直接注入记忆提供范例。其多数据集训练与「单独编码器隔离不同位姿质量」的做法，为混合渲染与真实数据训练提供参考。

局限与不足

数据与模型规模仍是瓶颈。当前基于 1B 参数量模型，当同一场景里超过约五个运动角色同时交互时，生成质量明显下降，出现外观不一致或回访时物体缺失，说明模型容量不足以同时追踪大量动态实体，需要更大模型与含更丰富动态交互的长视频数据支撑。其次，工作聚焦于「完整保存上下文并高效检索/查询」，记忆开销仍随生成时长线性增长，尚不构成真正可扩展的智能记忆；作者也指出上下文压缩与本文亮点正交，可先压缩表示再套用学习式查询，而压缩、摘要、更新、编辑与选择性遗忘仍是后续方向。此外，相机位姿估计依赖专门编码器，端到端训练以中等规模基座为实验主体，跨更大基座的整体收益仍需验证。