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

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.

• The paper introduces an answer-conditioned information gain reward that guides memory formation in chunk-wise long-context agents, leading to superior QA performance.
• The methodology computes per-token likelihood gains by comparing answer generation with and without memory, using normalized rewards to stabilize RL training.
• Empirical results demonstrate that InfoMem outperforms outcome-only and recall-augmented baselines across multiple benchmarks like CorpusQA and LongMemEval.

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

Introduction and Motivation

Long-context question answering (QA) tasks with LLMs remain challenging due to the necessity of identifying and retaining sparse, answer-relevant information within massive input contexts. Chunk-wise memory agent architectures sequentially process documents in manageable units, cumulatively updating a compact memory state. However, optimal memory-state formation is hindered by crude reward signals, most commonly sparse final-answer correctness or intermediate lexical overlap. Such rewards are insufficient for directly supervising whether the accumulated memory facilitates accurate answer generation.

InfoMem introduces a principled reward signal for chunk-wise memory agents, leveraging answer-conditioned information gain. The core mechanism compares the model's per-token log-likelihood of the ground-truth answer with and without access to the final memory, establishing whether the final memory genuinely supports answer generation. This reward is applied only to successful answer trajectories to avoid signal collapse and is normalized before being incorporated into the final policy optimization objective, addressing reward scale variance across prompts.

Figure 1: Schematic of InfoMem evaluating chunk-wise long-context RL via the final-memory answer-conditioned information-gain reward, applied on successful rollouts and normalized prior to GRPO reward composition.

Methodology

Theoretical Framework

InfoMem operationalizes the intuition that final memory $M$ should reduce epistemic uncertainty over the answer $Y$ given query $X$. The ideal metric is conditional mutual information $I(M; Y \mid X)$, but due to intractability in high-dimensional LLM latent spaces, InfoMem uses a single-sample surrogate:

$$r_{\mathrm{gain}}(x, M, y^*) = \frac{1}{|y^*|}\log P_\theta(y^* \mid x, M) - \frac{1}{|y^*|}\log P_\theta(y^* \mid x, \emptyset)$$

where the terms denote the per-token log-likelihood of the ground-truth answer under the presence or absence of the constructed memory.

Reward Implementation

• Trajectory Selection: The information-gain signal is computed only among successful trajectories (i.e., those yielding the correct answer), exploiting the property that among such rollouts, variances in $r_{\mathrm{gain}}$ reflect substantive differences in memory utility rather than confounding errors in earlier retrieval or generation stages.
• Normalization: To control for inter-prompt signal scale variance, the information-gain rewards across the successful trajectory set are standardized before being weighted and aggregated with the sparse outcome reward. This prevents reward dominance by easy or hard questions and regularizes optimization steps.
• Optimization: The final composed reward is then used for RL-based policy optimization (specifically, via GRPO), propagating the trajectory-level advantage to all tokens, including intermediate memory updates.

Diagnostic Analyses

Synthetic Evidence Discrimination

A synthetic evaluation demonstrates InfoMem's discriminative capacity. For SQuAD-derived QA, synthetic 'hallucinated' contexts are generated which preserve surface and topical characteristics but alter key facts, decoupling surface similarity and evidential support.

Figure 2: Example of synthetic hallucinated evidence constructed to mimic topic and style but obfuscate answer-relevant facts.

Among several evidence scoring schemes (embedding similarity, attention-based metrics, InfoMem's answer-conditioned log-likelihood gain), only the information gain achieves near-perfect mean reciprocal rank for the ground-truth context and exhibits a much higher standardized signal-to-noise ratio. This result substantiates that InfoMem's reward signal uniquely and stably captures answer-supportive memory content.

Supervision Side Ablation

Ablations compare the efficacy of restricting information-gain supervision to successful, failed, or all rollouts. Experiments demonstrate that restricting the signal to successful (i.e., correct-answer) rollouts provides the only stable and effective reward. Applying the signal to wrong or all rollouts induces instability and collapses validation performance, as failed rollouts entangle noise from faulty evidence extraction and answer prediction.

Figure 3: Top: Validation accuracy with varied information-gain supervision sides; only success-side supervision yields sustained performance. Bottom: Training accuracy without explicit outcome reward also underscores instability outside of success-only supervision.

Reward Conditioning and Exploitation Analysis

If the reward is conditioned on the query rather than on the answer (termed QueryPMI), models rapidly degenerate to copying/paraphrasing the query into memory, neglecting genuine evidence consolidation, and resulting in significant performance drops on all long-context tasks.

Figure 4: Fraction of rollouts with memory repetitions of the query increases rapidly under QueryPMI; InfoMem avoids this mode collapse.

Empirical Results

Benchmarks span multi-needle retrieval (MRCR-8needle), long-context corpus QA (CorpusQA), dialog memory (LongMemEval), and synthetic RULER QA. InfoMem achieves the highest accuracy on all benchmarks—outperforming outcome-only GRPO and the recall-augmented ReMemR1 RL baseline—with CorpusQA (19.5%), LongMemEval (12.8%), RULER synthetic QA (36.8%), and MRCR-8needle (0.279). Notably, outcome-only GRPO degrades on MRCR-8needle, confirming that outcome-only optimization fails in challenging retrieval tasks.

Learning curves further demonstrate that InfoMem achieves stable and accelerating training and validation performance, whereas alternatives plateau or degrade mid-training.

Implications and Future Directions

This work identifies three critical properties of effective final-memory rewards for chunk-wise long-context RL: (1) supervision is restricted to successful trajectories, (2) information-gain is normalized prior to reward composition, and (3) the reward is strictly answer-conditioned. These criteria distinguish useful, evidence-preserving memories from degenerate or spurious solutions.

InfoMem strengthens the direct optimization of non-parametric, actionable memory policies in long-context agents. In practical deployment, it provides a tractable, reward-regularized framework for tuning LLM driven agents on large-scale, noisy corpora with sparse supervision. Theoretically, it further aligns memory agent RL with information-theoretic objectives, bridging the gap between representation learning and optimal answer-inference under resource constraints.

Several avenues remain: scaling InfoMem with larger model backbones and broader datasets, adapting reward signals for intermediate or stepwise memory state shaping, and extending this reward framework to migration architectures such as retrieval-augmented generation and revisitable-memory agents.

Conclusion

InfoMem provides an efficacious, answer-conditioned information-gain reward for chunk-wise long-context memory agents, enforcing direct supervision on memory utility and yielding consistent improvements on diverse benchmarks relative to both outcome-only and lexical recall-based RL baselines. Its design principles delineate the necessary reward-shaping strategies for optimal memory-agent policy learning in high-noise, long-context environments. The approach augments the methodological toolkit for RL fine-tuning of LLM-based agents and motivates further research in information-theoretic reward design and scalable long-context architecture optimization.