[論文レビュー] Repeat After Me: Transformers are Better than State Space Models at Copying

Samy Jelassi, David Brandfonbrener+2 more|arXiv (Cornell University)|Feb 1, 2024

Simulation Techniques and ApplicationsDecision Sciences被引用数 3

ひとこと要約

この論文は、アテンションメカニズムによるn-gramの保存・取得が可能であるため、長大な入力系列からの情報のコピーと取得において、Transformersが一般化状態空間モデル（GSSMs）を上回ることを示している。理論的には、2層のTransformersは、そのサイズに指数的に比例する長さの系列をコピーできるが、GSSMsは固定サイズの潜在状態に起因して、その次元数を超えるビットを符号化できないという根本的な制限を受ける。

ABSTRACT

Transformers are the dominant architecture for sequence modeling, but there is growing interest in models that use a fixed-size latent state that does not depend on the sequence length, which we refer to as "generalized state space models" (GSSMs). In this paper we show that while GSSMs are promising in terms of inference-time efficiency, they are limited compared to transformer models on tasks that require copying from the input context. We start with a theoretical analysis of the simple task of string copying and prove that a two layer transformer can copy strings of exponential length while GSSMs are fundamentally limited by their fixed-size latent state. Empirically, we find that transformers outperform GSSMs in terms of efficiency and generalization on synthetic tasks that require copying the context. Finally, we evaluate pretrained large language models and find that transformer models dramatically outperform state space models at copying and retrieving information from context. Taken together, these results suggest a fundamental gap between transformers and GSSMs on tasks of practical interest.

研究の動機と目的

• GSSMsが効率的であるものの、入力文脈からの情報のコピーと取得において、Transformersに匹敵できるかどうかを調査すること。
• 固定サイズの潜在状態に起因するGSSMsの長系列表現における根本的な制限を分析すること。
• 合成コピー課題におけるTransformersとGSSMsの一般化性能および効率を評価すること。
• 事前学習モデル（例：Pythia 対 Mamba）のゼロショット設定における文脈の検索およびコピー性能をテストすること。
• 確率的パープレキシティ以外の文脈での言語モデルの能力に与えるアーキテクチャ的差異の実用的影響を理解すること。

提案手法

• GSSMsとTransformersの形式的定義を用いた文字列コピー課題の理論的分析により、GSSMsが潜在状態サイズを超える長さの系列をコピーできないことを証明する。
• アテンションヘッド数に指数的に比例する長さの系列をコピーできる2層のTransformersを構築し、アテンションによるn-gramの保存・取得を実現する。
• 訓練用に増加する長さの系列をコピーする合成データセットを用いた実験的評価を行い、訓練効率と長さ一般化性能を測定する。
• Transformersにおける訓練文脈窓を超えた一般化を向上させるために、ハード・アライボ（Hard-Alibi）位置エンコーディングを用いる。
• 文脈ベースの検索課題において、事前学習モデル（Pythia と Mamba）を評価し、さまざまな段落長におけるF1スコアを測定する。
• 訓練済みTransformersにおける学習済みアテンションパターンの分析により、コピーにn-gram検索メカニズムに依存していることを確認する。

(a) Copying: training efficiency.

実験結果

リサーチクエスチョン

• RQ1 Transformersモデルは、そのサイズに指数的に比例する長さの系列をコピーできるか？ もしそうなら、どのようにして達成しているか？
• RQ2 固定サイズの潜在状態に起因するGSSMsのコピー課題における根本的な表現的制限は何か？
• RQ3 合成コピー課題において、TransformersとGSSMsは訓練効率および長系列への一般化性能でどのように比較されるか？
• RQ4 事前学習済みTransformersは、ゼロショットの文脈検索およびコピー課題で、事前学習済みGSSMsを上回る性能を示すか？
• RQ5 GSSMsの固定メモリ容量は、長大な入力文脈を伴う自然言語処理タスクにおいて、どの程度性能を劣化させるか？

主な発見

• 2層のTransformersは、アテンションヘッド数に指数的に比例する長さの系列をコピーでき、n-gramの保存による文脈取得の強力なメカニズムを示している。
• GSSMsはコピー課題において根本的に制限を受ける：潜在状態サイズを超えるビット数を含む系列を正確にコピーすることはできない。
• 合成コピー課題において、TransformersはGSSMsよりも著しく高速に訓練され、長系列への一般化性能もはるかに優れている。Mambaは訓練長さを超えると性能が急激に低下する。
• 事前学習済みPythia（Transformersベース）モデルは、Mamba（GSSMベース）モデルよりもゼロショット文脈検索タスクで優れた性能を示し、段落長が延びるに従いMambaのF1スコアはより急激に低下する。
• Transformersにおける学習済みコピー機構は、アテンションに基づくn-gram検索に依存しており、理論的構築と一致する。
• MambaはPileデータセットにおいてPythiaよりも低いパープレキシティを達成しているが、文脈依存タスクでは性能が劣る。これは、訓練損失のみでは、検索やコピーの性能を予測できないことを示している。

(b) Copying: length generalization