大型语言模型的记忆层：提升事实知识，减少幻觉
随着企业在各种应用中不断采用大型语言模型（LLM），一个关键挑战是提高模型的事实知识并减少幻觉。Meta AI 的研究人员在最新论文中提出了“可扩展记忆层”，这可能是解决此问题的几种可能方案之一。
可扩展记忆层为 LLM 添加更多参数，以提高其学习能力，而无需额外的计算资源。这种架构对于那些可以为事实知识预留额外内存，但又希望模型更灵活的推理速度的应用非常有用。
传统的语言模型使用“密集层”在其参数中编码大量信息。在密集层中，所有参数都以其全部容量使用，并且在推理过程中大部分时间同时被激活。密集层可以学习复杂的函数，而增加其容量需要额外的计算和能源资源。
相比之下，对于简单的事实知识，使用具有联想记忆架构的更简单的层将更有效率和可解释。这就是记忆层的作用。它们使用简单的稀疏激活和键值查找机制来编码和检索知识。稀疏层比密集层占用更多内存，但一次只使用一小部分参数，这使得它们在计算方面效率更高。
记忆层已经存在多年，但在现代深度学习架构中很少使用。它们没有针对当前的硬件加速器进行优化。
当前最先进的 LLM 通常使用某种形式的“专家混合”（MoE）架构，该架构使用与记忆层略微相似的机制。MoE 模型由许多专门负责特定任务的较小的专家组件组成。在推理时，路由机制根据输入序列确定哪个专家被激活。PEER 是 Google DeepMind 最近开发的一种架构，它将 MoE 扩展到数百万个专家，从而对推理过程中被激活的参数进行更细粒度的控制。
记忆层在计算方面很轻，但在内存方面很重，这对当前的硬件和软件框架提出了特定的挑战。在他们的论文中，Meta 研究人员提出了几种修改方法来解决这些挑战，并使其能够大规模使用。
记忆层可以在多个 GPU 上并行存储知识，而不会减慢模型速度（来源：arXiv）
首先，研究人员配置了用于并行的记忆层，将它们分布在多个 GPU 上以存储数百万个键值对，而不会改变模型中的其他层。他们还实现了一个特殊的 CUDA 内核来处理高内存带宽操作。此外，他们开发了一种参数共享机制，该机制支持模型中多个记忆层之间共享一组记忆参数。这意味着用于查找的键和值在层之间共享。
这些修改使得在 LLM 中实现记忆层成为可能，而不会减慢模型速度。
“记忆层及其稀疏激活很好地补充了密集网络，为知识获取提供了更大的容量，同时在计算方面很轻，”研究人员写道。“它们可以有效地扩展，并为从业人员提供了一个有吸引力的新方向，可以权衡内存与计算。”
为了测试记忆层，研究人员通过用共享记忆层替换一个或多个密集层来修改 Llama 模型。他们在几个任务上比较了增强记忆的模型与密集 LLM 以及 MoE 和 PEER 模型，包括事实问答、科学和常识世界知识以及编码。
在事实问答任务中，一个在 1 万亿个 token 上训练的 13 亿参数记忆模型（实线）的性能接近一个 70 亿参数模型（虚线）的性能，因为它获得了更多记忆参数（来源：arxiv）
他们的研究结果表明，记忆模型在密集基线模型上有了显著的改进，并且可以与使用 2 倍到 4 倍计算资源的模型相媲美。它们还与具有相同计算预算和参数数量的 MoE 模型的性能相匹配。该模型的性能在需要事实知识的任务中尤其显著。例如，在事实问答中，一个具有 13 亿参数的记忆模型的性能接近 Llama-2-7B 的性能，后者在两倍的 token 和 10 倍的计算资源上进行了训练。
此外，研究人员发现，随着他们将实验从 1.34 亿参数扩展到 80 亿参数，记忆模型的优势保持一致。
“鉴于这些发现，我们强烈建议将记忆层集成到所有下一代 AI 架构中，”研究人员写道，同时补充说还有很大的改进空间。“特别是，我们希望能够开发出新的学习方法，以进一步提高这些层的有效性，从而减少遗忘、幻觉和持续学习。”
