LOCI 给世界模型装空间线性记忆:同预算重访 PSNR 高 0.89 dB

LOCI: Spatial Linear Memory for Streaming World Models

Ji Xia, Tingting Liao, Xuezhi Liang, Hao Li, Guangyi Liu

cs.CV

2026-10-01

把相机条件线性注意力记忆与 KV 缓存混进视频世界模型,同预算重访 PSNR 比 full softmax 高 0.89 dB,并能常量显存流式生成 300 秒。

这篇在解决什么

相机可控的视频世界模型有一个硬需求:走了很久回到看过的地方,原来的建筑、树木得原样画回来。这要求同时做到两件事,记住过去的观测,并在当前视角下取回对的那一条。

两条现有路线各有硬伤。全 softmax 注意力保留历史 KV,细节不丢,但显存随视频长度增长,论文里的同配方基线在 157 秒处 OOM;递归记忆(线性注意力、state-space)把历史压进固定大小状态,紧凑,但旧观测细节被压没了,取不回来。按视角选帧检索能省预算,但可能把后面才用得上的证据提前扔掉。MBZUAI 团队的 LOCI 两个都要,还让它们互相通信。

方法

基础是 Wan2.2-TI2V-5B 的 30 层视频 transformer,改成按 chunk 因果生成,每 chunk 5 个 latent 帧。15 层是混合块:chunk 内 softmax 注意力加一个递归线性注意力记忆;另外 15 层保留历史 softmax 注意力,直接读历史 KV。

历史访问分两种模式:全历史;或有界稀疏,1 个 sink 帧加最多 20 帧全景 bank、8 帧最近窗口和当前 chunk,合计不超过 34 个 latent 帧。bank 按视场角覆盖度更新,优先保留互补视角的早期观测;丢掉的帧不进档案,信息仍留在递归状态里。训练用 chunk 级 diffusion forcing 跑 5,000 步,递归分支只加 90.7M 参数(占 5.38B 的 1.7%)。

结果

MIND 基准给一段记忆视频,生成续段,与录像对比。LOCI 没在 MIND 上训练,zero-shot 在 MSE、PSNR、SSIM 上超过在 MIND 上训练的 GIM-World,LPIPS 略输(0.643 对 0.630),比自测组里最强的 HY-WorldPlay 高 1.82 dB:

方法PSNR↑SSIM↑LPIPS↓
GIM-World(在 MIND 训练,报告值)13.400.4140.630
HY-WorldPlay(8B,自测)12.540.3980.715
LOCI(zero-shot,全历史)14.360.4640.643
LOCI(有界稀疏)13.020.4090.679

隔离变量的对照(同 backbone、数据、配方、5,000 步)里,同等有界 KV 预算下混合结构比 full softmax 高 0.89 dB(95% 区间 0.64 到 1.15),50 段里 44 段更好,增益不来自存更多历史。留出 UE 轨迹重访时刻领先 0.62 dB;四条回到 60 秒前场景的轨迹,领先 0.68 到 1.93 dB。

代价:全历史模式等长峰值显存降约 30%,可生成长度从 157 秒延到 225 秒;有界模式跑满 300 秒,显存恒定 23.6 GiB(full softmax 同访问方式 27.6 GiB),速度持平,每视频秒 5.4 秒(H200)。

机制分析有实证:相机 yaw 从递归状态线性可解码,held-out R²=0.946;共同可见 token 只占历史约 0.1%,却拿到 11.1% 的历史注意力,富集 115 倍(full softmax 是 108 倍);推理时把读出置零,富集度下降,说明下游 query 确实在消费读出内容。

为什么重要

长视频流式生成卡在记忆上。LOCI 的结构工程上可落地:细节存 KV、概要存固定状态、相机几何做索引,组件(KDA、PRoPE、Wan backbone)全是现成的,增量参数 1.7%。对做世界模型、长视频生成和交互式仿真的人,常量显存加重访一致正好是流式落地的两个硬指标。增益不到 1 dB,是结构性渐进改进;受控实验做得干净,机制分析没有停在榜单层面。

局限与存疑

作者自列:留出轨迹全是 Unreal Engine 渲染,没有真实拍摄评测;所有模型绝对保真度都低,基准平均 PSNR 低于 15 dB,增益只能当相对值读;只测了一个 5B backbone 和 5,000 步短微调;逐 token 与逐 chunk 衰减的对比用了事后定义的扩展指标,作者自己标为探索性。

读下来另有几点。30 层每层都保留显式历史的相机注意力分支,递归路径只砍掉一半主注意力 KV,省 30% 不等于显存减半。MIND 上与外部模型的对比只是统一协议下的参考,各模型大小、采样步数、接收记忆段的方式都不同,作者也明说只有同配方 full softmax 那组隔离了变量。训练时递归状态从加噪特征构建,推理时从生成 chunk 提交,这处失配论文没有单独量化。

术语

原文与代码

相关论文

全部论文解读