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

This paper introduces the CogVSR dataset and a probing framework to identify sparse, functionally specialized attention heads in Vision-Language Models, demonstrating that targeting these heads through intervention significantly improves spatial reasoning capabilities.

The Gist

Imagine a Vision-Language Model (VLM) as a super-smart, multi-talented detective who can look at a photo and answer questions about it. You might ask, "Is the dog looking at the horse?"

For a long time, these detectives were great at naming things ("That's a dog!") but terrible at understanding where things are or how they relate to each other in space. They often got simple spatial questions wrong.

This paper, "Attention in Space," is like a team of neuroscientists putting the detective under a microscope to figure out why they struggle with space and how to fix it.

Here is the story of their discovery, broken down into simple parts:

1. The Problem: The Detective's "Brain Fog"

The researchers noticed that while the detective could describe a scene, it couldn't reliably tell if the dog was facing the horse. It's like having a detective who can list every item in a room but can't tell you if the chair is in front of or behind the table.

2. The New Tool: "CogVSR" (The Step-by-Step Recipe)

To understand the detective's brain, the researchers created a new dataset called CogVSR.

• The Analogy: Imagine you ask a human, "Is the dog facing the horse?" A human doesn't just guess. They break it down:
  1. What do I see? (A dog and a horse).
  2. Where are they? (Dog on the right, horse on the left).
  3. Which way is the dog looking? (Left).
  4. Does "left" point to the horse? (Yes).
  5. Conclusion: The statement is true.
• The Innovation: The researchers forced the AI to do this exact same thing. They broke complex questions into tiny, step-by-step "sub-questions," each requiring a specific mental skill (like "Spatial Perception" or "Relational Reasoning"). This is like giving the detective a checklist to follow.

3. The Discovery: Finding the "Specialist Cells"

Inside these AI models, there are thousands of tiny processing units called Attention Heads. Think of the model's brain as a massive office building with thousands of employees (the heads).

• The Finding: The researchers found that most employees are generalists, but a few are specialists.
• The "Space" Specialists: They discovered that there are specific employees whose only job is to understand space and geometry.
• The Bad News: These "Space Specialists" are extremely rare. In fact, there are far fewer of them than there are employees who just recognize objects or read text. It's like an office where you have 100 people who can read, but only 2 people who know how to navigate a map. This scarcity is why the AI struggles with spatial reasoning.

4. The Proof: The "Silence" Experiment

To prove these specialists were real, the researchers played a game of "Silence."

• The Experiment: They temporarily "muted" (turned off) the specific attention heads they thought were the space experts.
• The Result: The detective immediately went blind to space. It couldn't tell left from right anymore.
• The Control: When they muted random employees instead, the detective barely noticed. This proved that the "Space Specialists" were the only ones doing the heavy lifting for spatial tasks.

5. The Solution: Waking Up the Sleeping Giants

Since the AI has these specialists but they are too quiet or too few, the researchers asked: Can we wake them up?

• The Trick (Spatial Head Activation): They gave the detective a visual aid. Before asking the question, they drew boxes around the objects in the image (like highlighting the dog and the horse) and masked out the background noise.
• The Result: This forced the model to pay attention to the objects and their positions rather than just the general picture.
• The Outcome: Suddenly, the "Space Specialists" woke up and started working harder. The AI's accuracy on spatial tasks jumped by 10% or more, without needing to retrain the whole model from scratch.

The Big Picture

This paper teaches us two main things:

1. AI has a "Space Deficit": Current AI models are built with very few "brain cells" dedicated to understanding space, which is why they fail at these tasks.
2. We can fix it without rebuilding: By understanding how the AI thinks (interpreting its attention heads) and giving it better visual cues, we can unlock its hidden potential to understand the world around it.

In short: The researchers found the AI's "spatial brain," realized it was under-staffed, and figured out how to give it a caffeine shot so it could finally understand where things are in the world.

Technical Summary

1. Problem Statement

Despite significant advancements in Vision-Language Models (VLMs), spatial reasoning remains a critical bottleneck. While VLMs excel at tasks like image captioning and object detection, they frequently fail at simple spatial queries (e.g., "Is the dog facing the horse?").

• The Gap: Existing research has identified sparse attention heads responsible for visual grounding (aligning text to image regions). However, it remains unclear how VLMs internally coordinate these mechanisms to perform complex, multi-step spatial reasoning that requires integrating perception, relational logic, and decision-making.
• The Hypothesis: Similar to the human brain, which uses distinct regions for visual processing, spatial mapping, and relational reasoning, VLMs likely possess specialized attention heads dedicated to specific cognitive functions. The scarcity or misalignment of these "spatial heads" may explain current model failures.

2. Methodology

The authors propose a mechanistic interpretability framework centered on three core components:

