[論文レビュー] Scaling Transformer to 1M tokens and beyond with RMT

Aydar Bulatov, Yuri Kuratov+2 more|arXiv (Cornell University)|Apr 19, 2023

Topic Modeling被引用数 19

ひとこと要約

この論文は Recurrent Memory Transformer (RMT) を提示します。メモリ拡張型のセグメントレベル再帰アプローチで、トレーニング可能なメモリトークンを付与し、カリキュラム学習を用いることで、Transformerモデルが最大200万トークンを線形計算規模で処理できるようにします。

ABSTRACT

A major limitation for the broader scope of problems solvable by transformers is the quadratic scaling of computational complexity with input size. In this study, we investigate the recurrent memory augmentation of pre-trained transformer models to extend input context length while linearly scaling compute. Our approach demonstrates the capability to store information in memory for sequences of up to an unprecedented two million tokens while maintaining high retrieval accuracy. Experiments with language modeling tasks show perplexity improvement as the number of processed input segments increases. These results underscore the effectiveness of our method, which has significant potential to enhance long-term dependency handling in natural language understanding and generation tasks, as well as enable large-scale context processing for memory-intensive applications.

研究の動機と目的

• エンコーダーのみ/デコーダーのみのトランスフォーマーのプラグイン型ラッパーとして、メモリ拡張型セグメントレベル再発（RMT）を実証する。
• 推論時に線形計算かつ定常メモリで、極めて長い系列（最大2Mトークン）を処理できることを示す。
• 百万トークン規模の文脈にスケールするメモリ獲得/保持タスクを開発・ベンチマークし、メモリ操作の一般化を評価する。
• 長距離言語モデリングと形式推論タスクにおけるRMTの影響を調査し、分野横断的な実用的利点を評価する。

提案手法

• アーキテクチャを変更せず、事前学習済みトランスフォーマーにトークンベースのメモリモジュールを追加する。
• 長い入力を固定サイズのセグメントに分割し、セグメント内でのみ全体的なアテンションを実行して線形スケーリングを可能にする。
• セグメント間でメモリトークンの再発を用いて、メモリ出力が後続のセグメントに影響を与えるように訓練する。
• カリキュラム学習を用いて、タスク長を単一セグメントから複数セグメント文脈へと段階的に拡張する。
• 合成的な記憶化タスクを通じてメモリ操作を評価し、長距離言語モデリングおよび定理証明風生成へ実験を拡張する。

実験結果

リサーチクエスチョン

• RQ1 RMTは事前学習済みトランスフォーマーの有効な文脈長を、線形計算コストで/multi-million-token規模へ拡張できるか？
• RQ2 メモリ拡張型トランスフォーマーは、極めて長い系列全体で事実を記憶し、検索し、推論する能力がどの程度か？
• RQ3 長さが段階的に長くなるセグメントタスクで学習した場合、メモリ拡張モデルはより長い系列長に一般化するか？
• RQ4 長文言語モデリングと正式な定理証明生成における困惑度と予測品質へRMTの影響はどのようか？

主な発見

• 固定セグメントサイズに対してRMTは入力長と線形にスケールし、マルチセグメント入力に対して非再発モデルと比較してFLOPsを削減します（場合によっては約295×少なくなる）。
• メモリを用いると、事前学習済みのBERTバックボーンは最大2,000,000トークン（512トークンの4,096セグメント）にわたって情報を保存・検索できる。
• カリキュラム学習は安定性と一般化を向上させ、短いタスクで学習したモデルがはるかに長いタスクを解決できるようにする。
• 長距離言語モデリングでは、メモリを備えたRMTはベースラインより困惑度を改善し、セグメント境界での予測を安定化させるため、セグメントを跨ってメモリを保持する。
• RMTはアテンションパターン主導のメモリ操作を示し、極端に長い系列全体でのメモリ検索の一般化が可能であり、適切なタスクで2Mトークンを超えるスケーリングに技術的な intrinsic limitation がないことを示唆する。