bingreeky/MemGen

MemGen: Weaving Generative Latent Memory for Self-Evolving Agents

📚 What is MemGen?

MemGen is a research‑grade codebase that implements the MemGen model introduced in an ICLR‑2026 paper. The core idea is to give large language models a latent memory that lives inside the model’s own reasoning stream instead of in an external database. Two components work together:

  • Memory Trigger – decides, during generation, whether the model should pull from its latent memory.
  • Memory Weaver – creates compact “memory tokens” that summarise past experiences and are woven into the current context. Together they let an agent evolve its behaviour across interactions without changing the underlying model weights.

⚙️ How to get started

  1. Create the environment – either install the requirements.txt with pip or use the provided memgen.yml conda file.
  2. (Optional) Set up a retriever – the repo can hook into the Search‑R1 environment for external retrieval, but it’s not required for the core MemGen experiments.
  3. Pick a model & dataset – the README lists ready‑to‑use checkpoints (Qwen2.5‑1.5B‑Instruct and SmolLM3‑3B) for KodCode, GSM8K, and TriviaQA.
  4. Run the scripts – training and evaluation are wrapped in simple bash files, e.g. scripts/weaver_sft.sh for the Weaver module or scripts/trigger_train.sh for the Trigger module. Adjust LOAD_MODEL_PATH in scripts/eval.sh to point at the checkpoint you want to test.

🚀 What can you do with it?

  • Re‑produce the paper’s results on single‑turn and multi‑turn reasoning tasks.
  • Experiment with latent memory by swapping the Weaver or Trigger components, or by editing the configs/latent_memory/*.yaml files.
  • Compare against a vanilla baseline – the repo includes a commented‑out version of the standard generate function; swapping it in lets you see how much the latent memory improves performance.

📦 What’s inside?

Folder / File Purpose
memgen/model/modeling_memgen.py Core model class with the custom generate that handles memory tokens.
scripts/ Bash wrappers for training (*_sft.sh, *_grpo.sh, trigger_train.sh) and evaluation (*_eval.sh).
configs/latent_memory/ YAML configs for each dataset, controlling options like attention implementation and fast‑SFT mode.
requirements.txt / memgen.yml Dependency lists (PyTorch, FlashAttention‑2, HuggingFace Transformers, etc.).
assets/memgen.png Diagram of the Trigger‑Weaver pipeline.

📊 Evaluation & checkpoints

The table in the README lists HuggingFace links for each model‑dataset pair. Scripts such as scripts/eval/qwen2_5_gsm8k_sft_trigger.sh will download the checkpoint, run inference on the test split, and report standard metrics (e.g., accuracy for GSM8K, exact match for TriviaQA).

🛠️ Known limitations

  • Multi‑GPU training works with Distributed Data Parallel (DDP) but not with Fully Sharded Data Parallel (FSDP). The authors note this in the FAQ.
  • Multi‑turn GRPO scripts and checkpoints are still pending; only single‑turn SFT/GRPO are currently released.
  • Small formatting quirks (e.g., whether a prompt ends with \boxed{} followed by a period or a newline) can affect performance for the 1.5B model – the latest code fixes this by enforcing a consistent ChatML template.

📖 Citation

If you use MemGen in research, cite the ICLR‑2026 paper:

@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}
}

🙏 Acknowledgements

The project builds on the Search‑R1 retrieval environment, prior latent‑reasoning work (LatentSeek, SoftCoT, R3Mem), and the LAVIS code framework.

TL;DR: MemGen is a research codebase that adds a latent, internally‑generated memory to LLMs via a Trigger‑Weaver mechanism, enabling agents to improve over time without weight updates. It ships with ready‑to‑run training/evaluation scripts for several datasets and small‑to‑medium sized models.

Related

  • Project
    deepmodeling/dpgen

    DP‑GEN is a Python platform that automates the creation of deep‑learning interatomic potentials (Deep Potential models). It runs a concurrent learning loop—sampling structures, selecting the most informative ones for DFT calculations, and retraining the model—while handling job scheduling on HPC clusters and interfacing with many MD and quantum‑chemistry codes.

  • Dispatch
    Qwen 3.5‑397B‑A17B release: hybrid linear‑attention MoE model with 1 M token context and state‑of‑the‑art multimodal performance

    Qwen 3.5‑397B‑A17B is a 397 billion‑parameter multimodal model that activates only 17 billion parameters per token, delivering state‑of‑the‑art performance on language, coding, reasoning and vision tasks while being up to 19× faster than its predecessor.

  • Dispatch
    Qwen-AgentWorld release: language world model for seven domains and its impact on general agents

    Qwen releases Qwen‑AgentWorld, a language world model that simulates seven agent environments and improves general agents via controllable simulation and unified next‑state prediction.

  • Project
    aiming-lab/SimpleMem

    A unified memory stack for LLM agents that provides semantically lossless compression for text and multimodal data, featuring a self-evolving retrieval system to maximize recall while minimizing token usage.

  • Project
    deepgenteam/deepgen

    DeepGen 1.0 is a lightweight 5B parameter unified multimodal model that integrates image generation, editing, and reasoning into a single framework.