A. CogVSR: A Cognitively Grounded Benchmark

To analyze spatial reasoning granularly, the authors introduced CogVSR, a dataset that decomposes complex spatial questions into step-by-step sub-questions using a Chain-of-Thought (CoT) paradigm.

• Structure: 1,142 main questions broken down into 3,759 sub-questions.
• Cognitive Taxonomy: Each sub-question is annotated with one of eight cognitive functions:
  1. Spatial Perception: Understanding positions, orientations, and geometric relationships.
  2. Relational Reasoning: Comparing relationships between entities.
  3. High/Low-level Visual Perception: Object recognition and feature extraction.
  4. Language Information Extraction: Retrieving context.
  5. Knowledge Recall: Accessing prior knowledge.
  6. Math Reasoning: Counting and arithmetic.
  7. Decision Making: Selecting the final answer.
• Quality Control: A rigorous two-stage human verification pipeline ensures logical consistency and accurate functional labeling.

B. Probing Framework for Head Identification

The authors developed a framework to identify which attention heads are responsible for specific cognitive functions:

1. Feature Extraction: For each sub-question, they extract attention head activations (value vectors projected into the residual stream) across all layers. They focus on the top-$k$ most informative tokens in the generated answer.
2. Classification: A multi-label classifier (MLP) is trained to predict the cognitive function based on these head activations.
3. Importance Scoring: Using a gradient-based attribution method ($\text{gradient} \times \text{activation}$), they compute an importance score for every attention head regarding each cognitive function.

C. Intervention Strategies

To validate the causal role of identified heads, the authors performed two types of interventions:

• Negative Intervention (Ablation): Scaling down the output of identified "cognitive heads" to near-zero.
• Positive Intervention (Activation): Shifting the activation of heads along the direction of a specific function (calculated as the difference between correct and incorrect answer activations) to "boost" that capability.
• Spatial Head Activation (SHA): A novel method to activate latent spatial heads by providing explicit object priors (bounding boxes and masks) to reduce reliance on high-level visual cues, forcing the model to engage spatial reasoning mechanisms.

3. Key Contributions

1. CogVSR Dataset: The first benchmark to systematically disentangle spatial reasoning into interpretable cognitive sub-processes, enabling fine-grained analysis of VLM internals.
2. Discovery of Functional Specialization: Demonstrated that VLMs contain sparse, universal, and intrinsic attention heads dedicated to specific cognitive functions (e.g., specific heads for spatial perception vs. math reasoning).
3. Identification of Spatial Scarcity: Found that spatially specialized heads are significantly scarcer than those for other functions (like information extraction), providing a mechanistic explanation for why VLMs struggle with spatial tasks.
4. Activation Methods: Proposed Spatial Head Activation (SHA) and positive intervention strategies that improve spatial reasoning performance without retraining the model.

4. Experimental Results

The study was conducted across three major VLM families (Intern, Qwen, Llama) with varying scales (2B to 90B parameters).

• Sparsity and Universality: Heatmaps revealed that fewer than 9% of all attention heads have significant importance scores across the eight functions. This sparse organization is consistent across different model architectures and scales.
• Scarcity of Spatial Heads: Spatial perception and relational reasoning heads were found to be the least abundant compared to other functions, correlating with lower performance on spatial benchmarks.
• Intervention Impact:
  • Ablation: Masking identified cognitive heads caused a drastic performance drop (e.g., accuracy falling below 20% in some cases), whereas masking random heads had minimal effect.
  • Activation: Applying Spatial Head Activation (SHA) (using bounding boxes/masks) improved accuracy on spatial perception and relational reasoning tasks by ~10% for InternVL3-2B and ~5% for Llama3.2-90B-Vision.
  • Positive Intervention: Shifting activations toward spatial functions on downstream benchmarks (VSR, SpatialEval, 3DSRBench) consistently improved LLM-Judge accuracy.

5. Significance and Conclusion

This work provides a mechanistic explanation for the limitations of VLMs in spatial reasoning. By moving beyond token-level analysis to head-level functional mapping, the authors reveal that the bottleneck is not a lack of capacity, but an imbalance in the allocation of attention resources toward spatial functions.

• Theoretical Impact: It bridges the gap between cognitive science (human brain specialization) and deep learning, showing that VLMs possess analogous functional modules.
• Practical Impact: The proposed Spatial Head Activation (SHA) offers a training-free method to enhance spatial reasoning in existing models. This suggests that future VLM architectures should explicitly prioritize the development and allocation of spatially specialized attention heads to achieve robust spatial understanding.

In summary, the paper argues that improving VLM spatial reasoning requires not just more data, but a structural re-balancing of the model's internal attention mechanisms to better support the "scarcity" of spatial cognition.