摘要

本文介绍了一种增强大型语言模型（LLM）记忆能力的新方法，通过引入外部关联记忆来提高模型从长上下文中检索事实的能力。作为案例研究，本文测试了LARIMAR架构，这是一种最近提出的LLM解码器架构，通过外部关联记忆增强，用于处理包括passkey和needle-in-the-haystack测试在内的多个长上下文检索任务。实验表明，LARIMAR能够在测试时适应比训练时更长的上下文，同时保持解码器识别的记忆输出，且不增加GPU内存占用。

原理

LARIMAR架构的核心在于其外部记忆模块，该模块能够在测试时动态更新，以适应比训练时更长的上下文。外部记忆的结构类似于Kanerva Machine，但更新记忆的方式是通过找到线性系统的最小二乘解，而不是更新多元高斯后验分布。这种方法使得模型即使在训练时使用较短的上下文，也能在测试时泛化到更长的上下文，只要上下文的相关部分在分布内。

流程

在实验中，LARIMAR将上下文分为多个段落，每个段落被编码并写入记忆。在读取时，使用查询来计算读取键，从记忆中读取编码并传递给解码器。例如，在passkey测试中，模型能够从长达100万token的上下文中准确检索出passkey，而无需特定任务的训练。在needle-in-the-haystack测试中，模型同样展示了从长篇文本中检索特定信息的能力。

应用

LARIMAR的应用前景广泛，特别是在需要处理大量文本数据且需要高精度检索信息的场景，如法律文档分析、历史文献研究、大规模数据挖掘等。由于其能够在不增加GPU内存占用的前提下处理更长的上下文，LARIMAR为实现高效、可扩展的语言模型应用提供了新的可能性。