PaperlayerView exact revision on arXiv

Source-linked AI summary

CauScale: Neural Causal Discovery at Scale

Bo Peng, Sirui Chen, Jiaguo Tian, Yu Qiao, Chaochao Lu

arXiv:2602.08629v2cs.LGcs.AIstat.ML

TL;DR

Large-graph causal discovery is limited by the time and memory costs of existing methods. CauScale uses reduction, tied attention, and two-stream processing to address these bottlenecks, achieving strong accuracy, 500-node training, and large inference speedups. Its predictions remain sensitive to data quality, modeling assumptions, and distribution shift.

  • Problem

    Existing causal discovery methods face major time- and space-efficiency bottlenecks when scaling to large graphs.

  • Method

    CauScale is an amortized neural architecture combining a reduction unit, tied attention weights, and data and graph streams that preserve structural signals during compression.

  • Results

    CauScale scales training to 500-node graphs and inference to graphs with up to 1,000 nodes, achieving strong accuracy with 4×–13,000× speedups over existing approaches.

  • Takeaways & Limitations

    CauScale suggests a practical direction for pre-training efficient neural models for causal discovery at scale.

  • Takeaways & Limitations

    Predicted graphs can be sensitive to data quality, modeling assumptions, and distribution shift, so they should be treated as hypotheses requiring validation.

Abstract

from arXiv · show

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.

1. Introduction

Causal discovery matters across scientific and data-driven fields, but existing methods face severe time and space bottlenecks on large graphs. CauScale addresses these constraints with an efficient neural architecture and reports strong accuracy, scalability, and inference speed.

  • Causal structure learning identifies directed acyclic graphs whose edges represent direct causal relationships between variables.
  • Existing causal discovery methods become difficult to scale because conditional-independence testing can be time-prohibitive, while dataset-specific optimization remains computationally expensive.
  • CauScale combines a data stream, graph stream, reduction unit, tied attention weights, and data-graph blocks to improve efficiency while preserving relational signals.
  • 99.6% mAP is achieved on in-distribution data and 84.4% on out-of-distribution data.
  • 4× to 13,000× inference speedups are reported, while training scales to 500-node graphs where AVICI fails because of memory and space costs.
  • The paper presents CauScale as an early study of pre-training neural networks for efficient causal discovery at scale.

2. Related Work

Causal discovery methods differ according to whether they solve a new optimization problem for each dataset or learn a shared inference model. Related approaches include constraint-based, score-based, asymmetry-based, temporal, and amortized methods, each with scalability trade-offs.

  • Non-amortized methods solve an optimization or search problem independently for each dataset, resulting in high computational cost and limited scalability.
  • Constraint-based methods use conditional-independence tests and suffer exponential complexity as the number of variables grows.
  • Asymmetry-based methods exploit data-generating-process asymmetries, while Granger methods infer directional relationships from temporal precedence.
  • Amortized approaches learn a shared model that maps datasets directly to causal graphs, enabling zero-shot inference on unseen data.
  • Earlier amortized methods rely on high-dimensional embeddings that scale poorly with graph size, while SEA decomposes large graphs into subproblems.

3. Preliminary

Causal graphical models pair a joint distribution with a directed acyclic graph, whose edges encode direct causal influences. The preliminary formulation assumes causal sufficiency and distinguishes observational from perfect-intervention settings.

  • A causal graphical model consists of a joint distribution over variables and a directed acyclic graph containing those variables as nodes.
  • Each directed edge (i, j) encodes a direct causal influence from x_i to x_j.
  • The distribution factorizes over each variable conditioned on its parent set, written as p(x1, . . . , xn) = Qn j=1 p(xj | PAj).
  • Causal sufficiency is assumed, meaning no unobserved common causes jointly affect multiple variables in X.
  • Interventions replace a target variable’s conditional distribution; under perfect intervention, that variable is randomized independently of its parents.

4. CauScale

