Vector Quantization in the Brain: Grid-like Codes in World Models

TL;DR

GCQ addresses the challenge of building efficient world models by compressing observation–action sequences into discrete, grid-like codewords produced by continuous attractor neural networks. The method introduces an action-conditioned codebook and sequence template matching to jointly compress space and time, forming a cognitive map that enables long-horizon prediction, goal-directed planning, and inverse modeling. The paper highlights both practical performance gains over traditional two-stage models and theoretical implications for how grid-like codes could emerge in neural systems via fixed CANN-based codebooks mapped to experience. It offers a computational tool for sequence modeling and a neuroscience-inspired perspective on representation learning.

Abstract

We propose Grid-like Code Quantization (GCQ), a brain-inspired method for compressing observation-action sequences into discrete representations using grid-like patterns in attractor dynamics. Unlike conventional vector quantization approaches that operate on static inputs, GCQ performs spatiotemporal compression through an action-conditioned codebook, where codewords are derived from continuous attractor neural networks and dynamically selected based on actions. This enables GCQ to jointly compress space and time, serving as a unified world model. The resulting representation supports long-horizon prediction, goal-directed planning, and inverse modeling. Experiments across diverse tasks demonstrate GCQ's effectiveness in compact encoding and downstream performance. Our work offers both a computational tool for efficient sequence modeling and a theoretical perspective on the formation of grid-like codes in neural systems.

Figures (7)

• Figure 1: (A) Schematic of a CANN: Each green dot represents a neuron uniformly distributed on a torus. The neurons receive external input $I$. (B) Energy landscape of CANN dynamics: Each local minimum in the energy landscape corresponds to an attractor state, which manifests as a 2D Gaussian bump on the torus. (C) Template matching via CANN dynamics: The CANN inherently performs template matching between the external input $I$ and its attractor states. The input $I$ is matched to the attractor that maximizes their inner product. (D) Attractor transition: Under four distinct actions, the attractor initially at position $(0,0)$ stabilizes to four new attractor states. (E) Due to the periodic boundary conditions of the CANN, bump movements along the two axes naturally form grid-like patterns.
• Figure 2: Schematic of GCQ: The action-observation sequence is encoded by the encoder into a latent state composed of $m = 3$ codes. Through sequence template matching with the action-conditioned codebook, the decoder reconstructs the predicted observations. The gray arrows indicate the template matching process, and the gray dashed boxes represent the matching targets.
• Figure 3: (A) Schematic of long-horizon prediction. The figure illustrates the process of initializing with a sequence of length 3 and predicting two future observations. (B) Schematic of goal-directed planning. Some gray arrows are omitted for clarity. The blue arrows represent the mapping from actions in the latent space to real agent actions. Through iterative action generation, environment interaction, and observation, the agent continues until it outputs a no-op action $a_\emptyset$, indicating that the goal has been reached.
• Figure 4: (A)(B) Reconstruction FID and prediction FID for GCQ and VQ+UNet across different model sizes. (C) Prediction FID of GCQ varies with changes in the initialization length. (D) Prediction FID of GCQ (#Para:112M), VQ+UNet (96M) and VQ+Transformer (121M) changes as the prediction length increases. (E) Predictions on GSV dataset. The first patch in each row represents the trajectory drawn by the action. The first three rows correspond to movement actions, while the last two rows correspond to rotation actions. In practice, different actions are encoded within a single GCQ, enabling their use. For visualization convenience, they are plotted separately here. Rows 1, 2, and 4 show GCQ predictions with different initialization lengths (orange frames), predicting subsequent observations (blue frames). Rows 3 and 5 show VQ+UNet predictions under the same conditions. As the prediction length increases, the images become blurry.
• Figure 5: With the same setup as Fig. \ref{['fig4']}E, Rows 1--2: Prediction on the 2DMaze dataset. Rows 3--4: Prediction on the 3DShapes dataset.