[論文レビュー] How Transformers Learn Causal Structure with Gradient Descent

Eshaan Nichani, Alex Damian+1 more|arXiv (Cornell University)|Feb 22, 2024

Bayesian Modeling and Causal Inference被引用数 5

ひとこと要約

この論文は、簡略化された二層トランスフォーマーに対する勾配降下法が潜在的因果構造を学習し、因果グラフを最初のアテンション層にエンコードすることを示し、文脈内マルコフ連鎖設定でインダクションヘッドが現れることを示す。アテンションの勾配情報は相互情報量を反映する。

ABSTRACT

The incredible success of transformers on sequence modeling tasks can be largely attributed to the self-attention mechanism, which allows information to be transferred between different parts of a sequence. Self-attention allows transformers to encode causal structure which makes them particularly suitable for sequence modeling. However, the process by which transformers learn such causal structure via gradient-based training algorithms remains poorly understood. To better understand this process, we introduce an in-context learning task that requires learning latent causal structure. We prove that gradient descent on a simplified two-layer transformer learns to solve this task by encoding the latent causal graph in the first attention layer. The key insight of our proof is that the gradient of the attention matrix encodes the mutual information between tokens. As a consequence of the data processing inequality, the largest entries of this gradient correspond to edges in the latent causal graph. As a special case, when the sequences are generated from in-context Markov chains, we prove that transformers learn an induction head (Olsson et al., 2022). We confirm our theoretical findings by showing that transformers trained on our in-context learning task are able to recover a wide variety of causal structures.

研究の動機と目的

• 勾配に基づく学習がトランスフォーマーにおける因果構造を生み出す仕組みの理解を促す。
• 潜在的なグラフを固定するための因果構造を持つランダム系列タスクを導入する。
• 勾配降下法の下での注意機構のみを持つ二層トランスフォーマーの学習ダイナミクスを分析する。
• 注意行列の勾配が相互情報量を捉え、グラフのエッジを明らかにすることを示す。
• マルチヘッドアーキテクチャを介して木ではないグラフへの一般化を実証し、分布外の性能を評価する。

提案手法

• A^(1) および A^(2) に焦点を当てた簡略化された二層の分離型トランスフォーマーと縮約モデルを定義する。
• トークン位置上の潜在的DAGによって定義される因果構造を持つランダム系列タスクを構築する。
• 勾配降下法が潜在グラフをA^(1) にエンコードすることで回復することを証明する。
• 勾配がカイ二乗相互情報量の測度に対応し、データ処理不等式を介してエッジの回復を導くことを示す。
• 特殊ケース分析：文脈内マルコフ連鎖はインダクションヘッドを生み出す。
• 非木グラフをヘッド間で分散させるためのマルチヘッド拡張を提供し、経験的に検証する。

実験結果

リサーチクエスチョン

• RQ1 固定された因果グラフで生成されたデータから、トランスフォーマーの勾配降下法は潜在的因果構造を回復できるか？
• RQ2 訓練中、因果構造はトランスフォーマーの注意層内でどのようにエンコードされるか？
• RQ3 文脈内学習（例：マルコフ連鎖）の場面で、どのようなプリミティブ（例：インダクションヘッド）が現れるか？
• RQ4 潜在グラフが木構造でない場合のモデルの性能はどうなるか、そしてマルチヘッド設計で解決できるか？
• RQ5 訓練済みモデルは分布外の遷移に一般化するか？

主な発見

• 二層のディスエンタングルド トランスフォーマーに対する勾配降下法は、潜在的因果グラフを最初のアテンション層の隣接行列としてエンコードすることを学習する。
• 最初のアテンション層の勾配はトークン間のカイ二乗相互情報量に対応し、データ処理不等式が学習をグラフエッジへ集中させる。
• 文脈内マルコフ連結の特別ケースでは、遷移の文脈内推定を行うインダクションヘッドを発達させる。
• 因果グラフが木ではない場合、マルチヘッドトランスフォーマは潜在グラフをヘッドに分配して解決行動を達成できる。
• 実証的には、訓練されたトランスフォーマは提案タスク上で様々な因果構造を回復し、遷移に対する分布外一般化を示す。
• 理論的保証（定理1および定理2）は、特定の仮定の下で母集団損失の収束とOOD一般化を確立する。