InfoMem: Training Long-Context Memory Agents with Answer-Conditioned Information Gain

2026-06-02

AI Summary

This work addresses the challenge large language models face in effectively extracting answer-relevant information from extensive contexts. The authors propose InfoMem, a method that sequentially processes document segments and dynamically updates a compact memory representation through a reinforcement learning framework based on GRPO. The key innovation lies in an answer-conditional information gain reward mechanism that directly quantifies how well the final memory supports the generation of the ground-truth answer. To enhance training stability, the approach incorporates successful trajectory filtering and reward normalization. Under identical training budgets, InfoMem substantially outperforms existing reinforcement learning–based memory agent baselines and achieves superior performance on long-context understanding tasks.

Abstract

Long-context tasks require LLMs to identify and preserve answer-relevant information from large contexts. Chunk-wise memory agents address this issue by sequentially reading document chunks, updating a compact memory, and generating the final answer from the accumulated memory. However, existing RL-based chunk-wise agents either rely on sparse final-answer rewards or use lexical intermediate rewards for memory and retrieval actions. These signals supervise task success or local overlap, but do not directly evaluate whether the final memory supports the ground-truth answer. We propose InfoMem, a reward mechanism for training chunk-wise memory agents that evaluates final-memory utility using answer-conditioned information. InfoMem measures how much the final memory increases the model's per-token log-likelihood of the ground-truth answer. To stabilize RL optimization, InfoMem applies this signal only to successful trajectories and normalizes it before reward composition. Under the same GRPO framework and training budget, InfoMem improves long-context memory-agent performance over comparable memory-agent RL baselines. Analyses show that effective final-memory rewards should operate on successful trajectories, be normalized before reward composition, and be conditioned on the answer rather than the query. Our code is available at https://github.com/GenSouKa1/InfoMem.

Problem

• long-context tasks
• memory agents
• reinforcement learning
• reward mechanism
• information gain

Innovation

• answer-conditioned information gain
• long-context memory agents
• reinforcement learning
• memory utility reward
• information-theoretic reward