URaG: Unified Retrieval and Generation in Multimodal LLMs for Efficient Long Document Understanding

AI summary

Overview

• Research area: Multimodal large language models (MLLMs) for long multi-page document understanding, combining document retrieval with answer generation.
• Technical level: Intermediate. The paper assumes familiarity with Transformer attention, retrieval embeddings, LoRA fine-tuning, and multimodal benchmarks, though the central idea is explained intuitively.
• Scope: The paper proposes a single-model framework that performs evidence retrieval inside an MLLM's own early layers and then generates answers from only the retained pages, reporting state-of-the-art accuracy on several long-document benchmarks while cutting computation by 44–56%.

What This Paper Is About

Long documents overwhelm multimodal LLMs in two ways: most pages are irrelevant and interfere with the answer, and Transformer cost grows quadratically with sequence length. Prior fixes either compress all pages uniformly (losing fine visual detail) or bolt on a separate retriever (adding system complexity and blocking end-to-end training). URaG instead asks the MLLM to retrieve for itself: a small module reads the model's own early-layer hidden states, scores which pages matter, keeps the top-k pages, and discards the rest before the deeper layers generate the answer.

Key Contributions

1. An empirical study of how MLLMs read long documents. The authors visualize attention entropy and measure attention-based and embedding-based retrieval accuracy layer by layer across two representative MLLMs on two long-document benchmarks, showing a coarse-to-fine pattern: broad attention early, concentrated evidence focus later, and a re-check of all pages in the final two layers.
2. The URaG framework. A single MLLM that unifies retrieval and generation, using a lightweight cross-modal retrieval module (two linear projections with GELU activation) inserted at the sixth LLM layer to select the top-5 pages from the model's own hidden states.
3. Evidence that internal representations already localize evidence. Embedding-based retrieval from mid-level layers is reported to reach consistently high accuracy earlier and more stably than attention-based retrieval, motivating the design of the retrieval module.
4. Strong accuracy and efficiency results. State-of-the-art performance on MPDocVQA, DUDE, SlideVQA, LongDocURL, and MMLongBench-Doc, with 44–56% lower computation on longer inputs and a retrieval module that adds only 0.05–0.07% of total model parameters.

Main Findings

