💬 NLP

DocMemo: Dynamic Evidence Discovery via Probabilistic Memory-Guided Retrieval for Multi-Modal Document Understanding

DocMemo is a memory-guided framework that enhances long-document understanding by dynamically exploring evidence through a tri-level memory system and Bayesian belief updating, thereby overcoming the limitations of static retrieval and fragile cross-round state propagation in existing multi-modal document analysis methods.

Original authors: Hanshu Yao, Janfeng Zhong, Niu Lian, Jinpeng Wang

Published 2026-08-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Hanshu Yao, Janfeng Zhong, Niu Lian, Jinpeng Wang

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 solve a massive mystery, but the clues aren't hidden in a single notebook; they are scattered across a library containing hundreds of books, each with hundreds of pages. Some clues are written in text, some are hidden in complex charts, and others are tucked away in tiny diagrams. This is the world of "long-document understanding" for artificial intelligence. Currently, most AI models act like a student who tries to read the entire library at once. But since their brains (or "context windows") are too small to hold everything, they have to pick a few pages to read first. The problem is, if they pick the wrong pages in the beginning, they get stuck and can't fix their mistake. Other smarter AI models try to read a few pages, think, and then read more, but they often forget what they learned in the previous step, treating every new search as if it were their very first time. This paper, titled DocMemo, proposes a new way for AI to think: by giving it a "memory" that helps it remember what it has already learned, what it still needs to find, and how the pages in the library are connected, so it can hunt for clues more like a human detective and less like a confused robot.

The researchers behind DocMemo, Hanshu Yao and his team, realized that existing AI systems struggle because they are either too rigid or too forgetful. Some systems pick a fixed set of pages to read right at the start and stick to them, even if they miss the most important clue. Others try to search in rounds but fail to connect the dots between one round and the next, essentially starting over every time. To fix this, the team built a system that treats document reading as a dynamic exploration. Instead of just grabbing pages, DocMemo maintains a "tri-level memory" that acts like a detective's case file. First, it has a Document Schema Memory, which is like a map of the library's layout, knowing where different types of information (like tables or chapters) usually live. Second, it has a Page Belief Memory, which is a living list of "hunches" about which pages are likely to hold the answer. This list isn't static; it gets updated with every new piece of evidence, using a mathematical trick called "Bayesian updating" to get smarter about what to look for next. Finally, it has a Question Episodic Memory, which records the detective's own journey—the questions it asked, the dead ends it hit, and the refined questions it came up with along the way.

What makes DocMemo truly special is how it uses this memory to change its strategy on the fly. When the AI is unsure, it uses a method called "Thompson sampling" to balance between checking pages it is very confident about and exploring pages it is less sure about, just like a detective who checks the most likely suspect but also keeps an eye on the suspicious stranger. If the AI finds a relevant page, it doesn't just stop there; it uses "spatial proximity propagation" to assume that the answer might also be on the pages right next to it, because clues in long documents often cluster together. Furthermore, if the AI sees a page with a dense table or a complex chart, it doesn't just look at the whole page; it zooms in to read the fine details, a feature called "adaptive-granularity evidence access."

The team tested this new detective on three different sets of challenging documents, including academic papers and long reports with hundreds of pages. The results were promising: DocMemo consistently outperformed the best existing methods. On one specific benchmark called MMLongBench-Doc, it achieved an accuracy of 71.3%, beating the previous top performers. The researchers found that the system's ability to remember its past searches and update its "hunches" about page relevance was the key to its success. In fact, when they removed the memory components in their tests, the system's performance dropped significantly, proving that the memory is not just a nice-to-have feature but the engine driving the improvement. The study suggests that by giving AI a structured way to remember its exploration and update its beliefs dynamically, we can help it solve complex puzzles in massive documents much more effectively than before.

Technical Summary: DocMemo

Problem Statement

Long-document Visual Question Answering (DocVQA) requires models to locate sparse and heterogeneous evidence (text, tables, figures, layouts) across documents spanning hundreds of pages. The core challenge is not merely processing long inputs, but dynamically locating, updating, and integrating relevant pages under a limited evidence budget.

Existing systems face two primary limitations:

Static Retrieval: Single-turn methods fix a candidate page set before reasoning begins. If key evidence is missed initially, the system cannot recover, leading to rigid, one-shot pipelines.
Fragile Iterative Retrieval: While recent iterative methods allow multi-round acquisition, they lack a mechanism for propagating cross-round states. Information is typically maintained through context reconstruction, effectively treating each round as an independent retrieval event rather than a structured exploration with memory. This results in an inability to track dynamic changes in page relevance or refine search trajectories based on intermediate reasoning signals.

Methodology: DocMemo Framework

DocMemo proposes a memory-guided framework that formulates long-document reasoning as dynamic evidence exploration. It decouples memory from retrieval to enable continuous state refinement.

1. Tri-Level Document Memory

