CauScale: 대규모 신경망 기반 인과관계 발견

CauScale: Neural Causal Discovery at Scale

인과관계 발견은 과학적 인공지능 및 데이터 분석과 같은 데이터 기반 분야의 발전에 필수적이지만, 기존 방법은 대규모 그래프로 확장할 때 상당한 시간 및 공간 효율성 병목 현상을 겪습니다. 이러한 문제를 해결하기 위해, 우리는 효율적인 인과관계 발견을 위한 신경망 아키텍처인 CauScale을 제안합니다. CauScale은 최대 1000개의 노드를 가진 그래프에 대한 추론을 확장할 수 있도록 설계되었습니다. CauScale은 데이터 임베딩을 압축하는 감소 모듈을 통해 시간 효율성을 향상시키고, 축별 어텐션 맵을 유지하지 않고 공유된 어텐션 가중치를 사용하여 공간 효율성을 개선합니다. 높은 인과관계 발견 정확도를 유지하기 위해, CauScale은 두 개의 스트림을 채택합니다. 데이터 스트림은 고차원 데이터에서 관계 증거를 추출하고, 그래프 스트림은 통계적 그래프 사전 지식을 통합하고 중요한 구조적 신호를 보존합니다. CauScale은 기존 연구에서 공간 제약으로 인해 실패하는 500개의 노드를 가진 그래프에서도 성공적으로 확장됩니다. 다양한 그래프 크기와 인과 메커니즘을 가진 테스트 데이터에서, CauScale은 동일 분포 데이터에서 99.6%의 mAP를 달성하고, 이 분포 외부 데이터에서 84.4%의 mAP를 달성하며, 기존 방법보다 4~13,000배 더 빠른 추론 속도를 제공합니다. 프로젝트 페이지는 https://github.com/OpenCausaLab/CauScale 입니다.

Causal discovery is essential for advancing data-driven fields such as scientific AI and data analysis, yet existing approaches face significant time- and space-efficiency bottlenecks when scaling to large graphs. To address this challenge, we present CauScale, a neural architecture designed for efficient causal discovery that scales inference to graphs with up to 1000 nodes. CauScale improves time efficiency via a reduction unit that compresses data embeddings and improves space efficiency by adopting tied attention weights to avoid maintaining axis-specific attention maps. To keep high causal discovery accuracy, CauScale adopts a two-stream design: a data stream extracts relational evidence from high-dimensional observations, while a graph stream integrates statistical graph priors and preserves key structural signals. CauScale successfully scales to 500-node graphs during training, where prior work fails due to space limitations. Across testing data with varying graph scales and causal mechanisms, CauScale achieves 99.6% mAP on in-distribution data and 84.4% on out-of-distribution data, while delivering 4-13,000 times inference speedups over prior methods. Our project page is at https://github.com/OpenCausaLab/CauScale.