[論文レビュー] BayesDAG: Gradient-Based Posterior Inference for Causal Discovery

Yashas Annadani, Nick Pawlowski+4 more|arXiv (Cornell University)|Jul 26, 2023

Bayesian Modeling and Causal Inference被引用数 4

ひとこと要約

BayesDAG は、確率的勾配マルコフ連鎖モンテカルロ (SG-MCMC) と変分推論を活用することで、非線形構造的関数と有向無閉路グラフ (DAG) の両方におけるスケーラブルな後方推論を可能にする、画期的な勾配ベースのベイジアン因果発見フレームワークを提案する。合成および実世界のデータセットにおいて、DAG正則化を必要とせず、非線形関係を伴う100変数までスケーリング可能な状況で、因果発見および不確実性の定量化において最先端の性能を達成する。

ABSTRACT

Bayesian causal discovery aims to infer the posterior distribution over causal models from observed data, quantifying epistemic uncertainty and benefiting downstream tasks. However, computational challenges arise due to joint inference over combinatorial space of Directed Acyclic Graphs (DAGs) and nonlinear functions. Despite recent progress towards efficient posterior inference over DAGs, existing methods are either limited to variational inference on node permutation matrices for linear causal models, leading to compromised inference accuracy, or continuous relaxation of adjacency matrices constrained by a DAG regularizer, which cannot ensure resulting graphs are DAGs. In this work, we introduce a scalable Bayesian causal discovery framework based on a combination of stochastic gradient Markov Chain Monte Carlo (SG-MCMC) and Variational Inference (VI) that overcomes these limitations. Our approach directly samples DAGs from the posterior without requiring any DAG regularization, simultaneously draws function parameter samples and is applicable to both linear and nonlinear causal models. To enable our approach, we derive a novel equivalence to the permutation-based DAG learning, which opens up possibilities of using any relaxed gradient estimator defined over permutations. To our knowledge, this is the first framework applying gradient-based MCMC sampling for causal discovery. Empirical evaluation on synthetic and real-world datasets demonstrate our approach's effectiveness compared to state-of-the-art baselines.

研究の動機と目的

• 組み合わせ的 DAG 空間と連続的関数パラメータにおけるスケーラブルなベイジアン因果発見の課題に対処すること。
• DAG 正則化に依存する既存手法の限界（非巡回性を保証できないこと）や、変分推論の精度の問題を克服すること。
• 線形および非線形構造的因果モデルに適用可能なモデルに依存しないフレームワークの開発。
• 緩和や正則化を用いずに勾配ベースの MCMC を直接用いて DAG をサンプリングすることで、高品質な後方推論を実現すること。
• 合成および実世界のデータセットにおいて、100変数にまで及ぶ高次元の非線形因果モデルの有効性とスケーラビリティを実証すること。

提案手法

• DAG 上での同等のベイジアン推論を可能にするエッジ信念とノードポテンシャル $({\bm{W}}, {\bm{p}})$ の拡張空間を導入し、DAG 正則化を不要にする。
• この拡張空間と置換ベースの DAG 学習との間の同等性を導出し、置換に関する任意の緩和された勾配推定器の使用を可能にする。
• SG-MCMC と変分推論 (VI) を組み合わせたハイブリッドフレームワークを提案し、DAG と関数パラメータの両方の後方分布を同時にサンプリングする。
• 2番目の定式化において、隣接行列 $\tilde{\bm{W}}$ の連続的緩和を採用し、明示的な DAG 制約なしに勾配ベースの MCMC サンプリングを可能にする。
• ノードポテンシャル $\bm{p}$ とエッジ信念 $\bm{W}$ を結合する変分ネットワーク $\mu_{\phi}$ を採用し、サンプリング中に応答的な更新を可能にする。
• SG-MCMC における確率的勾配推定器を活用し、$d=100$ 変数を含む高次元設定へのスケーリングを実現する。

実験結果

リサーチクエスチョン

• RQ1 離散的 DAG 構造上でのベイジアン因果発見に勾配ベースの MCMC サンプリングを効果的に適用できるか？
• RQ2 エッジ信念とノードポテンシャルの拡張空間を用いた、緩和なしの DAG サンプリングが達成可能か？
• RQ3 提案フレームワークは、既存の変分推論および正則化ベース手法と比較して、後方分布の精度と不確実性定量化において優れているか？
• RQ4 100変数の高次元非線形因果モデルにスケーリング可能か？
• RQ5 変分ネットワークを介したノードポテンシャルとエッジ信念の結合が、サンプリングの効率性と精度を向上させるか？

主な発見

• BayesDAG は $d=20$ の非線形合成データセットにおいて、期待構造ハミング距離 (E-SHD)、エッジ F1、および負の対数尤度 (NLL) のすべての指標で、すべてのベースラインを上回る最先端の性能を達成する。
• $d=70$ および $d=100$ の非線形設定においても、E-SHD と NLL で優れた性能を維持しており、最良のベースラインと比較してエッジ F1 はわずかに低下するにとどまる。
• 変分ネットワークにおける $\bm{p}$ と $\bm{W}$ の良好な結合のおかげで、SG-MCMC+VI の定式化が、連続的緩和 ($\tilde{\bm{W}}$) を用いた SG-MCMC よりも優れた性能を示す。
• BayesDAG は $d=100$ 変数にまで効率的にスケーリング可能であり、1枚の40GB A100 GPU で448分の実行時間である。一方、DIBS や DDS はメモリ制約により失敗する。
• BGES や DDS は高速ではあるが、不確実性定量化において顕著に劣っており、線形モデルに限定されるか、完全にベイジアンではない。
• 実証的結果から、BayesDAG は特にデータが少ない状況でも、適切な不確実性キャリブレーションを伴う高精度な後方推論を提供することが確認された。