arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2402.14735cs.LGcs.ITmath.ITstat.ML
收藏

Transformer如何通过梯度下降学习因果结构

How Transformers Learn Causal Structure with Gradient Descent

Princeton University(普林斯顿大学)

Eshaan Nichani, Alex Damian, Jason D. Lee

AI总结:

该研究针对Transformer通过梯度下降学习因果结构的机制不明问题,提出上下文学习任务,证明两层Transformer的注意力矩阵梯度编码词元互信息、对应潜在因果图边,还验证了其可恢复多种因果结构。

摘要:

Transformer在序列建模任务上取得的巨大成功很大程度上可归功于自注意力机制,该机制支持信息在序列不同部分之间传递。自注意力让Transformer能够编码因果结构,这也使其尤其适用于序列建模。然而,Transformer通过基于梯度的训练算法学习此类因果结构的过程仍未得到充分理解。为了更深入地理解这一过程,我们提出了一项需要学习潜在因果结构的上下文学习任务。我们证明,在简化的两层Transformer上进行的梯度下降,会通过在第一注意力层中编码潜在因果图来学习解决该任务。我们证明的核心洞见是,注意力矩阵的梯度编码了词元之间的互信息。根据数据处理不等式,该梯度的最大元素对应潜在因果图中的边。作为特例,当序列由上下文马尔可夫链生成时,我们证明Transformer会学习到归纳头(induction head,Olsson等人,2022)。我们通过实验证实了理论发现:在我们的上下文学习任务上训练的Transformer能够恢复多种不同的因果结构。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