Data Construction

This repo uses a unified long-context synthesis pipeline to convert standard QA examples into single long documents by mixing gold evidence with many in-domain distractors. The goal is a controlled setting where the answer is always supported by gold evidence, while retrieval difficulty scales with context length.

---

Overview: Unified Long-Context Synthesis Pipeline

Each synthesized instance follows the same supervision format:

• Question: Q
• Answer: A
• Gold evidence set: D_gold (documents/paragraphs annotated by the source dataset)
• Distractor pool: D_dist (sampled from the same dataset corpus)

Construction Steps

For each original QA instance:

1. Build three text pools: Q, D_gold, and D_dist (in-domain distractors).
2. Shuffle candidate documents with a fixed random seed.
3. Insert each gold document exactly once (at the document level).
4. Append distractors until reaching a target token budget.
5. Truncate at document boundaries (avoid truncating inside a document whenever possible).

Why In-Domain Distractors?

We sample distractors from the same corpus to preserve domain/style match, making retrieval harder than mixing out-of-domain noise.

---

Training Data

We adopt a two-stage training scheme to reliably execute the PreThink–Retrieve–Write protocol under a bounded-memory constraint:

1. Cold-start SFT: distills protocol-following behavior (retrieve calls + concise memory writing) from a strong teacher.
2. RL alignment: improves long-horizon control decisions (when to retrieve, how to write memory, when to stop) under long-context inputs.

Cold-start SFT Data

We synthesize cold-start SFT data following the NIAH-style long-context QA construction used in MemAgent.

• Sources: HotpotQA, SQuAD, MuSiQue
• Samples per source: 4,096 (training split)
• Target context length: 32K tokens

Teacher Traces (Protocol-Consistent)

We use Qwen3-32B (thinking enabled) as a teacher to generate interaction traces that follow PreThink–Retrieve–Write:

• plan + emit structured retrieve calls,
• update bounded memory by writing compressed evidence,
• decide when to stop retrieving and answer.

Student Backbones (Distillation Targets)

We distill trajectories into the following student backbones:

• Qwen3-1.7B
• Qwen3-4B
• Qwen2.5-7B-Instruct

MuSiQue Question Decompositions (Teacher-Only)

MuSiQue provides optional multi-hop decompositions (sub-questions). We feed decompositions only to the teacher to elicit cleaner planning traces.
Students never observe decompositions, gold document IDs, or teacher-side annotations during training or inference.

Trajectory Filtering

To ensure supervision quality:

• Keep only traces whose final answers are correct under official dataset evaluation.
• Remove traces that would exceed the memory budget or require truncating the agent memory/state.

After filtering, we decompose successful traces into turns, yielding 29,717 single-turn dialogue instances for SFT.

---

RL Training Data

For RL training, we apply the same synthesis pipeline to extend HotpotQA instances to ~28K tokens:

• Base QA pairs: original HotpotQA question–answer pairs (unchanged)
• Context scaling: insert more in-domain distractors to increase retrieval difficulty
• Reward: Exact Match (EM) between generated answer and ground-truth answer

RL focuses on improving agent control decisions under long contexts (retrieval timing, memory management, early stopping) while keeping the environment consistent with SFT.

---

Evaluation Benchmark (Extreme Context Scaling)

To evaluate robustness under extreme context lengths, we create synthesized long-context benchmarks from:

• Tasks: HotpotQA, 2WikiMultihopQA, MuSiQue
• Comparison setting: synthetic SQuAD setting used in MemAgent (for direct comparison)
• Samples per task: 128 (test split)

Length Variants

For each fixed set of questions, we generate multiple variants at increasing target lengths, e.g.:

• ~32K / 28K
• ~64K / 56K
• ~128K / 112K
• ...
• up to ~1M / 896K tokens

Gold evidence is inserted once per instance (fixed seed), and distractors are sampled from the same corpus.
This ensures that differences across lengths reflect context scaling only (more distractors / longer inputs), not changes in questions or evidence.

Links

• Paper: https://arxiv.org/pdf/2602.02704
• GitHub: https://github.com/UCMP13753/InfMem

Citation

If you find this work useful, please consider citing our paper:

@misc{wang2026infmemlearningsystem2memory,
      title={InfMem: Learning System-2 Memory Control for Long-Context Agent}, 
      author={Xinyu Wang and Mingze Li and Peng Lu and Xiao-Wen Chang and Lifeng Shang and Jinping Li and Fei Mi and Prasanna Parthasarathi and Yufei Cui},
      year={2026},
      eprint={2602.02704},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2602.02704}, 
}