近日，全球社交巨头Meta分享了一项名为“Memory Layers”（记忆层）的创新技术，旨在显著改善预训练大模型在存储和查询数据方面的效率。这项技术的发布引发了人工智能领域的广泛关注，尤其是在计算算力日益成为模型训练和应用瓶颈的背景下。
随着Transformer架构在自然语言处理、图像生成和其他生成任务中的应用越来越广泛，研究者们发现，越来越多的模型参数需要消耗越来越大的算力进行存储和检索。Meta的“记忆层”技术应运而生，提供了一种新的高效查询机制。该机制通过比较查询键与两个较小集合中的键，快速找到最相关的键，无需遍历整个记忆层，从而有效降低了计算成本。
简而言之，就是可以在不增加算力的情况下显著增加大模型的参数。
Memory layers的核心算法
1、Product - Key Lookup
Product - Key Lookup是“记忆层”的核心算法之一，使用了一种分而治之的策略，将传统的单一键集合分解为两个较小的键集合，通过两个阶段的查找来减少必要的比较次数，从而提高查找效率。
首先，查询键被分割为两个子查询，每个子查询分别与两个半键集合进行比较。由于每个半键集合的大小只有原始键集合的平方根大小，因此这个阶段的计算量大幅减少。在第一阶段，每个半键集合中找到与子查询最相似的k个键，这个过程称为top-k查找。
在第二阶段，两个半键集合中找到的top-k键被合并，以确定最终的top-k键。这一步骤涉及到对两个半键集合中找到的键进行综合评分，以确定它们与原始查询键的整体相似度。需要考虑到两个半键集合中的键的组合，以找到最佳的匹配。
除了计算效率之外，Product-Key Lookup模块还优化了内存和带宽的使用。由于每个GPU只需要处理一半的键，因此内存的使用量减少了一半。由于每个GPU只需要返回与自己处理的键相关的值，所以内存带宽的需求也得到了优化。
Product-Key Lookup算法不仅提高了记忆层的查询效率，还为记忆层的应用开辟了新的可能性，使得记忆层可以被应用于更大规模的数据集和更复杂的任务中，包括大规模知识图谱的查询、长文本的语义检索等。
2、并行记忆层和共享记忆参数
并行记忆层主要是用于对硬件GPU的优化。在并行记忆层的设计中，每个计算单元只负责处理一部分数据，这样可以减少单个计算单元的负担，同时提高整体的处理速度。这种设计允许模型在保持单个计算单元负载合理的同时，处理更大规模的记忆层。使得模型可以扩展到数十亿甚至数百亿的参数，而不会受到单个计算单元性能的限制。
共享记忆参数则是另外一个重要优化方法，允许不同层的记忆层共享同一个参数集合。这种设计的优势在于，它减少了模型的总参数数量，同时提高了参数的利用率。当一个记忆层接收到输入后，它会先从共享记忆池中查找最相似的记忆单元，然后根据查询结果生成输出。由于所有记忆层都指向同一个记忆池，因此它们可以在不影响彼此的情况下同时进行操作。
研究人员还开发了一套动态调整策略来应对训练期间可能出现的变化，每当有新的键加入或旧有的键被更新时，系统会自动调整相应的子集，而无需对整个记忆池进行全面改造。这样的设计既简化了维护流程，又提高了系统的灵活性和适应性。
