[论文解读] BayesDAG: Gradient-Based Posterior Inference for Causal Discovery

Yashas Annadani, Nick Pawlowski+4 more|arXiv (Cornell University)|Jul 26, 2023

Bayesian Modeling and Causal Inference被引用 4

一句话总结

BayesDAG 提出了一种新颖的基于梯度的贝叶斯因果发现框架，通过利用随机梯度马尔可夫链蒙特卡洛（SG-MCMC）和变分推断，实现了对有向无环图（DAGs）和非线性结构函数的可扩展后验推断。该方法在合成数据集和真实世界数据集上均实现了最先进的因果发现性能与不确定性量化，可扩展至100个变量，且无需DAG正则化。

ABSTRACT

Bayesian causal discovery aims to infer the posterior distribution over causal models from observed data, quantifying epistemic uncertainty and benefiting downstream tasks. However, computational challenges arise due to joint inference over combinatorial space of Directed Acyclic Graphs (DAGs) and nonlinear functions. Despite recent progress towards efficient posterior inference over DAGs, existing methods are either limited to variational inference on node permutation matrices for linear causal models, leading to compromised inference accuracy, or continuous relaxation of adjacency matrices constrained by a DAG regularizer, which cannot ensure resulting graphs are DAGs. In this work, we introduce a scalable Bayesian causal discovery framework based on a combination of stochastic gradient Markov Chain Monte Carlo (SG-MCMC) and Variational Inference (VI) that overcomes these limitations. Our approach directly samples DAGs from the posterior without requiring any DAG regularization, simultaneously draws function parameter samples and is applicable to both linear and nonlinear causal models. To enable our approach, we derive a novel equivalence to the permutation-based DAG learning, which opens up possibilities of using any relaxed gradient estimator defined over permutations. To our knowledge, this is the first framework applying gradient-based MCMC sampling for causal discovery. Empirical evaluation on synthetic and real-world datasets demonstrate our approach's effectiveness compared to state-of-the-art baselines.

研究动机与目标

• 解决在组合性DAG空间与连续函数参数上的可扩展贝叶斯因果发现挑战。
• 克服现有方法依赖DAG正则化（无法保证无环性）或变分推断（可能面临精度问题）的局限性。
• 开发一种适用于线性和非线性结构因果模型的模型无关框架。
• 通过直接使用基于梯度的MCMC对DAG进行采样，无需松弛或正则化，实现高质量的后验推断。
• 在合成数据集和真实世界数据集上展示方法的有效性与可扩展性，包括最多100个变量。

提出的方法

• 引入边信念与节点势能的扩展空间 $({\bm{W}}, {\bm{p}})$，实现在无需DAG正则化的情况下对DAG进行等价贝叶斯推断。
• 推导出该扩展空间与基于排列的DAG学习之间的等价性，从而可使用任意松弛化的梯度估计器处理排列。
• 提出一种结合SG-MCMC与变分推断（VI）的混合框架，用于对DAG与函数参数进行联合后验采样。
• 在第二种公式中采用邻接矩阵 $\tilde{\bm{W}}$ 的连续松弛，实现无需显式DAG约束的基于梯度的MCMC采样。
• 使用变分网络 $\mu_{\phi}$ 联合节点势能 $\bm{p}$ 与边信念 $\bm{W}$，实现在采样过程中响应式更新。
• 在SG-MCMC中利用随机梯度估计器，实现对高维设置（包括 $d=100$ 个变量）的可扩展性。

实验结果

研究问题

• RQ1 基于梯度的MCMC采样能否有效应用于离散DAG结构上的贝叶斯因果发现？
• RQ2 能否通过边信念与节点势能的扩展空间实现无松弛的DAG采样方法？
• RQ3 所提出框架在后验精度与不确定性量化方面相较于现有变分方法与正则化方法表现如何？
• RQ4 该方法能否扩展至具有100个变量的高维非线性因果模型？
• RQ5 通过变分网络耦合节点势能与边信念，是否能提升采样效率与精度？

主要发现

• 在 $d=20$ 的合成非线性数据集上，BayesDAG 实现了最先进的性能，相较于所有基线方法在期望结构汉明距离（E-SHD）、边F1与负对数似然（NLL）方面均表现更优。
• 在 $d=70$ 与 $d=100$ 的非线性设置中，BayesDAG 在E-SHD与NLL方面保持卓越性能，仅在边F1上相比最佳基线略有下降。
• SG-MCMC+VI 框架优于采用连续松弛 ($\tilde{\bm{W}}$) 的SG-MCMC，原因在于变分网络中 $\bm{p}$ 与 $\bm{W}$ 的耦合更优。
• BayesDAG 高效扩展至 $d=100$ 个变量，单张40GB A100 GPU 上运行时间为448分钟，而DIBS与DDS因内存限制无法运行。
• BGES与DDS虽然更快，但在不确定性量化方面表现显著更差，且仅限于线性模型或非完全贝叶斯方法。
• 实证结果证实，BayesDAG 能够提供准确的后验推断并具备良好的不确定性校准，尤其在低数据场景下表现突出。