[论文解读] Memformer: A Memory-Augmented Transformer for Sequence Modeling

Qingyang Wu, Zhenzhong Lan+4 more|arXiv (Cornell University)|Oct 14, 2020

Topic Modeling被引用 7

一句话总结

Memformer 通过引入一种固定大小的外部动态记忆，构建了增强记忆的 Transformer 模型，实现了长序列建模的线性时间复杂度和恒定内存空间复杂度。通过采用具有遗忘机制的记忆写入器和读取器，以及一种新颖的记忆重播反向传播（MRBP）优化方法，其在自回归任务中实现了与 Transformer-XL 相当的性能，但内存使用量减少 8.1 倍，推理速度提升 3.2 倍。

ABSTRACT

Transformers have reached remarkable success in sequence modeling. However, these models have efficiency issues as they need to store all the history token-level representations as memory. We present Memformer, an efficient neural network for sequence modeling, that utilizes an external dynamic memory to encode and retrieve past information. Our model achieves linear time complexity and constant memory space complexity when processing long sequences. We also propose a new optimization scheme, memory replay back-propagation (MRBP), which promotes long-range back-propagation through time with a significantly reduced memory requirement. Experimental results show that Memformer has achieved comparable performance compared to the baselines by using 8.1x less memory space and 3.2x faster on inference. Analysis of the attention pattern shows that our external memory slots can encode and retain important information through timesteps.

研究动机与目标

• 解决标准 Transformer 在长序列建模中计算和内存复杂度呈二次增长的问题。
• 克服 RNN 和 Transformer-XL 等循环模型在保留长距离依赖关系方面的局限性。
• 为具有大内存表示的自回归模型提供高效训练与推理能力。
• 设计一种内存高效架构，在降低计算与内存成本的同时保持高性能。

提出的方法

• 引入一个外部动态记忆模块，用于存储过去标记的压缩表示，从而实现恒定的内存空间复杂度。
• 采用记忆写入器根据当前输入更新记忆槽，并通过遗忘机制控制过时信息的保留。
• 利用记忆读取器在解码过程中对记忆槽进行注意力计算，使模型能够检索相关的历史信息。
• 提出记忆重播反向传播（MRBP）优化方案，显著降低在时间反向传播过程中的内存消耗。
• 设计 Memformer 时采用独立的编码器与解码器堆叠结构，其中记忆仅在编码器中更新，并由解码器访问。
• 应用压缩机制在保持长距离依赖关系的同时减小内存大小，提升效率而不损失性能。

实验结果

研究问题

• RQ1 固定大小的外部记忆是否能在保持与 Transformer 相当性能的前提下，提升长序列建模的效率？
• RQ2 记忆写入器的注意力机制如何实现对长距离依赖关系的有效压缩与保留？
• RQ3 所提出的 MRBP 方法在具有大内存的循环模型中，对反向传播过程的内存消耗降低程度如何？
• RQ4 在长序列上，Memformer 是否在推理速度和内存效率方面优于 Transformer-XL 和 Compressive Transformer？
• RQ5 不同记忆大小和记忆槽类型如何影响模型随时间保留与检索信息的能力？

主要发现

• 当记忆大小为 1,024 时，Memformer 在性能上与 Transformer-XL 相当，但内存使用量减少 8.1 倍，推理速度提升 3.2 倍。
• 在记忆大小为 1,024 时，Memformer 在 WikiText-103 数据集上达到困惑度 20.53，优于使用 1,024 内存的 Transformer-XL 和 FLOPs 相近的 Compressive Transformer。
• 当记忆大小减少至 512 时性能显著下降，若完全移除记忆模块，困惑度上升至 25.57。
• 对记忆写入器注意力的分析显示，约 60–80% 的记忆槽在长时间内保持内容稳定，表明具备稳定的长期记忆保留能力。
• 类型为 m^{355} 的记忆槽对输入标记表现出高注意力，且遗忘偏置较大（3.20），表明其对动态或显著信息具有快速更新特性。
• 可视化结果证实，记忆槽能够学习关键元素（如命名实体和动词）的压缩表示，模拟人类记忆编码机制。