Memorilla 记忆压缩怎么用:16 个 token 装下整本文档

一句话结论:斯坦福开源的 Memorilla 把一整批文档压成 16 个记忆 token 挂到冻结模型上,FactKG 上拿到 84.33%,比把全部文档塞进上下文的 80.43% 还高,同时少用 108 倍的文档 token。

给大模型加长期记忆,这条路一直很贵。常规做法是把历史对话或长文档塞进上下文,代价随文档量线性上升,模型能记住多少取决于窗口有多长。斯坦福 SNAP 团队在 2026 年 10 月开源的 Memorilla 走了另一条路:把文档集合编码成固定数量的记忆 token,再让模型自己去读取其中与当前问题相关的部分。

这个仓库在 GitHub 上拿到 1136 个 star,Apache-2.0 协议,论文 arXiv 编号 2610.21375,10 月 15 日首次提交,11 月 3 日更新到第三版。它的价值不在于又一个记忆框架,而在于给出了一个可以自己复算的数字:在同等任务上,压缩记忆比全上下文更准。

16 个 token 到底装了什么

Memorilla 的工作方式是三步。第一步把文档集合编码成可复用的 embedding,每篇文档只编码一次,之后反复使用。第二步训练一个记忆模块,在每次提问时把检索到的相关信息放进模型的输入嵌入空间。第三步把这些信息压缩成固定数量的 token,通常是 16 个,直接交给一个冻结的 Qwen3-8B。

关键在于模型本身不需要重新训练。记忆模块是独立的部分,底座模型保持冻结,接入方只需要评估记忆模块本身的质量。对想给自己的模型加长期记忆的团队来说,这意味着不用承担全量重训的成本。

官方给的示例场景很具体。PersonaMem-v2 测试集里有一个用户,之前和助手有过 107 轮对话,总共约 32000 token,现在提出一个新问题。Memorilla 读完全部对话,交给冻结 Qwen3-8B 的只有 16 个记忆 token,比原始对话少约 2000 倍。这个例子需要一张 80 GB 显卡,首次运行下载约 25 GB 模型和 37 MB 测试集。

示例里的答案能说明压缩保留了什么。用户从没说过家里有糖尿病史,唯一线索藏在 107 轮对话的某一轮里,他问过母亲一侧的家族病史是否会比父亲一侧更容易。Memorilla 把这条线索和早餐问题连了起来,给出控糖又扛饿的建议;没有记忆的 Qwen3-8B 只能给出一份泛泛的早餐清单。

与全上下文正面比:更准且更省

这是全篇最关键的一组对照。

基准 Memorilla(16 token) 全上下文推理 差异
FactKG 事实型知识保持 84.33% 80.43% 高 3.90 个点
文档 token 消耗 基准量 108 倍 少用 108 倍
八项基准平均表现 优于 Mem0 与 MemoRAG — 五项领先
TextWorld 平均成功率 24.2% 18.6%(全历史上下文) 高 5.6 个点

准确率反超这件事需要解释一下。全部文档塞进上下文时,模型面对的是几千 token 的干扰信息,注意力被稀释,事实召回率会掉。把文档压缩成 16 个 token 之后,进入上下文的是与当前问题强相关的信息,模型看到的东西更干净。这就是压缩记忆和长上下文两条路线最核心的差别,一个在筛,一个在堆。

对照组的实现细节也值得记一笔。评测脚本里提供了四种基线:闭卷、top-5 文档的 RAG、全上下文,以及 Memorilla 本身。RAG 方案只取最相关的 5 篇文档,而全上下文取全部,这两种常见做法与 16 token 方案在同一套脚本下比较。

速度账:文档变多,延迟不变

Memorilla 的第二个结构性优势是成本不随文档量增长。解码器永远只读那 16 个记忆 token,模型侧的成本与文档数量无关,只有记忆模块那一侧会随集合增长。

场景 配置 响应耗时 倍数
单张 H100 文档 1 增长到 512 篇(约 128K token) 0.74 秒 基本持平
单张 H100 全部文档放进提示词 0.75 秒升至 19.9 秒 慢 27 倍
默认模式 每次提问重算全部文档 embedding 100 万文档 74 毫秒 基线
开启记忆缓存 写入时算好交叉注意力的键值 100 万文档 8.1 毫秒 约 9 倍提速
缓存加 torch.compile 同上 100 万文档 5.0 毫秒 约 15 倍提速
更大集合 400 万文档 33 毫秒 仍可接受

记忆缓存的代价是显存。缓存占用是 embedding 的四倍,400 万文档时需要 80 GiB,仍然装得进一张卡。缓存里的记忆 token 与默认模式在 bfloat16 舍入级别上一致,FactKG 准确率差距在 0.06 分以内,换句话说这个开关几乎不损失质量。

还有一组压缩比更极端的数据。论文里把 NarrativeQA 的书籍和电影剧本压缩超过 4000 倍。因为无论存多少文档,模型收到的永远是同样的 16 个 token,这套机制理论上可以扩展到百万 token 级别的素材量。

