[논문 리뷰] BayesDAG: Gradient-Based Posterior Inference for Causal Discovery

Yashas Annadani, Nick Pawlowski외 4명|arXiv (Cornell University)|2023. 07. 26.

Bayesian Modeling and Causal Inference인용 수 4

한 줄 요약

BayesDAG는 확률적 경사 하강 마르코프 체인 몬테 카를로(SG-MCMC)와 변분 추론을 활용하여 방향 비순환 그래프(DAGs)와 비선형 구조적 함수에 대한 확장 가능한 사후 추론을 가능하게 하는 혁신적인 경사 기반 베이지안 원인 발견 프레임워크를 제안한다. 이는 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와 함수 매개변수에 대한 공동 사후 샘플링을 수행한다.
• 두 번째 공식화에서 인접 행렬 $\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$ 비선형 설정에서 BayesDAG는 E-SHD와 NLL 측면에서 뛰어난 성능 유지를 보이며, 최고 기준선 대비 엣지 F1에서 약간의 감소만을 보인다.
• SG-MCMC+VI 공식화는 변분 네트워크 내에서 $\bm{p}$와 $\bm{W}$ 간의 더 나은 결합 덕분에 연속적 완화($\tilde{\bm{W}}$)를 사용하는 SG-MCMC보다 우수한 성능을 보인다.
• BayesDAG는 $d=100$ 개 변수까지 효율적으로 확장 가능하며, 단일 40GB A100 GPU에서 448분의 실행 시간을 기록한다. 반면 DIBS와 DDS는 메모리 제약으로 인해 실패한다.
• BGES와 DDS는 더 빠르지만 불확실성 정량화 측면에서 유의미하게 열등하며, 선형 모델에 국한되거나 완전히 베이지안이 아님을 확인한다.
• 실증 결과는 BayesDAG가 특히 데이터가 적은 환경에서 정확한 사후 추론과 적절한 불확실성 캘리브레이션을 제공함을 확인한다.