Meta 公司近日发布了一项突破性的研究成果，他们开发出一种新型的记忆层技术，可以显著提升大型语言模型（LLM）的事实准确性，并在参数规模上实现了前所未有的扩展。这项技术不仅挑战了传统神经网络的扩展方式，还为未来的 AI 架构设计提供了新的方向。
这项研究的核心在于利用可训练的键值查找机制，为模型增加额外的参数，而无需增加计算量（FLOPs）。这种方法的核心思想是，通过稀疏激活的记忆层来补充计算密集的前馈层，从而提供专门的存储和检索信息的能力。
与传统的稠密网络相比，记忆层在处理信息存储方面更具效率。例如，语言模型需要学习人名生日、国家首都等简单的关联信息，记忆层可以通过简单的键值查找机制实现，这种方式比使用前馈网络更高效。
该研究的主要贡献在于将记忆层的规模扩展到了前所未有的程度，达到了1280亿个参数。实验结果表明，在下游任务中，配备改进型记忆层的语言模型不仅优于计算量翻倍的稠密模型，在计算量和参数量匹配的情况下，也胜过混合专家模型。尤其在事实性任务上，性能提升更为显著。
Meta 的研究人员通过将 Transformer 网络中的一个或多个前馈网络（FFN）替换为记忆层来实现这一目标。这种替换方式在不同基础模型大小(从1.34亿到80亿参数)和记忆容量(高达1280亿参数)上都表现出了一致的优势。实验结果显示，记忆层可以将语言模型的事实准确性提高100%以上，同时在代码编写和一般知识方面也有显著提高。在许多情况下，配备记忆层的模型甚至可以达到需要4倍计算量的稠密模型的性能。
研究人员还对记忆层进行了多项改进，以克服其在规模化应用中的挑战:
采用乘积键查找机制:为了解决大规模记忆层中查询键检索的瓶颈，该研究采用了可训练的乘积量化键，从而避免了对每个查询键对进行比较。
并行化记忆层:为了在多 GPU 环境下实现记忆层的并行化，研究人员将嵌入查找和聚合操作分布在多个 GPU 上。
共享记忆机制:为了最大限度地共享参数，研究人员在所有记忆层之间使用共享的记忆参数池。
优化性能和稳定性:研究人员使用自定义的 CUDA 内核优化了 EmbeddingBag 操作，显著提高了内存带宽利用率。此外，还引入了带有 silu 非线性的输入相关门控机制，以提高训练性能和稳定性。
实验结果还揭示了以下关键发现:
记忆层的大小对性能有显著影响:随着记忆层大小的增加，事实性问答的性能持续提高。
多个记忆层优于单个记忆层:使用多个共享参数的记忆层可以提高性能，但过多的记忆层会降低性能。最佳的记忆层数量为三个。
记忆层能更快地学习事实:在训练初期，配备记忆层的模型性能提升更快，表明记忆层有助于模型更快地学习事实.
记忆层与稠密层互补:实验表明，稀疏的记忆层和稠密的前馈层都是必不可少的。
为了验证记忆层技术的有效性，研究人员在多个基准测试上进行了评估，包括:
事实性问题回答（NaturalQuestions， TriviaQA）
多跳问题回答（HotpotQA）
科学和常识性知识（MMLU， HellaSwag， OBQA， PIQA）
代码编写（HumanEval， MBPP）
结果显示，配备记忆层的模型在这些测试中均优于基线模型，尤其是在事实性问题回答上，性能提升最为明显。
Meta 的这项研究不仅为 AI 模型的扩展提供了新的思路，也为解决事实性问题和提高模型性能开辟了新的道路。研究人员认为，记忆层技术具有很强的可扩展性，未来有望在各种 AI 应用中得到广泛应用。他们还指出，虽然记忆层在硬件加速方面仍面临挑战，但相信通过持续的研究和优化，其性能可以与传统前馈网络相媲美甚至超越.
此外，Meta 的研究团队还希望通过新的学习方法进一步提升记忆层的性能，减少模型的遗忘、幻觉，并实现持续学习.
这项研究的发布无疑为 AI 领域注入了新的活力，也让我们对未来 AI 的发展充满了期待。
论文：https://arxiv.org/pdf/2412.09764