[论文解读] Memory Transformer

Mikhail Burtsev, Yuri Kuratov+2 more|arXiv (Cornell University)|Jun 20, 2020

Analog and Mixed-Signal Circuit Design被引用 6

一句话总结

本文提出 MemTransformer，一种通过引入可学习的 [mem] 令牌来存储全局和局部表征的内存增强型 Transformer 架构，以提升序列建模性能。该模型在机器翻译和语言建模任务中表现更优，注意力可视化显示其具备有效的内存读写与分块处理能力，且内存大小与模型准确率呈正相关。

ABSTRACT

Transformer-based models have achieved state-of-the-art results in many natural language processing tasks. The self-attention architecture allows transformer to combine information from all elements of a sequence into context-aware representations. However, information about the context is stored mostly in the same element-wise representations. This might limit the processing of properties related to the sequence as a whole more difficult. Adding trainable memory to selectively store local as well as global representations of a sequence is a promising direction to improve the Transformer model. Memory-augmented neural networks (MANNs) extend traditional neural architectures with general-purpose memory for representations. MANNs have demonstrated the capability to learn simple algorithms like Copy or Reverse and can be successfully trained via backpropagation on diverse tasks from question answering to language modeling outperforming RNNs and LSTMs of comparable complexity. In this work, we propose and study few extensions of the Transformer baseline (1) by adding memory tokens to store non-local representations, (2) creating memory bottleneck for the global information, (3) controlling memory update with dedicated layer. We evaluate these memory augmented Transformers and demonstrate that presence of memory positively correlates with the model performance for machine translation and language modelling tasks. Augmentation of pre-trained masked language model with memory tokens shows mixed results for tasks from GLUE benchmark. Visualization of attention patterns over the memory suggest that it improves the model's ability to process a global context.

研究动机与目标

• 为解决标准 Transformer 在访问全局序列信息时的局限性，即上下文信息分散存储于标记表征中。
• 探究在序列建模任务（如机器翻译和语言建模）中，添加可训练的通用内存是否能提升性能。
• 探索内存增强架构如何支持复杂内存操作，如复制、摘要和分块处理。
• 评估内存控制器在推理过程中内存大小变化下的鲁棒性与泛化能力。
• 评估内存增强技术在下游 GLUE 任务中对预训练模型（如 BERT 和 Transformer-XL）的可迁移性。

提出的方法

• 在输入序列开头引入 [mem] 令牌，作为全局和局部表征的通用内存存储。
• 提出三种变体：MemTransformer（基础内存增强）、MemCtrl（由专用 Transformer 层控制内存更新）和 MemBottleneck（仅保留内存到内存的注意力，无序列到序列的注意力）。
• 通过反向传播进行训练，采用标准目标函数，实现内存写入、读取及内存内操作的端到端学习。
• 采用自注意力机制，实现输入标记与 [mem] 令牌之间的注意力（读/写），以及 [mem] 令牌之间的注意力（内存内处理）。
• 通过注意力模式的可视化分析内存行为，识别出如对角线注意力（放大）、分块复制和垂直注意力（读取）等操作。
• 将内存增强应用于预训练模型（如 BERT、Transformer-XL），并在 GLUE 基准任务上进行微调。

实验结果

研究问题

• RQ1 在机器翻译和语言建模任务中，添加可学习的 [mem] 令牌是否能提升标准 Transformer 的性能？
• RQ2 MemTransformer 中的注意力模式是否显示出学习到的内存操作（如写入、读取、复制和汇总）的证据？
• RQ3 内存大小如何影响推理过程中的模型性能与鲁棒性，尤其是在内存大小被改变时？
• RQ4 内存增强是否能提升 BERT 和 Transformer-XL 等预训练模型在下游 GLUE 任务中的性能？
• RQ5 内存控制器是否能在不同内存大小下实现泛化？在推理过程中内存减少时，其性能是否呈渐进式下降？

主要发现

• [mem] 令牌的引入显著提升了 Transformer 架构在机器翻译和语言建模任务中的性能。
• 模型性能与内存大小正相关，表明更大的内存容量有助于表征学习。
• 注意力模式的可视化证实了内存读写操作及内存内处理（如分块复制和对角线放大）的存在。
• 内存损毁测试表明，预训练 MemTransformer 模型的性能严重依赖于内存的存在，证实了其功能性重要性。
• 内存控制器表现出鲁棒性，当推理过程中内存大小减小时性能呈渐进式下降，表明其具备良好的泛化能力。
• 在扩展后对内存增强模型进行微调可获得更优性能，表明内存增强与持续训练结合时具有显著有效性。