Sample, estimate, aggregate: A recipe for causal discovery foundation models

📄 arXiv: 2402.01929v3 📥 PDF

作者: Menghua Wu, Yujia Bao, Regina Barzilay, Tommi Jaakkola

分类: cs.LG, stat.ML

发布日期: 2024-02-02 (更新: 2025-03-28)

备注: Our code is available at https://github.com/rmwu/sea

期刊: Transactions on Machine Learning Research (03/2025)

💡 一句话要点

提出基于基础模型的因果发现方法以应对数据限制问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 因果发现 基础模型 生物数据 合成数据 机器学习 模型鲁棒性 因果推断

📋 核心要点

  1. 现有因果发现算法在处理大规模变量集时，容易受到数据限制和模型误设定的影响，导致结果不稳定。
  2. 本文提出了一种基于基础模型的监督学习方法，通过合成数据训练模型，预测因果图，克服了传统方法的局限性。
  3. 实验结果表明，该模型在生物和合成数据上表现出良好的泛化能力，能够在秒级内处理数百个变量的因果图。

📝 摘要（中文）

因果发现是从数据中推断因果结构的任务，能够揭示生物实验中的机制性洞察。然而，现有因果发现算法在变量集较大时容易受到误设定或数据有限的影响。为此，本文提出了一种受基础模型启发的方法，通过在大规模合成数据上训练的监督模型，从汇总统计中预测因果图。该方法能够在不同数据集上保持输出的相对一致性，并在生物和合成数据的实验中表现出良好的泛化能力，能够快速处理数百个变量的图。

🔬 方法详解

问题定义：本文旨在解决因果发现中因变量集庞大和数据稀缺导致的模型不稳定性问题。现有方法在数据有限或模型设定不当时，容易产生误导性结果。

核心思路：提出一种受基础模型启发的监督学习方法，利用大规模合成数据训练模型，从汇总统计中预测因果图。该方法基于观察到的发现算法输出错误在不同数据集间的相对一致性。

技术框架：整体架构包括数据生成、模型训练和因果图预测三个主要模块。首先生成合成数据，然后在此基础上训练模型，最后利用训练好的模型进行因果图的预测。

关键创新：最重要的创新在于通过合成数据训练的模型能够有效应对模型误设定和数据分布变化，且能在不同数据假设下灵活适应。

关键设计：模型设计中采用了特定的损失函数以增强鲁棒性，并在网络结构上进行了优化，以确保模型能够快速收敛并有效处理大规模数据。具体参数设置和网络结构细节在论文中有详细描述。

📊 实验亮点

实验结果显示，该模型在生物和合成数据集上均表现出优异的性能，能够在数秒内处理数百个变量的因果图。与传统方法相比，模型在处理数据稀缺和误设定时的鲁棒性显著提升，验证了其广泛的适用性和有效性。

🎯 应用场景

该研究的潜在应用领域包括生物医学研究、基因组学和药物开发等。通过准确推断因果关系，研究人员可以更好地理解生物过程，优化实验设计，并推动个性化医疗的发展。未来，该方法有望在其他领域的因果推断中发挥重要作用。

📄 摘要（原文）

Causal discovery, the task of inferring causal structure from data, has the potential to uncover mechanistic insights from biological experiments, especially those involving perturbations. However, causal discovery algorithms over larger sets of variables tend to be brittle against misspecification or when data are limited. For example, single-cell transcriptomics measures thousands of genes, but the nature of their relationships is not known, and there may be as few as tens of cells per intervention setting. To mitigate these challenges, we propose a foundation model-inspired approach: a supervised model trained on large-scale, synthetic data to predict causal graphs from summary statistics -- like the outputs of classical causal discovery algorithms run over subsets of variables and other statistical hints like inverse covariance. Our approach is enabled by the observation that typical errors in the outputs of a discovery algorithm remain comparable across datasets. Theoretically, we show that the model architecture is well-specified, in the sense that it can recover a causal graph consistent with graphs over subsets. Empirically, we train the model to be robust to misspecification and distribution shift using diverse datasets. Experiments on biological and synthetic data confirm that this model generalizes well beyond its training set, runs on graphs with hundreds of variables in seconds, and can be easily adapted to different underlying data assumptions.