[论文解读] Causal Discovery with Reinforcement Learning

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

Bayesian Modeling and Causal Inference参考文献 53被引用 87

一句话总结

本文使用神经编码器-解码器（类似 Transformer）生成有向无环图（DAG），并利用强化学习在满足无环性罚则的前提下搜索最大化分数型奖励的图，从而通过更灵活的评分函数和有限样本改进DAG发现。

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.

研究动机与目标

• 在无法进行对照实验时，激励因果结构学习以来自观测数据的学习。
• 提出一种神经组合优化方法来搜索最大化预定义分数的DAG。
• 通过罚项将无环性约束整合进来，并在适当的罚则下展示与硬性无环性的等价性。
• 在合成数据和真实数据集上展示相比传统基于分数的方法更有效的搜索效果与灵活性。

提出的方法

• 编码器-解码器网络从观测数据的随机子集生成DAG的二进制邻接矩阵。
• 编码器使用 Transformer 风格的自注意力编码器来捕捉变量之间的交互。
• 解码器计算编码器输出之间的成对分数 g_ij，并通过伯努利机制对样本化形成邻接矩阵的条目（i != j）。
• 基于分数的奖励将预定义的 DAG 分数（如 BIC）与无环性罚项 h(A) 及非DAG图的指示项结合起来。
• 通过 h(A) = trace(exp(A)) - d 的无环性约束，以及额外的指示罚项来保证输出为DAG。
• 使用带 critic 的 actor-critic 架构的策略梯度（REINFORCE）来训练网络以最大化期望奖励。
• 最终图形在训练过程中观测到的最佳得分DAG被采纳，可能会被剪枝以降低假阳性。

实验结果

研究问题

• RQ1 强化学习能否作为一种有效的搜索策略，在满足无环性约束的前提下，优化用于DAG发现的基于分数的目标？
• RQ2 基于RL的方法是否在搜索方面优于传统方法（GES、PC），并在提供对各种分数函数（如 BIC）和回归模型的灵活性方面有优势？
• RQ3 所提出的方法在有限样本下对线性高斯、LiNGAM 和非线性（二次）因果模型的表现如何？
• RQ4 在大规模图（边数多）下，无环性罚项对学习效率、DAG质量和边稀疏性的影响如何？

主要发现

• 在 LiNGAM 和线性高斯设置的12节点实验中，使用基于 BIC 分数的 RL（RL-BIC2）能够恢复真实的因果图。
• RL-BIC2 在 FDR、TPR 和 SHD 方面通常优于传统方法（GES、PC）以及某些神经基线，在线性和非线性数据模型上均有优势。
• 在更大图（d=30）且数据为 LiNGAM 时，RL-BIC2 实现了有竞争力的 FDR 和较高的 TPR，SHD 与替代方法相比接近或更好。
• 对于非线性二次模型，RL-BIC2 实现了非常高的 TPR（约0.98）和极低的 SHD，优于若干基线（NOTEARS 变体、ICA-LiNGAM、CAM、DAG-GNN、GraN-DAG）。