Dual-Grained Agent Memory and Shapley Context Attribution for Multimodal Agentic Learner

Jieke Wang; Tiancheng Shen; Yibo Yang; Ming-Hsuan Yang

cs.CVAug 24, 2026

Dual-Grained Agent Memory and Shapley Context Attribution for Multimodal Agentic Learner

Authors: Jieke Wang, Tiancheng Shen, Yibo Yang, Ming-Hsuan Yang

Organizations: UC Merced · Shanghai Jiao Tong University

Abstract

Frontier multimodal large language models (MLLMs) deliver impressive perception yet still falter on scientific and mathematical reasoning. Parameter-level adaptation is unavailable for closed-weight or on-device backbones, and stateless prompting forfeits any compounding benefit from problems already solved. We propose DG-Mem, a dual-grained agentic memory framework that augments a frozen MLLM with a non-parametric, externally stored memory built once from training-time rollouts and consulted read-only at test time. Motivated by the Complementary Learning Systems (CLS) account of human memory, DG-Mem factors its store into an instance-grounded exemplar memory and a category-level schema memory of IF-THEN rules, with a transient reflection store mediating their construction so that schemas are synthesized only from abstract reflections, never from exemplar text. Two design choices distinguish DG-Mem: an online concept categorizer that grows the category space incrementally during training rather than committing to a predefined taxonomy, and a Shapley context attribution procedure that decomposes correctness across the entire retrieved rule set and yields a per-rule utility that re-weights retrieval at test time. The pipeline introduces no gradient updates and is deployable on closed-weight or on-device backbones. Across MathVista, MMMU, and MMMU-Pro on four open-weight and proprietary backbones (Qwen3.5-27B, Qwen3.5-122B-A10B, GPT-5-Nano, Gemini-3-Flash), DG-Mem improves consistently over no-memory and competitive memory baselines.

Figures & tables

Method	Multimodal	Exemplar	Schema	Credit Assn.	Grad-Free
Mem0 [ 5 ]	$\times$	$\checkmark$	$\times$	$\times$	$\checkmark$
Memento [ 47 ]	$\times$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$
A-Mem [ 30 ]	$\times$	$\checkmark$	$\times$	$\times$	$\checkmark$
ExpeL [ 44 ]	$\times$	$\times$	$\checkmark$	$\times$	$\checkmark$
AWM [ 29 ]	$\times$	$\times$	$\checkmark$	$\times$	$\checkmark$
ReasoningBank [ 17 ]	$\times$	$\times$	$\checkmark$	$\times$	$\checkmark$

Table 1: Comparison of agentic memory methods across five design dimensions : whether the method (i) handles multimodal inputs (vision+text), (ii) maintains instance-grounded exemplars , (iii) maintains abstract schema rules , (iv) performs per-rule credit assignment beyond similarity-only retrieval, and (v) is gradient-free (no model parameter updates).

\checkmark

\times

indicates whether the method satisfies a given dimension. See the surrounding text for column definitions.

Panel of Figure 1 — Figure 1: Two-stage offline pipeline for memory training on a frozen MLLM backbone. Left (Stage 1 – Memory Construction): for each training problem $q$ , $\operatorname{MLLM}$ produces $N$ rollouts that an online classifier and three parallel reflectors turn into an exemplar memory $\mathcal{M}_{E}$ and a category-level schema memory $\mathcal{M}_{S}$ , with a transient reflection store $\mathcal{M}_{R}$ mediating the latter so that schema synthesis never reads exemplar text. Right (Stage 2 – Shapley Context Attribution): the same frozen $\operatorname{MLLM}$ is reused as an attribution device; for each problem we draw $B$ subsets of the active category’s rules, run a rollout per subset, and convert the binary rewards into per-rule Shapley contributions that are accumulated into a utility column on $\mathcal{M}_{S}$ and used to re-weight retrieval at test time. See § 3.1 and § 3.3 for full details.

Backbone	Method	MathVista*	MMMU*	MMMU-Pro*	Avg Imp.
Qwen3.5-122B-A10B	No-memory	83.6	85.8	77.9	-
	ViLoMem (offline)	85.2	86.4	80.0	+1.4
	ViLoMem (online)	88.4	87.4	80.3	+2.9
	DG-Mem (Ours)	88.7	89.2	86.1	+5.6
Qwen3.5-27B	No-memory	82.9	83.4	75.1	-
	ViLoMem (offline)	86.4	85.9	74.6	+1.8

