Memory Layers：Meta推出的记忆层增强语言模型

Memory Layers简介

Memory Layers是一种创新的神经网络组件，由Meta推出，它通过键值查找机制扩展模型参数而不增加计算量。这种层利用稀疏激活模式，与密集型前馈层形成互补，专门用于高效存储和检索信息。Memory Layers在大规模语言模型中显示出显著的性能提升，尤其是在处理事实性任务时，能够超越计算预算更高的密集模型和混合专家模型。

Memory Layers主要功能

1. 参数扩展：Memory Layers能够在不增加浮点运算（FLOPs）的情况下为模型增加额外的参数，从而扩展模型的容量。
2. 信息存储与检索：提供专门的机制来廉价地存储和检索信息，这对于语言模型来说尤其重要，因为它们需要记忆和使用大量的事实性信息。
3. 稀疏激活：与传统的密集层相比，Memory Layers采用稀疏激活模式，这意味着在任何给定时间只有一小部分参数被激活，从而减少了计算和内存带宽的需求。
4. 性能提升：在下游任务中，特别是事实性问答任务中，Memory Layers能够显著提高语言模型的准确性。
5. 与现有架构兼容：Memory Layers可以集成到现有的神经网络架构中，如Transformer，通过替换或补充前馈网络（FFN）层来增强模型性能。

Memory Layers技术原理

1. 键值查找机制：Memory Layers使用简单的键值查找机制，其中键和值都编码为嵌入向量，通过计算查询向量与键向量之间的相似度来检索相关的值向量。
2. 可训练参数：与注意力机制不同，Memory Layers中的键和值是可训练的参数，而不是激活函数的输出。
3. 稀疏性和规模：Memory Layers通常拥有大量的键值对，使得稀疏查找和更新成为必须，只有与查询最相似的顶部键和值参与输出计算。
4. 产品量化键（Product-key lookup）：为了解决大规模记忆中的查询-键检索瓶颈，Memory Layers采用了产品量化键技术，通过两个较小的键集合来有效地执行全局的top-k查找。
5. 并行化实现：Memory Layers通过在多个GPU上并行化嵌入查找和聚合操作，实现了大规模扩展。
6. 共享记忆参数池：在多层Memory Layers之间共享记忆参数池，以保持参数数量不变，最大化参数共享。
7. 性能和稳定性优化：通过优化PyTorch的EmbeddingBag操作，Memory Layers实现了更高效的CUDA内核，显著提高了内存带宽利用率，并优化了训练过程中的性能和稳定性。

Memory Layers应用场景

1. 事实性问答系统：用于构建能够准确回答基于事实的问题的系统，如TriviaQA和NaturalQuestions，提高模型的事实准确性。
2. 知识检索：在需要从大量数据中检索特定信息的应用中，Memory Layers可以快速定位和提取相关知识。
3. 语言模型预训练：在大规模语言模型的预训练阶段，Memory Layers有助于模型更有效地学习和记忆语言模式。
4. 多跳问答：在需要通过多个信息源或步骤来找到答案的复杂问答任务中，Memory Layers能够存储和检索中间信息。
5. 编程语言模型：在编程语言模型中，Memory Layers可以帮助模型记忆代码片段和编程概念，提高代码生成和理解的能力。
6. 常识推理：在需要模型具备世界知识和常识以进行推理的任务中，Memory Layers提供了一种有效的信息存储和检索机制。

Memory Layers项目入口

• 项目主页：https://ai.meta.com/research/publications/memory-layers-at-scale
• GitHub代码库：https://github.com/facebookresearch/memory
• arXiv技术论文：https://arxiv.org/pdf/2412.09764