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

Source: [InfoMem: Training Long-Context Memory Agents with Answer-Conditioned Information Gain](<https://www.sota2.com/research/paper/6a210666c10484bac23eff3b>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2026-06-04T05:00:05.280Z

Authors: Tiancheng Han, Yong Li, Wuzhou Yu, Qiaosheng Zhang, Wenqi Shao

Published: 2026-06-02T08:39:03.000Z

## 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.

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

- [Corpus-level evidence aggregation on CorpusQA 128K-token](<https://www.sota2.com/research/sota/corpus-level-evidence-aggregation-on-corpusqa-128k-token.md>): Corpus-level evidence aggregation; CorpusQA 128K-token; LLM-as-Judge Accuracy: 19.453.

- [Long-horizon Dialogue Memory Tracking on LongMemEval 115K-token S](<https://www.sota2.com/research/sota/long-horizon-dialogue-memory-tracking-on-longmemeval-115k-token-s.md>): Long-horizon dialogue memory and conversational information tracking; LongMemEval 115K-token S; LLM-as-Judge Accuracy: 12.8.

- [Multi-needle retrieval on MRCR-8needle 100-example](<https://www.sota2.com/research/sota/multi-needle-retrieval-on-mrcr-8needle-100-example.md>): Multi-needle retrieval; MRCR-8needle 100-example; Sequence Match Accuracy: 27.9.

- [Sparse retrieval-style question answering on RULER synthetic QA (standard)](<https://www.sota2.com/research/sota/sparse-retrieval-style-question-answering-on-ruler-synthetic-qa-standard.md>): Sparse retrieval-style question answering; RULER synthetic QA (standard); F1 Score: 0.36848.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2606.03329>)

[PDF](<https://arxiv.org/pdf/2606.03329.pdf>)
