[論文レビュー] Amortized Inference for Causal Structure Learning

Lars Lorch, Scott Sussex+3 more|arXiv (Cornell University)|May 25, 2022

Bayesian Modeling and Causal Inference被引用数 6

ひとこと要約

本稿では、観測データまたは介入データから直接因果グラフを予測する神経ネットワークベースの手法AVICIを提案する。これは、構造探索の計算コストを回避するものである。置換に同変なアーキテクチャを活用し、シミュレータで生成されたデータで訓練することで、より大きな未学習グラフに対しても頑健な一般化を達成し、特にゲノム分野で既存手法を上回る性能を発揮する。また、補正された不確実性推定値を提供する。

ABSTRACT

Inferring causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test. The resulting search is costly, and designing suitable scores or tests that capture prior knowledge is difficult. In this work, we propose to amortize causal structure learning. Rather than searching over structures, we train a variational inference model to directly predict the causal structure from observational or interventional data. This allows our inference model to acquire domain-specific inductive biases for causal discovery solely from data generated by a simulator, bypassing both the hand-engineering of suitable score functions and the search over graphs. The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training. On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain. Our code and models are publicly available at: https://github.com/larslorch/avici.

研究の動機と目的

• 因果構造学習におけるスコア関数や独立性検定の設計の計算コストの高さと、正確な設計の難しさに対処すること。
• 組合せ的探索空間が非効率的で構造的バイアスを有する伝統的な探索ベースの手法の限界を克服すること。
• シミュレータで生成されたデータからドメイン固有のインダクティブバイアスを学習することで、学習時に見られなかったより大きなグラフへの一般化を可能にすること。
• 因果予測の信頼性を高めるために、補正された不確実性推定値を提供すること。
• ノイズや分布シフトにさらされても失敗する既存手法が失敗するような挑戦的な分野、例えばゲノム分野において優れた性能を示すこと。

提案手法

• 本手法は、神経ネットワークを訓練して入力データから直接因果グラフを予測するアンモライズド変分推論フレームワーク（AVICI）を採用し、繰り返しの構造探索の必要性を排除する。
• 推論モデルは変数順序に関して置換に同変であるように設計されており、入力順序に依存しない性質を保ち、統計的効率性と一般化性能を向上させる。
• モデルは、データ生成プロセスのシミュレータで生成された合成データで訓練される。これにより、手作業で設計したスコア関数に依存せずにドメイン固有のインダクティブバイアスを学習できる。
• アーキテクチャは、変数間の依存関係をモデル化し、エッジ確率を予測するために注意メカニズムを備えたグラフニューラルネットワークを採用する。
• エッジ確率はエンドツーエンドで予測され、不確実性の定量化としきい値処理による点推定が可能になる。
• モデルのアーキテクチャと学習データに埋め込まれたインダクティブバイアスを活用することで、学習時に見られなかったより大きなグラフへの一般化が実現される。

実験結果

リサーチクエスチョン

• RQ1 合成データから直接因果構造を予測できる神経ネットワークを訓練できるか？ これは、グラフ上の組合せ的探索を回避するものである。
• RQ2 小さなグラフで訓練されたモデルが、はるかに大きな未学習グラフにどの程度一般化できるか？
• RQ3 手作業で設計したスコア関数に依存せずに、シミュレータで生成されたデータからドメイン固有のインダクティブバイアスを学習できるか？
• RQ4 モデルは、グラフ構造、メカニズム、ノイズ、問題サイズの分布シフトに対してどの程度頑健か？
• RQ5 特に医療やゲノム分野のような高リスク分野において、因果予測の補正された不確実性推定値を提供できるか？

主な発見

• AVICIは合成データおよびセミ合成的遺伝子発現データにおいて、既存手法を顕著に上回り、分布シフトや高ノイズ条件下でも優れた性能を示した。
• 遺伝子発現データセットでは、AVICIが唯一、妥当な因果構造を同定できた。これは、挑戦的な生物学的分野で最先端の性能を示している。
• 観測データのみを用いて、GrnドメインでAUROCが0.901（±0.04）かつAUPRCが0.656（±0.10）を達成した。これは、介入データがなくても強力な性能を示している。
• 期待不確実性誤差（ECE）は、非パラメトリックブートストラップ法を除き、すべてのベースラインより低く、良好に補正された不確実性推定値であることを示している。
• AVICIは、学習時に使用されたより小さなグラフよりも最大30変数のグラフに対しても一般化でき、強力なインダクティブバイアスの伝達を示した。
• SHDおよびF1指標において、AVICIはGESおよびGIESと同等の性能を示し、確率的指標ではDAG-GNN、LiNGAM、DCDIを上回った。