• Coarse-to-fine attention pattern confirmed. Attention entropy is high and attention-based retrieval accuracy low in the first 3 layers; entropy declines with fluctuations while retrieval accuracy rises in layers 3–20; deeper layers (20–34) show low entropy and consistently high retrieval accuracy; in the final two layers entropy rises again and retrieval accuracy drops slightly, which the authors read as the model revisiting all pages before answering.
• Embedding-based retrieval saturates earlier. Embedding-based retrieval reaches high accuracy around layer 12, earlier and more stably than attention-based retrieval, which is why the method uses embedding-based scoring.
• Retrieval results. URaG-7B reaches Top-1/Top-5 accuracy of 92.9/99.0 on SlideVQA, 68.3/86.0 on MMLongBench-Doc, 83.9/96.9 on DUDE, and 84.5/98.0 on MPDocVQA. URaG-3B reaches 92.1/98.9, 63.0/85.4, 83.0/97.0, and 84.4/98.0 respectively. For comparison, ColPali (3B) reaches 90.2/98.2, 60.3/80.2, 68.5/93.3, and 73.6/95.6.
• Generation results. URaG-7B scores 88.2 on MPDocVQA, 57.6 on DUDE, 72.1 on SlideVQA, and 52.2 on LongDocURL; URaG-3B scores 86.0, 54.1, 63.8, and 41.5. The corresponding Qwen2.5-VL baselines score 87.2/55.0/66.4/51.1 (7B) and 84.4/50.6/59.1/40.0 (3B).
• MMLongBench-Doc breakdown. URaG-7B attains 33.8 generalized accuracy and 32.8 F1 overall, versus 26.2/25.1 for Qwen2.5-VL 7B, 23.0/24.2 for SV-RAG (4B), and 21.0/22.6 for M3DocRAG (10B). URaG performs especially well on unanswerable questions (48.9 for URaG-3B, 43.5 for URaG-7B) and on visually intensive types such as charts and images.
• Strong on long inputs. The gains are most pronounced on SlideVQA and MMLongBench-Doc, whose inputs average 20 and 47.5 pages respectively.
• Beats a fine-tuned baseline without fine-tuning the backbone. URaG-3B without fine-tuning scores 62.1 on SlideVQA, 29.4 on MMLongBench-Doc, and 43.1 on LongDocURL, outperforming the fully supervised-fine-tuned baseline (61.9, 29.1, 37.3).
• Layer 6 is the sweet spot. Moving the retrieval module from layer 2 to 6 improves SlideVQA Top-1 from 89.0 to 92.1 and MMLongBench-Doc Top-5 from 82.0 to 85.4; going to layers 12 or 18 raises Top-1 slightly (93.1, 93.5) but Top-5 is already saturated at layer 6 and generation metrics decline (SlideVQA EM 62.9 at layer 12, 62.3 at layer 18 vs. 63.8 at layer 6).
• Both training stages matter. With no pretraining, SlideVQA Top-5 is 98.5 and MMLongBench-Doc Top-5 is 82.0; pretraining only gives 98.8 and 84.0; pretraining plus joint fine-tuning gives 98.9 and 85.4.
• Efficiency gains. On the 7B model, FLOPs fall from 415.9T to 232.8T at 20 pages (−44.0%), from 1246.4T to 574.9T at 60 pages (−53.9%), and from 2076.9T to 917.0T at 100 pages (−55.8%). On the 3B model, reductions are 34.8%, 42.5%, and 44.0%.
• Inference time and memory. At 100 input pages, per-question time drops from 32.07s to 18.74s (−41.57%) and peak GPU memory from 42.73 GB to 20.81 GB (−51.30%). At 20 pages, time drops 16.67% and memory 31.12%.
• Decoupled from a specific backbone. Re-implemented on InternVL2.5-4B, URaG-4B improves over its baseline on SlideVQA (51.9 vs. 45.2), MMLongBench-Doc (16.8 vs. 15.9), and LongDocURL (29.5 vs. 24.0).
• Retrieval behavior is interpretable. Patch-level similarity heatmaps for URaG-7B highlight query-relevant regions, such as people and a chalkboard for a counting question, and the word "landslides" for a definition question.

Methodology in Plain English

The approach has four pieces.

First, a diagnostic study. The authors run two existing MLLMs over long documents and record, at every Transformer layer, how attention is distributed across pages and how well that attention (or the hidden-state embeddings) can pick out the evidence page. This produces the coarse-to-fine picture that motivates everything else.

Second, a lightweight retrieval module. Given a document of n pages and a question, every page image is encoded into visual tokens and fed to the LLM together with the question. At the sixth layer, the module takes the hidden states, projects them down to 1024 and then 512 dimensions with two linear layers and GELU, and L2-normalizes the result. Visual features for each page and textual features for the question are separated out by position. Relevance is scored with contextualized late interaction (the ColBERT-style max-similarity sum), and the top-5 pages are kept while the rest are deleted directly from the hidden states.

Third, two-stage training. Stage one pretrains only the retrieval module with a contrastive-style retrieval loss while the rest of the model is frozen. Stage two adds LoRA adapters to both the LLM and the retrieval module and optimizes the sum of the retrieval loss and the generation (cross-entropy) loss with equal weight. Ground-truth evidence pages are always retained during training so information is never lost; remaining slots go to the highest-scoring pages.

Fourth, evaluation. Models are built on Qwen2.5-VL at 3B and 7B, trained for one epoch per stage with batch size 4, 8 gradient accumulation steps, AdamW, a learning rate of 1e-4 with 0.03 warm-up and cosine decay, and LoRA rank 32, alpha 64, dropout 0.1, on 4 NVIDIA A6000 GPUs. Training data come from MPDocVQA, DUDE, and SlideVQA. Retrieval is judged by Top-1 and Top-5 accuracy; generation by ANLS (MPDocVQA, DUDE), Exact Match (SlideVQA), Generalized Accuracy and F1 (MMLongBench-Doc), and Generalized Accuracy (LongDocURL).

