AWM: Answerable Working Memory for Long-Document VQA Agents

Dongzhuoran Zhou; Yuqicheng Zhu; Yule Liu; Zhen Yang; Rui Lu; Yuxiao Dong; Jie Tang; Evgeny Kharlamov

cs.CLAug 26, 2026

AWM: Answerable Working Memory for Long-Document VQA Agents

Authors: Dongzhuoran Zhou, Yuqicheng Zhu, Yule Liu, Zhen Yang, Rui Lu, Yuxiao Dong, Jie Tang, Evgeny Kharlamov

Organizations: University of Oslo · Bosch Center for AI · University of Stuttgart · Hong Kong University of Science and Technology (Guangzhou) ·

Abstract

Long-document visual question answering increasingly relies on VLM agents that retrieve candidate pages, inspect page images, write findings to working memory, and synthesize answers. Working memory should carry answer-supporting evidence across page inspections for later grounded answering, yet existing evaluation mainly checks final-answer correctness and evidence-page access. This creates a memory-quality blind spot: an agent may reach the right page and answer correctly while leaving behind memory too generic or incomplete to support answering once page context is removed. We introduce memory-only answerability, a diagnostic that asks whether a reader can answer from the question and terminal working memory alone. Building on this diagnostic, Answerable Working Memory (AWM) treats terminal working memory as an answerable evidence artifact, and AWM-GRPO incorporates this signal into the GRPO reward while preserving final-answer priority. Under GRPO, this reward assigns higher advantages to answer-correct trajectories whose terminal working memory remains answerable. On MMLongBench-Doc, even when gold evidence pages are provided, 42.5% of correct answers still cannot be answered from terminal working memory alone. AWM-GRPO improves final-answer accuracy over the RAG baseline by 8.1 and 11.9 points on MMLongBench-Doc and LongDocURL and reduces the memory-missing-correct rate by 2.7 points over answer-only GRPO.

Figures & tables

Figure 1: Combined agent-memory framework and memory-quality diagnostics. (a) The agent retrieves page images, writes source-linked findings to working memory, and outputs both a final answer and terminal working memory. (b) Standard evaluation scores the final answer, memory-only evaluation asks whether the question can be answered from terminal working memory alone, and the controlled comparison varies terminal working memory while holding evidence-page access fixed.

Final answer	Memory answerable	Case name
correct	yes	memory-supported correct
correct	no	memory-missing correct
wrong	yes	answering error
wrong	no	unresolved error

Table 1: Outcome cases induced by final-answer correctness and memory-only answerability.

Setting	$N$	Final-ans. acc. (%)	Memory-only acc. (%)	MMC / correct	$P_{\mathrm{mmc}}$ (%)
Full multi-turn	500	33.4	27.4	43 / 172	25.0
EP-given	500	36.2	30.2	77 / 181	42.5

Table 2: Memory-quality diagnostic for the Base AWM-Agent (Qwen3-VL-4B, no post-training) on a fixed 500-example answerable MMLongBench-Doc subset. Unanswerable questions are excluded. Accuracy follows the official mean per-question score; for MMC/correct and

P_{\mathrm{mmc}}

, we binarize the official score by treating a positive score as correct.

Memory condition	Final-ans. binary rate	Memory-only binary rate
Empty memory	2.8	2.6
Original agent memory (4B)	36.2	30.2
Improved memory (GPT-4o)	44.4	44.8

Table 3: Controlled memory comparison on a fixed 500-example MMLongBench-Doc EP-given answerable subset. The reader and scorer are fixed; only the terminal-memory condition changes. Improved memory is produced by GPT-4o.

Figure 2: Overview of AWM-GRPO. For each sampled trajectory, the agent produces both a final answer and terminal working memory. A frozen reader answers from the terminal working memory alone, and a frozen scorer evaluates the final answer and the memory-only answer. The two scores define the AWM reward, which is normalized within each GRPO group to update the policy.

Figure 3: Simulated GRPO advantage distributions. Each row samples synthetic groups of size

G{=}8

from a four-cell outcome mixture;

p_{11}

is the probability of memory-supported correct

(1,1)

. Left: answer-only reward collapses the four cells into answer-wrong and answer-correct regions. Right: AWM assigns a distinct raw reward to each of the four cells before group normalization.

Method	Backbone	Param	Paradigm	MMLongBench-Doc	LongDocURL
Open-source / RAG
Qwen2.5-VL ( Zheng et al., 2026 )	Qwen2.5-VL	7B	VLM	28.0	32.9
SV-RAG ( Chen et al., 2024 )	InternVL2	4B	RAG	23.0	–
M3DocRAG ( Cho et al., 2024 )	Qwen2-VL	7B	RAG	21.0	35.1
VisRAG ( Yu et al., 2025 )	MiniCPM-V	8B	RAG	18.8	41.9
MoLoRAG ( Wu et al., 2025 )	Qwen2.5-VL	7B	RAG	41.0	51.9

Table 4: Final-answer accuracy; higher is better. External rows follow their reported protocols and provide context only. Our five Qwen3-VL-4B rows use the official scoring protocol.

Variant	Final acc.	Mem acc.	$P_{\mathrm{mmc}}$ (%)
SFT	40.8	38.5	17.7
+ Answer-GRPO	43.2	42.5	19.9
+ AWM-GRPO	45.4	44.5	17.2

Table 5: Memory-quality analysis of the three AWM-Agent variants on 1,082 MMLongBench-Doc trajectories in the full multi-turn setting.

