Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes

TL;DR

This paper introduces VBG, a novel Bayesian method combining GFlowNets and Variational Bayes to jointly learn causal structures and mechanisms, providing uncertainty quantification and flexibility for non-linear models.

Contribution

It extends Bayesian causal structure learning with GFlowNets to jointly infer DAG structures and linear-Gaussian mechanisms, ensuring acyclicity and generalizing to non-linear causal models.

Findings

VBG is competitive in modeling posterior distributions over DAGs and mechanisms.

The method guarantees acyclic graph sampling.

It can generalize to non-linear causal mechanisms.

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.