CauScale combines data and graph streams to infer probabilistic causal graphs while reducing computation and attention-memory costs. Its architecture compresses observation embeddings, transfers relational evidence into the graph stream, and uses tied attention weights to improve scalability.

  • Overall Architecture: CauScale encodes observational data and a statistical graph prior into data- and graph-stream embeddings, then predicts a probabilistic adjacency matrix.The graph prior is defined as an inverse covariance matrix, and the output represents likelihoods of directed causal relations.
  • Data-Graph Block: Each data-graph block updates the data stream, summarizes it into a node-relation message, and injects that message into the graph stream.The data2graph layer extracts pairwise relational evidence and forms a relation matrix representing directed pairwise relationships between variables.
  • Reduction Unit: The reduction unit average-pools observation embeddings every k blocks, reducing the sample dimension by factor r while retaining variable-wise representations.This avoids naive subsampling and targets the typically much larger observation dimension; incomplete final chunks may be discarded for convenience.
  • Prediction Head: The prediction head applies a decomposed pairwise classifier and softmax probabilities over three states without explicitly enforcing acyclicity.The authors report comparable accuracy to the AVICI prediction head and empirically fewer cycles in decoded graphs.
  • Time Efficiency: 36.60% of baseline sample-axis compute and 48.13% of baseline node-axis compute remain when B=10, k=2, and r=2.The reduction changes the effective sample length geometrically across blocks, lowering the dominant axial-attention costs O(nm^2) and O(mn^2).
  • Space Efficiency: Tied attention weights reduce row-axis attention memory from O(RHC^2) to O(HC^2) and column-axis memory from O(CHR^2) to O(HR^2).Attention weights are shared across the target axis, avoiding separate axis-specific attention maps while preserving the output shape.

5. Experiment

CauScale is evaluated across causal discovery datasets, graph sizes, mechanisms, and robustness settings, with experiments emphasizing accuracy, efficiency, ablations, and generalization. It maintains strong accuracy while scaling to large graphs, substantially reducing inference time and memory demands, though performance is more sensitive to out-of-distribution mechanisms and severe latent confounding.

  • Experimental setup: CauScale is evaluated against constraint-based, score-based, FCM-based, pre-training-based, and statistical baselines using accuracy and efficiency metrics.The evaluation reports SHD, mAP, AUC, OA, cyclicity, AID metrics, inference time, and peak GPU memory.
  • Experimental setup: Training data span synthetic Erdős-Rényi and scale-free graphs with 10–500 nodes, multiple edge densities, linear and neural mechanisms, and synthetic single-cell GRNs.Synthetic graphs use additive or nonadditive Gaussian noise and 1,000 observations; GRN data are generated with SERGIO.
  • Accuracy: 99.6% mAP is achieved on linear synthetic graphs with n = 100, while 96.6% mAP is maintained on n = 1000 graphs unseen during training.On polynomial mechanisms, CauScale reaches 50.3% mAP versus 36.2% for SEA and 41.9% for SDCD; it also performs best across SERGIO-GRN metrics and settings.
  • Efficiency: 0.8288s inference time on n = 1000 graphs yields speedups of over 13,000× versus NOTEARS, 200× versus SEA-gies, and 4× versus AVICI.AVICI encounters an out-of-memory error on SERGIO-GRN at n = 100, whereas CauScale scales to n = 200 with 20,000 samples.
  • Cyclicity: CauScale produces zero cycles on in-distribution graph sizes and SERGIO-GRN, while n = 1000 graphs require cycle breaking and polynomial mechanisms remove 103 edges per graph on average.For other settings, cycle breaking removes at most 0.6 edges per graph on average and improves reported results.
  • Ablations and robustness: Ablations show that tied attention is six times faster than vanilla attention, average pooling outperforms alternative pooling strategies, and removing the graph stream causes the largest degradation.The model is also more sensitive to out-of-distribution noise and mechanisms than to graph structures, while remaining robust to mild latent confounding and becoming more sensitive to prior quality as graph size increases.

6. Conclusion

CauScale addresses time and memory bottlenecks in large-scale causal discovery, scaling training to 500-node graphs and inference to graphs with up to 1,000 nodes. It achieves strong accuracy and substantial speedups on synthetic and semi-synthetic benchmarks, while real-world observational evaluation remains future work.

  • CauScale combines a reduction unit, tied attention weights, and a two-stream design to improve efficiency while preserving structural signals.The reduction unit targets time efficiency, tied attention targets memory efficiency, and the two-stream design preserves structural information under compression.
  • 500-node training and 1,000-node inference demonstrate CauScale’s scalability on the evaluated benchmarks.
  • 4×–13,000× speedups and strong accuracy were achieved over existing approaches.
  • Evaluation on real-world observational datasets remains an important direction for future work.The method was evaluated on synthetic and semi-synthetic benchmarks, and its reliance on simulated training data may limit generalization when data-generating processes deviate from the assumed SCM framework.

Impact Statement

CauScale aims to make causal structure learning more computationally accessible by reducing runtime and memory requirements. However, predicted graphs remain sensitive to data quality, modeling assumptions, and distribution shift, so they require validation before high-stakes use.

  • CauScale reduces runtime and memory requirements for causal structure learning, supporting faster hypothesis generation in data-intensive scientific domains.
  • Predicted graphs should be treated as hypotheses and validated by domain experts and, where applicable, downstream experiments.The passage cautions that noisy or misspecified input data can produce spurious edges.

