张量记忆：用于长程 Transformer 的固定大小循环状态

Transformer 通过将空间和时间展平为长序列来处理图像和视频。虽然注意力和 KV 缓存保留了过去的特征，但其内存随序列长度增长且缺乏明确且持久的空间状态，这使得长程视频理解和遮挡相关推理变得困难。我们提出了张量记忆(Tensor Memory)，一种轻量级模块，通过在 Transformer 块中引入固定大小的循环 3D 记忆张量来实现增强。token 通过可微软写入机制写入到体素网格中，该写入机制在预测的连续 3D 位置附近以高斯加权体积形式存储内容；记忆通过高效的局部交互算子和门控循环动力学进行更新；token 通过连续采样和门控残差融合方式读取上下文。由于记忆张量具有恒定大小，张量记忆将状态容量与输入长度解耦，同时保持空间归纳偏置。我们在标准语言、图像和视频基准测试以及一个控制 toy 诊断套件上评估了该模块，后者旨在隔离持久状态何时有益；它与标准 Transformer 训练管线集成，可无需其他架构更改地附加或移除到现有块上。