Memory Papers
Memory Architecture

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

AuthorsTiancheng Han, Yong Li, Wuzhou Yu et al.

arXiv 20262026
arXivPDF

TL;DR

InfoMem uses answer-conditioned information-gain rewards on successful trajectories to shape final memory, boosting CorpusQA from 16.413% to 19.453% over Outcome-only GRPO.

SharePost on XLinkedIn

Read our summary here, or open the publisher PDF on the next tab.

THE PROBLEM

Outcome-only rewards cannot judge final memory utility

Existing chunk-wise memory agents rely on sparse final-answer rewards or lexical intermediate rewards that do not directly evaluate whether the final memory supports the ground-truth answer.

As a result, long-context question answering and corpus-level evidence aggregation can succeed or fail with indistinguishable rewards, leaving memory quality, answer support, and evidence selection under-supervised.

HOW IT WORKS

InfoMem — Answer-conditioned Information Gain for Final Memory

InfoMem introduces an answer-conditioned information-gain reward combined with Group Relative Policy Optimization (GRPO) and a chunk-wise long-context memory agent to directly score final-memory utility.

Think of InfoMem as checking whether the agent’s compact memory behaves like a high-bandwidth cache that actually makes the correct answer more probable, instead of just storing loosely related text.

This information-theoretic shaping lets InfoMem learn final memories that reduce uncertainty about the ground-truth answer, something a plain context window or outcome-only reward cannot target explicitly.

DIAGRAM

Chunk-wise Reading and Answer-conditioned Info Gain Flow

This diagram shows how InfoMem sequentially reads chunks, builds a final memory, and computes answer-conditioned information gain within a GRPO rollout.

DIAGRAM

Training and Evaluation Pipeline for Long-context Benchmarks

This diagram shows how InfoMem is trained on RULER-HotpotQA and evaluated on four long-context benchmarks under a shared chunk-wise memory framework.

PROCESS

How InfoMem Handles Chunk-wise Long-context Question Answering

  1. 01

    Chunk-wise Long-context Memory Agent

    InfoMem instantiates the chunk-wise long-context memory agent by dividing D into chunks {c1,...,cK} and updating Mt = πθ(Mt−1, ct, x) for each step.

  2. 02

    Final Memory and Answer Generation

    After processing all chunks, InfoMem obtains the final memory MK and generates the final answer ŷ = πθ(x, MK) conditioned on this compact memory.

  3. 03

    Answer-conditioned Information-gain Reward Definition

    InfoMem computes rgain(x, M, y*) as the per-token average log-likelihood difference of y* with final memory MK versus null memory ∅ under teacher forcing.

  4. 04

    Using Successful Trajectories as Positive Memory Signals

    InfoMem applies rgain only to successful rollouts S, normalizes rgain within S, and combines it with Routcome,i inside GRPO for policy optimization.

KEY CONTRIBUTIONS

Key Contributions

  • 01

    Information-theoretic Final-memory Utility

    InfoMem formulates final-memory utility via answer-conditioned information gain, where useful memories reduce uncertainty about the ground-truth answer y* given query x.

  • 02

    Answer-conditioned Information-gain Reward

    InfoMem introduces an answer-conditioned information-gain reward applied only to successful trajectories and normalized before reward composition within GRPO.

  • 03

    Improved Chunk-wise Long-context Memory Agents

    InfoMem improves chunk-wise long-context memory agents over Outcome-only GRPO and ReMemR1, reaching 19.453% on CorpusQA and 36.848% on RULER synthetic QA.

RESULTS

By the Numbers

CorpusQA

19.453%

+3.040 over Outcome-only GRPO

LongMemEval

12.800%

+2.800 over Outcome-only GRPO

MRCR-8needle

0.279%

+0.216 over Outcome-only GRPO

RULER synthetic QA

36.848%

+2.113 over Outcome-only GRPO

Table 2 reports performance on CorpusQA, LongMemEval, MRCR-8needle, and RULER synthetic QA, which test corpus-level aggregation, dialogue memory, multi-needle retrieval, and synthetic long-context QA. These results show that InfoMem consistently improves chunk-wise long-context memory agents over Outcome-only GRPO and ReMemR1 under the same GRPO framework and training budget.

BENCHMARK

By the Numbers

Table 2 reports performance on CorpusQA, LongMemEval, MRCR-8needle, and RULER synthetic QA, which test corpus-level aggregation, dialogue memory, multi-needle retrieval, and synthetic long-context QA. These results show that InfoMem consistently improves chunk-wise long-context memory agents over Outcome-only GRPO and ReMemR1 under the same GRPO framework and training budget.

BENCHMARK

Main evaluation results on long-context benchmarks

Scores (%) on CorpusQA for Qwen2.5-1.5B-Instruct variants with different RL objectives.

KEY INSIGHT

The Counterintuitive Finding

InfoMem without rgain normalization drops CorpusQA from 19.453% to 16.109%, nearly erasing the +3.040 gain over Outcome-only GRPO.

This is surprising because adding a seemingly informative reward signal should help, but InfoMem shows that unnormalized information gain can destabilize training and harm long-context retrieval.

WHY IT MATTERS

What this unlocks for the field

InfoMem unlocks a way to directly supervise final-memory quality using answer-conditioned information gain rather than only sparse correctness or lexical overlap.

With InfoMem, builders can train chunk-wise long-context agents whose compact memories are explicitly optimized to support the correct answer, making long-document reasoning and corpus-level aggregation more reliable under fixed context budgets.

~12 min read← Back to papers

Related papers

Memory Architecture

A Control Architecture for Training-Free Memory Use

Yanzhen Lu, Muchen Jiang et al.

· 2026

TAG routes low-confidence steps to uncertainty-based routing, filters them with guarded acceptance with rollback, chooses between bank selection across rule and exemplar memory, and prunes via evidence-based retirement inside a unified control loop. On SVAMP and ASDiv, TAG reaches 81.0% and 85.2% accuracy, improving over the 74.0% and 77.5% no-memory baselines while a compute-matched Retry baseline stays flat.

arXiv:2604.18206Read explainer

Questions about this paper?

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

Answers use this explainer on Memory Papers.

Checking…