A. Evaluation metrics.

The evaluation uses causal structure-learning metrics on synthetic and SERGIO-GRN graph data. The graph models include random, hub-dominated, and community-structured topologies.

  • Evaluation metrics: SHD measures the minimum edge insertions, deletions, and reversals needed to transform a predicted graph into the ground-truth graph.
  • Evaluation metrics: mAP is the area under the precision-recall curve over candidate edges, averaged across the graph.
  • Graph models: Synthetic and SERGIO-GRN data are generated from directed acyclic graph structures.
  • Graph models: Erdős-Rényi graphs add edges between node pairs with fixed probability p, producing approximately Poissonian degree distributions.
  • Graph models: Scale-Free graphs use preferential attachment, creating hubs and power-law degree distributions.
  • Graph models: Stochastic Block Models assign nodes to latent blocks with edge probabilities determined by block membership.

B.2. Synthetic Data Generation

Synthetic data generation varies graph size, topology, causal mechanism, noise, and interventions, while SERGIO-GRN supplies simulated gene-expression data. These configurations support evaluation across linear, nonlinear, additive, and nonadditive mechanisms.

  • Polynomial, linear, NN-Add, and NN mechanisms define target variables from parent variables and independent noise in different ways.NN-Add adds noise after the neural network, whereas NN concatenates noise with parent inputs.
  • Interventions randomize one node at a time, and larger graphs use fewer generated structures because generation becomes more computationally expensive.
  • SERGIO-GRN generates gene-expression data by sampling steady states of a stochastic dynamical system with Hill-function regulatory interactions.

C. Baseline Implementation Details

The baselines use official implementations or checkpoints with specified thresholds, transformations, and training configurations for comparative evaluation.

  • FCI uses the causal-learn implementation with Fisher-Z tests at α = 0.05.
  • NOTEARS applies a 0.3 threshold to estimated weight matrices before SHD computation.
  • DiffAN follows official hyperparameters, uses the nonapproximated residue setting, and transforms edge-existence p-values with −log10 for AUC and mAP.
  • SDCD uses its official implementation and defaults, runs on GPU, and applies a 0.5 discretization threshold for SHD.
  • AVICI uses official pretrained checkpoints because training on the datasets caused Out-of-Memory errors.
  • SEA uses a released checkpoint for synthetic Stage 1 and trains corresponding models for synthetic Stage 2 and SERGIO-GRN under default configurations.

D. CauScale Implementation Details

CauScale combines configurable neural components, staged training, FLOP accounting, and activation-memory analysis to support efficient causal discovery at large graph sizes.

  • The model uses 10 layers, 128-dimensional embeddings, and 16 attention heads selected through hyperparameter tuning.
  • Inputs are standardized variable-wise using each variable’s empirical mean and standard deviation.
  • Training Strategy: Training uses Adam at 1 × 10−4 on 8 GPUs with distributed data parallelism and two synthetic stages spanning 10–100 and 150–500 nodes.Stage 1 uses batch size 8 for 37 hours, while Stage 2 uses batch size 1 for 2.75 hours.
  • Architecture: The architecture encodes data and graph-prior inputs separately, alternates data-graph blocks with reduction units, and produces graph predictions from the graph stream.Every k blocks, the reduction unit pools the data-stream embedding along the sample dimension by factor r.
  • Prediction Head: The output head applies two linear layers to n(n −1)/2 concatenated edge-pair embeddings of dimension 2d.Its approximate cost is 4n2d2 + 6n2d.
  • Attention: Axial attention performs row-wise and column-wise self-attention followed by an FFN, while tied weights reduce attention-memory overhead.
  • Memory Analysis: During inference, persistent activation memory is (m + n)nd, while transient memory is governed by max(mnF, Hm2).The analysis states that only the current block’s live tensors occupy memory and that the peak occurs at b = 0.
  • AID Evaluation: CauScale achieves the best AID in most settings, except n = 1000 with sigmoid and polynomial mechanisms outside its training distribution.SDCD has the lowest scores in those two cases but produces cyclic predictions for 100% of graphs at n = 1000.

F.3. Effectiveness on observational datasets

The observational-data evaluation reports model performance under pure observational inputs, with results presented in the corresponding tables.

  • CauScale is evaluated on synthetic data using pure observational inputs, indicating robustness without intervention data.
  • Table 5 reports AID results, where lower values indicate better performance and cyclic predictions undergo confidence-guided greedy cycle breaking.
  • Table 6 provides full CauScale results with standard deviations, and each setting uses five independent graphs.
  • Table 7 reports model performance under pure observational data.
Loading 2602.08629v2…