[论文解读] Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes

Mizu Nishikawa-Toomey, Tristan Deleu+3 more|arXiv (Cornell University)|Nov 4, 2022

Bayesian Modeling and Causal Inference被引用 5

一句话总结

该论文提出了一种名为变分贝叶斯-DAG-GFlowNet（VBG）的新方法，通过变分贝叶斯和GFlowNets联合学习因果DAG结构和线性高斯机制参数的后验分布。通过将DAG-GFlowNet扩展以同时建模结构和机制参数，VBG保证了无环性，支持灵活的参数化，并在模拟数据上实现了具有竞争力的性能，其不确定性估计和后验近似优于基线方法。

ABSTRACT

Bayesian causal structure learning aims to learn a posterior distribution over directed acyclic graphs (DAGs), and the mechanisms that define the relationship between parent and child variables. By taking a Bayesian approach, it is possible to reason about the uncertainty of the causal model. The notion of modelling the uncertainty over models is particularly crucial for causal structure learning since the model could be unidentifiable when given only a finite amount of observational data. In this paper, we introduce a novel method to jointly learn the structure and mechanisms of the causal model using Variational Bayes, which we call Variational Bayes-DAG-GFlowNet (VBG). We extend the method of Bayesian causal structure learning using GFlowNets to learn not only the posterior distribution over the structure, but also the parameters of a linear-Gaussian model. Our results on simulated data suggest that VBG is competitive against several baselines in modelling the posterior over DAGs and mechanisms, while offering several advantages over existing methods, including the guarantee to sample acyclic graphs, and the flexibility to generalize to non-linear causal mechanisms.

研究动机与目标

• 解决在贝叶斯因果结构学习中，联合学习因果结构与机制参数在不确定性下的挑战。
• 克服现有方法的局限性，这些方法或无法保证无环性，或在机制参数化方面缺乏灵活性，或无法从后验分布中采样。
• 通过提供一种原则性且可扩展的后验近似方法，实现因果发现中的主动学习，涵盖DAG和机制参数的后验分布。
• 通过变分贝叶斯方法扩展DAG-GFlowNet的能力——此前该方法仅限于结构学习——并引入机制参数推断。
• 提供一个统一的框架，通过机制分布的灵活参数化，支持非线性机制的泛化。

提出的方法

• 提出联合后验近似 $P(G, \theta \mid \mathcal{D}) \approx q_{\phi}(G) q_{\lambda}(\theta \mid G)$，使用变分贝叶斯，其中 $G$ 为DAG结构，$\theta$ 为机制参数。
• 使用GFlowNet对DAG上的变分分布 $q_{\phi}(G)$ 进行建模，其奖励函数 $\widetilde{R}(G)$ 由ELBO导出，包含期望似然和KL散度项。
• 定义奖励函数 $\log \widetilde{R}(G) = \mathbb{E}_{\theta \sim q_{\lambda}}[\log P(\mathcal{D} \mid \theta, G)] - \mathrm{KL}(q_{\lambda}(\theta \mid G) \| P(\theta \mid G)) + \log P(G)$，以指导GFlowNet的训练。
• 通过交替优化 $\phi$（GFlowNet参数）和 $\lambda$（机制参数分布参数）来执行ELBO的坐标上升，其中 $q_{\lambda}(\theta \mid G)$ 使用变分推断进行更新。
• 通过GFlowNet架构的内在结构约束，从构造上保证无环性，该架构仅采样有效DAG。
• 利用共轭先验下机制参数的闭式边缘化（例如线性高斯模型），高效计算 $P(\mathcal{D} \mid G)$ 和 $P(\theta \mid G)$。

实验结果

研究问题

• RQ1 我们能否使用一种可扩展且可微的框架，联合学习因果DAG与线性高斯机制参数的后验分布？
• RQ2 所提出的VBG方法在训练和采样过程中是否能保证无环性，而不同于以往的方法？
• RQ3 VBG在对结构和机制的不确定性建模方面，与现有基线方法相比表现如何？
• RQ4 通过机制分布的灵活参数化，该方法能否泛化至非线性因果机制？
• RQ5 将变分贝叶斯与GFlowNets结合是否能提升后验近似的质量，并支持因果发现中的主动学习？

主要发现

• VBG成功建模了DAG与机制参数的联合后验分布，在模拟数据上后验近似和结构恢复方面表现具有竞争力。
• 由于GFlowNet的结构约束，该方法在所有采样图中均保证无环性，克服了DiBS和VCN等方法的关键局限。
• VBG支持机制参数 $\theta$ 的灵活参数化，可泛化至非线性机制，而BCD-Nets则受限于特定参数化形式。
• 该方法在训练后支持无界后验采样，而DiBS在训练后不具备此能力。
• 实证评估表明，VBG在结构汉明距离和真实后验与近似后验之间的KL散度等指标上优于基线方法。
• 该方法具有可扩展性，适用于主动学习，因其能对结构和机制提供校准良好的不确定性估计，从而指导信息丰富的干预。