💬 NLP

Attention Grounded Enhancement for Visual Document Retrieval

The paper proposes AGREE, a framework that leverages cross-modal attention from multimodal large language models as proxy supervision to guide visual document retrievers in learning fine-grained, region-level relevance, thereby significantly improving performance on complex, non-extractive queries compared to global-supervision-only baselines.

Original authors: Wanqing Cui, Wei Huang, Yazhi Guo, Yibo Hu, Meiguang Jin, Junfeng Ma, Keping Bi

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Wanqing Cui, Wei Huang, Yazhi Guo, Yibo Hu, Meiguang Jin, Junfeng Ma, Keping Bi

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 you are trying to find a specific page in a massive, messy library of documents. Some pages are just text, but many are complex "visual documents" filled with charts, tables, photos, and text all mixed together.

The Problem: The "Gist" vs. The "Details"
Current search engines for these documents are like a librarian who only reads the summary on the back of a book. They can tell you, "Yes, this book is about the topic you asked about," but they don't know which specific paragraph or which chart holds the answer.

Because they only look at the "big picture" (global relevance), they often get tricked. If you ask a tricky question that requires connecting two ideas that aren't next to each other (like linking the word "bottleneck" to a hidden symbol in a diagram), the librarian might miss it entirely. They rely too much on finding exact keyword matches, like a dog chasing a squeaky toy, rather than understanding the actual meaning.

The Solution: AGREE (The "Super-Teacher" Librarian)
The authors of this paper propose a new training method called AGREE (Attention-Grounded REtriever Enhancement).

Think of AGREE as hiring a super-smart, hyper-observant teacher (a Multimodal Large Language Model, or MLLM) to train the librarian.

Here is how the training works, step-by-step:

The Teacher's "Gaze": When the teacher sees a question and a document, they don't just say "Yes, this is relevant." They point their finger at the exact spots on the page that matter.
- Analogy: Imagine the teacher is using a laser pointer. If you ask, "Where is the hidden symbol?" the teacher doesn't just nod; they shine the laser on the tiny symbol, even if it's small, and also on the nearby text that explains it. They highlight both the obvious matches and the subtle, hidden connections.
The "Heat Map" Lesson: The teacher creates a "heat map" (a visual guide showing where they are looking). This map tells the librarian: "Pay attention to this specific corner of the chart, and this specific word in the table."
The Training: The librarian (the search engine) is then forced to learn from this heat map. Instead of just learning "Book A is a match," the librarian learns, "To find the answer, I must look specifically at the top-right corner and the bottom-left text."
The Result: The librarian stops guessing based on the whole book and starts understanding why a page is relevant. They learn to spot the "invisible" clues that connect the question to the answer.

Why This Matters
The paper tested this on a very difficult benchmark called ViDoRe V2, which is full of tricky questions that require reasoning, not just keyword matching.

Before AGREE: The librarian was like a student who memorized the table of contents but couldn't find the specific facts inside.
After AGREE: The librarian became a detective. They could find the answer even when the question used different words than the document (e.g., asking about "gay customers" and finding the answer under "LGBT" in the text).

The Key Takeaway
The paper claims that by using this "teacher's gaze" (attention maps) to guide the search engine, the system becomes much better at finding the right information. It doesn't just know that a document is relevant; it knows where the relevance is hiding.

In simple terms: AGREE teaches the search engine to stop skimming the cover and start reading the fine print, using a super-smart AI teacher to show it exactly where to look.

The paper notes that this method works significantly better than previous methods, especially for complex questions, and the code is available for others to try. It does not claim to be used for medical diagnosis or other specific real-world applications beyond document search and retrieval.

Technical Summary: Attention Grounded Enhancement for Visual Document Retrieval (AGREE)

Problem Statement

Visual document retrieval (VDR) aims to retrieve document pages semantically relevant to a given text query. Unlike traditional image retrieval, VDR involves heterogeneous content (text, tables, figures, layouts) with high information density. While recent state-of-the-art (SOTA) methods, such as ColQwen, utilize screenshot-based encoding with fine-grained late interaction (MaxSim) to capture token-level and patch-level alignments, they face a critical limitation: training supervision is coarse.

Current retrievers are trained using only global binary relevance labels (relevant vs. non-relevant). These labels indicate if a page matches a query but provide no insight into which specific regions support the match. Consequently, retrievers tend to rely on surface-level cues (e.g., exact keyword matching) and struggle to capture implicit semantic connections or non-extractive queries (e.g., reasoning about "autoencoder" vs. "encoder-decoder" structures). This hinders performance on complex, real-world retrieval tasks where evidence is distributed across visual and textual elements.

Methodology: AGREE Framework

