技术摘要：张量记忆：面向长程 Transformer 的固定大小循环状态

1. 问题陈述

Transformer 通过将空间和时间维度展平为长令牌序列来处理图像和视频。虽然注意力机制和 KV 缓存保留了过去的特征，但随着序列视界的增长，它们存在两个主要局限性：

1. 内存扩展：实际推理依赖于存储过去的键和值，其规模与令牌数量呈线性增长。这对于长视频和文档而言，在计算和内存方面变得不可行。
2. 缺乏持久状态：基础 Transformer 缺乏显式的持久状态。它们关注过去的特征，但并未维护一种结构化的表示，以便在输入变化时进行更新和查询。这导致其在遮挡、观测缺失的情况下表现脆弱，并且需要不断从日益增长的历史中重建全局上下文，而非维持一个紧凑的持久状态。

现有的方法，如稀疏注意力、历史压缩或检索增强上下文，虽然往往能扩展上下文长度，但通常将信息存储为非结构化的令牌向量或槽位集合。它们并未显式保留空间结构，也未维护适合视频理解等基于状态任务的紧凑“世界状态”。

2. 方法论：张量记忆

作者提出了张量记忆（Tensor Memory），这是一个轻量级模块，通过引入固定大小的循环 3D 记忆张量来增强标准 Transformer 块。该模块在保持空间归纳偏置的同时，将状态容量与输入长度解耦。

核心架构

记忆被定义为一对体积（隐藏状态 $h_t$ 和单元状态 $c_t$），其维度为 $C \times D \times H \times W$，其中空间维度是恒定的，与输入长度无关。该模块按每个令牌块（chunk）的周期运行：

1. 分块（Chunking）：输入令牌被分组为块。从块的第一个令牌派生“读取查询”，而“写入摘要”则是整个块的学习投影。
2. 连续寻址：一个共享的坐标预测器（使用绑定的权重）在体素网格内生成用于读取（$\mu_{read}$）和写入（$\mu_{write}$）的连续 3D 坐标。
3. 读取（连续采样）：模型使用三线性插值从 $\mu_{read}$ 处的隐藏状态 $h_t$ 中读取。检索到的特征被投影回令牌维度，并通过门控残差连接融合到令牌流中。一个学习的标量门 $\gamma$ 控制这种融合的强度，允许模型在特定任务中记忆路径无益时抑制该路径。
4. 写入（可微软写入）：模型通过预测内容向量和扩散参数 $\sigma$ 来写入信息。它将此内容作为围绕预测的 3D 位置 $\mu_{write}$ 的高斯加权体积进行沉积。这种“软写入”避免了硬体素分配，允许梯度塑造信息的放置位置。
5. 循环更新：写入体积与当前隐藏状态拼接。一个轻量级的、因式分解的 3D 算子（深度 3D 卷积后接逐点混合）计算 ConvLSTM 风格的门（$i_t, f_t, o_t, g_t$）。这些门更新单元状态和隐藏状态，使信息能够持久存在、局部传播，并在观测缺失时仍然可用。

关键设计选择

• 固定大小状态：记忆张量的大小是恒定的，防止了与 KV 缓存相关的线性内存增长。
• 空间结构：3D 体素网格显式保留了空间关系，不同于非结构化的令牌槽位。
• 可微软写入：使用高斯加权体积允许平滑、连续的更新，而非离散、硬性的分配。
• 门控融合：学习的门 $\gamma$ 确保该模块可以附加到预训练模型上而不会破坏训练稳定性；如果记忆不必要，该门将抑制该路径。

3. 主要贡献

• 新颖的记忆机制：引入了一种固定大小的循环 3D 体素状态，令牌可以通过连续坐标向其写入和从中读取，充当显式的空间状态。
• 集成：该模块干净地集成到标准 Transformer 训练流程中，可以附加到或从现有块中移除，而无需架构更改。
• 诊断套件：开发了一个受控的“玩具”诊断套件，旨在隔离持久状态有益的场景（如坐标绑定、遮挡、地图构建）与不有益的场景。

4. 实验结果

作者在语言、图像和视频领域的标准基准测试以及诊断套件上评估了张量记忆。

• 语言建模（WikiText-2）：张量记忆取得了最大的绝对增益，将测试困惑度从基础 Transformer 的 130.30 降低至 100.07。其表现优于包括 Transformer-XL、局部注意力和线性注意力在内的基线模型。
• 图像块重建（CUB-200-2011）：在掩码块重建任务中，与基础 ViT、局部 ViT 和 Registers ViT 相比，张量记忆变体（空间 ViT）在掩码比率（25%、50%、75%）下均显示出 PSNR 和 SSIM 的持续改进。
• 视频动作识别（UCF-101）：该模型在 Top-1（89.21% vs 82.12%）和 Top-5 准确率上均优于联合时空 Transformer 基线和强大的“注册令牌”基线。作者指出，这是由于每步体素循环导致训练挂钟时间显著增加。
• 玩具诊断：
  • 坐标绑定：随着写入次数的增加，张量记忆保持了高准确率（75–91%），而基于注意力的基线模型则崩溃至接近随机性能（7–14%）。
  • 无害控制：在不需要记忆的任务中，学习的门 $\sigma(\gamma)$ 保持在初始化附近（抑制记忆路径），所有模型均达到 100% 的准确率。
  • 地图构建：在需要随时间保留信息的长程网格导航任务中，张量记忆显著优于基线模型。

5. 意义与主张

该论文声称，张量记忆通过为 Transformer 提供显式、紧凑且结构化的状态，提供了一种处理长程推理的实用方法。

• 容量与长度解耦：通过使用固定大小的张量，模型避免了 KV 缓存的线性内存扩展，使其在理论上对长序列更具可扩展性。
• 对遮挡的鲁棒性：持久状态允许模型维护关于场景的“信念状态”，使其比仅依赖从令牌流重建上下文的模型更能抵抗观测缺失或遮挡。
• 通用性：虽然其动机源于长程视频，但该机制已被证明对语言（积累全局上下文）和图像（结构化空间聚合）均有益。
• 未来工作的基础：作者将当前的实现视为基础。他们建议未来的方向包括更高维度的记忆（例如，用于时间感知空间记忆的 4D 记忆）、更丰富的动态（例如，光流扭曲）以及稀疏更新以提高效率。

作者对当前的局限性保持谦逊，承认密集的体素更新是吞吐量的瓶颈，并且如果任务需要存储大量不同的细节，固定大小的状态可能会饱和。然而，他们断言，添加一个小的结构化状态是实现更强大的长程感知和推理系统的可行途径。