URaG: Unified Retrieval and Generation in Multimodal LLMs for Efficient Long Document Understanding (2511.10552v1)

Abstract: Recent multimodal LLMs (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.

• The paper proposes the URaG framework that embeds a lightweight cross-modal retrieval module in early transformer layers to effectively select relevant pages.
• It demonstrates a coarse-to-fine reasoning strategy where early layers capture broad context and deeper layers focus on evidence-rich content, improving retrieval accuracy and reducing FLOPs.
• URaG achieves state-of-the-art retrieval performance and efficiency, with up to 56% FLOP reduction and minimal parameter overhead across multiple long-document benchmarks.

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

Motivation and Problem Setting

The proliferation of multimodal LLMs (MLLMs) has substantially improved the understanding of single-page documents; however, scaling these models to multi-page, long-document scenarios presents persistent challenges. Processing extended inputs, particularly when much of the content is irrelevant to the task, introduces both information interference and a quadratic increase in computational cost inherent to Transformer-based architectures. Existing solutions fall into two categories: token compression, which attenuates informative content, and external retriever integration, which raises system complexity and precludes end-to-end optimization.

Figure 1: Comparison of different methods for long document understanding. Unified retrieval and generation in MLLMs leverages early-layer features for evidence retrieval, enabling efficient, accurate long-document reasoning.

Empirical Analysis: Coarse-to-Fine Reasoning in MLLMs

To ground their approach, the authors analyze the intrinsic reasoning patterns of MLLMs during long-document processing. Layerwise attention distribution was examined, revealing a human-like coarse-to-fine strategy: early layers allocate attention broadly, while deeper layers concentrate on evidence-rich pages.

Figure 2: (a) Attention entropy (b) Attention-based retrieval accuracy (c) Embedding-based retrieval accuracy confirm progressive focus on relevant information across layers.

Attention entropy drops and retrieval accuracy increases in deeper layers, substantiating the hypothesis that semantic discrimination emerges in mid-level representations. Embedding-based retrieval accuracy is notably higher and more stable than attention-based retrieval, reaching maturity around layer 12, thus informing the design of the framework’s retrieval module.

URaG Framework: Architecture and Methodology

URaG (Unified Retrieval and Generation) is predicated on embedding a lightweight cross-modal retrieval module into the early layers of an MLLM. The design leverages the evidence localization capabilities of MLLMs, transforming early hidden states into actionable selectors that retain the top-$k$ most relevant pages with respect to the user query.

Figure 3: Architecture overview of URaG, showing unified retrieval-generation within a single MLLM enabled by the cross-modal retrieval module.

Cross-Modal Retrieval Module

The retrieval module comprises two linear projections with GELU activations, applied to early-layer hidden states (typically layer 6). After dimensionality reduction and normalization, visual features for each page and the textual feature for the query are extracted. Contextualized late interaction computes the relevance:

$$s_{q,v^{(p)}} = \sum_{i \in [|E_q|]} \max_{j \in [|E_v^{(p)}|]} E_{q_i} \cdot E_{v_j}^{(p) T}$$

Top-$k$ pages are retained for further processing, while irrelevant ones are discarded, effectively shrinking the visual token sequence considered by the deeper layers.

Training Strategy

A two-stage training approach is adopted:

1. Retrieval Pretraining: Only retrieval module parameters are updated using a contrastive retrieval loss, with model backbone frozen.
2. Joint Fine-tuning: Low-rank adaptation (LoRA) is applied to both LLM and retrieval module, jointly optimized with retrieval and generation losses.

This protocol enables efficient evidence selection and robust answer generation, avoiding the pitfalls of modular, independently trained retrievers.

Experimental Results

URaG demonstrates superior retrieval and generation performance on multiple long-document benchmarks, including MPDocVQA, DUDE, SlideVQA, LongDocURL, and MMLongBench-Doc. Noteworthy numerical findings:

• Computational Overhead: URaG reduces FLOPs by 44–56% relative to non-retrieval baselines, with the retrieval module accounting for <0.1% of total parameters.
• Accuracy: State-of-the-art Top-1 and Top-5 retrieval accuracy across datasets; significant improvement in question types requiring fine-grained visual-semantic alignment.
• Parameter and Memory Efficiency: Decreased inference time and GPU memory consumption as input length scales.

URaG’s performance is robust even without backbone fine-tuning, outperforming fully-finetuned baselines by inserting only the pretrained retrieval module.

Figure 4: Replicated analysis of attention entropy and retrieval accuracy on MMLongBench-Doc, reinforcing the coarse-to-fine inductive bias.

Figure 5: Patch-level similarity heatmap for “How many people are standing in front of the chalkboard in the photograph?”, illustrating fine-grained evidence localization.

Figure 6: Patch-level similarity heatmap for “What are landslides?”, indicating precise semantic query-image alignment.

Figure 7: Qualitative results from URaG, demonstrating effective extraction of key content from long, diverse documents.

Ablation, Scalability, and Limitations

• Module Placement: Inserting the retrieval module at layer 6 is optimal; deeper insertions yield minor gains in Top-1 retrieval but do not improve overall comprehension.
• Training Stages: Retrieval pretraining and joint fine-tuning are both necessary; skipping either degrades retrieval and answer accuracy.
• Fixed Top-$k$ Selection: The current design uses $k=5$; scenarios requiring more or fewer pages would benefit from an adaptive retrieval mechanism.
• Generalizability: The method retains effectiveness when transplanted to alternative MLLM backbones (e.g., InternVL2.5-4B), demonstrating portability.

Theoretical and Practical Implications

By tightly integrating retrieval and generation within the MLLM backbone, URaG recapitulates human reading behaviors and introduces a scalable, unified paradigm for long-document reasoning. The approach eliminates the need for external retrievers, facilitating end-to-end training and simplifying deployment pipelines. Its computational thrift renders it especially attractive for real-world document-centric applications in legal, financial, and business domains.

The coarse-to-fine attention pattern elucidated through empirical analysis may inform future research into dynamic, confidence-based retrieval strategies, enabling adaptive $k$ selection and further efficiency gains. Additionally, the latent semantic alignment surfaced by mid-layer embeddings could inspire new approaches for multimodal fusion and evidence localization across broader domains.

Conclusion

URaG introduces a unified retrieval-generation methodology for MLLMs, driven by insights into the models' hierarchical attention patterns. Its cross-modal retrieval module, operating in early transformer layers, efficiently retains relevant content, minimizing resource expenditure and maximizing accuracy. With strong experimental results and negligible additional parameter cost, URaG points to a compelling direction for scalable, multimodal long-document understanding and sets a blueprint for future investigation into joint retrieval-generation architectures.