[논문 리뷰] Ordering-Based Causal Discovery with Reinforcement Learning

Xiaoqiang Wang, Yali Du외 5명|arXiv (Cornell University)|2021. 05. 14.

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)으로 변수 순서 탐색을 재구성한다.
• 디코더가 단계별로 순서를 생성하는 인코더-디코더 신경망 아키텍처를 사용해 순서 생성 과정을 구현한다.
• 변수 선택 및 데이터 적합도 기반의 점수 함수에서 유도된 점수 기반 평가 메커니즘을 설계해 각 생성된 순서의 품질을 평가한다.
• 안정적이고 효율적인 정책 업데이트를 가능하게 하기 위해 Proximal Policy Optimization(PPO)을 사용해 강화학습 에이전트를 훈련시킨다.
• 최종 순서를 인과적 DAG로 변환하기 위해 변수 선택(예: CAM 프루닝)을 적용하여 사이클성과 구조적 타당성을 보장한다.
• 학습을 가속화하기 위해 미리 학습된 모델을 통합함으로써 수렴 시간을 단축시키되 성능에 영향을 주지 않는다.

실험 결과

연구 질문

• 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시간 제한 내에 수렴하지 못하는 데에 비해 훨씬 빠르게 수렴한다.
• 미리 학습된 모델을 사용함으로써 훈련 시간이 크게 단축되었으며, 30노드 그래프에서는 약 700회 반복, 10노드 그래프에서는 약 5000회 반복 내에 CORL-2가 수렴한다.
• 실제 Sachs 데이터셋(11노드)에서 CORL-1, CORL-2, RL-BIC2는 SHD 11을 기록하여 CAM(SHD 12)과 NOTEARS(SHD 19)를 능가한다.
• 이 방법은 다양한 그래프 유형(ER1, ER4)과 데이터 분포에서 뛰어난 성능을 유지하며, 구조적 복잡성과 표본 크기의 변화에 대해 강건함을 입증한다.