[论文解读] Ordering-Based Causal Discovery with Reinforcement Learning

Xiaoqiang Wang, Yali Du+5 more|arXiv (Cornell University)|May 14, 2021

Bayesian Modeling and Causal Inference参考文献 23被引用 7

一句话总结

本文提出CORL，一种基于强化学习的因果发现方法，将问题重新表述为通过多步马尔可夫决策过程学习变量顺序。采用编码器-解码器架构生成顺序，并结合定制奖励的强化学习，CORL在合成数据和真实数据上均达到最先进性能，优于现有基于强化学习的方法，在150节点线性图上与NOTEARS等基于梯度的方法表现相当或更优。

ABSTRACT

It is a long-standing question to discover causal relations among a set of variables in many empirical sciences. Recently, Reinforcement Learning (RL) has achieved promising results in causal discovery from observational data. However, searching the space of directed graphs and enforcing acyclicity by implicit penalties tend to be inefficient and restrict the existing RL-based method to small scale problems. In this work, we propose a novel RL-based approach for causal discovery, by incorporating RL into the ordering-based paradigm. Specifically, we formulate the ordering search problem as a multi-step Markov decision process, implement the ordering generating process with an encoder-decoder architecture, and finally use RL to optimize the proposed model based on the reward mechanisms designed for~each ordering. A generated ordering would then be processed using variable selection to obtain the final causal graph. We analyze the consistency and computational complexity of the proposed method, and empirically show that a pretrained model can be exploited to accelerate training. Experimental results on both synthetic and real data sets shows that the proposed method achieves a much improved performance over existing RL-based method.

研究动机与目标

• 解决现有基于强化学习（RL）的因果发现方法在有向无环图（DAG）空间中搜索时存在的效率低下和可扩展性受限问题。
• 通过将搜索空间从庞大的DAG空间转移到更小的变量顺序空间，克服高计算成本和大规模DAG空间探索的不可行性。
• 利用现代强化学习的强大搜索能力，同时通过基于顺序的结构学习确保无环性。
• 通过在强化学习训练流程中引入预训练模型，提升样本效率并加快收敛速度。
• 在线性和非线性因果模型上均实现优异性能，包括真实世界生物数据。

提出的方法

• 将变量顺序搜索建模为多步马尔可夫决策过程（MDP），其中每一步对应选择顺序中的下一个变量。
• 采用编码器-解码器神经网络架构实现顺序生成过程，解码器逐步生成顺序。
• 设计一种奖励机制，基于变量选择和数据拟合的评分函数评估每种生成顺序的质量。
• 使用近端策略优化（PPO）训练强化学习智能体，实现稳定且高效的策略更新。
• 应用变量选择（如CAM剪枝）将最终顺序转换为因果DAG，确保无环性和结构有效性。
• 集成预训练模型以加速训练，减少收敛时间而不损失性能。

实验结果

研究问题

• RQ1 通过从DAG空间转向变量顺序空间，能否使基于强化学习的因果发现方法更高效且更具可扩展性？
• RQ2 在大规模线性数据上，学习顺序的基于强化学习方法与NOTEARS等基于梯度的方法相比表现如何？
• RQ3 预训练模型是否能显著加速基于强化学习的因果发现训练过程，同时不损害最终精度？
• RQ4 在具有高斯过程结构方程的非线性、可识别因果模型中，所提方法表现如何？
• RQ5 在真实世界生物数据（如Sachs数据集）上，该方法的实证性能如何？

主要发现

• CORL-2在150节点线性图上优于最先进基于梯度的方法NOTEARS，结构汉明距离（SHD）更低。
• 在30节点非线性数据（高斯过程模型）上，CORL-2在所有基于强化学习的方法中SHD最优，并在ER1和ER4图上与CAM表现相当或更优。
• 在100节点线性图上，CORL-1和CORL-2约7小时完成收敛，显著快于RL-BIC2，后者在15小时时限内未能收敛。
• 使用预训练模型可显著缩短训练时间，CORL-2在30节点图上约700次迭代内收敛，在10节点图上约5000次迭代内收敛。
• 在真实Sachs数据集（11个节点）上，CORL-1、CORL-2和RL-BIC2均实现最优SHD（11），优于CAM（SHD 12）和NOTEARS（SHD 19）。
• 该方法在不同图类型（ER1、ER4）和数据分布下均保持优异性能，展现出对结构复杂性和样本量变化的鲁棒性。