Arrow: A Foundation Model for Causal Discovery
作者: Ryan Thompson, He Zhao, Daniel M. Steinberg, Edwin V. Bonilla
分类: cs.LG
发布日期: 2026-05-08
💡 一句话要点
提出Arrow基础模型,通过骨架与拓扑排序分解实现零样本因果发现
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 因果发现 基础模型 Transformer架构 零样本学习 拓扑排序 结构学习 表格数据分析
📋 核心要点
- 传统因果发现方法通常针对特定数据集进行昂贵的迭代优化,难以在处理新数据时实现快速的零样本推理。
- Arrow通过将DAG分解为无向骨架与拓扑排序,利用Transformer架构进行上下文建模,从根本上保证了图结构的无环性。
- 实验表明,该模型在多种合成与真实数据集上表现优异,推理效率远高于传统方法,验证了预训练模型在因果发现领域的潜力。
📝 摘要(中文)
本文介绍了Arrow,一种用于观测表格数据零样本因果发现的基础模型。Arrow将有向无环图(DAG)分解为无向骨架和拓扑排序,从而在构建过程中保证了无环性。给定新数据集,该模型利用基于Transformer的架构对观测内部及观测间的变量进行上下文建模,进而预测骨架边概率和节点排序分数,共同定义因果图。Arrow在具有基准图的合成数据集上进行监督训练,采用了由骨架-排序分解导出的端到端可微有向边复合似然函数。训练分布涵盖了多样的图族、函数形式、噪声模型和数据集形态。在分布内及分布外的合成、半合成和真实数据集上,Arrow在显著降低推理成本的同时,匹配或超越了现有的因果发现方法。研究结果表明,在大规模多样化合成数据上进行预训练,能够产生快速、准确且可复用于新数据集的零样本因果发现模型。
🔬 方法详解
问题定义:因果发现旨在从观测数据中推断变量间的因果关系。现有方法多依赖于针对单个数据集的搜索或约束优化,计算开销大且泛化能力弱,难以应对大规模或多样化的实际数据场景。
核心思路:Arrow的核心思想是将复杂的DAG学习问题转化为两个子任务的联合预测:预测变量间的无向骨架(确定连接关系)和预测节点的拓扑排序(确定因果方向)。这种分解策略天然规避了搜索空间中的环路问题。
技术框架:模型采用Transformer架构,将观测数据作为输入,通过自注意力机制捕捉变量间的复杂依赖关系。输出层分为两部分:一是预测边存在的概率矩阵,二是预测每个节点的拓扑排序分数,两者结合生成最终的因果图。
关键创新:引入了端到端可微的“有向边复合似然”损失函数,该函数直接作用于骨架-排序分解结构,使得模型能够通过监督学习在海量合成数据上进行预训练,从而获得强大的零样本迁移能力。
关键设计:训练数据覆盖了广泛的图结构、函数形式(线性/非线性)及噪声分布,确保模型具备极强的鲁棒性。通过这种大规模预训练,模型在推理阶段仅需一次前向传播即可输出因果图,极大提升了计算效率。
🖼️ 关键图片
📊 实验亮点
Arrow在多种基准测试中展现了卓越性能,不仅在准确度上匹配甚至超越了传统的基于搜索或约束的因果发现算法,更在推理速度上实现了数量级的提升。其在分布外(OOD)数据集上的稳健表现,证明了通过大规模合成数据预训练构建通用因果发现模型的可行性与高效性。
🎯 应用场景
Arrow适用于需要快速从观测数据中提取因果关系的领域,如生物医学研究中的基因调控网络推断、经济学中的宏观指标因果分析、以及工业生产中的故障根因溯源。其零样本特性使其在处理缺乏先验知识的新数据集时具有极高的实用价值,能够显著降低因果分析的门槛与时间成本。
📄 摘要(原文)
We introduce Arrow, a foundation model for zero-shot causal discovery on observational tabular data. Arrow factorizes a directed acyclic graph into an undirected skeleton and a topological order, guaranteeing acyclicity by construction. Given a new dataset, it uses a transformer-based architecture to contextualize variables within and across observations, then predicts skeleton edge probabilities and node order scores that together define a graph. Arrow is trained in a supervised fashion on synthetic datasets with ground-truth graphs, using an end-to-end differentiable directed edge composite likelihood induced by the skeleton-order factorization. The training distribution spans diverse graph families, functional forms, noise models, and dataset shapes. Across in- and out-of-distribution synthetic, semi-synthetic, and real datasets, Arrow matches or outperforms existing causal discovery methods at substantially lower inference cost than competitive alternatives. Our results demonstrate that large-scale pretraining on diverse synthetic data can yield zero-shot causal discovery models that are fast, accurate, and reusable on new datasets.