Learning Causal Graphs at Scale: A Foundation Model Approach

• URL: http://arxiv.org/abs/2506.18285v1
• Date: Mon, 23 Jun 2025 04:41:02 GMT
• Title: Learning Causal Graphs at Scale: A Foundation Model Approach
• Authors: Naiyu Yin, Tian Gao, Yue Yu,
• Abstract summary: We propose Attention-DAG (ADAG), a novel attention-mechanism-based architecture for learning multiple linear Structural Equation Models (SEMs)<n>ADAG learns the mapping from observed data to both graph structure and parameters via a nonlinear attention-based kernel.<n>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.
• Score: 28.966180222166766
• License: http://creativecommons.org/licenses/by/4.0/
• 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, in this work we leverage the recent success of linear transformers and develop a foundation model approach for discovering multiple order-consistent 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 estimation of the underlying linear SEMs. By formulating the learning process across multiple tasks 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 learning 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 specifically designed for DAG learning, representing a step toward more efficient and generalizable down-stream applications in causal discovery.

Related papers

• Anomaly Detection and Generation with Diffusion Models: A Survey [51.61574868316922]
  Anomaly detection (AD) plays a pivotal role across diverse domains, including cybersecurity, finance, healthcare, and industrial manufacturing.<n>Recent advancements in deep learning, specifically diffusion models (DMs), have sparked significant interest.<n>This survey aims to guide researchers and practitioners in leveraging DMs for innovative AD solutions across diverse applications.
  arXiv  Detail & Related papers  (2025-06-11T03:29:18Z)
• KARE-RAG: Knowledge-Aware Refinement and Enhancement for RAG [63.82127103851471]
  Retrieval-Augmented Generation (RAG) enables large language models to access broader knowledge sources.<n>We demonstrate that enhancing generative models' capacity to process noisy content is equally critical for robust performance.<n>We present KARE-RAG, which improves knowledge utilization through three key innovations.
  arXiv  Detail & Related papers  (2025-06-03T06:31:17Z)
• Transfer Learning with Foundational Models for Time Series Forecasting using Low-Rank Adaptations [0.0]
  This study proposes the methodology LLIAM, a straightforward adaptation of a kind of FM, Large Language Models, for the Time Series Forecasting task.<n>A comparison was made between the performance of LLIAM and different state-of-the-art DL algorithms, including Recurrent Neural Networks and Temporal Convolutional Networks, as well as a LLM-based method, TimeLLM.<n>The outcomes of this investigation demonstrate the efficacy of LLIAM, highlighting that this straightforward and general approach can attain competent results without the necessity for applying complex modifications.
  arXiv  Detail & Related papers  (2024-10-15T12:14:01Z)
• MTLSO: A Multi-Task Learning Approach for Logic Synthesis Optimization [19.13500546022262]
  MTLSO is a Multi-Task Learning approach for Logic Synthesis Optimization. We introduce an auxiliary task of binary multi-label graph classification alongside the primary regression task. We also employ a hierarchical graph representation learning strategy to improve the model's capacity for learning expressive graph-level representations.
  arXiv  Detail & Related papers  (2024-09-09T21:20:36Z)
• Data-Centric Long-Tailed Image Recognition [49.90107582624604]
  Long-tail models exhibit a strong demand for high-quality data. Data-centric approaches aim to enhance both the quantity and quality of data to improve model performance. There is currently a lack of research into the underlying mechanisms explaining the effectiveness of information augmentation.
  arXiv  Detail & Related papers  (2023-11-03T06:34:37Z)
• Discovering Dynamic Causal Space for DAG Structure Learning [64.763763417533]
  We propose a dynamic causal space for DAG structure learning, coined CASPER. It integrates the graph structure into the score function as a new measure in the causal space to faithfully reflect the causal distance between estimated and ground truth DAG.
  arXiv  Detail & Related papers  (2023-06-05T12:20:40Z)
• End-to-End Meta-Bayesian Optimisation with Transformer Neural Processes [52.818579746354665]
  This paper proposes the first end-to-end differentiable meta-BO framework that generalises neural processes to learn acquisition functions via transformer architectures. We enable this end-to-end framework with reinforcement learning (RL) to tackle the lack of labelled acquisition data.
  arXiv  Detail & Related papers  (2023-05-25T10:58:46Z)
• DAGs with No Curl: An Efficient DAG Structure Learning Approach [62.885572432958504]
  Recently directed acyclic graph (DAG) structure learning is formulated as a constrained continuous optimization problem with continuous acyclicity constraints. We propose a novel learning framework to model and learn the weighted adjacency matrices in the DAG space directly. We show that our method provides comparable accuracy but better efficiency than baseline DAG structure learning methods on both linear and generalized structural equation models.
  arXiv  Detail & Related papers  (2021-06-14T07:11:36Z)
• On the Role of Sparsity and DAG Constraints for Learning Linear DAGs [16.97675762810828]
  We study the role of sparsity and DAG constraints for learning DAG models in the linear Gaussian and non-Gaussian cases. We propose a likelihood-based score function, and show that one only has to apply soft sparsity and DAG constraints to learn a DAG equivalent to the ground truth DAG.
  arXiv  Detail & Related papers  (2020-06-17T23:43:23Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.