[论文解读] Recurrent Memory Transformer

Aydar Bulatov, Yuri Kuratov+1 more|arXiv (Cornell University)|Jul 14, 2022

Topic Modeling被引用 25

一句话总结

本文提出了一种记忆增强型分段循环Transformer——循环记忆Transformer（RMT），通过在输入中注入专用的记忆标记并实现跨分段的循环机制，提升了长序列建模能力。RMT在复制、反转和关联检索等长上下文任务上优于Transformer-XL，同时在语言建模任务中达到与之相当的性能，但内存使用量最多减少10倍，展现出卓越的记忆效率与长期依赖学习能力，且无需对Transformer架构进行任何修改。

ABSTRACT

Transformer-based models show their effectiveness across multiple domains and tasks. The self-attention allows to combine information from all sequence elements into context-aware representations. However, global and local information has to be stored mostly in the same element-wise representations. Moreover, the length of an input sequence is limited by quadratic computational complexity of self-attention. In this work, we propose and study a memory-augmented segment-level recurrent Transformer (RMT). Memory allows to store and process local and global information as well as to pass information between segments of the long sequence with the help of recurrence. We implement a memory mechanism with no changes to Transformer model by adding special memory tokens to the input or output sequence. Then the model is trained to control both memory operations and sequence representations processing. Results of experiments show that RMT performs on par with the Transformer-XL on language modeling for smaller memory sizes and outperforms it for tasks that require longer sequence processing. We show that adding memory tokens to Tr-XL is able to improve its performance. This makes Recurrent Memory Transformer a promising architecture for applications that require learning of long-term dependencies and general purpose in memory processing, such as algorithmic tasks and reasoning.

研究动机与目标

• 为解决标准Transformer在处理长序列时因自注意力计算复杂度为二次方而带来的局限性，以及全局信息表征模糊的问题。
• 在不修改Transformer架构的前提下，提升序列建模任务中长期依赖与全局上下文的建模能力。
• 探究专用记忆标记与分段级循环机制在提升记忆效率与长上下文任务性能方面的有效性。
• 评估记忆增强型Transformer是否能在算法推理与语言建模任务中超越或匹配如Transformer-XL等最先进模型的性能。
• 分析RMT中注意力机制如何学习使用记忆，特别是通过可解释的读/写模式。

提出的方法

• RMT在输入序列中加入特殊记忆标记，作为全局与局部信息的专用存储，使其与输入标记表示解耦。
• 模型以分段方式处理序列，将前一enco段的记忆隐藏状态传递至当前段，实现循环机制，支持长上下文建模。
• 记忆操作（读取、写入、重写）通过注意力机制端到端学习，训练过程中梯度可反向传播至记忆状态。
• 该方法无需修改Transformer架构，仅在输入与输出序列层面进行调整。
• RMT同时学习控制记忆操作与序列表示处理，实现紧凑且上下文感知的表示。
• 该方法与现有模型兼容：RMT可与Transformer-XL的缓存机制结合，进一步提升性能。

实验结果

研究问题

• RQ1 记忆增强型、循环式Transformer架构是否能在长序列建模上超越标准Transformer与Transformer-XL？
• RQ2 使用专用记忆标记是否能减少输入标记与全局上下文之间的信息混杂，从而在算法任务中取得更优性能？
• RQ3 RMT在显著减少内存使用量的前提下，能在多大程度上实现与Transformer-XL相当的语言建模性能？
• RQ4 RMT中的注意力模式与Transformer-XL相比有何不同，特别是在记忆访问与信息流动方面？
• RQ5 RMT是否能与Transformer-XL中的现有缓存机制有效结合，以进一步提升性能？

主要发现

• 当序列长度超过某一阈值，尤其是在超过2个分段时，RMT在复制、反转和关联检索等长上下文算法任务上优于Transformer-XL。
• 在语言建模任务中，RMT性能与Transformer-XL相当，但内存使用量最多减少10倍，展现出卓越的记忆效率。
• 注意力图可视化显示，RMT能够将整个分段信息压缩至记忆标记中，避免了Transformer-XL中常见的表示混杂现象。
• 在4个分段的任务中，当内存大小为6时，Transformer-XL在反转任务上的平均准确率仅为0.8，这是由于表示混杂所致；而RMT通过隔离内存存储，实现了1.0的准确率。
• RMT能够通过从读取记忆到写入记忆的重写机制，更长时间保留近期分段的信息，从而提升多分段序列任务的性能。
• 将RMT与Transformer-XL的缓存机制结合，可在语言建模任务中进一步提升性能，表明两种方法具有互补优势。