[论文解读] Memformer: The Memory-Augmented Transformer

Qingyang Wu, Zhenzhong Lan+2 more|arXiv (Cornell University)|May 4, 2021

Topic Modeling参考文献 15被引用 12

一句话总结

Memformer 引入了一种增强记忆的 Transformer 架构，通过使用统一的记忆机制和一种新颖的优化方案——记忆重放反向传播（Memory Replay Back-Propagation），实现了长序列建模的 O(n) 时间复杂度和 O(1) 的空间复杂度。它在 WikiText-103 上优于现有的长距离模型（如 Transformer-XL 和 Compressive Transformer），在长序列上的性能和可扩展性表现更优。

ABSTRACT

Transformer models have obtained remarkable accomplishments in various NLP tasks. However, these models have efficiency issues on long sequences, as the complexity of their self-attention module scales quadratically with the sequence length. To remedy the limitation, we present Memformer, a novel language model that utilizes a single unified memory to encode and retrieve past information. It includes a new optimization scheme, Memory Replay Back-Propagation, which promotes long-range back-propagation through time with a significantly reduced memory requirement. Memformer achieves O(n) time complexity and O(1) space complexity in processing long sequences, meaning that the model can handle an infinite length sequence during inference. Our model is also compatible with other self-supervised tasks to further improve the performance on language modeling. Experimental results show that Memformer outperforms the previous long-range sequence models on WikiText-103, including Transformer-XL and Compressive Transformer.

研究动机与目标

• 解决 Transformer 在长序列上自注意力机制的二次方时间复杂度问题。
• 实现在极低内存消耗下对长序列进行高效的反向传播时间传播。
• 设计一种统一的记忆机制，以有效编码和检索过去信息。
• 实现处理任意长序列的 O(n) 时间复杂度和 O(1) 空间复杂度。
• 提升语言建模及其他自监督任务的性能。

提出的方法

• 引入单一统一的记忆模块，用于在长序列中存储和检索上下文信息。
• 采用记忆重放反向传播（Memory Replay Back-Propagation），一种新颖的优化方案，可显著降低反向传播时间过程中的内存需求。
• 使用增强记忆的注意力机制，使其同时关注输入标记和存储的记忆状态。
• 应用可学习的记忆更新机制，以动态编码长距离依赖关系。
• 通过记忆重放实现端到端训练，确保在长序列中梯度流动的完整性。
• 保持与自监督预训练目标的兼容性，以提升语言建模性能。

实验结果

研究问题

• RQ1 统一记忆机制是否能降低 Transformer 在长序列上自注意力机制的计算和内存开销？
• RQ2 记忆重放反向传播是否能在极小内存占用下实现长序列上的稳定且高效的训练？
• RQ3 Memformer 是否能在保持或提升长上下文任务性能的同时，实现线性时间复杂度和常数空间复杂度？
• RQ4 Memformer 在标准基准测试中与 SOTA 长上下文模型（如 Transformer-XL 和 Compressive Transformer）相比表现如何？
• RQ5 Memformer 在语言建模任务的自监督预训练中能获得多大程度的性能增益？

主要发现

• Memformer 实现了 O(n) 时间复杂度和 O(1) 空间复杂度，使推理阶段能够高效处理任意长序列。
• 该模型在 WikiText-103 基准测试中优于 Transformer-XL 和 Compressive Transformer，展现出更优的语言建模性能。
• 记忆重放反向传播显著降低了内存使用量，同时实现了长序列上的稳定训练。
• 统一记忆机制在不增加计算成本的前提下，有效捕捉了长距离依赖关系。
• 当扩展到其他自监督预训练任务时，Memformer 仍保持强劲性能。
• 与现有长上下文 Transformer 变体相比，该模型展现出显著的可扩展性和效率优势。