装一次要什么代价

上手门槛比想象的低,环境要求只有 Linux、CUDA 显卡和 Python 3.11,用 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,评测不需要。如果只想评测,可以用 PyTorch 自带的注意力实现替代:

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

训练成本要提前知道。三阶段训练在 4 张 H100 80GB 上约需 19 小时,下载约 310 GiB 的数据与 embedding。在一张卡上评测全部八个基准约 15 分钟。权重在 Hugging Face 的 memorilla/Memorilla-Qwen3-8B,分为 stage1、stage2、stage3 三阶段,外加从 stage1 派生的 PersonalizationV4 和 PersonaMem-v2 两套。全部使用 Qwen3-8B 解码器与 Qwen3-Embedding-4B 编码器。

想用自己的数据也无需改代码。每个数据划分放一个 parquet 文件,行里带 collection_id、question、answer、documents 四列,指定数据目录即可。single_task 配置默认从 stage2 检查点热启动,传 –init_memory 空字符串则从头训练。

智能体持续学习这条路

Memorilla 另一个值得注意的用法是让智能体从自己的历史里持续学习。这部分用 GRPO 在 TextWorld 文字冒险游戏上训练记忆模块,基于 SkyRL 的一个分支实现,并在仓库的 rl 目录下有完整文档。

论文里描述的现象很直观。只把完整历史放在上下文里,智能体会反复在已经探索过的房间之间游荡,直到撞上 50 轮上限;加入 8 个训练过的记忆 token 之后,它能直接抵达目标。平均成功率从 18.6% 提到 24.2%。仓库还附了一段回放动画,展示一个保留测试的游戏在没有记忆和有记忆两种情况下的行为差异。

随项目一起发布的 PersonalizationV4 基准规模不大但设计精巧。149 个虚构用户,每人有过 200 轮简短对话,每个问题把用户置于一个新场景,问他最可能怎么做或偏好什么。问题被刻意设计成需要至少组合两条用户事实,而这两条事实只在对话里出现过,从未被直接陈述。数据在 Hugging Face 上以 CC BY 4.0 发布。

几条使用上的实际提醒

安全上有几处需要留意。仓库会从 Hugging Face Hub 和本地路径加载模型与检查点,只加载信任的来源。解码器的自定义模型代码只在传入 –trust_remote_code 时才会运行,默认的 Qwen3 系列不需要这个开关。记忆检查点用 torch.load 的 weights_only 模式加载。

版本也值得记一笔。仓库锁定的 vLLM 0.9.1、transformers 4.53.3、PyTorch 2.7.0 都存在公开的安全公告,涉及恶意构造的模型配置和检查点。这些问题只在加载不可信来源时才会实际发生。

内存缓存默认关闭,需要显式传入 memory_cache 与 compile_memory 参数才生效。评测时给 evaluate.py 加上这两个开关即可。

常见问题

16 个记忆 token 会不会丢信息

会丢一部分,但丢的不是当前问题需要的那部分。Memorilla 学的是把与提问相关的信息放进记忆 token,所以对某个具体问题而言它保留的恰好是相关部分。证据是 FactKG 上的 84.33% 高于全上下文的 80.43%,以及八项基准里五项优于 Mem0 和 MemoRAG。真正的信息损失出现在「用户曾经说过但当前问题完全无关」的细节上,这类内容本来就不影响当前回答。

和 RAG 相比该选哪个

RAG 取回的是文档原文片段,直接塞进提示词,延迟随文档量增长;Memorilla 压缩成固定数量 token,延迟基本与文档量无关,代价是需要训练记忆模块。仓库把 RAG top-5 列为正式基线,可以直接跑对比。数据量大、文档还在持续增长、关心延迟的场景更偏向 Memorilla;文档本身需要模型逐字引用的场景,RAG 更直接。

有没有比 16 更少的可能

仓库里的 TextWorld 实验用了 8 个训练过的记忆 token,说明数量是可调的。核心机制是模型收到的 token 数固定,与存入多少文档无关,这也意味着把数量进一步压小在架构上可行,只是要在召回率上付出代价。官方默认给的是 16 个。

小结

Memorilla 提供了一个可以自己复算的对照。压缩记忆和长上下文这两条路线在同一个基准上正面相遇,前者用 108 倍的 token 节省换来了 3.90 个点的准确率提升,同时在文档量增长 512 倍时把延迟稳定在 0.74 秒。

这套思路对做智能体产品的团队有直接含义。当历史记录是持续增长的,把全部历史塞进上下文这条路迟早撞上延迟墙,先把知识压缩成固定预算再交给模型,是一条已被实测验证过的替代方案。底座模型保持冻结这一点尤其关键,它把采用门槛从「重新训练一个模型」降到了「训练一个记忆模块」。

相关阅读

全部文章归档与模型专题