Variant	Subset	Final acc.	Mem acc.	$P_{\mathrm{mmc}}$ (%)
+ Answer-GRPO	MM EP	45.4	41.2	19.1
+ AWM-GRPO	MM EP	48.0	43.5	16.4

Table 6: Memory-quality analysis of Answer-GRPO and AWM-GRPO on 733 answerable MMLongBench-Doc examples in the EP-given setting.

Figure 4: AWM-GRPO memory-only accuracy over training steps on MMLongBench-Doc . Memory-only accuracy increases from 38.8 at step 40 to 43.5 at step 280.

		Final acc.		$P_{\mathrm{mmc}}$ (%)
Source	$N$	Ans	AWM	Ans	AWM
Pure-text	135	46.1	58.5	22.8	23.1
Table	143	36.7	39.5	15.2	23.9
Chart	107	45.7	46.0	21.4	17.5
Figure	187	41.0	43.9	20.3	16.2
Layout	32	75.9	74.3	4.8	15.8

Table 7: Evidence-source comparison on MMLongBench-Doc in the full multi-turn setting. “Ans” denotes Answer-GRPO and “AWM” denotes AWM-GRPO; bold marks the better value within each source.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Memory condition	Correct / N	Rate (%)	95% Wilson CI (%)
Empty memory	14 / 500	2.8	[1.7, 4.6]
Original agent memory (4B)	181 / 500	36.2	[32.1, 40.5]
Improved memory (GPT-4o)	222 / 500	44.4	[40.1, 48.8]

Appendix

Table 8: Final-answer positive-score rate in the controlled memory intervention. Intervals are 95% Wilson confidence intervals.

Original outcome	N	Original memory correct	Improved memory correct	Change (count; pp)
Memory-supported correct	104	104	88	$-16$ ; $-15.4$
Memory-missing correct	77	0	48	$+48$ ; $+62.3$
Answering error	47	47	29	$-18$ ; $-38.3$
Unresolved error	272	0	59	$+59$ ; $+21.7$
Total	500	151	224	$+73$ ; $+14.6$

Appendix

Table 9: Memory-only correctness after grouping examples by the original agent’s outcome. Change is measured within each fixed row.

Long-Document VQA

See all 17 papers

DocTrace: Towards Traceable Long Document VQA via Hierarchical Evidence Graph Reasoning
Aug 4, 2026Le Xiang, Zhicheng Guan, Hong Chen +5Visual Question Answering Long-Document VQA
Baidu Basic Model Research and Development Department, Baidu Inc. · Tsinghua Shenzhen International Graduate School, Tsinghua University
VELD-QA: A Large-Scale Dataset for Visual Evidence Grounding and Question Answering in Long Documents
Oct 14, 2026Zhivar Sourati, Marianne Menglin Liu, Rongguang Wang +12Long-Document VQA Document QA
University of Southern California · Oracle AI · Turing Enterprise Inc +1
Doc-V*:Coarse-to-Fine Interactive Visual Reasoning for Multi-Page Document VQA
Apr 15, 2026Yuanlei Zheng, Pei Fu, Hang Li +9Document QA Long-Document VQA
School of Software Engineering, Huazhong University of Science and Technology · MiLM Plus, Xiaomi Inc. · School of Computer Science and Technology, Huazhong University of Science and Technology +1

Aug 4, 2026J/K move · Enter open · S save

DocTrace: Towards Traceable Long Document VQA via Hierarchical Evidence Graph Reasoning

Le Xiang, Zhicheng Guan, Hong Chen, Xiaocong Lin+4

Baidu Basic Model Research and Development Department, Baidu Inc. · Tsinghua Shenzhen International Graduate School, Tsinghua University

Visual Question Answering Long-Document VQA Multimodal Large Language Models VLM Reasoning

Open paper AskPDF arXiv

Long Document Visual Question Answering (LongDocVQA) requires Multimodal Large Language Models (MLLMs) to locate, integrate, and reason over heterogeneous document elements distributed across multiple pages. Existing approaches, including end-to-end MLLMs, retrieval-augmented generation (RAG) pipelines, and document agents, often lack explicit mechanisms to represent and verify how grounded evidence is progressively composed during reasoning, limiting both answer accuracy and traceability. In this paper, we cast LongDocVQA as an explicit evidence graph reasoning problem rather than implicit answer prediction. To this end, we propose DocTrace, a hierarchical framework that progressively performs evidence localization, structured document parsing, and evidence graph reasoning to enable explicit evidence provenance. To effectively learn these capabilities, we develop a two-stage training framework: joint Supervised Fine-Tuning (SFT) first initializes evidence localization and graph reasoning abilities, followed by task-specific Group Relative Policy Optimization (GRPO) with dedicated rewards to further optimize these capabilities. Extensive experiments on MMLongBench-Doc, LongDocURL, and SlideVQA demonstrate that DocTrace consistently outperforms both existing open-source baselines and proprietary MLLMs. Compared with the Qwen3-VL-8B-Instruct backbone, DocTrace achieves absolute improvements of 14.4, 11.3, and 11.7 points on the three benchmarks, respectively. Beyond competitive performance, DocTrace constructs traceable evidence graphs with explicit node-level provenance, enabling transparent and verifiable reasoning for long document understanding.

Abstract

Long-Document VQA

Explore similar work

DocTrace: Towards Traceable Long Document VQA via Hierarchical Evidence Graph Reasoning

VELD-QA: A Large-Scale Dataset for Visual Evidence Grounding and Question Answering in Long Documents

Doc-V*:Coarse-to-Fine Interactive Visual Reasoning for Multi-Page Document VQA