Goal-Oriented Sequential Bayesian Experimental Design for Causal Learning

TL;DR

GO-CBED is a goal-oriented Bayesian experimental design framework that efficiently targets specific causal questions through sequential interventions, outperforming traditional methods especially with limited resources.

Contribution

We introduce a novel goal-oriented Bayesian framework with a variational estimator and transformer policy for targeted causal experimental design.

Findings

Outperforms existing baselines in causal reasoning tasks

Effective with limited experimental budgets

Handles complex causal mechanisms

Abstract

We present GO-CBED, a goal-oriented Bayesian framework for sequential causal experimental design. Unlike conventional approaches that select interventions aimed at inferring the full causal model, GO-CBED directly maximizes the expected information gain (EIG) on user-specified causal quantities of interest, enabling more targeted and efficient experimentation. The framework is both non-myopic, optimizing over entire intervention sequences, and goal-oriented, targeting only model aspects relevant to the causal query. To address the intractability of exact EIG computation, we introduce a variational lower bound estimator, optimized jointly through a transformer-based policy network and normalizing flow-based variational posteriors. The resulting policy enables real-time decision-making via an amortized network. We demonstrate that GO-CBED consistently outperforms existing baselines across various causal reasoning and discovery tasks-including synthetic structural causal models and semi-synthetic gene regulatory networks-particularly in settings with limited experimental budgets and complex causal mechanisms. Our results highlight the benefits of aligning experimental design objectives with specific research goals and of forward-looking sequential planning.