InfoMem: Training Long-Context Memory Agents with Answer-Conditioned Information Gain
The paper introduces InfoMem, a new reward mechanism designed for training long-context memory agents in artificial intelligence. It focuses on improving the evaluation of final memory utility by using answer-conditioned information. The proposed method shows enhanced performance over existing reinforcement learning baselines in managing long-context tasks.
- ▪InfoMem evaluates final-memory utility based on how much it increases the model's likelihood of the ground-truth answer.
- ▪The mechanism applies rewards only to successful trajectories and normalizes them before composition.
- ▪InfoMem improves the performance of chunk-wise memory agents under the same training conditions compared to existing methods.
Opening excerpt (first ~120 words) tap to expand
Computer Science > Artificial Intelligence arXiv:2606.03329 (cs) [Submitted on 2 Jun 2026] Title:InfoMem: Training Long-Context Memory Agents with Answer-Conditioned Information Gain Authors:Tiancheng Han, Yong Li, Wuzhou Yu, Qiaosheng Zhang, Wenqi Shao View a PDF of the paper titled InfoMem: Training Long-Context Memory Agents with Answer-Conditioned Information Gain, by Tiancheng Han and 4 other authors View PDF HTML (experimental) 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.
…
Excerpt limited to ~120 words for fair-use compliance. The full article is at arXiv cs.AI.
Discussion
0 commentsMore from arXiv cs.AI
