💻 computer science

Direct Visual Grounding by Directing Attention of Visual Tokens

This paper proposes a novel KL attention loss (KLAL) that directly supervises the attention of visual tokens to relevant language tokens in Vision Language Models, thereby improving performance on visual grounding tasks by addressing the insufficient attention signals provided by standard next-token prediction.

Original authors: Parsa Esmaeilkhani, Longin Jan Latecki

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Parsa Esmaeilkhani, Longin Jan Latecki

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers can look at a picture and talk about it, just like you and I chat about a photo on our phones. This is the realm of Vision-Language Models (VLMs), a branch of artificial intelligence that tries to bridge the gap between what a camera sees and what a human says. Think of these models as super-smart translators that take an image, break it down into tiny digital puzzle pieces called "visual tokens," and mix them with words to answer questions like "What color is the cat?" or "How many apples are there?"

For a long time, scientists thought these models were getting really good at this. But there was a nagging mystery: sometimes, even when the answer was right, the computer seemed to be ignoring the picture entirely. It was as if a student taking a test was reading the question, glancing at the diagram, and then answering based entirely on what they remembered from a textbook, completely forgetting to look at the specific lines and shapes right in front of them. The big question for researchers was: Why does the computer's "brain" stop paying attention to the visual clues when it's time to give the final answer? This paper dives into that exact problem, trying to figure out if we can teach these AI models to actually look at the right parts of the image when they speak.

The Problem: The "Ghost" in the Machine

Meet the Vision-Language Model. You can think of it as a very chatty robot that has read a million books and looked at a million photos. When you ask it a question about an image, it breaks the image down into a grid of tiny squares (visual tokens) and mixes them with the words you typed.

Here's the weird part: The paper found that in the final moments of thinking, right before the robot types its answer, it often stops looking at the picture. It's like a chef who has all the ingredients on the counter but, when it's time to plate the dish, decides to ignore the food and just guess what it should taste like based on a recipe they memorized. The researchers discovered that the "attention" the computer gives to the most important parts of the image (like the intersection of two lines or a specific object) is almost zero. It's as if the visual clues are ghosts—present in the system, but invisible to the part of the brain that writes the answer.

The paper argues that the standard way these models are trained (called "Next-Token Prediction") is like a teacher who only grades the final essay but never checks if the student actually looked at the reference photos. The computer learns to guess the next word correctly, but it doesn't learn why it's the right word or which part of the image made it right.

The Solution: A "Laser Pointer" for Attention

To fix this, the authors came up with a clever new training trick called KL Attention Loss (KLAL).

Imagine you are teaching a dog to fetch a ball. In the old way, you'd just say "Good dog!" when it brings the ball back, but the dog might have just been running around randomly and happened to grab the ball by luck. In the new way proposed by this paper, you use a laser pointer. Every time the dog looks at the ball, you shine the light on it and say, "Yes! Look right there!"

In the computer world, the "laser pointer" is a special mathematical signal that forces the model to focus its attention on the specific pixels of the image that matter for the answer.

If the question is "How many times do these lines cross?", the laser pointer shines on the exact spots where the lines intersect.
If the question is "Is this node connected to that one?", the laser pointer traces the path of the line connecting them.

The model is then punished (via a loss function) if it looks anywhere else. It's like telling the robot: "You can't just guess the answer; you have to prove you're looking at the right spot in the picture."

What They Found: From Guessing to Seeing

The researchers tested this idea on some tricky tasks that require precise visual reasoning, like counting how many times two squiggly lines cross each other or tracing a path through a complex graph to see if two points are connected. They even made up a new game called "Line Tracing" where the computer has to follow a path through a maze of dots.

Here is what happened when they turned on the "laser pointer" (KLAL):

Better Answers: The models got significantly better at the tasks. For example, on the "Line Intersection" task, one model improved its accuracy from about 47% to over 70% when they added this new training method. On a "Pointing" task (where the model has to say exactly where an object is), the improvement was massive, jumping from a near-random guess to getting it right nearly half the time.
Better Focus: The most exciting part wasn't just the score; it was how the model thought. When the researchers looked at the model's "attention maps" (which show where the computer is looking), they saw a dramatic change. Before the new training, the model was looking at random parts of the image. After the training, the model's attention was tightly locked onto the intersection points or the specific object, just like a human would look.
Stronger Memory: The paper also found that this training didn't just change where the model looked; it actually made the computer's internal memory of the image stronger. The digital "weights" of the important image parts became heavier and more distinct, meaning the model truly "understood" the visual data better, not just the words.

The Verdict

The paper suggests that the reason these AI models sometimes fail at simple visual tasks isn't because they aren't smart enough, but because they aren't being taught to look properly. By adding a direct "laser pointer" signal during training, they forced the models to connect their words directly to the pixels they represent.

The results show that this method works across different types of models, from open-source ones to some of the biggest commercial AI systems. Even the most advanced commercial models struggled with these specific geometric tasks until this method was applied. The authors conclude that if we want AI to truly understand the world, we can't just let it guess; we have to teach it to pay attention to the right things, just like a teacher guiding a student's eyes to the right part of a diagram.