Table 2: Top-1 accuracy (%) on held-out test splits across multimodal reasoning benchmarks. Each backbone is evaluated under four conditions: No-memory (same backbone and decoding configuration with retrieval disabled), ViLoMem (offline) (logical/visual error memory frozen after training), ViLoMem (online) (memory continues to grow on the test split), and DG-Mem (Ours). Avg Imp. reports the per-backbone mean improvement over the No-memory baseline across the three benchmarks. * on each benchmark name indicates that we report on a fixed held-out test partition rather than the full benchmark. Best result per backbone is in bold .

Datasets	MathVista*	MMMU*	MMMU-Pro*
Baseline	86.2	82.4	69.0
– w/o $\mathcal{M}_{E}$ (stage-1 schema only)	88.4 +2.2	82.1 -0.3	73.4 +4.4
– w/o $\mathcal{M}_{S}$ (stage-1 exemplar only)	87.6 +1.4	84.0 +1.6	73.4 +4.4
– w/o utility ( $\alpha{=}1$ , stage-1 both memory)	91.0 +4.8	84.0 +1.6	75.0 +6.0
Full method	91.6 +5.4	84.7 +2.3	77.1 +8.1

Table 3: Component ablations on Gemini-3-Flash. Each row disables a single mechanism while leaving the rest of the pipeline intact. w/o

\mathcal{M}_{E}

retains the category-level schema memory but removes the exemplar store, so retrieval returns only IF-THEN rules. w/o

\mathcal{M}_{S}

keeps exemplars but disables schema retrieval, so the prompt is augmented only with nearest-neighbor solved instances. w/o utility sets the retrieval interpolation

\alpha{=}1

, ranking rules by embedding similarity alone and discarding the Stage-2 Shapley utility column. Colored deltas ( green / red ) report the absolute change in accuracy relative to the No-memory Baseline row.

Test question	Retrieved exemplar ( $\mathcal{M}_{E}$ )	Retrieved schema rule ( $\mathcal{M}_{S}$ )	Baseline $\to$ Ours
MathVista #90. “Subtract all yellow matte blocks. Subtract all tiny brown cylinders. How many objects are left?”	Solved template: “Subtract all red things. Subtract all tiny matte balls. How many objects.” Strategy: “Catalog every object’s color, size, material, and shape, then translate the natural-language descriptors to those attribute categories before subtracting.”	Visual Counting and Attribute Analysis: “Perform an exhaustive initial census of all objects, assign a count of zero to any missing categories, and calculate the cardinality of each subset independently before applying the subtraction.”	Gold: 5 Baseline: 4 (skipped enumeration; “There are 4 objects left.”) Ours: 5 (lists 7 objects, removes 2)
MathVista #13. “How many objects are preferred by more than 90% of people in at least one category?” (grouped bar chart)	Solved template: “How many algorithms have accuracy higher than 9 in at least one dataset?” Verification check: “Re-evaluate any bar that appears to end exactly on a midpoint or grid line to confirm whether it meets a ‘strictly greater than’ or ‘strictly less than’ condition.”	Data Analysis and Interpretation: “When identifying threshold-based counts, systematically compute the metric for every distinct item rather than relying on visual estimation of distances or assumed ordering.”	Gold: 0 Baseline: 2 (visually estimated bars near the 90% line) Ours: 0 (computes per-bar height, applies strict $>$ )
MathVista #792. “What is the age gap between these two people?” (photo of P. Hammond and M. Kaljurand)	Solved template: same age-gap question structure on a different pair. Verification check: “Cross-verify the identity of the less prominent individual against the primary subject’s known associates or event context, as misidentification of this figure is the most frequent source of error.”	Arithmetic: Age Problems and Data Extraction: “Subtract the earlier birth year from the later birth year, ensuring the values belong to the exact subjects specified in the query.”	Gold: 7 Baseline: 3 (used wrong birth year for Kaljurand: 1958) Ours: 7 (corrects to 1962, then $1962{-}1955$ )

Figure 2: Qualitative cases where DG-Mem corrects the No-memory backbone. Each row is a held-out test instance for which both retrieved grains are populated and the same baseline error is corrected under category-only , exemplar-only , and both-memory retrieval, ruling out the alternative explanation that any prompt perturbation flips the answer. The retrieved schema rule names the procedure the baseline omitted (an exhaustive census; per-item threshold computation; grounding the arithmetic in the correct entities), and the retrieved exemplar supplies a solved template whose verification check directly addresses the failure mode (boundary-bar evaluation; identity disambiguation of the less prominent subject).

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Benchmark	Train (Stage 1 & 2)	Test (held out)	Total
MathVista (MINI)	750	250	1,000
MMMU (DEV+VAL)	788	262	1,050
MMMU-Pro	1,298	432	1,730

Appendix

Table 4: Train/test partitions per benchmark. The same split file is used by every backbone and by every method (No-memory, ViLoMem, DG-Mem) so that comparisons are head-to-head.

