[论文解读] Large Memory Layers with Product Keys

Guillaume Lample, Alexandre Sablayrolles+3 more|HAL (Le Centre pour la Communication Scientifique Directe)|Jul 10, 2019

Topic Modeling被引用 50

一句话总结

引入一种基于产品键的可扩展内存增强层，能够实现非常大的内存容量并进行精确最近邻搜索，集成到 Transformer 中用于大规模语言建模，并在精度和推理速度方面优于更深层基线。

ABSTRACT

This paper introduces a structured memory which can be easily integrated into a neural network. The memory is very large by design and significantly increases the capacity of the architecture, by up to a billion parameters with a negligible computational overhead. Its design and access pattern is based on product keys, which enable fast and exact nearest neighbor search. The ability to increase the number of parameters while keeping the same computational budget lets the overall system strike a better trade-off between prediction accuracy and computation efficiency both at training and test time. This memory layer allows us to tackle very large scale language modeling tasks. In our experiments we consider a dataset with up to 30 billion words, and we plug our memory layer in a state-of-the-art transformer-based architecture. In particular, we found that a memory augmented model with only 12 layers outperforms a baseline transformer model with 24 layers, while being twice faster at inference time. We release our code for reproducibility purposes.

研究动机与目标

• 通过引入一个大规模、可训练的内存层，在不成比例增加计算量的情况下提升模型容量。
• 在可扩展的键值内存中提供基于产品键的精确最近邻搜索。
• 在大型语言建模任务中展示内存层在 Transformer 内的有效性。
• 证明记忆增强模型在计算量更低或相当的情况下能超越更深的基线。
• 分析内存大小、网络中的放置位置以及规范化对性能和使用的影响。

提出的方法

• 定义一个具有查询网络、双子键产品键空间和一个值查找表的内存。
• 使用由两个码本的外积形成的产品键来隐式定义一个庞大的键集合。
• 通过查询的内积相似度选择前-k个键，然后对相应值进行加权求和（稀疏读取）。
• 采用多头记忆注意力，每个头部有自己的查询网络和子键，但共享值。
• 通过稀疏内存更新和查询网络的批量归一化训练，以改善键的覆盖和使用。
• 在选定的 Transformer 层中用内存层替换前馈网络（FFN），允许加性残差连接（x → x + PKM(x)）。

实验结果

研究问题

• RQ1 产品键记忆（PKM）是否在保持精确搜索的同时提供可扩展的容量且开销极小？
• RQ2 内存大小、头部数量、k-NN 和内存放置如何影响大型 Transformer 的困惑度、内存使用和速度？
• RQ3 在性能、使用和效率方面，产品键记忆与扁平键记忆相比如何？
• RQ4 哪些训练配置（批归一化、稀疏更新）能够在大规模语言建模数据集上优化内存使用和模型性能？

主要发现

• 一个具有 12 层且带有产品键的记忆在 CC-News 上的表现超过同维数的 24 层 Transformer，推理时间大约是其一半。
• 增加内存尺寸通常降低困惑度，当批量归一化改善键使用时，内存尺寸为 1M 槽的效果尤为显著。
• 在可比尺寸下，产品键相比扁平键提供更高的内存使用率、较低的困惑度和更快的推理。
• 内存放置位置很重要；在中间 Transformer 层（大约第4–5层）插入内存可获得最佳性能。
• 4 个记忆头，每头 32 个键，总共 512^2 槽，在速度与困惑度之间取得有利权衡；相比扁平键，产品键降低了参数量，并在推理时具有更好的扩展性。
• 在查询网络中的批量归一化在大内存时显著改善键的使用和困惑度。