Direct Visual Grounding by Directing Attention of Visual Tokens

TL;DR

This work addresses the misalignment between visual tokens and language tokens in vision-language models, where visually relevant regions receive weak attention during answer generation. It introduces KL Attention Loss ($L_{KLAL}$), an auxiliary loss that directly supervises the attention distribution of visual tokens to the final answer, using ground-truth attention maps $P(\\mathcal{S})$ derived automatically from task geometry or grounding annotations; the total objective is $L_{total} = L_{NTP} + \lambda L_{KLAL}$ with $\\lambda=1$. Empirically, KLAL+NTP yields consistent gains across geometric reasoning tasks (Line Intersection, Line Tracing), pointing and grounding datasets (Grid Patch, PixMo-Points), and the REC benchmark RefCOCO, while also enhancing attention interpretability and the embedding quality of visual tokens. The approach is model-agnostic, requires no task-specific heads, and demonstrates strong cross-dataset transfer, suggesting practical impact for robust visual grounding in real-world VLM applications. The paper also introduces a Line Tracing dataset to stress-test connectivity reasoning in VLMs and shows that even commercial VLMs benefit from targeted attention supervision.

Abstract

Vision Language Models (VLMs) mix visual tokens and text tokens. A puzzling issue is the fact that visual tokens most related to the query receive little to no attention in the final layers of the LLM module of VLMs from the answer tokens, where all tokens are treated equally, in particular, visual and language tokens in the LLM attention layers. This fact may result in wrong answers to visual questions, as our experimental results confirm. It appears that the standard next-token prediction (NTP) loss provides an insufficient signal for directing attention to visual tokens. We hypothesize that a more direct supervision of the attention of visual tokens to corresponding language tokens in the LLM module of VLMs will lead to improved performance on visual tasks. To demonstrate that this is indeed the case, we propose a novel loss function that directly supervises the attention of visual tokens. It directly grounds the answer language tokens in images by directing their attention to the relevant visual tokens. This is achieved by aligning the attention distribution of visual tokens to ground truth attention maps with KL divergence. The ground truth attention maps are obtained from task geometry in synthetic cases or from standard grounding annotations (e.g., bounding boxes or point annotations) in real images, and are used inside the LLM for attention supervision without requiring new labels. The obtained KL attention loss (KLAL) when combined with NTP encourages VLMs to attend to relevant visual tokens while generating answer tokens. This results in notable improvements across geometric tasks, pointing, and referring expression comprehension on both synthetic and real-world data, as demonstrated by our experiments. We also introduce a new dataset to evaluate the line tracing abilities of VLMs. Surprisingly, even commercial VLMs do not perform well on this task.

Figures (5)

• Figure 1: Processing flow in VLMs
• Figure 2: Visualization of the ground truth map for the attention of answer tokens. Yellow highlights the target patches with the highest mass of the distribution, while green indicates the surrounding neighbors. The other patches have a fixed low value.
• Figure 3: Attention maps show the attention of the last answer token to visual tokens. The results in the first two rows are from LLava-v1.5, and the last two rows show results from Qwen2.5-VL. First column shows input images and tasks. Second column shows the attention maps of out-of-the-box models. Third column shows the NTP-finetuned model. Fourth column shows the attention maps of models finetuned with the proposed NTP+KLAL. The red text indicates wrong answers and the green the correct ones. In the third row, the green box indicates the ground truth patch, and the red box denotes the predicted patch. Not only are the answers produced by NTP + KLAL correct, but the attention maps are also dramatically improved and become much more interpretable.
• Figure 4: The bar plots illustrate the ratio of the average attention of target tokens w.r.t the answer token to the average attention of all visual tokens computed over all Line Intersection test images.
• Figure 5: The first row shows attention maps from the last output token of LLava-v1.5, while the second row presents a visualization derived from the magnitudes of visual token embeddings. Both types of visualizations reveal similar patterns, as enhanced attention through NTP + KLAL corresponds to increased focus of high-norm tokens on the target regions.