Hyperparameter	Value	Description
$\alpha$	0.8	Similarity / utility interpolation in Eq. ( 5 )
Category similarity threshold	0.40	Min cosine similarity for category retrieval
Schema rule similarity threshold	0.35	Min cosine similarity for rule retrieval
$K_{\max}$ (rules per category)	2	Top- $k$ rules retrieved per active category

Appendix

Table 5: Test-time retrieval and decoding hyperparameters.

Datasets	MathVista*	MMMU*	MMMU-Pro*
Baseline	64.8	67.6	48.2
– w/o $\mathcal{M}_{S}$ (stage-1 exemplar only)	68.8 +4.0	72.5 +4.9	55.6 +7.4
– w/o $\mathcal{M}_{E}$ (stage-2 schema only)	67.2 +2.4	74.0 +6.4	58.1 +9.9
Full method	74.4 +9.6	79.4 +11.8	64.4 +16.2

Appendix

Table 6: Component ablations on GPT-5-Nano. Each row disables a single mechanism while leaving the rest of the pipeline intact.

Datasets	MathVista*	MMMU*	MMMU-Pro*
Baseline	83.6	85.8	77.9
– w/o $\mathcal{M}_{S}$ (stage-1 exemplar only)	87.2 +3.6	85.8 +0	80.9 +3.0
– w/o $\mathcal{M}_{E}$ (stage-2 schema only)	86.4 +2.8	86.5 +0.7	81.3 +3.4
Full method	88.7 +5.1	89.2 +3.4	86.1 +8.2

Appendix

Table 7: Component ablations on Qwen3.5-122B-A10B. Each row disables a single mechanism while leaving the rest of the pipeline intact.

Memory-Augmented VLMs

Memory-augmented vision-language models use persistent or retrieved context to improve reasoning across long inputs and interactions.

See all 160 papers

MemTrain: Self-Supervised Context Memory Training
Jun 2, 2026Ziheng Li, Xingrun Xing, Haoqing Wang +2LLM Agent Memory Persistent Memory for Language Models
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Samsung Research, Beijing, China
Mem- $π$ : Adaptive Memory through Learning When and What to Generate
May 20, 2026Xiaoqiang Wang, Chao Wang, Hadi Nekoei +5LLM Agent Memory RL for LLM Agents
1ServiceNow AI Research · 2Mila – Quebec AI Institute · 3Université de Montréal +3
Memory as a Controlled Process: Learned Adaptive Memory Management for LLM Agents
Jul 15, 2026Eric Hanchen Jiang, Zhi Zhang, Yuchen Wu +11Persistent Memory for Language Models Memory Management for LLM Agents
University of California Los Angeles · University of Washington · Northwestern University

Jun 2, 2026J/K move · Enter open · S save

MemTrain: Self-Supervised Context Memory Training

Ziheng Li, Xingrun Xing, Haoqing Wang, Zhi-Hong Deng+1

State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Samsung Research, Beijing, China

LLM Agent Memory Persistent Memory for Language Models Long-Horizon LLM Agents Memory-Augmented Language Model Agents Memory Management for LLM Agents

Open paper AskPDF arXiv

Memory is an indispensable capability for long-horizon LLM agents, enabling them to preserve and utilize information accumulated across extended interactions. Existing memory-agent approaches are typically trained end-to-end with reinforcement learning on downstream tasks. However, collecting high-quality annotated problems for memory-intensive scenarios is costly, and the resulting training data often lack sufficient diversity to cover general memory behaviors. In this work, we propose MemTrain, a self-supervised training framework for generally enhancing the context-memory capability of LLM agents for more effective downstream post-training. MemTrain introduces two coupled proxy tasks over unlabeled Wikipedia corpora: (1) an end-to-end masked reconstruction objective, which requires the model to recover masked entities after multiple rounds of memory updates, thereby encouraging memory maintenance from the final outcome perspective; and (2) an intermediate memory recall objective, which requires the model to reconstruct masked historical information using intermediate memory states, encouraging faithful compression and memory completeness throughout the interaction process. The two objectives are jointly optimized using GRPO. Extensive experiments on long-text QA and search-based QA benchmarks demonstrate that MemTrain consistently improves downstream memory-intensive reasoning performance across different models, achieving gains of up to 17.67 points over direct task-specific post-training.

Abstract

Memory-Augmented VLMs

Explore similar work

MemTrain: Self-Supervised Context Memory Training

Mem-πππ: Adaptive Memory through Learning When and What to Generate

Memory as a Controlled Process: Learned Adaptive Memory Management for LLM Agents

Mem- $π$ : Adaptive Memory through Learning When and What to Generate