To address the lack of fine-grained supervision without incurring the cost of manual annotation, the authors propose AGREE (Attention-Grounded REtriever Enhancement). AGREE leverages the cross-modal attention mechanisms of pre-trained Multimodal Large Language Models (MLLMs) as a proxy for fine-grained supervision. The framework operates in three key stages:

1. MLLM Attention Annotation

AGREE employs a frozen, pre-trained MLLM (specifically Qwen2.5-VL) to generate query-conditioned, patch-level attention maps.

Extraction Strategy: The authors compare two strategies: Answer-Token Attention (extracting attention from the token preceding the answer generation) and Query-Token Attention (averaging attention from all query tokens to image patches). Experiments show that Query-Token Attention yields better alignment with human judgments, covering both explicit and implicit matches more robustly.
Normalization: To emphasize semantic relevance over global aggregation, the query-specific attention is normalized against a general image description attention.
Aggregation: Attention maps are aggregated across all layers of the MLLM to form a unified supervision signal.

2. Spatial-Preserving Attention Downsampling

A resolution mismatch exists between the high-resolution inputs required for MLLM reasoning (up to 2048 patches) and the lower-resolution inputs used by efficient retrievers (e.g., 768 or 1024 patches).

Solution: AGREE employs adaptive max pooling to downsample the high-resolution attention maps. Unlike average pooling, which dilutes salient peaks, max pooling preserves the highest attention values within each window, ensuring that key regions highlighted by the MLLM remain prominent in the lower-resolution grid used by the retriever.

3. Attention-Guided Training

The retriever is trained with a dual-objective function:

Global Loss ( $L_{global}$ ): A standard contrastive loss using in-batch hard negatives to maintain page-level relevance.
Local Alignment Loss ( $L_{local}$ ): Applied only to positive query-page pairs, this loss encourages the retriever's patch-level similarity scores to mimic the downsampled MLLM attention map.
- The authors evaluate three alignment strategies: Kullback–Leibler (KL) divergence, Top-K salient contrastive, and Cosine Similarity.
- Cosine Similarity is found to be the most effective, as it focuses on the directional agreement of score patterns (prioritizing salient regions) without being overly penalized by noise in low-salience background areas or requiring manual threshold tuning.

The total loss is defined as $L_{total} = L_{global} + \lambda \cdot L_{local}$ . Crucially, this fine-grained supervision is used only during training; at inference, the retriever operates independently without requiring the MLLM or external attention signals.

Key Contributions

Novel Training Framework: AGREE is the first framework to distill query-conditioned MLLM attention into late-interaction visual document retrievers, providing scalable, fine-grained supervision for implicit relevance matching without manual annotation.
Performance Gains: Extensive experiments demonstrate that AGREE significantly improves retrieval performance on challenging, non-extractive tasks, proving effective across different backbone models (PaliGemma and Qwen2.5-VL).
Interpretability and Reasoning: The framework shifts retrieval behavior from surface-level keyword matching to semantically grounded, rationale-aware alignment, enabling models to capture implicit query-document correspondences.

Experimental Results

The authors evaluated AGREE on the ViDoRe V2 benchmark, a challenging out-of-domain dataset requiring semantic reasoning beyond keyword overlap.

Performance: Compared to the global-supervision-only baseline (ColQwen2.5), AGREE achieved:
- +12.82% improvement in average nDCG@1 (from 0.5481 to 0.6184).
- +5.03% improvement in average nDCG@5 (from 0.5859 to 0.6154).
Robustness: On the easier ViDoRe V1 benchmark (driven by explicit textual overlap), AGREE remained competitive, confirming that fine-grained supervision does not degrade performance on explicit tasks.
Qualitative Analysis:
- Coverage: AGREE-trained retrievers showed significantly higher coverage of human-annotated relevant regions (both explicit and implicit) compared to baselines.
- Case Studies: Visualizations demonstrated that AGREE successfully highlights implicit matches (e.g., linking "gay" to "LGBT" or identifying answer regions not containing the query keyword) where baselines failed or highlighted only surface-level matches.

Significance and Claims

The paper claims that the primary limitation of current visual document retrieval is not the encoding capability but the coarseness of the training signal. By treating the MLLM not as a model to be compressed, but as an annotator that synthesizes missing fine-grained labels, AGREE effectively transfers the teacher's intrinsic reasoning capabilities to the retriever.

The authors emphasize that this approach allows retrievers to generalize beyond simple extractive patterns found in training data. The success of AGREE hinges on the quality of the proxy labels (attention maps) and the balance between global and local objectives. The work suggests that future improvements in visual document retrieval may rely less on larger training corpora and more on precise, scalable supervision mechanisms that guide models to understand why a document is relevant, leading to more interpretable and accurate retrieval systems.