Skip to content

AttriMem: Attribution-Guided Process Feedback for Agent Memory Learning

🕒 Published (v1): 2026-07-23 09:35 UTC · Source: Arxiv · link

Why this paper was selected

Attribution-guided memory learning for agents; addresses what to store/compress/discard—core harness problem

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

AttriMem introduces an attribution-guided process-feedback framework for training LLM agent memory-construction policies via RL. Rather than relying on sparse outcome-level rewards, it uses token-level attribution (via ContextCite) to estimate each intermediate memory action's causal contribution to the final answer and converts these scores into dense local process rewards. On long-horizon dialogue QA, it consistently outperforms retrieval-based, heuristic, and RL-based baselines across three benchmarks.

Problem

Existing RL-based memory-learning methods for LLM agents suffer from a fine-grained credit-assignment bottleneck: outcome-only rewards indicate whether the task succeeded but cannot identify which specific memory contents (at token or span granularity) caused success or failure. Redistributing final rewards (e.g., Mem-T) or deriving action-level QA rewards (e.g., MemBuilder) still cannot localize credit within individual memory operations. Constructing ground-truth process signals is prohibitively difficult because intermediate memory decisions have no unique correct target—appropriate credit depends on the agent's uncertain reasoning trajectory.

Method

AttriMem augments GRPO-based RL training with attribution-derived process rewards alongside the global outcome reward. The combined per-step advantage is:

\[\hat{A}_{i,t} = \hat{A}^{\text{out}}_i + \lambda \hat{A}^{\text{proc}}_{i,t}\]

For each memory action at step \(t\) that produces textual update \(\Delta M_t\), AttriMem applies ContextCite (Cohen-Wang et al., 2024): it randomly ablates 32 subsets of memory tokens, fits a lightweight linear model to the resulting answer-score changes, and estimates per-token attribution \(\phi((\Delta M_t)^{(j)}, \hat{y}; c)\) as the counterfactual drop in final-answer score when that token is removed. The process reward for step \(t\) is:

\[r_t = \frac{1}{\sqrt{l_t}} \sum_{j=1}^{l_t} \phi\!\left((\Delta M_t)^{(j)},\, \hat{y};\, c\right)\]

where the \(\sqrt{l_t}\) normalization mitigates length bias. The 32 ablated variants are processed in a single batched forward pass, adding roughly one forward-pass of latency. The base model is Qwen3-4B with SFT warmup (3,000 steps) followed by GRPO RL (400 steps). Memory is structured into four typed stores: core (user profile), episodic (timestamped events), semantic (entity facts), and procedural (workflows).

Key Contributions

  • Identifies the fine-grained credit-assignment bottleneck as the central failure mode of existing RL-based memory-learning methods.
  • Formally characterizes why intermediate process ground-truth is unattainable: the appropriate credit is trajectory-dependent and cannot be specified in advance.
  • Proposes AttriMem: the first framework to derive token-level process rewards from causal attribution of final-answer contribution back to individual memory operations.
  • Demonstrates that token-level attribution rewards improve both in-domain performance and zero-shot cross-domain generalization, and stabilize RL optimization relative to outcome-only and action-level reward baselines.

Results

  • LoCoMo (out-of-domain): AttriMem (SFT+RLTok) achieves 82.48, outperforming MemBuilder (SFT+RL) at 80.01 by +2.47 and the best heuristic baseline (LightMem, 80.53) by +1.95.
  • LongMemEval (in-domain): 83.25, outperforming MemBuilder at 81.75 by +1.50.
  • PerLTQA (out-of-domain): 84.49, outperforming MemBuilder at 83.74 by +0.75.
  • AttriMem achieves best performance across all three benchmarks among all compared methods.
  • Generalizes across answer models: with Claude 4.5 Sonnet (82.48 / 83.25 / 84.49), GPT-4.1 (80.41 / 82.00 / 82.18), and Qwen3-4B variants, showing the memory policy transfers to unseen readers.
  • Qualitative analysis shows AttriMem's Step-400 memory retains precise factual spans (e.g., "3 books/2 months," "20 bottles saved," "bus pass replaced twice") that MemBuilder distorts or omits.

Limitations

  • Attribution overhead: ContextCite requires 32 ablated context variants per training step; although parallelizable to approximately one additional batched forward pass, this increases training cost.
  • Evaluation is confined to long-horizon dialogue QA; generalization to other agentic task types (code execution, tool use, web navigation) is not demonstrated.
  • The \(\lambda\) hyperparameter weighting process vs. outcome advantage requires tuning and may be task-sensitive.
  • Only 400 RL steps are used; the behavior under longer training or larger models is not analyzed.

Relevance to Agentic AI / LLM Agents

Credit assignment across long action sequences is a foundational problem in training any LLM agent that takes sequences of intermediate decisions (memory writes, tool calls, planning steps) before receiving a final reward. AttriMem's insight—that token-level attribution from a final response can serve as surrogate process supervision without requiring labeled intermediate targets—is directly transferable to other agentic settings where intermediate ground truth is unavailable. The framework addresses a structural weakness in GRPO-style RLVR training, which has become the dominant paradigm for agent fine-tuning, and the four-type memory architecture it trains over is representative of production-grade agent memory systems. For researchers tracking memory-augmented and long-horizon LLM agents, this work establishes a concrete technical path beyond outcome-only RL.