🤖 AI Summary
Existing recurrent memory agents suffer from sharply degraded memory retention in ultra-long contexts due to their single memory block architecture, leading to unreliable performance. This work proposes a training-agnostic Multi-Head Memory (MHM) framework that partitions memory into multiple independent heads and employs a staged “select-update” mechanism—updating only one head per step while shielding the others from overwriting—thereby shifting the burden of memory retention from agent behavior to architectural design. The MHM-LRU variant achieves uniform head utilization with zero additional token overhead. Across context lengths from 100K to 1M tokens, the approach substantially improves both memory retention and end-to-end accuracy, boosting retention on RULER-HQA (896K tokens) from under 30% to 73.96%, with consistent gains observed across diverse model families, scales, and tasks.
📝 Abstract
Recurrent memory agents extend LLMs to arbitrarily long contexts by iteratively consolidating input into a fixed-size memory window. Despite their scalability, these agents exhibit a well-documented reliability problem: end-to-end performance degrades systematically as context length grows. We diagnose this failure by decomposing performance into two factors--memory capture and memory retention--and quantitatively confirm that retention is the dominant bottleneck. Retention collapses because existing designs maintain memory as a monolithic text block, forcing every update to risk overwriting previously retained content. Motivated by this diagnosis, we propose Multi-Head Recurrent Memory (MHM), a general, training-free framework that partitions memory into independent heads governed by a stage-wise select-then-update strategy. At each step, exactly one head is selected for update while the remaining heads are structurally shielded from overwriting, shifting the burden of retention from model behavior to architectural design. As a lightweight instantiation, we introduce Least-Recently-Updated MHM (MHM-LRU), which guarantees uniform head utilization with zero additional token overhead. Extensive experiments on long-context benchmarks show that MHM-LRU substantially improves both retention and end-to-end accuracy across the 100K--1M token range, where baselines degrade sharply. On RULER-HQA at 896K tokens, MHM-LRU improves the memory retention rate from less than 30% to 73.96%. These gains generalize across model families, scales, and task types, positioning architectural optimization as a practical and cost-efficient path toward reliable long-context recurrent memory.