Mechanistic Attention Guidance for Agent Memory Refinement

2026-07-20

AI Summary

Current agent memory systems predominantly rely on textual outputs for optimization, making error attribution difficult and prone to hallucination. This work introduces, for the first time, a mechanistic attention signal to analyze retrieval-head attention patterns, constructing a context-utilization matrix that reveals how memory segments are actually employed during task execution. Leveraging this insight, the method enables segment-level memory correction and enhancement through attention aggregation, context-utilization matrix construction, targeted memory updates, and re-execution validation. Evaluated on interactive decision-making benchmarks, the approach significantly outperforms text-only baselines, simultaneously improving both task performance and memory efficiency.

Abstract

Existing self-evolving memory systems mainly improve agent memory based on textual outputs, such as task trajectories and reflections. However, this text-based paradigm rarely incorporates internal mechanistic signals, leaving how retrieved memory is actually utilized during task execution underexplored. This limitation can lead to unreliable error attribution and hallucinated memory modifications. In this work, we show that retrieval-head attention provides a mechanistic signal for revealing segment-level memory utilization. By aggregating attention over memory segments and decision steps, we construct a context utilization matrix that exposes recurring memory-use patterns and indicates corresponding refinement strategies. Building on this observation, we propose Attention-Guided Memory Refinement (AGMR), a framework that uses utilization patterns revealed by attention to guide targeted segment-level memory updates. AGMR corrects or enhances memory for failed executions, simplifies memory for successful executions, and verifies each update through re-execution. Experiments on interactive decision-making benchmarks show that AGMR improves both task performance and memory efficiency over text-only memory refinement baselines. Code is available at https://anonymous.4open.science/r/AGMR_code-3262/

Problem

• agent memory
• mechanistic signals
• memory utilization
• error attribution
• hallucinated memory

Innovation

• mechanistic attention
• memory refinement
• attention-guided learning
• agent memory
• context utilization matrix