Attention Grounded Enhancement for Visual Document Retrieval

TL;DR

This work addresses the challenge of non-extractive, region-level matching in visual document retrieval, where traditional global relevance signals fail to reveal which content drives a match. It introduces AGREE, a training framework that leverages cross-modal attention from multimodal large language models to generate fine-grained, query-conditioned supervision for retrievers, combining a global contrastive objective with a local alignment loss that grounds patch-level relevance to MLLM attentions. AGREE employs three key components: (i) MLLM attention annotation to produce patch-level signals, (ii) spatial-preserving attention downsampling to align with the retriever's patch grid, and (iii) attention-guided retriever training with a dual objective and several local-loss options (KL, Top-K, Cosine). Empirical results on ViDoRe V2 (and ViDoRe V1) show significant gains in early-ranked retrieval and improved interpretability of similarity maps, demonstrating the method’s effectiveness for implicit and non-extractive queries and its potential to improve grounding in real-world visual document retrieval tasks. The approach is efficient to train, requires no manual annotation, and is broadly compatible with existing backbones, suggesting strong practical impact for retrieval-augmented generation and document understanding systems.

Abstract

Visual document retrieval requires understanding heterogeneous and multi-modal content to satisfy information needs. Recent advances use screenshot-based document encoding with fine-grained late interaction, significantly improving retrieval performance. However, retrievers are still trained with coarse global relevance labels, without revealing which regions support the match. As a result, retrievers tend to rely on surface-level cues and struggle to capture implicit semantic connections, hindering their ability to handle non-extractive queries. To alleviate this problem, we propose a \textbf{A}ttention-\textbf{G}rounded \textbf{RE}triever \textbf{E}nhancement (AGREE) framework. AGREE leverages cross-modal attention from multimodal large language models as proxy local supervision to guide the identification of relevant document regions. During training, AGREE combines local signals with the global signals to jointly optimize the retriever, enabling it to learn not only whether documents match, but also which content drives relevance. Experiments on the challenging ViDoRe V2 benchmark show that AGREE significantly outperforms the global-supervision-only baseline. Quantitative and qualitative analyses further demonstrate that AGREE promotes deeper alignment between query terms and document regions, moving beyond surface-level matching toward more accurate and interpretable retrieval. Our code is available at: https://anonymous.4open.science/r/AGREE-2025.

Figures (8)

• Figure 1: The similarity map of ColQwen2.5 (left), versus the query-to-image attention map from Qwen2.5-VL (right).
• Figure 2: Overview of the AGREE training framework.
• Figure 3: Human annotation (left) and top-3% high attention areas using "query-token attention" (others) given the query "What % of the US workforce are not women?".
• Figure 4: (a) Results on ViDoRe V1 using PaliGemma and Qwen2.5-VL as backbones; (b) Results on ViDoRe V2 with different local alignment losses $\mathcal{L}_{\text{local}}$.
• Figure 5: Coverage of human-annotated matching areas by top-k% attention regions with models of different sizes and different attention extraction strategies.