[論文レビュー] Memorizing Transformers

Yuhuai Wu, Markus N. Rabe+2 more|arXiv (Cornell University)|Mar 16, 2022

Topic Modeling被引用数 39

ひとこと要約

この論文は、デコーダーのみの Transformer に対して、大規模で微分不能な外部メモリを近似的な kNN アクセスを介してメモリ化・検索することで、過去の入力を記憶し、取得する能力を付与し、長大なコンテキストデータセットにおける言語モデリングを改善します。

ABSTRACT

Language models typically need to be trained or finetuned in order to acquire new knowledge, which involves updating their weights. We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately. In this work, we extend language models with the ability to memorize the internal representations of past inputs. We demonstrate that an approximate kNN lookup into a non-differentiable memory of recent (key, value) pairs improves language modeling across various benchmarks and tasks, including generic webtext (C4), math papers (arXiv), books (PG-19), code (Github), as well as formal theorems (Isabelle). We show that the performance steadily improves when we increase the size of memory up to 262K tokens. On benchmarks including code and mathematics, we find that the model is capable of making use of newly defined functions and theorems during test time.

研究の動機と目的

• 推論時にウェイトを更新するのではなく、過去の入力を記憶して新しい知識を獲得するよう、言語モデルの拡張を動機づける。
• 外部ストアから正確な記憶を取り出す、単純でスケーラブルな kNN 増強アテンション機構を提案する。
• 記憶サイズが複数の長大なコンテキスト領域とモデル規模にわたって困惑度を改善することを示す。

提案手法

• 前のステップから生成されたキー/バリューで更新される、(key, value) の非微分可能な外部メモリを追加する。
• 1 つのアテンション層で各クエリに対してメモリへの近似的な kNN ルックアップを行い、外部メモリアテンションと局所密結合アテンションを学習ゲートで組み合わせる。
• 継続的な文脈のために Transformer-XL 型のキャッシュを使用し、局所的な文脈を維持するスライディング因果マスクを適用する。
• 訓練ステップ間の分布シフトによる陳腐化を緩和するため、キーとクエリを正規化する。
• 近似的な kNN 実装を用いて検索品質と計算効率のバランスを取り、数万トークン規模のメモリを実現する。

実験結果

リサーチクエスチョン

• RQ1 kNN を介してアクセスされる大規模な外部メモリは、長文テキストとコードデータセットの困惑度を改善するのか？
• RQ2 メモリサイズは性能にどう影響するのか、収穫逓減点はあるのか？
• RQ3 事前学習済みモデルはファインチューニングを通じて外部メモリを効果的に活用できるのか、それともゼロからの訓練が必要か？
• RQ4 モデルがメモリから検索する際に現れるパターン（例: 関数名、定義、補題など）はどのようなものか？

主な発見

• 外部メモリは、長文データセットとアーキテクチャ群（例: 長いウェブテキスト、書籍、コード、形式的証明）全体で困惑度を一貫して改善する。
• メモリサイズの増大は非常に大規模なメモリ規模まで性能向上をもたらし、コードや数学タスクに顕著な利点がある。
• ゲーティング機構は外部メモリを優先するようヘッドに偏りを生じさせ、非微分可能なメモリ勾配なしで長距離検索を効果的に可能にする。
• モデルは定義、補題、関数名を思い出すために memory を活用することができ、形式的証明やコードに対して検索のような挙動を示す。
• 事前学習モデルを memory 活用にファインチューニングすると、 memorized 訓練 regime へのギャップを迅速に縮小でき、実用的な適用性を示す。