snap-stanford/memorilla
机器学习把大型文档集合压缩成少量记忆 token 交给冻结 LLM 的记忆模块,含训练、评估代码与检查点。
项目介绍
Memorilla 是一个面向冻结 LLM 的记忆系统:它读取整个文档集合(例如用户过去的对话、一本长书),把与问题相关的信息压缩成固定数量的记忆 token(发布配置中为 16 个),再交给不做改动的基座模型。由于 token 预算不随语料增长,同一解码器可以处理几篇文档,也可以处理数百篇文档。
仓库提供训练与评估代码、参考检查点和基准测试脚本。论文报告的结果是:在 FactKG 上,使用记忆 token 的准确率为 84.33%,高于把全部文档塞进上下文的 80.43%,同时 token 用量少 108 倍;在 TextWorld 上,记忆模块通过奖励学习后,平均成功率从 18.6% 提升到 24.2%。这些是作者给出的结果,仓库可用于复现或把方法迁移到自有数据。
目标用户是研究长上下文检索、个性化或智能体记忆的研究者和工程师,前提是具备 CUDA GPU。
快速使用
安装要求 Linux、CUDA GPU 和 Python 3.11,官方脚本在 80 GB H100 上运行过 安装步骤。使用 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,评估不需要它;训练时也可以加 --attn_implementation sdpa 改用 PyTorch 注意力 可选项。
随后可在 PersonaMem-v2 的一个测试样本上运行推理示例,需要一张 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]
model = MemoryVLLM.from_pretrained("memorilla/Memorilla-Qwen3-8B", "pmv2")
answer = model.answer(row["question"], row["documents"], system_prompt=PERSONA_PROMPT)
print(answer)
预期结果是打印出基于压缩记忆生成的回答,而不是基于完整对话历史生成的回答。
若要从头训练,可依次运行三阶段配置(4 张 H100 约 19 小时,三阶段合计下载约 310 GiB 数据);也可以从官方检查点继续,先用 hf download memorilla/Memorilla-Qwen3-8B --local-dir checkpoints 下载 训练与评估 检查点。单张 GPU 评估一个检查点的全部 8 个基准约需 15 分钟。
使用自有数据
数据为 parquet 文件,每行对应一个问题,列为 collection_id、question、answer、documents,放在 my_data/data/my_task/ 下;首次使用时会计算嵌入。single_task 配方默认从 Stage 2 检查点热启动,传 --init_memory "" 则从头训练 数据指南。
使用约束
- 硬件:Linux 加 CUDA GPU,文档中的配方按 80 GB H100 级别显卡编写;完整训练需要 4 张 GPU。
- 记忆模块与训练时使用的 Qwen3-8B 解码器和 Qwen3-Embedding-4B 编码器绑定,发布检查点即基于这组默认配置。
- 训练与完整基准套件涉及大规模数据和嵌入下载(三阶段约 310 GiB);离线环境和多卡问题可参考排障文档。
- PersonalizationV4 以 CC BY 4.0 发布,Memorilla-Data 各配置沿用其来源数据集许可证;代码为 Apache-2.0 许可证。
HotRepo 点评
仓库在 2026-11-10 进入 Python 日榜,当天最好名次为第 15,新增 62 颗 Star;2026-11-06 的 Hacker News 讨论帖热度不高 HN 讨论。比起名次,更值得注意的是定位:基座模型保持冻结,把无界的文档集合压缩进固定 token 预算,并配套论文、检查点和自发布的个性化基准。具备所需 GPU 的读者可以直接验证方法;硬件不足的读者仍可把论文和数据格式文档作为长上下文记忆设计的参考。
本文由 AI 根据项目 README 生成,更新于 ,以项目官方文档为准。