[论文解读] Amortized Inference for Causal Structure Learning

Lars Lorch, Scott Sussex+3 more|arXiv (Cornell University)|May 25, 2022

Bayesian Modeling and Causal Inference被引用 6

一句话总结

本文提出 AVICI，一种基于神经网络的近似因果结构学习方法，可直接从观测数据或干预数据中预测因果图，避免了对图结构进行昂贵的搜索。通过利用排列等变架构并在模拟器生成的数据上进行训练，AVICI 在更大、未见过的图上实现了稳健的泛化性能，且在基因组学领域尤其优于现有方法，同时提供校准的不确定性估计。

ABSTRACT

Inferring causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test. The resulting search is costly, and designing suitable scores or tests that capture prior knowledge is difficult. In this work, we propose to amortize causal structure learning. Rather than searching over structures, we train a variational inference model to directly predict the causal structure from observational or interventional data. This allows our inference model to acquire domain-specific inductive biases for causal discovery solely from data generated by a simulator, bypassing both the hand-engineering of suitable score functions and the search over graphs. The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training. On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain. Our code and models are publicly available at: https://github.com/larslorch/avici.

研究动机与目标

• 解决因果结构学习中设计精确评分函数或独立性检验所面临的高计算成本和困难问题。
• 克服传统基于搜索的方法因组合爆炸的搜索空间和结构偏差而导致的局限性。
• 通过从模拟器生成的数据中学习特定领域的归纳偏置，实现对训练时未见的更大图的泛化。
• 为因果预测提供校准的不确定性估计，提升在实际应用中的可靠性。
• 在基因组学等挑战性领域展示优越性能，其中现有方法在噪声和分布偏移下表现不佳。

提出的方法

• 该方法采用近似变分推理框架（AVICI），通过训练神经网络直接从输入数据预测因果图，从而消除对迭代结构搜索的需求。
• 推理模型设计为对变量顺序具有排列等变性，确保对输入顺序的不变性，提升统计效率和泛化能力。
• 模型在由数据生成过程模拟器生成的合成数据上进行训练，使其能够学习特定领域的归纳偏置，而无需手工设计的评分函数。
• 架构采用带有注意力机制的图神经网络，以建模变量间的依赖关系并预测边的概率。
• 边的概率以端到端方式预测，支持不确定性量化和阈值处理以获得点估计。
• 通过模型架构和训练数据中编码的归纳偏置，实现对训练时未见的更大图的泛化。

实验结果

研究问题

• RQ1 是否可以训练神经网络直接从数据预测因果结构，从而绕过图结构的组合搜索？
• RQ2 在小图上训练的模型在泛化到显著更大的未见图时表现如何？
• RQ3 模型能否从模拟器生成的数据中学习特定领域的归纳偏置，而无需依赖手工设计的评分函数？
• RQ4 模型在图结构、机制、噪声和问题规模发生分布偏移时的鲁棒性如何？
• RQ5 模型能否为因果预测提供校准的不确定性估计，特别是在基因组学等高风险领域？

主要发现

• 在合成数据和半合成基因表达数据上，AVICI 显著优于现有方法，尤其在分布偏移和高噪声条件下表现突出。
• 在基因表达数据集上，AVICI 是唯一能够推断出合理因果结构的方法，展示了在具有挑战性的生物学领域中的最先进性能。
• 仅使用观测数据，AVICI 在 Grn 数据集上实现了 AUROC 0.901（±0.04）和 AUPRC 0.656（±0.10）的优异表现，表明其在无干预数据时也具备强大性能。
• 模型的期望校准误差（ECE）低于所有基线方法，仅非参数自助法除外，表明其不确定性估计具有良好的校准性。
• AVICI 成功泛化至最多含 30 个变量的图，即使训练时仅使用较小图，显示出强大的归纳偏置迁移能力。
• 在 SHD 和 F1 指标上，AVICI 与 GES 和 GIES 表现相当，且在概率性指标上优于 DAG-GNN、LiNGAM 和 DCDI。