ahmedehabb/Memory-R2

Memory-R2 7B — Memory Manager

The trained memory-management policy from Memory-R2: Fair Credit Assignment for Long-Horizon Memory-Augmented LLM Agents (arXiv:2605.21768). This is the paper's main contribution and deployed "champion" (32sess_champion_v2, LoGo-GRPO curriculum, global step 5).

It is a Qwen2.5-7B-Instruct model fine-tuned with LoGo-GRPO (turn-level + token-level credit assignment) via a curriculum of 8 → 16 → 32-session rollouts on the LoCoMo long-horizon dialogue dataset. Given a running conversation, it decides what to INSERT / UPDATE / DELETE in an external memory store.

This model only manages memory — it does not answer questions. A separate answer agent reads the memory store this model produces and generates answers; it can be any instruction-tuned LLM. Our own SFT+RL-trained answer agent is released separately at [ahmedehabb/Memory-R2-answer-agent](https://huggingface.co/ahmedehabb/Memory-R2-answer-agent).

Headline results (tab:main)

This memory manager is held constant; only the paired answer agent changes:

See the paper's tab:different-answer-agent for more pairings (untrained Qwen-7B, etc.) — the memory manager is not tied to any one answer agent.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

memory_manager = AutoModelForCausalLM.from_pretrained("ahmedehabb/Memory-R2", torch_dtype="auto", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("ahmedehabb/Memory-R2")

Full inference code and the memory-store protocol are in the project repository (see the paper for the official release).

Training

• —Base model: Qwen/Qwen2.5-7B-Instruct
• —Algorithm: LoGo-GRPO (turn-level + token-level advantage), curriculum-trained 8-session → 16-session → 32-session
• —Reward: per-session cumulative F1 against gold QA + a memory-compression penalty (λ=0.3)
• —Judge for reward/logging during training: GPT-OSS-120B

Citation

@misc{yan2026memoryr2faircreditassignment,
      title={Memory-R2: Fair Credit Assignment for Long-Horizon Memory-Augmented LLM Agents},
      author={Sikuan Yan and Ahmed Bahloul and Ercong Nie and Susanna Schwarzmann and Riccardo Trivisonno and Volker Tresp and Yunpu Ma},
      year={2026},
      eprint={2605.21768},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.21768},
}