The system maintains a structured memory state $\mathcal{M}^{(t)} = (\mathcal{M}_{schema}, \mathcal{M}^{(t)}_{belief}, \mathcal{M}^{(t)}_{epi})$ across reasoning rounds:

Document Schema Memory ( $\mathcal{M}_{schema}$ ): An offline, query-independent prior constructed from document type, topic descriptions, and a structural index of page ranges and topic labels. It serves as a fixed navigation guide.
Page Belief Memory ( $\mathcal{M}^{(t)}_{belief}$ ): A dynamic state modeling the relevance confidence of each page using a Beta distribution $Beta(\alpha_i, \beta_i)$ . $\alpha_i$ and $\beta_i$ accumulate evidence for relevance and irrelevance, respectively. This memory evolves with cross-round feedback.
Question Episodic Memory ( $\mathcal{M}^{(t)}_{epi}$ ): A query-specific record of intermediate findings, unmet information needs, and query refinement trajectories. It allows the system to narrow the search space along an existing evidence chain rather than restarting from the original query.

2. Bayesian Page Belief Updating & Thompson Sampling

DocMemo treats page selection as a confidence updating process:

Bayesian Update: After each reasoning round, the system identifies useful ( $U_t$ ) and irrelevant ( $V_t$ ) pages. Beliefs are updated via Beta-Bernoulli conjugacy: $\alpha_i \leftarrow \alpha_i + \mathbb{I}[i \in U_t]$ and $\beta_i \leftarrow \beta_i + \mathbb{I}[i \in V_t]$ .
Spatial Propagation: Motivated by the locality of evidence in long documents, positive feedback on a page $i$ is propagated to neighboring pages $j$ (within radius $r$ ) with a distance decay factor $\gamma$ . Negative feedback is not propagated to avoid suppressing potentially relevant nearby pages.
Thompson Sampling: To balance exploitation and exploration, the system samples a relevance estimate $\theta^{(t)}_i$ from the posterior Beta distribution for each page.
Candidate Scoring: The final score combines current visual similarity $c^{(t)}_i$ with the sampled belief $\theta^{(t)}_i$ : $score^{(t)}_i = (1 - \lambda_t)c^{(t)}_i + \lambda_t\theta^{(t)}_i$ . The weight $\lambda_t$ increases over iterations, shifting reliance from raw visual similarity to accumulated memory confidence.

3. Adaptive-Granularity Evidence Access

To handle information-dense regions (e.g., tables), the system supplements full-page context with fine-grained visual regions (high-resolution crops) during reasoning, ensuring detailed layout understanding.

Key Contributions

Memory-Guided Framework: Introduction of DocMemo, which organizes document exploration via a tri-level memory structure separating structural priors, dynamic page beliefs, and query-specific reasoning traces.
Probabilistic State Modeling: Implementation of Bayesian page belief updating with Thompson sampling and spatial contiguity propagation, enabling uncertainty-aware and locality-aware dynamic retrieval across rounds.
Empirical Validation: Demonstration of state-of-the-art performance on three long-document DocVQA benchmarks, with ablation studies confirming the necessity of structured memory and dynamic belief updating.

Experimental Results

DocMemo was evaluated on MMLongBench-Doc, LongDocURL, and PaperTab.

Overall Performance: DocMemo achieved an average accuracy of 77.6% across the three benchmarks, outperforming strong baselines including proprietary models (e.g., GPT-4o, Claude-4-Sonnet) and open-source agentic systems (e.g., SimpleDoc, DocLens).
- On MMLongBench-Doc, it reached 71.3% accuracy (vs. 67.4% for the next best agentic baseline, AgenticOCR).
- On LongDocURL, it achieved 81.1% (vs. 72.3% for SimpleDoc).
- On PaperTab, it achieved 80.4% (vs. 65.4% for SimpleDoc).
Fine-Grained Analysis: The model showed significant gains in table-heavy (TAB) and unanswerable (UNA) categories, suggesting improved evidence localization and better handling of unanswerable queries through accumulated cross-round states.
Efficiency: DocMemo reduced the number of retrieval iterations required to achieve high accuracy compared to baselines (0.41× iterations for 1.18× accuracy gain), resulting in a 2.40× overall efficiency improvement.
Ablation Studies: Removing the tri-level memory or Bayesian updating caused significant accuracy drops (from 71.3% to ~68.5%), validating that dynamic belief revision and structured memory are critical for robust reasoning.

Significance

The paper claims that the key to improving long-document reasoning lies in explicitly modeling and continuously updating document exploration states. By transforming iterative retrieval from simple context accumulation into a structured evidence exploration driven by persistent document knowledge, dynamic page states, and query-local experience, DocMemo addresses the limitations of static and stateless iterative approaches. The framework demonstrates that separating stable document priors from rapidly changing relevance beliefs allows agents to recover from early retrieval errors and adaptively refine their search trajectories.