DCD-PFN: A Decoupling-Aware Foundation Model for Causal Discovery

2026-06-19

AI Summary

This work addresses the limitations of traditional causal discovery methods, which suffer from poor performance and low computational efficiency in highly nonlinear and noisy systems, as well as the inability of existing PFN-based tabular foundation models to support explicit structural causal discovery. The authors propose DCD-PFN, the first approach to integrate a disentanglement mechanism into the PFN framework. By pretraining on sample-level disentangled weights to identify Markov boundaries, DCD-PFN enables zero-shot discovery of local causal relationships and efficiently reconstructs the global causal graph through parallel integration of local results. The method demonstrates robust zero-shot generalization across diverse scenarios, significantly improving both the accuracy and computational efficiency of causal graph reconstruction while maintaining interpretability and theoretical grounding.

Abstract

Causal discovery is critical for understanding complex data-generating mechanisms, yet traditional algorithms often struggle with highly non-linear and noisy systems, or suffer from severe computational bottlenecks. Recent tabular foundation models based on Prior-Data Fitted Networks (PFNs) have demonstrated remarkable zero-shot inference capabilities, but their potential for explicit structural causal discovery remains underexplored. To bridge this gap, we propose DCD-PFN, a decoupling-aware foundation model for causal discovery. Instead of directly amortizing global graph reconstruction, DCD-PFN focuses on local causal discovery through a decoupling-based paradigm. Through pre-training on diverse synthetic Structural Causal Models (SCMs), the model learns sample-wise decoupling weights that enable Markov boundary (MB) identification. Furthermore, by leveraging parallelized local discovery, DCD-PFN efficiently reconstructs global causal graphs while remaining grounded in the theoretical foundations of decoupling-based causal discovery. Experiments demonstrate that our foundation model achieves robust zero-shot generalization.

Problem

• causal discovery
• non-linear systems
• noisy systems
• computational bottlenecks
• structural causal models

Innovation

• decoupling-aware
• causal discovery
• Prior-Data Fitted Networks
• Markov boundary
• zero-shot generalization