[論文レビュー] Recurrent Memory Transformer

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

Topic Modeling被引用数 25

ひとこと要約

本稿では、入力に専用のメモリトークンを組み込み、セグメント間で再帰を可能にする、メモリ拡張型でセグメントレベルの再帰的Transformer、Recurrent Memory 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は、グローバルおよびローカル情報の専用のストレージとして機能する特別なメモリトークンを入力系列に追加し、これにより入力トークンの表現とは独立して情報を保持する。
• モデルは系列をセグメントに分割して処理し、前のセグメントのメモリの隠れ状態を現在のセグメントに引き渡すことで、再帰的処理と長文脈モデリングを実現する。
• 読み取り、書き込み、再書き込みのメモリ操作は、すべてのバックプロパゲーションがメモリ状態を経由して勾配が流れ込むように、エンドツーエンドで学習される。
• この手法はTransformerアーキテクチャに変更を加えない。変更は入力および出力系列レベルでのみ行われる。
• RMTは、メモリ操作と系列表現処理の両方を制御するように訓練され、コンパクトでコンテキストに適応した表現を可能にする。
• このアプローチは既存モデルと互換性がある：RMTはTransformer-XLのキャッシュと組み合わせることで、さらなる性能向上が可能である。

実験結果

リサーチクエスチョン

• RQ1 メモリ拡張型で再帰的なTransformerアーキテクチャは、標準のTransformerやTransformer-XLを上回る長文脈モデリングを実現できるか？
• RQ2 専用のメモリトークンの使用により、入力トークンとグローバルコンテキストの情報が混合するのを軽減し、アルゴリズム的タスクでより良い性能が得られるか？
• RQ3 RMTは、顕著に小さいメモリサイズで、Transformer-XLと同等の言語モデリング性能を達成できるか、その程度はどの程度か？
• RQ4 RMTの注意パターンは、特にメモリアクセスと情報フローにおいて、Transformer-XLとはどのように異なるか？
• RQ5 RMTは、Transformer-XLのキャッシュ機構と効果的に組み合わせられるか、性能向上に寄与するか？

主な発見

• RMTは、2つのセグメントを超える長さの系列において、コピー、リバース、関連検索といった長文脈アルゴリズム的タスクで、Transformer-XLを上回る性能を示した。
• 言語モデリングタスクでは、RMTはTransformer-XLと同等の性能を達成したが、最大10倍の小さなメモリサイズで実現した。これは、優れたメモリ効率を示している。
• 注意マップの可視化結果から、RMTはセグメント全体をメモリトークンに圧縮する能力を学習しており、Transformer-XLで見られる表現の混合を回避している。
• 4つのセグメントを持つタスクでは、メモリサイズ6のTransformer-XLはリバースタスクで平均正解率0.8にとどまり、表現の混合が原因であるのに対し、RMTはメモリストレージを分離することで1.0の正解率を達成した。
• RMTが読み取りメモリから書き込みメモリに再書き込みする能力により、最近のセグメント情報の保持期間が延び、複数セグメントの系列処理で性能が向上した。
• RMTをTransformer-XLのキャッシュと組み合わせることで、言語モデリングの性能が向上した。これは、両者のアプローチが相補的な強みを示していることを示している。