bingreeky/MemGen

📚 什么是 MemGen？

MemGen 是一个研究级代码库，实现了 ICLR-2026 论文中介绍的 MemGen 模型。核心思想是给大型语言模型一个 潜在记忆，它存在于模型自身的推理流中，而不是外部数据库中。两个组件协同工作：

记忆触发器 – 在生成过程中决定模型是否应从其潜在记忆中提取。
记忆编织器 – 创建紧凑的“记忆令牌”，总结过去的经验，并将其编织到当前上下文中。它们共同让代理在交互中进化其行为，而无需改变底层模型权重。

⚙️ 如何开始

创建环境 – 使用 pip 安装 requirements.txt 或使用提供的 memgen.yml conda 文件。
（可选）设置检索器 – 仓库可以连接到 Search‑R1 环境进行外部检索，但核心 MemGen 实验不需要它。
选择模型和数据集 – README 列出了现成的检查点（Qwen2.5‑1.5B‑Instruct 和 SmolLM3‑3B），用于 KodCode、GSM8K 和 TriviaQA。
运行脚本 – 训练和评估封装在简单的 bash 文件中，例如 scripts/weaver_sft.sh 用于 Weaver 模块，或 scripts/trigger_train.sh 用于 Trigger 模块。调整 scripts/eval.sh 中的 LOAD_MODEL_PATH 以指向要测试的检查点。

🚀 你能用它做什么？

重现论文结果 在单轮和多轮推理任务上。
实验潜在记忆 通过替换 Weaver 或 Trigger 组件，或编辑 configs/latent_memory/*.yaml 文件。
与普通基线比较 – 仓库包含标准 generate 函数的注释版本；替换它可以让您看到潜在记忆提高了多少性能。

📦 里面有什么？

文件夹 / 文件	用途
`memgen/model/modeling_memgen.py`	核心模型类，带有处理记忆令牌的自定义 `generate`。
`scripts/`	训练（`_sft.sh`、`_grpo.sh`、`trigger_train.sh`）和评估（`*_eval.sh`）的 Bash 包装器。
`configs/latent_memory/`	每个数据集的 YAML 配置，控制注意力实现和快速 SFT 模式等选项。
`requirements.txt` / `memgen.yml`	依赖列表（PyTorch、FlashAttention‑2、HuggingFace Transformers 等）。
`assets/memgen.png`	Trigger‑Weaver 管道的示意图。

📊 评估和检查点

README 中的表格列出了每个模型‑数据集对的 HuggingFace 链接。诸如 scripts/eval/qwen2_5_gsm8k_sft_trigger.sh 之类的脚本将下载检查点，在测试集上运行推理，并报告标准指标（例如，GSM8K 的准确率，TriviaQA 的精确匹配）。

🛠️ 已知限制

多 GPU 训练 适用于分布式数据并行（DDP），但不适用于完全分片数据并行（FSDP）。作者在 FAQ 中提到了这一点。
多轮 GRPO 脚本和检查点仍在等待中；目前仅发布了单轮 SFT/GRPO。
小的格式问题（例如，提示是否以 \boxed{} 后跟句点或换行符结尾）可能会影响 1.5B 模型的性能 – 最新代码通过强制一致的 ChatML 模板修复了此问题。

📖 引用

如果您在研究中使用 MemGen，请引用 ICLR‑2026 论文：

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

🙏 致谢

该项目基于 Search‑R1 检索环境、先前的潜在推理工作（LatentSeek、SoftCoT、R3Mem）和 LAVIS 代码框架。

TL;DR： MemGen 是一个研究代码库，通过 Trigger‑Weaver 机制为 LLM 添加潜在、内部生成的记忆，使代理无需权重更新即可随时间改进。它附带针对多个数据集和小到中型模型的现成训练/评估脚本。