Transformers Provably Learn Directed Acyclic Graphs via Kernel-Guided Mutual Information

TL;DR

This paper introduces a new information-theoretic metric, KG-MI, enabling transformers to learn and recover complex DAG structures with provable guarantees, extending previous tree-based results to more general graphs.

Contribution

The work proposes KG-MI combined with multi-head attention to provably learn DAGs, providing convergence proofs and structure recovery guarantees for transformer models.

Findings

Gradient ascent on transformers converges to the global optimum for DAG sequences.

Attention scores at convergence reflect the true DAG structure.

Experimental results support theoretical claims.

Abstract

Uncovering hidden graph structures underlying real-world data is a critical challenge with broad applications across scientific domains. Recently, transformer-based models leveraging the attention mechanism have demonstrated strong empirical success in capturing complex dependencies within graphs. However, the theoretical understanding of their training dynamics has been limited to tree-like graphs, where each node depends on a single parent. Extending provable guarantees to more general directed acyclic graphs (DAGs) -- which involve multiple parents per node -- remains challenging, primarily due to the difficulty in designing training objectives that enable different attention heads to separately learn multiple different parent relationships. In this work, we address this problem by introducing a novel information-theoretic metric: the kernel-guided mutual information (KG-MI), based on the $f$-divergence. Our objective combines KG-MI with a multi-head attention framework, where each head is associated with a distinct marginal transition kernel to model diverse parent-child dependencies effectively. We prove that, given sequences generated by a $K$-parent DAG, training a single-layer, multi-head transformer via gradient ascent converges to the global optimum in polynomial time. Furthermore, we characterize the attention score patterns at convergence. In addition, when particularizing the $f$-divergence to the KL divergence, the learned attention scores accurately reflect the ground-truth adjacency matrix, thereby provably recovering the underlying graph structure. Experimental results validate our theoretical findings.