梦晨 发自 凹非寺
量子位 | 公众号 QbitAI
在实验中，仅有2.4B参数的Memory3模型不仅打败了许多7B-13B的模型，在专业领域任务如医学上的表现也超过了传统的RAG方法，同时推理速度更快，“幻觉”问题也更少。
目前相关论文已上传到arXiv，并引起学术界关注。
知识按使用频率分类
这一方法受人脑记忆原理启发，独立于存储在模型参数中的隐性知识和推理时的短期工作工作记忆，给大模型添加了显式记忆。
具体来说，人类的记忆大致可以分为三部分:
显式记忆：可以主动回忆的长期记忆，比如读过的文章。获取显式记忆很容易，但提取时需要一定的回忆过程。
隐式记忆：无意识使用的长期记忆，比如骑自行车的技能。获取隐式记忆需要大量重复练习，但使用时毫不费力。
外部信息：存在大脑之外的信息，如考试时的备考资料。获取和使用都很轻松，但遇到新问题时作用有限。
可以看出，三种记忆形式在获取和使用的效率上形成了鲜明的互补。人脑会根据知识的使用频率，巧妙地在它们之间分配存储位置，从而最小化整体开销。
反观大模型，目前主要依赖在参数中以隐式记忆的形式来存储知识，这导致两个问题：
知识分配效率低：无论一个知识使用得多频繁，都一视同仁塞进参数里，导致大量冷知识占用了宝贵的参数空间。
知识提取效率低：每次使用知识，都得动用大量参数参与计算。
目前在训练阶段，团队将大模型比作显式记忆能力受损的患者，靠学习如何系鞋带一样的大量重复练习才能背下一点知识，消耗大量的数据和能量。
在推理阶段，大模型又好像一个人每写一个单词时都要回忆起毕生所学的一切，就很不合理。
基于以上思路，团队按照知识的预期使用频率（横轴）计算了读写成本（纵轴），阴影区域表示给定记忆格式的最小成本区域。
结果发现，把常用知识塞进模型参数里成本最低，但容量有限；不常用的知识直接检索效率最高，但每次读取都要重新编码，成本高；而显式记忆则是个平衡点，对于使用次数中等的大部分知识最划算。
记忆电路理论
团队进一步在论文中提记忆电路理论，在大模型语境下重新定义知识和记忆，以确定哪些知识更适合存储为显式记忆，以及什么样的模型架构适合读写显式记忆。
通过分析一些已知的大模型内部机制，如事实问答、搜索复制粘贴等，团队认为大模型中的每条知识都可以表示为一个输入-输出关系，加上实现这个关系的内部电路（circuit）。
电路指计算图中的一个子图，由一些注意力头和MLP神经元组成，这些电路的输入输出具有一定的语义关联。大模型的知识可进一步分为两类:
具体知识（specific knowledge）：电路的输入和输出都具有明确的语义，如常识、常见短语等。
抽象知识（abstract knowledge）：电路的输出语义可变，如搜索、复制、粘贴，需要通过输入推理出输出。
接下来，作者引入可分离知识（separable knowledge）的概念：如果一个知识可以仅通过文本实现而不必内置到模型参数里，那它就是可分离的。
可模仿知识（imitable knowledge）是可分离知识的一个特例，可以直接用描述这条知识自身的文本去“教会”另一个不具备这条知识的大模型，无需通过参数来编码。
一个核心结论是，具体知识都是可模仿的，因此也是可分离的，都可转化为显式记忆。论文从理论上给出了（非形式化）证明。
团队进一步把具体知识按使用次数分成“无关紧要”、专业知识和常见短语三个等级，不同等级按照读写成本分别适合三种不同的记忆格式。
拥有显式记忆的大模型Memory3
那么如何实现显式记忆呢？
以注意力层的key-value向量作为显式记忆的载体，在推理之前，Memory3模型将所有引用文本转换为显式记忆，并将它们保存在硬盘或非易失性内存设备上。
在推理时，模型会查询与当前上下文最相关的一些显式记忆，将它们并入注意力机制中，与上下文的key-value向量一起计算注意力分数，生成下一个token。
然而，海量文本转化成的显式记忆不仅需要更多的磁盘空间，而且在推理过程中还会占用GPU内存，从而损害LLM生成的吞吐量。
为此，Memory3采取了多维度压缩优化策略：
layer维度：只有前半部分的注意力层（记忆层）产生和存取显式记忆，后半部分仍然是普通的注意力层。
head维度：每层只有少部分head（如1/5）负责处理显式记忆的key-value，其他head保持原样。
token维度：对于每个head，只选取参考文本中最相关的少量token（如8个），提取其key-value作为显式记忆。
最后再进一步用向量量化（vector quantization）压缩每个key和value向量到更短的表示。
多级压缩的组合，使得显式记忆的规模从45.9TB压缩到4.02TB，压缩到一个GPU集群通常配备的存储容量之内。
另外，团队在显式记忆的读写上还有一些值得注意的细节设计：
推理时为了避免不同文本片段重复检索显式记忆，Memory3每隔64个token做一次检索，中间共享检索结果。
频繁调用显式记忆会产生IO开销。为此，Memory3在内存中维护了一个固定大小的缓存，存储最近访问过的显式记忆。
对于参考文本，模型使用不同的输入符号（“<s>Reference:”）将其与普通文本区分开，避免干扰文本理解。
对于显式记忆中的不同文本片段，模型为其分配了同一区间的位置编码，保留局部上下文。这种”平行”位置编码避免了长文本中间部分被忽略的问题。
最终训练出来的Memory3模型，在HuggingFace排行榜上的评测结果如下，显式记忆将平均分数提高了2.51%。
相比之下Llama2-7B和13B之间的分数差异为4.91%，而13B模型的非嵌入参数数量接近7B模型的两倍。
因此，可以说显式记忆可以将“有效模型大小”提高了2.51/4.91≈51.1%。如果用Qwen-1.8B和4B来做参考，计算结果相似，“有效模型大小”提高49.4%。
在幻觉评估上，Memory3避免了将文本压缩到模型参数中可能会导致的信息丢失，表现的比大部分模型要好。
论文中还详细报告了从数据到训练、微调和对齐过程的具体设置，感兴趣的可以查看原文。
论文地址：
https://arxiv.org/abs/2407.01178
参考链接：
[1]https://x.com/rohanpaul_ai/status/1809782336021537094
