This repo is the official implementation of [ICLR 2026] MemGen: Weaving Generative Latent Memory for Self-Evolving Agents.
Inspired by the human brain’s ability to dynamically integrate memory and reasoning, MemGen introduces a novel framework that empowers AI agents to evolve through experience—without relying on rigid parameter updates or external databases.
Unlike traditional approaches, MemGen generates latent memory tokens directly within the model’s reasoning stream. It features:
- A Memory Trigger that decides when to recall memory.
- A Memory Weaver that synthesizes past experiences into compact, latent sequences—seamlessly enriching ongoing reasoning.
A: DDP is supported, but FSDP is not currently supported. Thank you for your understanding.
A: TriviaQA multi-turn Weaver GRPO training and evaluation are available for the SmolLM3-3B Reasoner with either a Qwen2.5-1.5B-Instruct or SmolLM3-3B Weaver. See the scripts and checkpoints in the Quick Evaluation table below. Additional multi-turn environments, including AlfWorld, are planned for a future release.
A: In the previous version, single-turn training did not use the ChatML template (for both the baseline and MemGen), which led to lower performance. In addition, we identified a small but impactful formatting issue: for the 1.5B model, whether the prompt ends with \boxed{} followed by . or \n significantly affects performance. In particular, appending \n after “Put your answer within \boxed{}” can noticeably degrade results compared with appending .. While surprising, this behavior was consistent in our tests. The updated codebase consistently applies the ChatML template across all datasets and resolves these formatting inconsistencies. We still observe stable performance gains from MemGen under this unified setup.
We apologize for any inconvenience caused by the earlier version.
Create and activate the MemGen environment:
Option 1: Install via requirements.txt
conda create -n memgen python=3.10
conda activate memgen
pip install -r requirements.txt
Option 2: Install via memgen.yml
conda env create -f memgen.yml
conda activate memgen
Option 3: Set Up Search Environment
Please follow the instructions in the Search-R1 to configure the retriever environment.
Below are several MemGen models based on Qwen2.5-1.5B-Instruct and SmolLM3-3B across multiple datasets. We are currently in the process of carefully validating additional checkpoints to ensure they are fully reproducible and can be released in a clean, one-click setup. We appreciate your patience as we complete this verification process.
| model | dataset | mode | link | eval_script | train_script |
|---|---|---|---|---|---|
| Qwen2.5-1.5B-Instruct | KodCode | weaver-sft + trigger | huggingface link | scripts/eval/qwen2_5_kodcode_sft_trigger.sh |
scripts/train/qwen2_5_kodcode_sft_trigger.sh |
| Qwen2.5-1.5B-Instruct | KodCode | weaver-grpo | huggingface link | scripts/eval/qwen2_5_kodcode_grpo.sh |
scripts/train/qwen2_5_kodcode_grpo.sh |
| Qwen2.5-1.5B-Instruct | GSM8K | weaver-sft + trigger | huggingface link | scripts/eval/qwen2_5_gsm8k_sft_trigger.sh |
scripts/train/qwen2_5_gsm8k_sft_trigger.sh |
| Qwen2.5-1.5B-Instruct | GSM8K | weaver-grpo | huggingface link | scripts/eval/qwen2_5_gsm8k_grpo.sh |
scripts/train/qwen2_5_gsm8k_grpo.sh |
| Qwen2.5-1.5B-Instruct | TriviaQA | weaver-sft | huggingface link | scripts/eval/qwen2_5_triviaqa.sh |
scripts/train/qwen2_5_triviaqa.sh |
| SmolLM3-3B | KodCode | weaver-sft | huggingface link | scripts/eval/smollm_kodcode.sh |
scripts/train/smollm_kodcode.sh |
| SmolLM3-3B | GSM8K | weaver-sft + trigger | huggingface link | scripts/eval/smollm_gsm8k_sft_trigger.sh |
scripts/train/smollm_gsm8k_sft_trigger.sh |
| SmolLM3-3B (Weaver: Qwen2.5-1.5B-Instruct) | TriviaQA | weaver-sft | huggingface link | scripts/eval/smollm_triviaqa_qwen_weaver_sft.sh |
scripts/train/smollm_triviaqa_qwen_weaver_sft.sh |
| SmolLM3-3B (Weaver: Qwen2.5-1.5B-Instruct) | TriviaQA | weaver-grpo | huggingface link | scripts/eval/smollm_triviaqa_qwen_weaver_grpo.sh |
scripts/train/smollm_triviaqa_qwen_weaver_grpo.sh |
| SmolLM3-3B (Weaver: SmolLM3-3B) | TriviaQA | weaver-sft | huggingface link | scripts/eval/smollm_triviaqa_smollm_weaver_sft.sh |
scripts/train/smollm_triviaqa_smollm_weaver_sft.sh |
| SmolLM3-3B (Weaver: SmolLM3-3B) | TriviaQA | weaver-grpo | huggingface link | scripts/eval/smollm_triviaqa_smollm_weaver_grpo.sh |
scripts/train/smollm_triviaqa_smollm_weaver_grpo.sh |
If you prefer to evaluate the vanilla model instead of MemGen, simply modify memgen/model/modeling_memgen.py by replacing the active generate function with the commented-out alternative generate implementation directly above it, and then run the standard evaluation script.
MemGen consists of two modules: Weaver and Trigger. We follow a two-stage training approach, training each module separately.
If you would like to reproduce results for a specific dataset + model, please refer to the table above. If the corresponding checkpoint is not yet available, we kindly ask for your patience as we are actively preparing more comprehensive releases.
-
Train the Weaver model (SFT or GRPO)
bash scripts/weaver_sft.sh bash scripts/weaver_grpo.sh
-
Evaluate the Weaver model
Before running, make sure to updateLOAD_MODEL_PATHinscripts/eval.shto point to the trained checkpoint:<weaver_dir>bash scripts/eval.sh
- Train the Trigger model
bash scripts/trigger_train.sh
- Evaluate the Trigger model
Before running, make sure to updateLOAD_MODEL_PATHinscripts/eval.shto point to the trained checkpoint:<trigger_dir>bash scripts/eval.sh
Set in configs/latent_memory/<dataset>.yaml or override with --options (e.g. model.fast_sft True):
model.attn_implementation:flash_attention_2(default),sdpaoreager.model.fast_sft: batch Weaver SFT examples instead of processing them one by one (about 2x faster; same loss and gradients up to floating-point summation order). Off by default.
If you find this repository helpful, a citation to our paper would be greatly appreciated:
@article{zhang2025memgen,
title={MemGen: Weaving Generative Latent Memory for Self-Evolving Agents},
author={Zhang, Guibin and Fu, Muxin and Yan, Shuicheng},
journal={arXiv preprint arXiv:2509.24704},
year={2025}
}
- We sincerely thank Search-R1 for open-sourcing their search web environment.
- We sincerely thank the previous latent reasoning works such as LatentSeek, SoftCoT, R3Mem and so on.
- We also extend our heartfelt thanks to LAVIS for their code framework design.