How Transformers Learn Causal Structure with Gradient Descent

要約

シーケンス モデリング タスクにおけるトランスフォーマーの驚異的な成功は、シーケンスの異なる部分間で情報を転送できるセルフ アテンション メカニズムに大きく起因すると考えられます。
セルフアテンションにより、トランスフォーマーは因果構造をエンコードできるため、シーケンス モデリングに特に適しています。
しかし、変換器が勾配ベースのトレーニング アルゴリズムを介してそのような因果構造を学習するプロセスは、依然としてよく理解されていません。
このプロセスをより深く理解するために、潜在的な因果構造の学習を必要とするコンテキスト内学習タスクを導入します。
単純化された 2 層変換器の勾配降下法が、最初の注目層の潜在因果グラフをエンコードすることによってこのタスクを解決することを学習することを証明します。
私たちの証明の重要な洞察は、注意行列の勾配がトークン間の相互情報をエンコードしているということです。
データ処理の不平等の結果として、この勾配の最大のエントリは潜在因果グラフのエッジに対応します。
特殊なケースとして、シーケンスがコンテキスト内のマルコフ連鎖から生成される場合、変換器が誘導頭部を学習することを証明します (Olsson et al., 2022)。
コンテキスト内学習タスクで訓練されたトランスフォーマーがさまざまな因果構造を回復できることを示すことで、理論的発見を確認します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google