Transformers Provably Learn Directed Acyclic Graphs via Kernel-Guided Mutual Information
This work addresses a theoretical gap for researchers in machine learning and graph learning, providing provable guarantees for transformers on general DAGs, which is incremental but important for advancing understanding.
The paper tackles the problem of extending provable guarantees for transformer-based models to learn directed acyclic graphs (DAGs) with multiple parents per node, and proves that training a multi-head transformer with a novel kernel-guided mutual information objective converges to the global optimum in polynomial time, accurately recovering the underlying graph structure.
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.