[Paper Review] Causal Discovery from Incomplete Data using An Encoder and Reinforcement Learning

Xiaoshui Huang, Fujin Zhu+2 more|arXiv (Cornell University)|Jun 9, 2020

Bayesian Modeling and Causal Inference24 references4 citations

TL;DR

This paper proposes a reinforcement learning (RL)-based causal discovery method that integrates a deep encoder for robust feature representation from incomplete data, enabling end-to-end learning of causal structures despite missing values. The approach outperforms direct imputation-combination methods, achieving up to 43.2% improvement in TPR on nonlinear models with missing data.

ABSTRACT

Discovering causal structure among a set of variables is a fundamental problem in many domains. However, state-of-the-art methods seldom consider the possibility that the observational data has missing values (incomplete data), which is ubiquitous in many real-world situations. The missing value will significantly impair the performance and even make the causal discovery algorithms fail. In this paper, we propose an approach to discover causal structures from incomplete data by using a novel encoder and reinforcement learning (RL). The encoder is designed for missing data imputation as well as feature extraction. In particular, it learns to encode the currently available information (with missing values) into a robust feature representation which is then used to determine where to search the best graph. The encoder is integrated into a RL framework that can be optimized using the actor-critic algorithm. Our method takes the incomplete observational data as input and generates a causal structure graph. Experimental results on synthetic and real data demonstrate that our method can robustly generate causal structures from incomplete data. Compared with the direct combination of data imputation and causal discovery methods, our method performs generally better and can even obtain a performance gain as much as 43.2%.

Motivation & Objective

• To address the challenge of causal discovery in observational data with missing values, which commonly degrades performance of existing methods.
• To develop a unified framework that jointly handles data imputation and causal structure learning, rather than treating them as separate stages.
• To imitate human-like reasoning on incomplete data by learning robust feature representations from non-missing values.
• To enable end-to-end optimization of causal graph discovery using reinforcement learning on incomplete data.
• To improve the robustness and accuracy of causal discovery in nonlinear and real-world datasets with missingness.

Proposed method

• A custom neural network encoder is designed to extract robust feature representations from incomplete observational data by focusing on available non-missing values.
• The encoder's output is used as input to a reinforcement learning agent that searches for the optimal DAG in a continuous space.
• The RL agent uses an actor-critic algorithm to optimize the policy based on rewards derived from graph score functions.
• The entire system is trained end-to-end, jointly optimizing the encoder and RL policy for causal structure discovery.
• The method integrates data imputation and feature extraction into a single differentiable framework, avoiding standalone imputation steps.
• The approach is evaluated using score functions like BIC and likelihood-based criteria within the RL framework.

Experimental results

Research questions

• RQ1 Can a deep encoder effectively extract meaningful representations from incomplete observational data for causal discovery?
• RQ2 Does integrating data imputation and causal structure learning within a unified RL framework improve performance compared to sequential imputation and discovery?
• RQ3 How robust is the proposed method in recovering true causal structures when data contain up to 20% missing values?
• RQ4 Can the method outperform state-of-the-art causal discovery algorithms on nonlinear models with missing data?
• RQ5 Does the end-to-end RL framework with a learned encoder generalize well to real-world biological datasets with missing values?

Key findings

• On synthetic nonlinear data with 20% missing values, the method achieved a 31% higher TPR and 6 lower SHD compared to GAIN+NOTEARS.
• On nonlinear Gaussian process models with missing data, the method achieved a 43.2% performance gain in TPR and reduced SHD by 19 points compared to GAIN+RL+BIC2.
• The method outperformed GAIN+NOTEARS and GAIN+RL+BIC2 on real-world T-cell signaling data, achieving a 13 SHD and 41.1% TPR, compared to 18 SHD and 11.8% TPR for the next best method.
• The low FDR and high TPR in real data suggest the method effectively reduces false positives while recovering true edges.
• The results demonstrate that end-to-end integration of encoding and RL leads to significantly better causal discovery than standalone imputation followed by structure learning.
• The method maintains robust performance across both synthetic and real-world datasets with missing values, indicating strong generalization.