MaxIM improves memory summaries for AI with better reward signals
MaxIM: Maximally Informative Incremental Summarization via Reinforcement Learning
Artificial IntelligenceMachine Learning
Summary
AI systems often need to remember important facts from long interactions but struggle because they get feedback only rarely and late. The authors developed MaxIM, a method that teaches AI to make short memory summaries that keep the most useful information for future tasks. It does this by using an information-based way to score memory steps more often, helping the AI learn which details to keep. Tests on reading scientific papers and dialogue tasks show MaxIM keeps facts better and helps AI perform tasks more reliably without needing lots of human feedback.
What this means in practice
- •For conversational ai teams: Create chatbot systems that maintain useful memory summaries for longer and diverse conversations without extensive human feedback.
- •For document understanding developers: Build AI systems that read and summarize lengthy technical documents by incrementally storing task-relevant information efficiently.
Authors
Jihwan Jeong, Guy Tennenholtz, Yinlam Chow, Chih-Wei Hsu, Craig Boutilier
Abstract
AI systems must often maintain bounded memory of unbounded interaction histories. However, optimizing the memory representation, and its updates, for downstream task utility is difficult for standard reinforcement learning (RL) due to severe credit assignment bottlenecks induced by sparse, delayed rewards. We introduce MaxIM (Maximally Informative Incremental Memory), a framework that formulates incremental summarization as a capacity-constrained, information-theoretic reward-shaping problem. Grounded in value suboptimality bounds, we show that bounding sequential information loss motivates maximizing the predictive log-likelihood of future, task-relevant targets, a variational lower bound on predictive sufficiency. We introduce a recursive consistency condition to ensure summaries are grounded by maximizing mutual information (MI) with the observed history. We mitigate credit assignment issues with a parameterized asymmetric critic that computes step-wise variational MI bounds for dense reward shaping. Empirical evaluation on long-horizon scientific reading and multi-turn dialogue benchmarks shows that MaxIM outperforms RL, supervised fine-tuning, and prompt-focused baselines on factual grounding, comprehensive autorater metrics, and downstream task utility across strict token budgets and extended horizons, without the need for expensive human preference labels.