CauScale: Neural Causal Discovery at Scale

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$\times$–13,000$\times$ inference speedups over prior methods.

Understanding cause and effect from data is one of the most important problems in science. When doctors want to know if a drug actually cures a disease — not just correlates with recovery — or when biologists want to understand which genes control others, they need causal discovery: the ability to automatically reconstruct cause-and-effect relationships from observations. Existing methods, however, are too slow or require too much computer memory to work at the scale of modern datasets with thousands of variables. We built CauScale, a neural network that learns to identify causal relationships efficiently. The core idea is to progressively compress the data as it flows through the network — like zooming out on a map once the important details have already been captured — rather than processing everything at full resolution throughout. This keeps computation manageable without discarding the signals that reveal causation. CauScale is up to 13,000 times faster than existing methods and can handle training on graphs with up to 5,00 variables, a scale where prior approaches run out of memory, making large graph causal discovery easier.