Interpreting Attention Mechanisms in Vision-Language Models for Spatial Reasoning

This repository provides code and instructions for analyzing and visualizing attention mechanisms in Vision-Language Models (VLMs) with a focus on spatial relational reasoning. We conduct experiments on both multimodal (vision + language) and text-only settings to study cross-attention, self-attention, and head/layer specialization.

---

Installation

We evaluate two models:

• LLaVA-v1.5-7B (Hugging Face)
• Qwen2-VL-7B-Instruct (Hugging Face)

Due to differences in dependencies, we recommend setting up separate virtual environments for the two models.

Installation for LLaVA

1. Download the modified LLaVA code from this repository and place it under models/llava/.

cd ..
git clone https://github.com/zjysteven/VLM-Visualizer.git
cp -r VLM-Visualizer/models/llava/ VLM-reasoning/models/
cd VLM-reasoning/

Note: This version is adapted from LLaVA v1.5 to support exporting attention maps from the CLIP vision encoder.

2. Modify the following line in models/llava/model/multimodal_encoder/clip_encoder.py (around line 30):

- self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name, device_map=device_map)
+ self.vision_tower = CLIPVisionModel.from_pretrained(
+     self.vision_tower_name,
+     device_map=device_map,
+     attn_implementation="eager",
+     torch_dtype=torch.bfloat16
+ )

3. Follow the instructions in env_setup_llava.bash to create and activate the LLaVA virtual environment.

Installation for Qwen2-VL

Here we describe how to set up a minimal Python environment for running Qwen2-VL attention analysis using the 🤗 Transformers library.

Requirements

• Python 3.9+
• CUDA-enabled GPU (recommended)
• PyTorch compatible with your CUDA version

Install Dependencies

Qwen2-VL requires a recent version of transformers.

pip install "transformers>=4.57.0"

Install additional commonly used dependencies:

pip install torch numpy pillow opencv-python matplotlib

⚠️ Make sure your installed PyTorch version matches your CUDA setup. Refer to: https://pytorch.org/get-started/locally/

---

Data Preparation

Download the following two datasets:

• ShapeWorld-based Dataset: https://polybox.ethz.ch/index.php/s/6gN7q5LqbpczGdJ, visit https://github.com/YimingZhao-art/ShapeWorld to reproduce the dataset.
• Fixed-position Counterfactual Dataset: https://polybox.ethz.ch/index.php/s/pjmxGEPZnEXQTjR, or run bash generate_corner_datasets.sh to reproduce the dataset.

After downloading, place them under the data/ directory. The directory structure should look like:

data/spatial_twoshapes/agreement/relational/shard0/world-0.png
data/dataset_topleft/images/pair_00000_control.png

---

Running Experiments

In the LLaVA Environment

1. Cross-attention and self-attention analysis:

In vlm_atten_analysis_llava.py, set the data storage path at line 581. By default, it is BASE = f"data/spatial_twoshapes/agreement/relational/shard{shard_id}".

python vlm_atten_analysis_llava.py

2. Text-only attention analysis:

In text_only_llm_analysis_llava.py, set the data storage path at line 36. By default, it is CAPTION_BASE = "/home/maqima/VLM-Visualizer/data/spatial_twoshapes/agreement/relational/".

python text_only_llm_analysis_llava.py

In the Qwen2-VL Environment

1. Cross-attention and self-attention analysis: In vlm_atten_analysis_qwen2.py, set the data storage path at line 20. By default, it is BASE_ROOT = "data/spatial_twoshapes/agreement/relational".

python vlm_atten_analysis_qwen2.py

2. Text-only attention analysis: In text_only_llm_analysis_qwen2.py, set the data storage path at line 32. By default, it is BASE_DIR = "data/spatial_twoshapes/agreement/relational".

python text_only_llm_analysis_qwen2.py

3. Head and layer specialization analysis:

bash run_corner_experiments_qwen2.sh

---

Acknowledgements

• VLM-Visualizer (GitHub) A visualization toolkit for inspecting attention maps in LLaVA.

• LLaVA (GitHub) The official implementation of the LLaVA model.

• attention (GitHub) Code for attention aggregation in large language models, which this project heavily builds upon.