In short, the paper proves that when you force an AI to look at the evidence before it speaks, it stops hallucinating and starts actually seeing.

Technical Summary: Direct Visual Grounding by Directing Attention of Visual Tokens

Problem Statement
Vision-Language Models (VLMs) integrate visual and textual tokens within Large Language Model (LLM) modules, yet they frequently fail at tasks requiring precise visual reasoning, such as geometric inference, object counting, and spatial relationship verification. A critical issue identified is that in the final layers of the LLM, answer tokens allocate little to no attention to the visual tokens most relevant to the query. Instead, the model often relies on language priors, leading to hallucinations or incorrect answers even when visual information is present. The standard Next-Token Prediction (NTP) loss provides an insufficient signal to direct the model's attention toward specific visual regions, resulting in a scenario where visual tokens are effectively ignored or treated equally with irrelevant tokens.

Methodology
The authors propose a novel training framework that introduces KL Attention Loss (KLAL) to directly supervise the attention mechanism of VLMs. The core hypothesis is that explicitly aligning the attention distribution of language answer tokens with ground truth (GT) visual regions improves performance.

Attention Supervision: The method computes the attention distribution of visual tokens with respect to the final generated answer token. This distribution is treated as a probability vector.
Ground Truth Construction: GT attention maps are constructed automatically without requiring new manual labels:
- Synthetic Tasks: Maps are derived from task geometry (e.g., patches containing line intersections or paths connecting nodes).
- Real Images: Maps are projected from standard annotations (bounding boxes or point annotations) using smoothing functions.
Loss Function: The proposed KLAL minimizes the Kullback-Leibler (KL) divergence between the model's predicted attention distribution and the constructed GT distribution across all LLM layers.
Training Objective: The total loss combines the standard NTP loss with the KLAL: $L_{total} = L_{NTP} + \lambda L_{KLAL}$ . This approach is model-agnostic, requires no architectural modifications, and does not necessitate additional task-specific heads (e.g., for segmentation or localization).

Key Contributions

Novel Loss Function: Introduction of KLAL, which directly grounds answer language tokens to relevant visual tokens by aligning attention distributions via KL divergence.
Empirical Evidence: Demonstration that improving attention to relevant image regions directly correlates with increased performance on visual tasks, offering better explainability through attention maps.
Embedding Enhancement: Discovery that the method improves the deep embeddings of visual tokens, causing high-norm tokens to concentrate on image regions most relevant to the correct answer.
New Dataset: Creation of a synthetic Line Tracing dataset to evaluate VLMs' ability to trace paths and determine connectivity between nodes in complex graphs, a task where even commercial models struggle.
Comprehensive Evaluation: Extensive benchmarking of open-source and commercial VLMs on geometric reasoning, pointing, and referring expression comprehension (REC) tasks.

Experimental Results
The method was evaluated on five datasets: Line Intersection, Line Tracing, Grid Patch pointing, PixMo-Points, and RefCOCO.

Geometric Tasks: On the Line Intersection and Line Tracing tasks, adding KLAL to NTP finetuning significantly outperformed base models and NTP-only finetuning. For instance, Qwen2.5-VL improved from 47.62% to 70.23% on Line Intersection and from 49.62% to 62.21% on Line Tracing.
Pointing Tasks: The approach yielded substantial gains in coordinate localization. Qwen2.5-VL achieved 44.90% accuracy on Grid Patch and 35.77% on PixMo-Points with NTP+KLAL, surpassing both base models and other state-of-the-art baselines like GPT-4o and Gemini-2.0 Flash.
Generalization: Cross-dataset transfer experiments showed that NTP+KLAL enabled effective transfer between pointing datasets, whereas NTP alone failed to generalize.
Real-World Grounding: On the RefCOCO dataset, the method provided measurable accuracy improvements (e.g., 91.45% vs. 90.05% on validation) even for models that already performed strongly.
Attention Visualization: Attention maps confirmed that models trained with KLAL focused significantly more on target regions (e.g., intersection points) compared to base models, where attention was diffuse or misaligned.

Significance and Claims
The paper claims that the primary limitation in current VLM performance on visual tasks is the lack of direct supervision for visual attention within the LLM module. By drawing an analogy to early childhood learning—where direct supervision links visual dominance to word learning—the authors argue that their approach transforms passive exposure into active learning. The significance of this work lies in demonstrating that:

Standard NTP is insufficient for grounding visual tokens.
Direct attention supervision via KL divergence is a simple, effective, and model-agnostic solution that improves both answer accuracy and the interpretability of model decisions.
Even advanced commercial VLMs exhibit fundamental weaknesses in geometric and connectivity reasoning, which can be mitigated through this targeted attention supervision.

The Problem: The "Ghost" in the Machine

The Solution: A "Laser Pointer" for Attention

What They Found: From Guessing to Seeing

The Verdict

Technical Summary: Direct Visual Grounding by Directing Attention of Visual Tokens

More like this