[Paper Review] Causal Discovery with Reinforcement Learning

Shengyu Zhu, Ignavier Ng+1 more|arXiv (Cornell University)|Jun 11, 2019

Bayesian Modeling and Causal Inference53 references87 citations

TL;DR

The paper uses a neural encoder-decoder (Transformer-like) to generate DAGs and employs reinforcement learning to search for graphs that maximize a score-based reward under acyclicity penalties, improving DAG discovery with flexible score functions and finite samples.

ABSTRACT

Discovering causal structure among a set of variables is a fundamental problem in many empirical sciences. Traditional score-based casual discovery methods rely on various local heuristics to search for a Directed Acyclic Graph (DAG) according to a predefined score function. While these methods, e.g., greedy equivalence search, may have attractive results with infinite samples and certain model assumptions, they are usually less satisfactory in practice due to finite data and possible violation of assumptions. Motivated by recent advances in neural combinatorial optimization, we propose to use Reinforcement Learning (RL) to search for the DAG with the best scoring. Our encoder-decoder model takes observable data as input and generates graph adjacency matrices that are used to compute rewards. The reward incorporates both the predefined score function and two penalty terms for enforcing acyclicity. In contrast with typical RL applications where the goal is to learn a policy, we use RL as a search strategy and our final output would be the graph, among all graphs generated during training, that achieves the best reward. We conduct experiments on both synthetic and real datasets, and show that the proposed approach not only has an improved search ability but also allows a flexible score function under the acyclicity constraint.

Motivation & Objective

• Motivate causal structure learning from observational data when controlled experiments are infeasible.
• Propose a neural combinatorial optimization approach to search for DAGs that maximize a predefined score.
• Integrate acyclicity constraints via penalty terms and show equivalence to hard acyclicity under appropriate penalties.
• Demonstrate improved search effectiveness and flexibility over traditional score-based methods across synthetic and real datasets.

Proposed method

• Encoder-decoder network generates a binary adjacency matrix for a DAG from randomized subsets of observed data.
• Encoder uses a Transformer-style self-attention encoder to capture interactions among variables.
• Decoder computes pairwise scores g_ij between encoder outputs and samples them through a Bernoulli mechanism to form adjacency entries (i != j).
• Score-based reward combines a predefined DAG score (e.g., BIC) with acyclicity penalties h(A) and an indicator for non-DAG graphs.
• Acyclicity is enforced using h(A) = trace(exp(A)) - d and an additional indicator penalty to guarantee DAG outputs.
• Policy gradient with an actor-critic setup (REINFORCE with a critic) is used to train the network to maximize the expected reward.
• The final graph is taken as the best-scoring DAG observed during training and may be pruned to reduce false positives.

Experimental results

Research questions

• RQ1 Can reinforcement learning be effectively used as a search strategy to optimize a score-based objective for DAG discovery under acyclicity constraints?
• RQ2 Does an RL-based approach provide improved search over traditional methods (GES, PC) while offering flexibility to use various score functions (e.g., BIC) and regression models?
• RQ3 How does the proposed method perform on linear-Gaussian, LiNGAM, and nonlinear (quadratic) causal models under finite samples?
• RQ4 What is the effect of acyclicity penalties on learning efficiency, DAG quality, and edge sparsity in large graphs?

Key findings

• RL with BIC-based scoring (RL-BIC2) recovers true causal graphs in LiNGAM and linear-Gaussian settings in 12-node experiments.
• RL-BIC2 generally outperforms traditional methods (GES, PC) and some neural baselines on both linear and nonlinear data models in terms of FDR, TPR, and SHD.
• On larger graphs (d=30) with LiNGAM data, RL-BIC2 achieves competitive FDR and high TPR, with SHD comparable to or better than alternatives.
• For nonlinear quadratic models, RL-BIC2 achieves very high TPR (≈0.98) and extremely low SHD, outperforming several baselines (NOTEARS variants, ICA-LiNGAM, CAM, DAG-GNN, GraN-DAG).