Why This Matters

Research impact. The paper reframes retrieval for long documents as something an MLLM already does internally rather than something bolted on from outside. The layer-wise analysis is a reusable diagnostic, and the result that a pretrained retrieval module inserted into a frozen MLLM beats a fully fine-tuned baseline challenges the assumption that long-context gains require expensive supervised tuning. It also offers a middle path between token compression (which discards detail) and external retrievers (which cannot be trained end to end).

Real-world applications.

• Contract and compliance review, where relevant clauses sit among hundreds of pages of boilerplate.
• Financial and analyst report processing, where answers are spread across charts, tables, and text.
• Slide-deck and presentation question answering, the setting behind the SlideVQA results.
• Enterprise document search and question answering over scanned, layout-rich PDF collections, where the 41.57% inference-time and 51.30% memory savings at 100 pages directly reduce serving cost.

Industry relevance. The gains are reported in FLOPs, wall-clock time, and GPU memory, not just accuracy, which matters for deployment economics. Because the retrieval module adds only 0.05–0.07% of parameters and the framework works on both Qwen2.5-VL and InternVL2.5 backbones, it is a lightweight add-on rather than a new model architecture. The absence of an external retriever also removes a separate component to train, version, and maintain.

Future Directions

• Adaptive top-k. The authors identify the fixed top-5 selection as the main limitation: when evidence spans more than five pages, essential information can be dropped; when it spans fewer, redundant pages are carried. They propose adjusting k based on query difficulty or similarity confidence.
• Better handling of cross-page and unanswerable questions. URaG's weakest results are on multi-page evidence (16.9 for URaG-7B on MMLongBench-Doc cross-page questions) even though unanswerable questions do well, suggesting the retrieval step itself is the bottleneck for dispersed evidence.
• Reducing the remaining gap to full fine-tuning. Fine-tuning the MLLM backbone degrades LongDocURL performance (37.3 vs. 40.0 for the untuned baseline), which the authors attribute to overfitting or domain mismatch; how to fine-tune without that regression is unresolved.
• Extending the analysis beyond the studied models and benchmarks. The diagnostic covers two MLLMs and two benchmark subsets, and the appendix notes that documents longer than 10 pages were cropped to a continuous 10-page span for analysis, leaving the layer dynamics of much longer contexts only partially explored.

Target Audience

This paper suits researchers and engineers working on multimodal LLMs, long-context modeling, and document AI who want a concrete, efficiency-oriented alternative to token compression and external retrieval pipelines. It is also useful for practitioners deploying document question answering at scale, since it reports FLOPs, latency, and memory alongside accuracy, and for students studying how attention behaves across Transformer depth. Readers should be comfortable with attention mechanisms, embedding similarity, and LoRA-style fine-tuning to get the most from the methodology section.

Authors’ abstract

Recent multimodal large language models (MLLMs) still struggle with long document understanding due to two fundamental challenges: information interference from abundant irrelevant content, and the quadratic computational cost of Transformer-based architectures. Existing approaches primarily fall into two categories: token compression, which sacrifices fine-grained details; and introducing external retrievers, which increase system complexity and prevent end-to-end optimization. To address these issues, we conduct an in-depth analysis and observe that MLLMs exhibit a human-like coarse-to-fine reasoning pattern: early Transformer layers attend broadly across the document, while deeper layers focus on relevant evidence pages. Motivated by this insight, we posit that the inherent evidence localization capabilities of MLLMs can be explicitly leveraged to perform retrieval during the reasoning process, facilitating efficient long document understanding. To this end, we propose URaG, a simple-yet-effective framework that Unifies Retrieval and Generation within a single MLLM. URaG introduces a lightweight cross-modal retrieval module that converts the early Transformer layers into an efficient evidence selector, identifying and preserving the most relevant pages while discarding irrelevant content. This design enables the deeper layers to concentrate computational resources on pertinent information, improving both accuracy and efficiency. Extensive experiments demonstrate that URaG achieves state-of-the-art performance while reducing computational overhead by 44-56%. The code is available at https://github.com/shi-yx/URaG.