🤖 AI Summary
This study addresses the challenge of simultaneously recovering directed causal structures and identifying latent confounders from purely observational data. To this end, it proposes FoundCause—the first end-to-end model within an amortized causal discovery framework that explicitly accounts for latent-variable confounding and missing data. Key innovations include a permutation-invariant Transformer encoder integrating dual attention over both samples and variables, a statistical-feature-guided conditional attention mechanism, a learnable latent-variable token module, and a triangle-refinement decoder for inferring higher-order structural relationships. Evaluated across 15 real-world datasets, FoundCause substantially outperforms 15 baseline methods, achieving a 9.6% increase in F₁ score, a 1.2% gain in AUROC, and an 18.9% reduction in structural Hamming distance, all with only a single forward pass for inference.
📝 Abstract
Causal discovery from observational data remains challenging due to the need to recover directed structure and latent confounding without interventions. We propose FoundCause, an amortized causal discovery model trained entirely on synthetic data that maps datasets directly to causal graphs in a single forward pass. By learning from large collections of simulated structural causal models, FoundCause captures transferable statistical patterns that generalize beyond individual datasets. The architecture incorporates several key inductive biases for causal discovery. It uses a permutation-invariant transformer encoder with alternating attention over samples and variables to jointly model cross-variable dependence and per-variable distributions. Pairwise statistical features derived from classical asymmetry measures are injected through statistics-conditioned attention, guiding the model toward known causal signals. A factorized decoder separates edge existence from direction, while a triangular refinement module enables reasoning over higher-order causal motifs such as chains and colliders. In addition, a dedicated confounder module based on learnable latent tokens explicitly models hidden common causes, and the model explicitly handles missing data via its masked input representation. To our knowledge, FoundCause is the first amortized causal discovery approach to explicitly model latent confounding. FoundCause outperforms 11 classical non-amortized methods (e.g., PC, GES, NOTEARS-style optimization) and 4 amortized causal discovery methods on 15 real-world datasets, achieving +9.6% improvement in $F_1$, +1.2% in AUROC, and an 18.9% reduction in structural Hamming distance relative to the strongest non-amortized methods, while performing inference in a single forward pass.