Attention in Space: Functional Roles of VLM Heads for Spatial Reasoning

TL;DR

This paper investigates how attention heads in Vision-Language Models contribute to spatial reasoning, introducing a new dataset and framework to analyze and improve the models' understanding of space.

Contribution

It presents CogVSR, a dataset for decomposing spatial reasoning, and a probing framework to identify and enhance spatially specialized attention heads in VLMs.

Findings

Spatially specialized heads are fewer than other functions.

Removing functional heads degrades spatial reasoning performance.

Activating spatial heads improves model accuracy.

Abstract

Despite remarkable advances in large Vision-Language Models (VLMs), spatial reasoning remains a persistent challenge. In this work, we investigate how attention heads within VLMs contribute to spatial reasoning by analyzing their functional roles through a mechanistic interpretability lens. We introduce CogVSR, a dataset that decomposes complex spatial reasoning questions into step-by-step subquestions designed to simulate human-like reasoning via a chain-of-thought paradigm, with each subquestion linked to specific cognitive functions such as spatial perception or relational reasoning. Building on CogVSR, we develop a probing framework to identify and characterize attention heads specialized for these functions. Our analysis across diverse VLM families reveals that these functional heads are universally sparse, vary in number and distribution across functions. Notably, spatially specialized heads are fewer than those for other cognitive functions, highlighting their scarcity. We propose methods to activate latent spatial heads, improving spatial understanding. Intervention experiments further demonstrate their critical role in spatial reasoning: removing functional heads leads to performance degradation, while emphasizing them enhances accuracy. This study provides new interpretability driven insights into how VLMs attend to space and paves the way for enhancing complex spatial reasoning in multimodal models.