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.

翻译：Transformer在序列建模任务中取得的惊人成功，很大程度上可归因于其自注意力机制——该机制允许信息在序列的不同部分之间传递。自注意力使Transformer能够编码因果结构，这使其特别适合序列建模。然而，Transformer通过基于梯度的训练算法学习这种因果结构的过程仍未被充分理解。为深入理解这一过程，我们引入了一项需要学习潜在因果结构的上下文学习任务。我们证明，在简化的双层Transformer上应用梯度下降时，模型通过将潜在因果图编码到第一注意力层中来求解该任务。我们证明的核心洞见在于：注意力矩阵的梯度编码了词元之间的互信息。根据数据处理不等式，该梯度中最大的项对应潜在因果图中的边。作为特例，当序列由上下文马尔可夫链生成时，我们证明Transformer会学习一个归纳头（Olsson等人，2022）。通过在上下文学习任务上训练Transformer，我们验证了其能够恢复多种因果结构，从而证实了理论发现。