MosaicMem:用于可控视频世界模型的混合空间记忆 MosaicMem: Hybrid Spatial Memory for Controllable Video World Models
融合显式3D与隐式注意力的Patch级混合空间记忆机制
前置知识
显式空间记忆(Explicit Spatial Memory)
通过外部3D估计器(如深度估计网络)将2D观测信息提升到3D空间,构建几何缓存(如点云或3D高斯),当相机重新访问某区域时,将缓存的3D结构投影到查询视角以条件化生成。代表方法包括GEN3C(点云缓存)和VMem(surfel索引)。其核心思想是通过显式的几何重建来保证跨视角一致性,但对动态物体的处理能力较差,且长时间序列中跨视角对齐误差会累积。
MosaicMem的核心设计是融合显式和隐式记忆的优势,因此理解显式记忆的工作原理、优势和局限是理解本文动机的关键前提。
隐式空间记忆(Implicit Spatial Memory)
不构建显式3D结构,而是将世界状态存储在模型的隐式表征中,通常通过反馈已带姿态的帧并依赖注意力机制进行检索。代表方法包括Context-as-Memory(CaM)和RTFM。这种方法灵活,能处理动态变化和非刚性变形,但存在空间漂移问题——即使提供完全准确的相机位姿,生成的视频仍会表现出不精确的自运动,且基于帧的表征高度冗余,消耗大量上下文窗口。
理解隐式记忆的局限性(空间漂移、上下文冗余)是理解MosaicMem为何选择Patch级表征作为中间方案的核心动机。
RoPE(旋转位置编码)
Rotary Position Embedding,一种将位置信息通过旋转矩阵编码到Query和Key向量中的位置编码方法。在DiT(Diffusion Transformer)架构中,RoPE用于编码视频token的时空坐标。MosaicMem提出了Warped RoPE机制,将源视角的RoPE坐标通过3D重投影变换到目标视角,从而建立检索到的记忆Patch与当前生成区域之间的精确空间对应关系。
Warped RoPE是MosaicMem实现Patch级空间对齐的关键技术组件,理解标准RoPE的工作原理才能理解这种坐标变换的意义。
PRoPE(投影位置编码)
Projective Positional Encoding,一种通过将相机截锥体几何信息直接注入自注意力层来实现相机控制的机制。给定每帧的投影矩阵 $\tilde{P}_i$,PRoPE通过投影变换 $\tilde{P}_{i_1} \tilde{P}_{i_2}^{-1}$ 编码两个视角之间的完整相对关系,并通过GTA风格的变换注意力应用。在本文中,PRoPE被用作DiT架构中相机条件化的接口,显著提升了视角可控性。
PRoPE是MosaicMem实现精细相机控制的独立组件,论文通过消融实验证明了其与MosaicMem记忆机制的互补性。
Flow Matching
一种生成建模框架,通过学习从高斯噪声 $X^0 \sim \mathcal{N}(0, I)$ 到目标分布 $X^1$ 的连续概率流ODE来实现采样。给定连续流时间 $\lambda \in [0,1]$,模型学习神经向量场 $u_\theta$ 来传输视频状态。MosaicMem基于文本+图像到视频(TI2V)的Flow Matching模型Wan 2.2构建,将空间记忆 $\mathcal{M}$ 作为额外条件引入生成过程。
理解Flow Matching是理解MosaicMem生成框架的基础,特别是空间记忆如何作为条件信号融入扩散生成过程。
3D VAE时空压缩
视频扩散模型中使用的3D变分自编码器,对视频在空间和时间维度上进行压缩。MosaicMem使用的Wan 2.2模型采用4倍时间压缩,即一个隐空间帧索引 $\ell$ 对应4个原始帧 $\{4\ell+k\}_{k=0}^3$。这种压缩在提高计算效率的同时引入了时空歧义,降低了DiT中RoPE坐标的等效分辨率,是MosaicMem需要设计Warped RoPE和Warped Latent两种对齐机制的技术原因。
理解VAE的时空压缩特性是理解MosaicMem对齐机制设计动机的关键,也是理解PRoPE如何处理时间压缩的技术基础。
研究动机
视频扩散模型正从生成短小合理的片段走向世界模拟器,需要在相机运动、重访和3D结构干预下保持一致性,但现有空间记忆机制存在根本性缺陷。显式空间记忆(如GEN3C、VMem)依赖外部3D估计构建几何缓存,在静态场景中能保证几何一致性,但难以维护和更新包含多个独立运动物体的连贯缓存,生成的动态分数仅为1.18-1.41(以平均光流幅度衡量),远低于真实场景需求。隐式空间记忆(如Context-as-Memory、WorldMem)将状态存储在模型隐表示中,能处理动态变化,但存在严重的空间漂移问题:即使提供完全准确的相机位姿,RotErr仍高达4.65°-5.87°,TransErr达0.43-0.49,且基于帧的表征高度冗余,在有限上下文窗口下限制了持久性。这两种范式都无法同时满足长时程、动态场景下精确空间记忆的需求。
本文的目标是本文的目标是设计一种混合空间记忆机制,融合显式和隐式记忆的互补优势:既利用显式3D估计实现精确的空间定位和目标化检索,又保留隐式记忆通过模型原生注意力机制条件化生成的能力,从而在长时程视频生成中同时实现精确的相机运动控制、动态物体建模和跨视角一致性。具体而言,该机制需要支持分钟级导航视频生成、基于记忆的场景编辑和自回归视频生成。
与已有工作不同的是,本文的独特切入角度在于发现了Patch这一被前人忽略的中间表征粒度。显式方法以3D图元(点、体素、高斯)为基本存储单元,隐式方法以整帧为基本存储单元,而MosaicMem以Patch为基本单元,形成了一个从粗到细的表征谱系中的新位置。这种设计的精妙之处在于:Patch粒度既能通过3D提升实现可靠的局部化(继承显式记忆的几何优势),又能通过模型原生的token拼接和注意力机制进行检索(继承隐式记忆的灵活性)。更关键的是,Patch级表征天然支持稀疏检索——由于视频的高度时间冗余性,从不同时空位置分布的记忆Patch往往足以重建整个序列,且与首帧重叠的区域无需额外记忆,大幅减少了所需的条件token数量。
核心方法
MosaicMem的整体思路可以用'镶嵌拼贴'来直觉理解:就像马赛克艺术将小块瓷砖拼成完整画面一样,MosaicMem将视频中观察到的场景区域切割成Patch,通过3D提升将每个Patch定位到世界空间中,当相机移动到新视角并重新访问某区域时,从记忆空间中检索空间对齐的Patch,将它们像镶嵌拼图一样拼贴到目标帧上。具体技术路线分为三个阶段:(1) 记忆构建——对输入视频的每个Patch,使用现成的3D估计器(Depth Anything V3)推断深度和相机信息,将Patch提升到3D空间;(2) 记忆检索与对齐——当需要生成新视角时,将记忆Patch通过重投影变换到目标视角,并通过Warped RoPE或Warped Latent机制与生成token对齐;(3) 条件化生成——将对齐后的记忆Patch展平并拼接到token序列中,通过DiT的原生注意力机制条件化生成,同时PRoPE模块提供精细的相机运动控制。整个流程建立在Wan 2.2(5B参数的TI2V DiT)之上,使用Flow Matching框架进行训练。
MosaicMem的核心创新是提出Patch-and-Compose接口,将显式和隐式记忆范式在Patch粒度上统一。与显式记忆的本质区别在于:显式方法将Patch投影后直接渲染为条件信号(如ControlNet分支或通道拼接),生成器本质上是视频修复;而MosaicMem将检索到的Patch作为上下文token注入DiT,让模型通过注意力机制自主决定是依赖空间记忆进行一致重建,还是根据文本提示合成新的动态内容。与隐式记忆的本质区别在于:隐式方法以整帧为存储单元,检索是基于注意力的软匹配;而MosaicMem以Patch为单元,检索由3D几何引导的硬匹配驱动,通过显式的重投影实现精确的空间定位。这种混合设计的关键优势是选择性持久性——保留应该持久的内容(静态场景结构),同时允许模型自由演变应该变化的内容(动态物体、光照变化)。
方法步骤详情
MosaicMem的完整方法流程包含以下步骤:(1) Patch提取与3D提升:给定输入视频,对每个Patch $P$ 使用现成3D估计器(Depth Anything V3)推断深度图 $D$ 和关联的相机内参/外参 $(K_i, T_i)$,将2D Patch提升到3D世界空间。(2) 记忆存储:将提升后的Patch连同其3D位置、深度和相机参数存入记忆空间 $\mathcal{M}$。(3) 视角查询:当相机移动到新视角需要生成内容时,根据目标相机位姿从记忆空间中检索空间对齐的Patch集合。(4) 空间对齐——Warped RoPE:对检索到的记忆Patch $P$,将其原始RoPE坐标 $(u,v)$ 通过反投影到3D世界空间,再重投影到目标相机 $(K_j, T_j)$,得到新坐标 $(u', v') = \Pi(K_j T_j (u,v,D))$,其中 $\Pi(\cdot)$ 为透视除法。3D RoPE坐标定义为 $(j, u', v')$,保留重投影坐标的分数部分并在更高分辨率下采样RoPE以保持精度。(5) 空间对齐——Warped Latent:利用重投影坐标 $(u', v')$ 对源隐空间表征进行空间重采样,通过可微分双线性网格采样在分数坐标处获取warped latent patch。两种对齐机制通过混合训练策略组合使用。(6) 条件化生成:将warped后的记忆Patch展平并拼接到token序列中作为条件输入DiT,同时PRoPE模块通过投影矩阵注入相机控制信号——每帧使用4个子相机矩阵 $\{\tilde{P}_{\ell,k}\}_{k=0}^3$ 处理4倍时间压缩,通过分块对角矩阵 $D_t^{PRoPE}$ 编码完整的相对相机几何。(7) 训练:基于Wan 2.2 5B模型微调,使用AdamW优化器,学习率 $1 \times 10^{-5}$,在8×H100 GPU集群上训练250k步,有效批量大小64,推理时使用50步去噪。
技术新颖性
MosaicMem的技术新颖性体现在多个层面。首先,在概念层面,Patch作为基本记忆单元是一个全新的设计选择,填补了显式(3D图元)和隐式(整帧)之间的表征空白,这不是简单的工程折中,而是基于视频高度时间冗余性的深刻观察——分布式稀疏Patch足以重建完整序列。其次,Warped RoPE是一种新的位置编码机制,不同于传统RoPE的固定坐标,它将源视角的RoPE坐标通过3D重投影动态变换到目标视角,在DiT的注意力层中建立像素级精确的跨视角对应关系,这在视频生成领域是首创。第三,Warped Latent提供了互补的对齐方案,直接在特征空间进行空间变换而非修改位置编码,在自回归生成中表现更稳定——消融实验表明Warped RoPE在极端情况下可能导致边界区域物体重复生成,而混合训练两种策略获得最鲁棒的效果。第四,MosaicMem展示了训练-free的注入能力——即使不进行任何微调,仅将MosaicMem作为上下文条件,Wan 2.2也能将记忆Patch投射到正确的时空位置并生成有意义的内容。
实验结果
论文通过四个维度的实验全面验证了MosaicMem的有效性。在空间记忆评估中(Table 1),MosaicMem在相机控制方面实现了RotErr 0.51°、TransErr 0.06,相比显式记忆最优基线SEVA(RotErr 1.42°、TransErr 0.12)分别降低了64%和50%,相比隐式记忆基线CaM(RotErr 4.65°、TransErr 0.43)降低了89%和86%。在视觉质量方面,MosaicMem的FID为65.67、FVD为232.95,相比显式记忆最优的SEVA(FID 74.67、FVD 301.77)分别改善了12%和23%。在记忆检索一致性方面,MosaicMem的SSIM达0.75、PSNR达23.57、LPIPS为0.11,全面优于所有基线。最具说服力的是动态分数:MosaicMem达到2.58,是显式记忆最优(1.41)的1.83倍、隐式记忆(1.72)的1.50倍,证明其成功融合了两种范式的优势。消融实验揭示了各组件的贡献:PRoPE alone的RotErr为4.91°,MosaicMem alone为0.79°,完整模型仅0.51%,说明记忆检索和相机控制模块具有强互补性。长时程视频生成实验(Section 4.3)展示了MosaicMem能生成2分钟的导航视频,保持跨时间的一致性,而CaM在长序列中逐渐累积伪影直至生成崩溃。记忆操作实验(Section 4.4)展示了MosaicMem支持场景拼接和空间编辑,能创建不同视觉风格但几何连续的复合场景。在自回归生成实验中(Table 2),MosaicMem在16 FPS实时生成下实现了总质量分81.11(RELIC为79.08)、RotErr 0.89°(RELIC为4.99°),在所有指标上全面超越Matrix-Game和RELIC。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 空间记忆评估——相机控制精度 | RotErr (°) | 0.51 | SEVA 1.42 / CaM 4.65 | 相比显式最优降低64%,相比隐式最优降低89% |
| 空间记忆评估——平移误差 | TransErr | 0.06 | SEVA 0.12 / CaM 0.43 | 相比显式最优降低50%,相比隐式最优降低86% |
| 空间记忆评估——视觉质量 | FID | 65.67 | SEVA 74.67 / CaM 85.32 | 相比显式最优降低12%,相比隐式最优降低23% |
| 空间记忆评估——视频质量 | FVD | 232.95 | SEVA 301.77 / CaM 392.11 | 相比显式最优降低23%,相比隐式最优降低41% |
| 空间记忆评估——检索一致性 | SSIM | 0.75 | SEVA 0.66 / CaM 0.49 | 相比显式最优提升14%,相比隐式最优提升53% |
| 空间记忆评估——动态建模 | Dynamic Score | 2.58 | VWM 1.41 / CaM 1.72 | 相比显式最优提升83%,相比隐式最优提升50% |
| 自回归生成——总质量分 | Quality Score | 81.11 | RELIC 79.08 / Matrix-Game 75.11 | 相比RELIC提升2.6% |
| 自回归生成——相机控制 | RotErr (°) | 0.89 | RELIC 4.99 / Matrix-Game 5.32 | 相比RELIC降低82% |
局限与改进
尽管MosaicMem取得了显著改进,论文也暴露了若干局限性。首先,MosaicMem依赖现成的3D估计器(Depth Anything V3)进行深度推断和相机位姿估计,这意味着其空间记忆质量受制于上游3D估计的精度——在深度估计失败或相机位姿不准确的区域,Patch提升和重投影都会产生误差。其次,Warped RoPE和Warped Latent两种对齐机制各有优缺点且存在互补的失败模式:Warped RoPE在自回归生成中可能导致边界区域物体重复生成,而Warped Latent在相机运动极慢时可能引入类似问题,论文通过混合训练策略缓解但未彻底解决。第三,论文的长时程生成实验仅展示了2分钟视频,尚未验证更长时间尺度(如5-10分钟)下的记忆稳定性和计算开销增长趋势。第四,MosaicMem-World基准虽然引入了重访轨迹,但主要基于Unreal Engine 5合成场景和游戏环境,与真实世界的复杂性存在差距——真实世界数据仅作为辅助来源。第五,论文未充分讨论MosaicMem的计算开销:虽然声称Patch级表征减少了上下文冗余,但未给出与整帧隐式记忆方法的推理速度、显存占用和延迟的直接对比。最后,记忆操作能力(场景拼接、编辑)目前仅展示了最简单的实现,缺乏对复杂空间编辑操作(如物体移动、场景变形)的系统评估。
独立分析的弱点
MosaicMem存在几个值得深入分析的弱点。(1) 对3D估计器的强依赖性:整个方法链的可靠性建立在Depth Anything V3的深度和位姿估计之上,但该估计器本身在弱纹理、透明/反射表面、遮挡等场景中表现不稳定,论文未讨论这种上游误差如何传播到下游记忆质量和生成结果中。改进方向包括引入多视角几何验证、不确定性感知的记忆加权,或联合优化3D估计和记忆构建。(2) Patch粒度固定:论文未讨论Patch大小的选择对性能的影响——过大的Patch可能无法精确定位动态物体边界,过小的Patch可能增加检索开销和对齐难度。改进方向是设计自适应Patch分割策略,根据场景复杂度动态调整粒度。(3) 记忆空间管理策略缺失:论文未详细说明记忆空间的增长控制机制——在长时程生成中,记忆空间可能无限增长导致检索效率下降。改进方向包括基于重要性的记忆淘汰策略、分层记忆索引或记忆压缩。(4) 训练数据偏差:MosaicMem-World主要基于合成数据构建,真实世界数据占比有限,可能导致模型在真实场景中的泛化能力不足。(5) 动态物体建模虽优于显式方法,但论文未展示多个独立运动物体同时存在时的处理能力——这是显式记忆的核心瓶颈,论文声称MosaicMem能处理但缺乏定量验证。
未来方向
基于MosaicMem的成果,未来研究可沿以下方向拓展。(1) 端到端联合优化:当前3D估计、记忆构建和视频生成是独立模块,未来可探索端到端训练,让3D估计器针对记忆检索任务进行优化,或让记忆模块与生成模型联合学习。(2) 自适应记忆粒度:设计根据场景语义和几何复杂度自动调整Patch大小和密度的机制,例如在动态物体边界使用更细粒度的Patch,在静态区域使用更大Patch以减少冗余。(3) 实时交互式世界模拟:MosaicMem的记忆操作能力(场景拼接、编辑)为交互式世界模拟器提供了基础,未来可探索实时用户交互下的记忆更新和场景演变,实现类似Genie 3的实时交互体验。(4) 多模态记忆扩展:当前记忆仅包含视觉信息,未来可将语言、动作、物理属性等多模态信息与Patch关联,支持更丰富的条件化生成和世界理解。(5) 大规模世界模型中的记忆系统:随着视频世界模型向更大规模发展,MosaicMem的Patch级稀疏检索机制可扩展为分布式记忆系统,支持跨场景、跨任务的记忆共享和复用。
复现评估
从复现角度来看,MosaicMem提供了相对良好的可复现基础。论文基于开源的Wan 2.2 5B模型构建,3D估计使用开源的Depth Anything V3,Flow Matching训练框架也有公开实现。训练配置明确:AdamW优化器、学习率 $1 \times 10^{-5}$、250k步、有效批量64、8×H100 GPU集群,这些参数清晰可复现。然而存在几个复现障碍:(1) MosaicMem-World数据集的构建涉及Unreal Engine 5场景和商业游戏(Cyberpunk 2077),完整数据集的获取和重现存在许可和技术门槛;(2) PRoPE的具体实现细节(特别是4倍时间压缩的相机矩阵展开)需要对照原始PRoPE论文实现;(3) Warped RoPE的高分辨率采样和分数坐标保留的实现细节未完全公开;(4) 自回归版本Mosaic Forcing基于Causal Forcing和Rolling Forcing的蒸馏流程,需要额外的训练步骤。总体而言,论文的核心方法具有可复现性,但完整的基准数据集和端到端流水线可能需要显著的工程努力。
论文图表
左图展示了三种空间记忆机制的对比:显式记忆(3D点云/高斯缓存)、隐式记忆(帧级隐表示)和MosaicMem(Patch级混合记忆),配以可视化示意。右图A展示了MosaicMem相比隐式记忆实现更准确的相机运动控制,右图B展示了MosaicMem相比显式记忆能生成文本驱动的动态内容(如骑士骑马、汽车驶过花园),而显式记忆生成的场景保持静态。
这张图是全文的核心概念图,直观展示了MosaicMem在显式和隐式记忆之间的定位,以及其双重优势——精确相机控制和动态内容生成,是理解论文动机和贡献的最佳入口。