[論文レビュー] 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を提案する。エンコーダ・デコーダアーキテクチャを用いて順序を生成し、特化した報酬を用いた強化学習により、合成データおよび実データの両方で最先端の性能を達成し、既存の強化学習ベース手法を上回り、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.

研究の動機と目的

• 有向無閉路グラフ（DAG）の探索空間を探索する既存の強化学習（RL）ベースの因果探索手法の非効率性とスケーラビリティの制限を解消すること。
• DAGの探索空間が非常に大きいため計算コストが高く、実行不可能となる問題を、より小さい変数順序の空間への探索に移行することで克服すること。
• 現代の強化学習の強力な探索能力を活用しながら、順序に基づく構造学習によりサイクルのない構造を保証すること。
• 事前学習モデルを強化学習の学習パイプラインに統合することで、サンプル効率と収束速度を向上させること。
• 線形および非線形の因果モデル、特に生物学的実データを含む、優れた性能を示すこと。

提案手法

• 各ステップが順序における次の変数の選択に対応する多段階のマルコフ決定過程（MDP）として、変数順序探索を定式化する。
• デコーダが段階的に順序を生成するエンコーダ・デコーダニューラルネットワークアーキテクチャを用いて、順序生成プロセスを実装する。
• 変数選択とデータ適合性に基づくスコア関数から導出される報酬機構を設計し、各生成順序の質を評価する。
• 安定的かつ効率的な方策更新を可能にするために、プロキシマルポリシー最適化（PPO）を用いて強化学習エージェントを訓練する。
• 最終的な順序を因果的DAGに変換するために、変数選択（例：CAMのプルーニング）を適用し、サイクルのない構造的妥当性を保証する。
• 訓練を加速するために事前学習モデルを統合し、性能に影響を与えることなく収束時間を短縮する。

実験結果

リサーチクエスチョン

• RQ1 DAG空間から変数順序空間に探索を移行することで、強化学習ベースの因果探索の効率性とスケーラビリティを向上させられるか？
• RQ2 大規模な線形データにおいて、順序を学習する強化学習ベース手法は、NOTEARSのような勾配ベース手法と比較してどのように性能を発揮するか？
• RQ3 事前学習モデルを用いることで、強化学習ベースの因果探索における訓練を顕著に加速できるか、最終的な精度に影響を与えないか？
• RQ4 ガウス過程構造方程式を用いた非線形で同定可能な因果モデル（例：GP-SEM）において、提案手法はどのように性能を発揮するか？
• RQ5 生物学的実データ（例：Sachsデータセット）において、本手法の経験的性能はいかがなっているか？

主な発見

• CORL-2は150ノードの線形グラフにおいて、最先端の勾配ベース手法であるNOTEARSを上回り、構造的ハミング距離（SHD）が低くなる。
• ガウス過程モデルを用いた30ノードの非線形データにおいて、CORL-2はすべての強化学習ベース手法の中で最良のSHDを達成し、ER1およびER4グラフではCAMと同等またはそれを上回る。
• 100ノードの線形グラフにおいて、CORL-1およびCORL-2は約7時間で収束するが、RL-BIC2は15時間の時間制限内に収束しなかった。
• 事前学習モデルの使用により訓練時間が顕著に短縮され、30ノードのグラフでは約700イテレーション、10ノードのグラフでは約5000イテレーションでCORL-2が収束した。
• 11ノードの実際のSachsデータセットにおいて、CORL-1、CORL-2、RL-BIC2はSHDが11という最良の値を達成し、CAM（SHD 12）およびNOTEARS（SHD 19）を上回った。
• 本手法はER1やER4などの異なるグラフタイプやデータ分布においても強力な性能を維持しており、構造的複雑さやサンプルサイズの変動に対しても頑健であることが示された。