AttriMem: Attribution-Guided Process Feedback for Agent Memory Learning

2026-07-23

AI Summary

Current memory construction methods for large language model agents suffer from limited generalization and misalignment with downstream task objectives due to the absence of fine-grained credit assignment. This work proposes an attribution-guided process feedback framework that, for the first time, incorporates token-level attribution signals to jointly leverage global outcome rewards and local contribution rewards. By employing reinforcement learning, the framework dynamically optimizes strategies for memory retrieval, storage, and updating. This approach overcomes the conventional reliance of reinforcement learning on coarse-grained rewards and demonstrates significant performance gains over retrieval-based, heuristic, and existing RL methods in long-horizon conversational question answering. Moreover, it exhibits strong cross-benchmark and cross-answer-model generalization while substantially improving training stability.

Abstract

Effective memory is crucial for LLM agents, yet constructing it effectively remains challenging. A memory-construction policy decides what information to extract, store, update, compress, or discard as interactions accumulate. Heuristic memory methods rely on subjective, task-specific rules, which can misalign with downstream objectives and limit cross-task adaptability. RL-based methods, by contrast, learn from task feedback but mainly use outcome- or module-level rewards. These coarse signals indicate task success but cannot identify which intermediate memory contents support the final answer, creating a fine-grained credit-assignment bottleneck. However, constructing such process feedback is prohibitively difficult because intermediate memory decisions lack unique ground-truth targets, while the appropriate credit varies with the agent's uncertain reasoning trajectory and therefore cannot be specified in advance. We propose AttriMem, an attribution-guided process-feedback framework for learning memory-construction policies with RL. AttriMem augments the global outcome reward with local rewards derived from token-level contributions to the final answer. Experiments on long-horizon dialogue question answering show that AttriMem outperforms retrieval-based, heuristic, and RL-based baselines, generalizes across benchmarks and answer models, stabilizes RL optimization.

Problem

• memory construction
• credit assignment
• process feedback
• LLM agents
• reinforcement learning

Innovation

• attribution-guided feedback
• memory construction
• reinforcement learning
• token-level contribution
• process reward