Technical Summary: Tensor Memory: Fixed-Size Recurrent State for Long-Horizon Transformers

1. Problem Statement

Transformers process images and videos by flattening spatial and temporal dimensions into long token sequences. While attention mechanisms and KV caching preserve past features, they suffer from two primary limitations as the sequence horizon grows:

1. Memory Scaling: Practical inference relies on storing past keys and values, which scales linearly with the number of tokens. This becomes computationally and memory-prohibitive for long videos and documents.
2. Lack of Persistent State: Vanilla Transformers lack an explicit, persistent state. They attend to past features but do not maintain a structured representation that can be updated and queried as the input changes. This leads to fragility under occlusion, missing observations, and the need to repeatedly reconstruct global context from an ever-growing history rather than maintaining a compact, persistent state.

Existing approaches, such as sparse attention, history compression, or retrieval-augmented contexts, often extend context length but store information as unstructured collections of token vectors or slots. They do not explicitly preserve spatial structure or maintain a compact "world state" suitable for state-based tasks like video understanding.

2. Methodology: Tensor Memory

The authors propose Tensor Memory, a lightweight module that augments standard Transformer blocks with a fixed-size recurrent 3D memory tensor. This module decouples state capacity from input length while preserving a spatial inductive bias.

Core Architecture

The memory is defined as a pair of volumes (hidden state $h_t$ and cell state $c_t$) with dimensions $C \times D \times H \times W$, where the spatial dimensions are constant regardless of input length. The module operates in a cycle per token chunk:

1. Chunking: Input tokens are grouped into chunks. A "read query" is derived from the first token of the chunk, and a "write summary" is a learned projection of the entire chunk.
2. Continuous Addressing: A shared coordinate predictor (using tied weights) generates continuous 3D coordinates for reading ($\mu_{read}$) and writing ($\mu_{write}$) within the voxel grid.
3. Reading (Continuous Sampling): The model reads from the hidden state $h_t$ at $\mu_{read}$ using trilinear interpolation. The retrieved feature is projected back to the token dimension and fused into the token stream via a gated residual connection. A learned scalar gate $\gamma$ controls the strength of this fusion, allowing the model to suppress the memory path if it is not beneficial for a specific task.
4. Writing (Differentiable Soft Write): The model writes information by predicting a content vector and a spread parameter $\sigma$. It deposits this content as a Gaussian-weighted volume around the predicted 3D location $\mu_{write}$. This "soft write" avoids hard voxel assignment, allowing gradients to shape where information is placed.
5. Recurrent Update: The write volume is concatenated with the current hidden state. A lightweight, factorized 3D operator (depthwise 3D convolutions followed by a pointwise mix) computes ConvLSTM-style gates ($i_t, f_t, o_t, g_t$). These gates update the cell and hidden states, enabling information to persist, propagate locally, and remain available even when observations are missing.

Key Design Choices

• Fixed-Size State: The memory tensor size is constant, preventing the linear memory growth associated with KV caching.
• Spatial Structure: The 3D voxel grid explicitly preserves spatial relationships, unlike unstructured token slots.
• Differentiable Soft Writes: Using Gaussian-weighted volumes allows for smooth, continuous updates rather than discrete, hard assignments.
• Gated Fusion: The learned gate $\gamma$ ensures the module can be attached to pretrained models without destabilizing training; if memory is unnecessary, the gate suppresses the path.

3. Key Contributions

• Novel Memory Mechanism: Introduction of a fixed-size, recurrent 3D voxel state that tokens can write to and read from via continuous coordinates, acting as an explicit spatial state.
• Integration: The module integrates cleanly into standard Transformer training pipelines and can be attached to or removed from existing blocks without architectural changes.
• Diagnostic Suite: Development of a controlled "toy" diagnostic suite designed to isolate scenarios where persistent state is beneficial (e.g., coordinate binding, occlusion, map building) versus scenarios where it is not.

4. Experimental Results

The authors evaluated Tensor Memory on standard benchmarks across language, image, and video domains, as well as the diagnostic suite.

• Language Modeling (WikiText-2): Tensor Memory achieved the largest absolute gains, reducing test perplexity from 130.30 (Base Transformer) to 100.07. It outperformed baselines including Transformer-XL, Local attention, and Linear attention.
• Image Patch Reconstruction (CUB-200-2011): In a masked patch reconstruction task, the Tensor Memory variant (Spatial ViT) showed consistent improvements in PSNR and SSIM across mask ratios (25%, 50%, 75%) compared to Base ViT, Local ViT, and Registers ViT.
• Video Action Recognition (UCF-101): The model outperformed both a joint spatiotemporal Transformer baseline and a strong "register token" baseline on Top-1 (89.21% vs 82.12%) and Top-5 accuracy. The authors note this came at the cost of significantly higher training wall-clock time due to the per-step voxel recurrence.
• Toy Diagnostics:
  • Coordinate Binding: Tensor Memory maintained high accuracy (75–91%) as the number of writes increased, whereas attention-based baselines collapsed to near-random performance (7–14%).
  • No-Harm Control: On tasks requiring no memory, the learned gate $\sigma(\gamma)$ remained near its initialization (suppressing the memory path), and all models achieved 100% accuracy.
  • Map Building: Tensor Memory significantly outperformed baselines in long-horizon grid navigation tasks where information must be retained over time.

5. Significance and Claims

The paper claims that Tensor Memory offers a practical way to handle long-horizon reasoning by providing Transformers with an explicit, compact, and structured state.

• Decoupling Capacity from Length: By using a fixed-size tensor, the model avoids the linear memory scaling of KV caches, making it theoretically more scalable for long sequences.
• Robustness to Occlusion: The persistent state allows the model to maintain a "belief state" about the scene, making it more robust to missing observations or occlusion compared to models that rely solely on reconstructing context from a token stream.
• Versatility: While motivated by long-horizon video, the mechanism is shown to be beneficial for language (accumulating global context) and images (structured spatial aggregation).
• Foundation for Future Work: The authors view the current implementation as a foundation. They suggest future directions include higher-dimensional memories (e.g., 4D for time-aware spatial memory), richer dynamics (e.g., optical flow warping), and sparse updates to improve efficiency.

The authors remain modest about the current limitations, acknowledging that the dense voxel updates are a throughput bottleneck and that the fixed-size state may saturate if a task requires storing a vast number of distinct details. However, they assert that adding a small structured state is a viable path toward more capable long-horizon perception and reasoning systems.