❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发非常感兴趣，我会每日跟你分享最新的 AI 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！

🥦 微信公众号｜搜一搜：蚝油菜花 🥦

🚀 快速阅读

1. 功能扩展：Memory Layers 通过键值查找机制扩展模型参数，不增加计算负担。
2. 性能提升：显著提升模型在事实性任务中的准确性，超越传统密集模型。
3. 技术原理：基于稀疏激活和产品量化键，实现高效信息存储与检索。

正文（附运行示例）

Memory Layers 是什么

Memory Layers 是 Meta 推出的一种通过可训练的键值查找机制为模型增加额外参数的技术。它不会增加浮点运算次数（FLOPs），而是通过稀疏激活的方式补充计算密集型的前馈层，提供专门的容量来廉价地存储和检索信息。

在大规模语言模型中，Memory Layers 显示出显著的实用性，尤其是在处理事实性任务时，能够显著提高模型的性能。通过替换 Transformer 层中的一个或多个前馈网络（FFN）为记忆层，Memory Layers 在不显著增加计算成本的情况下，显著提升了模型的事实准确性和知识获取能力。

Memory Layers 的主要功能

• 参数扩展：支持模型在不增加计算负担（即不增加 FLOPs）的情况下增加额外的参数，扩展模型的容量。
• 信息存储与检索：提供专门的机制存储和检索信息。
• 补充密集层：与计算密集型的前馈层（feed-forward layers）互补，基于稀疏激活模式提供信息存储和检索的专用容量。
• 提高事实准确性：在事实性任务中，显著提高语言模型的准确性，让模型更准确地记忆和使用训练数据中的事实。
• 性能提升：在计算和参数匹配的情况下，让模型性能超过传统的密集模型和混合专家模型。

Memory Layers 的技术原理

• 键值查找机制：基于简单的键值查找机制工作，其中键和值都编码为嵌入向量。给定查询向量 q，及一组键 K 和值 V，输出是值的一个软组合，权重根据 q 与对应键的相似度确定。
• 稀疏激活：与传统的注意力机制不同，Memory Layers 中的键和值是可训练的参数，而不是激活值。由于键和值的数量通常很大，Memory Layers 需要进行稀疏查找和更新，只有与查询最相似的 top-k 键和对应的值参与输出计算。
• 产品键查找：为解决大规模记忆中查询-键检索的瓶颈，Memory Layers 采用产品量化键（product-quantized keys），用两组较小的键有效地执行全键集的 top-k 查找。
• 并行化实现：由于 Memory Layers 对内存的需求很大，在多个 GPU 上并行化嵌入查找和聚合实现大规模扩展。
• 共享记忆参数：用共享的记忆参数池，在不增加参数数量的情况下，基于多个记忆层增加性能。

如何运行 Memory Layers

1. 克隆仓库并创建环境

git clone https://github.com/facebookresearch/memory
cd memory

bash setup/create_env.sh
# 或者如果你有 SLURM 集群访问权限
sbatch setup/create_env.sh

2. 激活环境

conda activate lingua_<date>

3. 下载并准备数据

python setup/download_prepare_hf_data.py fineweb_edu <MEMORY> --data_dir ./data --seed 42

4. 下载分词器

python setup/download_tokenizer.py llama3 <SAVE_PATH> --api_key <HUGGINGFACE_TOKEN>

5. 启动调试任务

# stool 代表 SLURM 工具！
python -m lingua.stool script=apps.main.train config=apps/main/configs/pkplus_373m_1024k.yaml nodes=1 partition=<partition>
# 如果你想在本地运行，可以使用 torchrun
torchrun --nproc-per-node 8 -m apps.main.train config=apps/main/configs/pkplus_373m_1024k.yaml
# 或者你也可以在 1 个 GPU 上运行
python -m apps.main.train config=apps/main/configs/pkplus_373m_1024k.yaml

资源

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

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发非常感兴趣，我会每日跟你分享最新的 AI 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！

🥦 微信公众号｜搜一搜：蚝油菜花 🥦