Linear Mechanisms for Spatiotemporal Reasoning in Vision Language Models

TL;DR

The paper proposes spatiotemporal IDs as a simple, linear mechanism for spatiotemporal reasoning in Vision Language Models. Through causal interventions, it demonstrates that these IDs mediate intermediate and final beliefs, and shows how they can be extracted and used to diagnose and improve SoTA VLMs. The approach extends to video models, uncovering temporal IDs that parallel spatial IDs, and offers a Spatial Loss to guide finetuning for stronger, more interpretable models. Overall, the work advances interpretability and principled design of more aligned VLMs by revealing a linear, causally influential internal reasoning pathway.

Abstract

Spatio-temporal reasoning is a remarkable capability of Vision Language Models (VLMs), but the underlying mechanisms of such abilities remain largely opaque. We postulate that visual/geometrical and textual representations of spatial structure must be combined at some point in VLM computations. We search for such confluence, and ask whether the identified representation can causally explain aspects of input-output model behavior through a linear model. We show empirically that VLMs encode object locations by linearly binding \textit{spatial IDs} to textual activations, then perform reasoning via language tokens. Through rigorous causal interventions we demonstrate that these IDs, which are ubiquitous across the model, can systematically mediate model beliefs at intermediate VLM layers. Additionally, we find that spatial IDs serve as a diagnostic tool for identifying limitations in existing VLMs, and as a valuable learning signal. We extend our analysis to video VLMs and identify an analogous linear temporal ID mechanism. By characterizing our proposed spatiotemporal ID mechanism, we elucidate a previously underexplored internal reasoning process in VLMs, toward improved interpretability and the principled design of more aligned and capable models. We release our code for reproducibility: https://github.com/Raphoo/linear-mech-vlms.

Figures (31)

• Figure 1: Hypothesis for spatiotemporal visual reasoning. The VLM linearly binds spatiotemporal localization to object word activations in early layers. Subsequent linguistic reasoning about the object is informed by its location in space and time per the spatiotemporal ID.
• Figure 2: Results from Targeted Intervention (§ \ref{['sec:causality']}). Median binary belief swap due to spatial ID steering is 64.4%, and 29.5% for noise. Spatial IDs have 43.6% above-chance influence on average. We conclude that spatial IDs mediate models' beliefs about objects' locations in space.
• Figure 3: Mirror swapping experiment (§ \ref{['sec information flow']}). Activations from case 1 and 2 are partially swapped at a select layer, in one of three arrangements. Computations continue normally after this point.
• Figure 4: Ratio change in log probability for logits "left" and "right" from mirror swap (A) and attribute swap (B) interventions. (A) shows distinct binary belief swaps, where text tokens have an influence after middle layers. Image patches stop having an influence after that point, and object word tokens only have an influence in these middle layers. The control, (B), is noisy.
• Figure 5: Spatial IDs in a grid. Color and saturation of markers represent the location of the object when spatial ID was extracted. x and y axes are coefficients of ID projections onto $h_L$ and $v_L$. L, R represent "left", "right" textual activations.