![DocMemo Logo](figures/logo.png)

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

[![arXiv](https://img.shields.io/badge/arXiv-2608.07067-b31b1b.svg)](https://arxiv.org/abs/2608.07067)

![DocMemo Method](figures/method.png)

DocMemo is a memory-guided framework for multimodal question answering over long documents. It overcomes the limitations of static page retrieval and unstructured cross-round context accumulation by organizing document exploration into three complementary memory components:

- **Document Schema Memory** — captures query-independent document structure, topics, and global semantics.
- **Page Belief Memory** — maintains dynamically updated relevance estimates for candidate pages through Bayesian belief updating, Thompson sampling, and spatial proximity propagation.
- **Question Episodic Memory** — records query-specific evidence, intermediate findings, and reasoning trajectories across retrieval rounds.

This structured design enables DocMemo to iteratively refine its search strategy, recover from early retrieval errors, and integrate sparse evidence distributed across long documents. DocMemo further supports adaptive-granularity evidence access, supplementing page-level retrieval with fine-grained visual information when needed. Experiments across multiple long-document understanding benchmarks demonstrate state-of-the-art performance over strong static and iterative retrieval baselines.

## 0. Install

Create an environment and install the dependencies.

```bash
conda create -n docmemo python=3.11 -y
conda activate docmemo
cd /path/to/DocMemo
python -m pip install -r requirements.txt
```

## 1. Prepare data

Create the data directory and place the dataset PDFs in `documents` (using MMLongBench as an example).

```bash
mkdir -p data/MMLongBench/
```

## 2. Generate page embeddings

Generate one embedding file for each PDF before starting vLLM.

```bash
bash scripts/generate_embeddings.sh \
  --documents-dir data/MMLongBench/documents \
  --output-dir data/MMLongBench/embeddings \
  --model vidore/colqwen2.5-v0.2 \
  --device cuda \
  --image-dpi 150 \
  --batch-size 2
```

Replace the model ID with `/path/to/colqwen2.5-v0.2` when using local weights.

## 3. Start vLLM

Start an OpenAI-compatible vision-language model service in another terminal.

```bash
conda activate vllm
vllm serve /path/to/vision-language-model \
  --served-model-name docmemo-model \
  --host 127.0.0.1 \
  --port 8000 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 32768 \
  --trust-remote-code
```

## 4. Generate page summaries

Generate one visual summary for each PDF page.

```bash
bash scripts/generate_summaries.sh \
  --documents-dir data/MMLongBench/documents \
  --output-dir data/MMLongBench/summaries \
  --model docmemo-model \
  --base-url http://localhost:8000/v1 \
  --image-dpi 150 \
  --max-tokens 2048 \
  --n-jobs 1
```

## 5. Generate document memory

Generate document-level memory from the page summaries.

```bash
bash scripts/generate_doc_memory.sh \
  --summaries-dir data/MMLongBench/summaries \
  --model docmemo-model \
  --base-url http://localhost:8000/v1 \
  --max-tokens 2048 \
  --n-jobs 1
```

## 6. Run document QA

Run iterative retrieval and multimodal question answering.

```bash
bash scripts/run_docmemo.sh \
  --input-file data/MMLongBench/samples.json \
  --documents-dir data/MMLongBench/documents \
  --embeddings-dir data/MMLongBench/embeddings \
  --summaries-dir data/MMLongBench/summaries \
  --output-file data/MMLongBench/results.json \
  --colqwen-model vidore/colqwen2.5-v0.2 \
  --embedding-device cpu \
  --retrieval-model docmemo-model \
  --qa-model docmemo-model \
  --retrieval-base-url http://localhost:8000/v1 \
  --qa-base-url http://localhost:8000/v1 \
  --max-rounds 3 \
  --max-page-retrieval 30 \
  --max-pages 9 \
  --image-dpi 150 \
  --n-jobs 1
```

### 6.1 Experimental Results

The figure below highlights the main experimental results of DocMemo across representative benchmarks.

![DocMemo Method](figures/results.png)

## 7. Citation

If you find our work useful, please consider citing:

```bibtex
@article{yao2026docmemo,
  title={DocMemo: Dynamic Evidence Discovery via Probabilistic Memory-Guided Retrieval for Multi-Modal Document Understanding},
  author={Yao, Hanshu and Zhong, Janfeng and Lian, Niu and Wang, Jinpeng},
  journal={arXiv preprint arXiv:2608.07067},
  year={2026}
}