Meta开源“记忆层”，重塑Transformer架构大模型

11月3日凌晨三点，全球社交巨头Meta分享了一个创新研究——Memory layers（记忆层）。

目前，Transformer架构的预训练大模型在存储、查询数据时，随着参数的变大对算力的需求呈指数级增长。“记忆层”则提出了新的高效查询机制替代了传统的查询方法，通过比较查询键与两个较小集合中的键，可以快速找到最相关的键，而无需遍历模型的整个记忆层。

换言之，可以在不增加算力的情况下显著增加大模型的参数。在传统的键值查找中，每个查询都需要与记忆层中的每个键进行比较，以找到最匹配的值。该方法在键的数量较少时是可行的，但随着记忆层规模的增长，这种暴力搜索的方式变得非常低效，需要消耗巨大算力和时间。

Product - Key Lookup是“记忆层”的核心算法之一，使用了一种分而治之的策略，将传统的单一键集合分解为两个较小的键集合，通过两个阶段的查找来减少必要的比较次数，从而提高查找效率。

除了计算效率之外，Product-Key Lookup模块还优化了内存和带宽的使用。由于每个GPU只需要处理一半的键，因此内存的使用量减少了一半。由于每个GPU只需要返回与自己处理的键相关的值，所以内存带宽的需求也得到了优化。

Product-Key Lookup算法不仅提高了记忆层的查询效率，还为记忆层的应用开辟了新的可能性，使得记忆层可以被应用于更大规模的数据集和更复杂的任务中，包括大规模知识图谱的查询、长文本的语义检索等。

并行记忆层则主要是用于对硬件GPU的优化。在传统的Transformer架构模型中，随着模型规模的增加，计算和内存需求也随之增长。特别是在处理大规模数据集时，单一的计算单元很难满足这种需求。并行记忆层通过在多个GPU之间分配任务，有效解决这一难题。在并行记忆层的设计中，每个计算单元只负责处理一部分数据，这样可以减少单个计算单元的负担，同时提高整体的处理速度。

共享记忆参数则是另外一个重要优化方法，允许不同层的记忆层共享同一个参数集合。这种设计的优势在于，它减少了模型的总参数数量，同时提高了参数的利用率。

为了应对训练期间可能出现的变化，研究人员开发了一套动态调整策略。每当有新的键加入或旧有的键被更新时，系统会自动调整相应的子集，而无需对整个记忆池进行全面改造。这样的设计既简化了维护流程，又提高了系统的灵活性和适应性。