Policy improves memory use in long document understanding
Learning What to Remember: Long-horizon Counterfactual Memory Optimization
Computation and LanguageMachine Learning
Summary
Memory in language models helps remember important information over long talks or texts, but figuring out what to keep is tricky. The authors created a way to measure how much each memory update actually helps later tasks, so the model learns which memories are worth keeping. They tested this on extracting details from documents and found their method keeps better facts while using much less memory. Their approach also works well when applied to different topics without extra training.
What this means in practice
- •For natural language processing teams: Enhance document understanding systems by learning which information to retain for improved extraction accuracy with smaller memory needs.
- •For enterprise chatbot developers: Build chatbots that remember important user details over long conversations without bloating memory, improving response relevance.$Commercial implications: Enables more efficient persistent memory handling in commercial chatbots, reducing costs and improving engagement with fewer resources.
Authors
Jiaming Tang, Mingyan Liu, Armin Sarabi
Abstract
Persistent textual memory allows language models to carry information across long interactions, but learning what to remember is fundamentally a credit-assignment problem. A memory rewrite may only become useful many steps later, while much of the observed utility may be inherited from information already stored before the rewrite. We introduce Memory Gain Policy Optimization (MGPO), which isolates the incremental value of each memory rewrite by crediting it for its marginal contribution to current and future downstream utility. This turns delayed memory utility into a direct learning signal for optimizing what information should persist. We study MGPO on document-level information extraction, where structured supervision makes the effects of individual memory updates directly measurable. MGPO improves extraction while reducing average memory length by nearly 80% relative to the initial memory policy before optimization. The learned memory policy also supports reuse and transfer across domains, downstream models without further training. These results show that effective memory learning depends not only on preserving useful information, but on identifying which memory updates create lasting incremental value.