CauScale: 大規模ニューラル因果探索

要旨

因果関係発見は、科学AIやデータ分析といったデータ駆動分野の発展に不可欠である。しかし、既存手法は大規模グラフへのスケーリングにおいて、時間効率と空間効率の重大なボトルネックに直面している。この課題に対処するため、我々は最大1000ノードまでのグラフで推論をスケール可能な、効率的な因果関係発見のために設計されたニューラルネットワークアーキテクチャ「CauScale」を提案する。CauScaleは、データ埋め込みを圧縮する縮小ユニットにより時間効率を向上させ、軸固有の注意マップの維持を回避するための結合注意重みを採用することで空間効率を改善する。高い因果発見精度を維持するため、CauScaleは二重ストリーム設計を採用している：データストリームは高次元観測データから関係的証拠を抽出し、グラフストリームは統計的グラフ事前分布を統合し、重要な構造信号を保持する。CauScaleは、空間制限により従来手法が失敗する訓練段階において、500ノードのグラフへのスケーリングに成功した。様々なグラフ規模と因果メカニズムを持つテストデータにおいて、CauScaleはインディストリビューションデータで99.6%、アウトオブディストリビューションデータで84.4%のmAPを達成し、従来手法に比べて4〜13,000倍の推論速度向上を実現した。プロジェクトページはhttps://github.com/OpenCausaLab/CauScale にある。

English

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.

CauScale: 大規模ニューラル因果探索

CauScale: Neural Causal Discovery at Scale

要旨

Support