“Large Memory Layers with Product Keys” 介绍了一种可集成到神经网络的结构化内存，通过产品键设计实现大规模且高效的存储与检索，在语言建模任务中提升性能并兼顾计算效率，为神经网络架构优化提供了新方法。

  1. 研究背景：随着数据增多神经网络规模扩大，但计算复杂度剧增。已有研究尝试在不增加计算复杂度的情况下提升网络容量，如条件计算模型、记忆增强神经网络等，但各有局限。本文提出一种新的键值内存层来解决相关问题，并聚焦于语言建模任务，将其集成到Transformer架构中。

  2. 可学习的产品键内存

    • 内存设计：由查询网络、含两组子键的键选择模块和值查找表组成。查询网络将输入映射到潜在空间，键选择模块利用产品键结构高效找到与查询最相似的键，通过加权求和得到输出。产品键由两个向量码本的外积定义，能大幅减少比较次数。
    • 复杂度：相比普通键的扁平表示，产品键在搜索相似键时计算量显著降低，复杂度为。
    • 多头内存注意力：采用多头机制，每个头独立计算查询选择键，共享相同的值，增加了模型的表达能力和键的使用效率。

  3. 实验

    • 数据集与评估指标：在从 Common Crawl 提取的 300 亿字语料上进行大规模语言建模实验，以测试集困惑度衡量模型性能，同时用内存使用率和 KL 散度评估内存使用情况。
    • 训练细节：基于 Transformer 架构，使用 Adam 优化器训练，在不同层数和维度的模型中穿插内存层，设置多头、选择键数量和内存槽数量等参数，并采用特定技术提高训练速度和减少内存使用。
    • 结果：增加内存比增加层数更有利于降低困惑度，如 12 层带单个内存的模型性能优于 24 层无内存模型，且推理速度更快。增加内存层数可进一步提升性能，在大模型中添加内存对推理时间影响可忽略不计。
    • 消融研究：内存越大困惑度越低；查询网络中批归一化可提高大内存的使用率和性能；内存置于 Transformer 中间层效果最佳；增加头数或 k - NN 数量可提升性能和内存使用率；产品键相比普通键在内存使用、困惑度和速度上表现更优。

  4. 结论：提出的内存层以极小计算开销显著提升神经网络容量，其基于键的分解和稀疏读写机制，在大规模语言建模实验中表现出色，如 12 层模型达到 24 层 BERT - large 模型性能且运行时间减半。