这个仓库提供了训练和评估 Memorilla 的代码。Memorilla 是一种记忆系统，它从整个文档集合中收集信息，将其压缩成固定数量、与查询相关的 memory token，供冻结的 LLM 使用。代码基于我们的论文 Memorilla: Latent Semantic Memory for LLMs。

冻结的 LLM 只能使用其上下文窗口容纳的内容，因此 Memorilla 会读取整个文档集合（例如用户过去的聊天记录或一本长篇书籍），然后只交给 LLM 16 个 memory token，这些 token 携带了每个问题所需的信息。在如此小的预算下，它相较于 Mem0 和 MemoRAG 等方法取得了最佳的平均性能；在 FactKG 上，它甚至比把全部文档都交给 LLM 更准确（84.33% vs. 80.43%），同时使用的 token 少了 108 倍。此外，Memorilla 还能帮助 agent 从自身历史中持续学习，将其在 TextWorld 游戏上的平均成功率从 18.6% 提升到 24.2%。

安装 快速开始 训练与评估 Checkpoints 使用你自己的数据 PersonalizationV4 TextWorld 上的强化学习 文档 开发 引用 许可证

安装

Memorilla 需要 Linux、CUDA GPU 和 Python 3.11；各配方在 NVIDIA H100 80GB GPU 上运行。使用 uv 安装：

$ git clone https://github.com/snap-stanford/memorilla.git && cd memorilla $ uv venv --python 3.11 && source .venv/bin/activate $ uv pip install -e . $ uv pip install --no-build-isolation flash-attn==2.7.4.post1

FlashAttention-2 与可选组件

最后一条命令安装 FlashAttention-2，训练默认使用它。评估不需要它，训练也可以改用 PyTorch 的注意力实现：

$ scripts/train.sh configs/stage1_enwiki.yaml --attn_implementation sdpa

可选组件：

$ uv pip install -e ".[dev]" # tests, linters and pre-commit $ uv pip install -e ".[wandb]" # Weights & Biases logging $ uv pip install -e ".[personalizationv4]" # OpenAI client for the PersonalizationV4 pipeline $ uv pip install -e ".[rl]" # TextWorld, for the RL data tools

快速开始

以下是一个来自 PersonaMem-v2 测试集的示例。用户此前与助手进行了 107 次对话，总计约 32,000 个 token，现在提出了一个新问题。Memorilla 读取所有对话，然后仅以 16 个 memory token 的形式交给冻结的 Qwen3-8B，大约少了 2,000 倍。你需要一块 80 GB GPU，首次运行会下载约 25 GB 的模型和 37 MB 的测试集。

from memorilla.benchmarks import PERSONA_PROMPT from memorilla.data import ensure_data, load_split from memorilla.vllm import MemoryVLLM row = load_split("pmv2", "test", ensure_data("pmv2", ["test"]))[73] # one user's question and past conversations model = MemoryVLLM.from_pretrained("memorilla/Memorilla-Qwen3-8B", "pmv2") # Qwen3-8B, the memory module and the encoder answer = model.answer(row["question"], row["documents"], system_prompt=PERSONA_PROMPT) print(f"Question: {row['question']}\nMemorilla: {answer}")

输出：

Question: What are some easy breakfast ideas that will keep me full all morning? Memorilla: Since you’ve got a family history of type 2 diabetes, you might enjoy breakfasts that balance protein, healthy fats, and fiber to keep you full without spiking blood sugar—like scrambled eggs with avocado and a side of roasted veggies, or Greek yogurt with chia seeds and berries. These options give you steady energy while keeping your blood sugar in check.

用户从未说过家族有糖尿病史。唯一的线索出现在 107 次对话中的某一次，用户当时询问母系家族的糖尿病史是否比父系家族史更能提高患病风险。Memorilla 将这条线索与一个看似无关的早餐问题联系起来，而没有记忆的 Qwen3-8B 给出的则是一份通用的早餐清单。

这只是一个例子。由于无论存储多少文档，LLM 始终接收相同的 16 个 memory token，该方法有可能扩展到数百万 token。在论文中，Memorilla 将 NarrativeQA 中的书籍和电影剧本压缩了 4,000 倍以上。在一块 H100 上，将集合从 1 篇文档增长到 512 篇文档（约 128K token），Memorilla 的响应时间保持在 0.74 秒，而将所有文档放入 prompt 会使模型从 0.75 秒减慢到 19.9 秒，慢了 27 倍。

Checkpoints 下列出的其他 checkpoint 以相同方式加载。

以下配方从头训练记忆模块，并在基准上对其进行评估。数据和预计算的 embedding 会在首次使用时从 memorilla/Memorilla-Data 下载。

Train the three stages with the Qwen3-8B decoder on 4 GPUs.

$ scripts/train.sh configs/stage1_enwiki.yaml $ scripts/train.sh configs/stage2_mixture.yaml $ scripts/train.sh configs/stage3_multitask.yaml

Continue on the personalisation benchmarks.

$ scripts/train.sh configs/personalization_pv4.yaml $ scripts/train.sh configs/personalization_pmv2.yaml

Evaluate a checkpoint on every benchmark (or name one, e.g. factkg).

$ scripts/evaluate.sh runs/stage3/epoch-04 all $ scripts/evaluate.sh runs/personalization_pv4/epoch-04 pv4

Text-only baselines: closed book, RAG over the top-5 documents, full context.

$ scripts/baselines.sh closed_book all $ scripts/baselines.sh rag all --top_k 5 $ scripts/baselines.sh full_context triviaqa

在 4 块 H100 80GB GPU 上，三个阶段大约需要 19 小时，并下载约 310 GiB 的数据和 embedding；在一块 GPU 上评估一个 checkpoint 在所有八个基准上的表现大约需要 15 分钟。训练和评估指南描述了主要选项，每个脚本的 --help 列出了全部选项。

训练好的记忆模块位于 Hugging Face 的 memorilla/Memorilla-Qwen3-8B。所有模块均使用 Qwen3-8B 解码器和 Qwen3-Embedding-4B 编码器训练，这也是 train.py 和 evaluate.py 的默认配置。

$ hf download memorilla/Memorilla-Qwen3-8B --local-dir checkpoints $ scripts/evaluate.sh checkpoints/pv4 pv4 # evaluate a checkpoint

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）