Towards a Foundation Model Approach for Causal Graph Learning

ICLR 2026 Conference Submission21821 Authors

Keywords: Causal Discovery; Inverse Problems; Attention Mechanism

TL;DR: We propose a foundation model approach using attention mechanisms for DAG learning tasks, enabling accurate and efficient causal discovery, especially in low-data regimes.

Abstract: Due to its human-interpretability and invariance properties, Directed Acyclic Graph (DAG) has been a foundational tool across various areas of AI research. %, leading to significant advancements. However, DAG learning remains highly challenging, due to its super-exponential growth in computational cost and identifiability issues, particularly in small-sample regimes. To address these two challenges, we leverage the recent success of transformers and develop a foundation model approach for discovering multiple DAGs across tasks. In particular, we propose Attention-DAG (ADAG), a novel attention-mechanism-based architecture for learning multiple linear Structural Equation Models (SEMs). ADAG learns the mapping from observed data to both graph structure and parameters via a nonlinear attention-based kernel, enabling efficient multi-task generalization of the underlying linear SEMs. By formulating the learning process across multiple domains as a continuous optimization problem, the pre-trained ADAG model captures the common structural properties as a shared low-dimensional prior, thereby reducing the ill-posedness of downstream DAG tasks in small-sample regimes. We evaluate our proposed approach on benchmark synthetic datasets and find that ADAG achieves substantial improvements in both DAG learning accuracy and zero-shot inference efficiency. To the best of our knowledge, this is the first practical approach for pre-training a foundation model for unsupervised DAG learning, representing a step toward more efficient and generalizable down-stream applications in causal discovery.

Primary Area: causal reasoning

Submission Number: 21821