PaperlayerView exact revision on arXiv

Source-linked AI summary

Amortized Inference for Causal Structure Learning

Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, Bernhard Schölkopf

arXiv:2205.12934v4cs.LGstat.ML

TL;DR

Causal discovery is costly because graph spaces are huge and suitable inductive biases are difficult to encode in scores or tests. AVICI trains a variational inference model on simulator-generated data to predict causal structures directly. It reports strong performance and robust generalization on synthetic and semisynthetic data, including gene expression settings.

  • Problem

    Causal graph search is costly as the graph space grows superexponentially, and existing approaches struggle to specify suitable domain-specific inductive biases.

    Source: §1 · p. 1

  • Method

    AVICI trains a variational inference model on simulator-generated graph-and-data samples to predict causal structures, using a forward-KL objective and permutation-aware architecture.

    Source: §1 · p. 2, §3.2 · p. 3, §4.2 · p. 4

  • Results

    AVICI significantly outperforms existing algorithms on synthetic and semisynthetic gene expression data and shows robust behavior under shifts in graphs, mechanisms, noise, and problem sizes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that AVICI solves complex-domain structure-learning problems that are intractable for existing approaches, with potential extension to other difficult domains.

    Source: §7 · p. 10

    The acyclicity constraint is not guaranteed to be satisfied by the variational family, and AVICI provides no theoretical performance guarantees.

    Source: §7 · p. 10

Abstract

from arXiv · show

Inferring causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test. The resulting search is costly, and designing suitable scores or tests that capture prior knowledge is difficult. In this work, we propose to amortize causal structure learning. Rather than searching over structures, we train a variational inference model to directly predict the causal structure from observational or interventional data. This allows our inference model to acquire domain-specific inductive biases for causal discovery solely from data generated by a simulator, bypassing both the hand-engineering of suitable score functions and the search over graphs. The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training. On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain. Our code and models are publicly available at: https://github.com/larslorch/avici.

1 Introduction

Causal structure learning is difficult because the graph space grows superexponentially, making search costly and realistic inductive biases hard to specify. AVICI replaces graph search with simulator-trained inference, and reports strong performance and robustness across synthetic and semisynthetic settings.

1 Introduction

  • The space of potential causal graphs grows superexponentially with the number of variables, making structure inference a difficult inverse problem.Standard approaches search graphs by maximizing a score or testing conditional independences.Source: §1 · p. 1
  • Score-based methods rely on difficult-to-verify assumptions, while constraint-based methods lack domain-specific bias and may identify exponentially large equivalence classes.Searching directed graphs may also introduce unwanted bias and artifacts.Source: §1 · p. 1
  • Intractable graph search can impose hard constraints, such as node degree, that limit the suitability of search in real-world domains.Source: §1 · p. 2
  • AVICI trains an inference model to predict causal structures from datasets generated by a domain simulator, bypassing graph search and hand-engineered scores.Its permutation-aware architecture generalizes to substantially larger problem instances than those seen during training.Source: §1 · p. 2
  • AVICI significantly outperforms existing causal discovery algorithms on synthetic and semisynthetic gene expression data, often by a large margin.The models also show calibrated uncertainties and robust behavior under shifts in graphs, mechanisms, noise, and problem sizes.Source: §1 · p. 2

2 Background and Related Work

The paper defines causal structure through direct effects, while assuming causal sufficiency. It situates AVICI among graph-search and neural approaches, emphasizing its variational objective and generalization to larger problems.

2.1 Causal Structure

  • A causal structure is a directed graph whose edges represent direct causal effects; causal sufficiency assumes the variables include all common causal parents.The paper notes that its approach does not require one particular formalization and can also be applied with acyclicity constraints.Source: §2.1 · p. 2

2.2 Related Work

  • Earlier neural approaches mainly modeled nonlinear mechanisms or inferred one dataset’s structure; prior amortized work studied narrower settings, while AVICI generalizes to larger problems and realistic genomics data.Classical methods search graphs using likelihoods or conditional-independence tests, and some use assumed noise properties.Source: §2.2 · p. 2

3 AVICI: Amortized Variational Inference for Causal Discovery

AVICI amortizes causal discovery by learning to predict a variational posterior from data sampled from a domain distribution. Its forward-KL objective uses simulator samples rather than requiring a tractable data likelihood.

3.1 Variational Objective

  • AVICI trains an inference model to predict variational parameters by minimizing the expected forward KL from the posterior over causal structures to a variational distribution.The objective is made tractable using samples from the joint distribution of graphs and data generated by a simulator.Source: §3.1 · p. 3
  • The variational objective is a lower bound on the mutual information between causal structures and observations, and is tight when the expected posterior-to-approximation KL is zero.Source: §3.1 · p. 3

3.2 Likelihood-Free Inference using the Forward KL

  • The forward-KL objective supports inference in complex domains without an explicit likelihood and typically yields more reliable uncertainty estimates than reverse KL.It requires an explicit variational density and samples from the true data-generating distribution.Source: §3.2 · p. 3
  • Reverse-KL variational inference requires a tractable marginal likelihood, which can restrict methods to conjugate linear-Gaussian or categorical mechanisms with zero measurement error.Likelihoods may also be sensitive to measurement scale when noise scale is not learned jointly.Source: §3.2 · p. 3

4 Inference Model

The inference model maps datasets to variational parameters for causal graphs, using a factorized family and a permutation-aware architecture. Optional acyclicity constraints and dual optimization support domains where acyclic graphs are justified.

4 Inference Model

  • The inference-model section introduces the variational family, architecture, optional acyclicity constraint, and training procedure for model parameters.Source: §4 · p. 3

4.1 Variational Family

  • The inference model maps a dataset to a d-by-d matrix of variational parameters, and each sample may include indicators for variables intervened upon.Other settings, including unknown intervention targets or incomplete measurements, can be encoded analogously.Source: §4.1 · p. 4

4.2 Model Architecture

  • The model processes an n×d×k representation with shared attention layers alternating over variables and observations, preserving the task’s permutation symmetries.Max-pooling over observations and inner products between variable embeddings produce edge probabilities while maintaining invariance over samples and equivariance over variables.Source: §4.2 · p. 4, Figure 1 · p. 4

4.3 Acyclicity

  • When an acyclic prior is justified, the optimization can constrain the expected predicted graph to be acyclic, although the variational family itself cannot enforce acyclicity.The acyclicity function is zero exactly when predicted edge probabilities induce an acyclic graph; power iteration approximates the relevant spectral radius.Source: §4.3 · p. 5

4.4 Optimization

  • Training uses stochastic optimization with samples from the domain’s data-generating process and solves the acyclicity-constrained program through a dual formulation.Without an acyclicity constraint, training uses primal updates with λ = 0.Source: §4.4 · p. 5

5 Experimental Setup

The experiments evaluate AVICI across simulated causal domains, including substantial shifts in graphs, mechanisms, and noise. They specify evaluation metrics and train one inference model for each domain.

5 Experimental Setup

  • The evaluation uses simulated data with known causal structure to control task aspects, alongside an additional real-world proteomics evaluation in Appendix E.Real-world datasets with ground-truth causal structures are scarce, and believed ground truths may be incomplete or change with expert knowledge.Source: §5 · p. 5

5.1 Domains and Simulated Components

  • Three domains cover linear SCMs, nonlinear random-Fourier-feature SCMs, and semisynthetic gene regulatory networks, with test distributions designed to shift graphs, mechanisms, and noise.For interventional data, half the samples are observational and half are single-variable interventions.Source: §5.1 · pp. 5–6
  • Training uses Erdős-Rényi and scale-free graphs, while shifted tests add small-world, community, and geometric graphs, plus yeast and E. coli network subgraphs for GRNs.The GRN subgraphs preserve structural patterns such as motifs and modularity.Source: §5.1 · p. 6

5.2 Evaluation Metrics

  • Structure predictions are assessed with SHD, SID, and edge precision, recall, and F1 at a 0.5 threshold; uncertainty is assessed with AUPRC, AUROC, and ECE.SHD measures graph edit distance, while SID measures closeness in interventional adjustment sets.Source: §5.2 · pp. 6–7

5.3 Inference Model Configuration

  • The authors train one inference model per domain on datasets with 2–50 variables and 200 samples, sometimes including 50 interventional samples; SCM models use an acyclicity constraint.The three models share architecture and optimization hyperparameters except for dropout rate.Source: §5.3 · p. 7

6 Experimental Results

Across fully out-of-distribution benchmarks, AVICI outperforms most baselines; it also shows robust behavior across distribution shifts and unseen problem sizes. Its edge probabilities are well calibrated, and ablations identify the value of attending across both data axes.

6.1 Out-Of-Distribution Generalization

  • LINEAR performance hardly degrades under distribution shift, whereas GRN performance degrades more in out-of-distribution scenarios; large-degree training examples help explain better predictions on some GRN shifts.Acyclicity is perfectly satisfied in LINEAR and RFF, and AUPRC and AUROC are less affected by shift than thresholded point-estimate metrics.Source: §6.1 · pp. 7–8
  • On unseen problem sizes, AUPRC improves with more observations but with diminishing gains, while declining smoothly as variable count increases, including beyond training settings of n = 200 and d ≤ 50.The reported size study varies n and d on unseen in-distribution data.Source: §6.1 · p. 8

6.2 Benchmarking

  • On fully out-of-distribution data, AVICI trained on LINEAR outperforms all baselines with observational or interventional data, while only GIES achieves comparable accuracy; on RFF, AVICI surpasses GraN-DAG and DCDI.The benchmark uses n = 1000 observations and d = 30 variables under shifts in graphs, mechanisms, and noise.Source: §6.2 · p. 8, Table 1 · p. 8
  • AVICI's marginal edge probabilities are the most calibrated by ECE in each domain, and its calibration plots closely follow the perfect-calibration line.Baselines often show overconfident predictions, with calibration curves below the diagonal.Source: §6.2 · p. 9, Figure 4 · p. 9

6.3 Ablations

  • Ablations show a significant performance drop when attention uses only the variable axis and aggregates observations once through pooling after 2L self-attention layers.Attention only over observations does not process variable interactions until predicting θ.Source: §6.3 · p. 9, Table 2 · p. 10
  • An alternative relational-network parameterization shows no statistically significant performance difference from the simpler model, so AVICI uses fewer parameters and less memory.The alternative applies an MLP to each edge pair using an additional learned vector and matrices.Source: §6.3 · p. 9
  • Causal discovery performance increases monotonically with network depth and training time; even L = 4 or shorter training matches most baselines, and the main model has 4.2 × 10^6 parameters.The model occupies 17.0 MB at f32 precision, and a forward pass takes a few seconds on CPU.Source: §6.3 · p. 10, Table 2 · p. 10

7 Discussion

AVICI amortizes causal structure learning through variational inference, using simulators to encode inductive biases for difficult domains. The authors report robust distribution-shift behavior and identify trade-offs and directions for future work.

7 Discussion

  • AVICI addresses complex causal structure learning domains that existing approaches find intractable, and its approach can benefit from ongoing development of domain simulators.The authors suggest extension to settings without causal sufficiency and fine-tuning on labeled real-world data as future work.Source: §7 · p. 10
  • Acyclicity is induced empirically but is not guaranteed, and AVICI offers no theoretical performance guarantees.Some classical methods have infinite-sample guarantees under specific data-generating assumptions.Source: §7 · p. 10
  • Robustness to distribution shifts suggests pretrained models could be useful out of the box beyond the studied domains; fine-tuning on labeled real-world datasets is a proposed future direction.The authors make their code and models publicly available.Source: §7 · p. 10

A Domain Specification and Simulation

The appendix specifies training and evaluation data-generating processes, including random graph models, extracted biological networks, and simulated causal mechanisms. For gene-expression data, it describes SERGIO simulation and the preprocessing used before causal discovery.

A Domain Specification and Simulation

  • Tables 3 and 4 summarize training and out-of-distribution process parameters for LINEAR, RFF, and GRN domains.Parameter configurations are sampled across option combinations, and out-of-distribution graph classes are sampled in equal proportions.Source: A Domain Specification and Simulation · pp. 15–16

A.1.1 Random graph models

  • Evaluation graphs include Erdős-Rényi and scale-free training graphs, with Watts-Strogatz, stochastic block, and geometric random graph models among the described graph classes.Undirected random graphs are oriented using the upper-triangular half of the adjacency matrix; graph-model classes are sampled equally for evaluation datasets.Source: §A.1.1 · p. 17

A.1.2 Subgraph Extraction from Real-World Networks

  • For GRN evaluation, subgraphs come from E. coli and S. cerevisiae networks using an extraction procedure designed to preserve source-network properties.The procedure greedily adds neighboring nodes to maximize modularity, with added randomness from the top 20 percent of candidate graphs; effect signs are assigned using known E. coli signs or a fitted Beta distribution for S. cerevisiae.Source: §A.1.2 · p. 17

A.2.1 Structural Causal Models

  • LINEAR and RFF domains use structural causal models with additive noise that may be heteroscedastic, and variable-wise standardization is applied before inference.LINEAR mechanisms are affine, while RFF mechanisms approximate Gaussian-process draws using random Fourier features.Source: §A.2.1 · p. 18

A.2.2 Single-Cell Gene Expression Data

  • SERGIO generates single-cell gene-expression data in two stages: biological expression snapshots across five to ten cell types, followed by technical measurement noise.The noise includes outlier genes, varying library sizes, dropouts, and UMI count sampling; knockout interventions are simulated by setting knocked-out genes’ production rates to zero.Source: §A.2.2 · pp. 18–19
  • Gene-expression data are log2-CPM normalized, then nonzero values are shifted by their minimum and scaled by the overall standard deviation before causal discovery.This preprocessing retains the original sparsity pattern and within-sample relative ordering.Source: §A.2.2 · p. 19

B Evaluation Metrics

The section defines graph-accuracy and edge-probability metrics, noting limitations of structural scores and the role of calibration. It also describes how training batches and buffered simulated data are managed during optimization.

B Evaluation Metrics

  • SHD measures graph-edit distance, while SID measures closeness through valid adjustment sets; neither captures edge-level precision and recall well.SID can reward dense predictions, while an empty graph can achieve competitive SHD for sparse graphs.Source: B Evaluation Metrics · p. 20
  • F1 combines edge precision and recall, penalizing both empty and dense predictions that perform poorly on either measure.The authors report test-case means and standard errors for precision, recall, and F1.Source: B Evaluation Metrics · p. 20
  • AUPRC and AUROC assess edge-probability rankings across thresholds; AUPRC is more suitable when comparing sparse graphs across different variable counts.AUROC increasingly discounts accuracy on the shrinking proportion of true edges as d grows in sparse graphs with O(d) edges.Source: B Evaluation Metrics · p. 20
  • ECE summarizes calibration as a weighted average deviation between predicted edge probabilities and empirical edge frequencies across probability bins.Calibration plots aggregate statistics over test cases and use 10 bins; ECE does not measure prediction accuracy.Source: B Evaluation Metrics · pp. 20–21

C.1 Optimization

  • Training samples variable counts from 2 to 50 while fixing observations at 200, with batch sizes adjusted to fit GPU memory and higher dimensions oversampled.The loss is divided by d^2 to keep loss and gradient scales approximately equal across variable counts.Source: §C.1 · p. 21
  • A continuously refreshed buffer stores 200 simulated graph-data pairs for each variable count, and training samples minibatches from it.Stored datasets contain four times the training observation count and are subsampled to 200 observations when used.Source: §C.1 · p. 21

C.2 Architecture

AVICI uses a Transformer encoder with normalization and residual sublayers, plus position-wise mappings that produce representations for edge predictions. The appendix also states that architecture and optimization specifications were selected through held-out in-distribution experimentation.

C.2 Architecture

  • The encoder has eight layers with four residual sublayers each, multi-head attention with eight heads, and feedforward modules with ReLU activations.Layer normalization is applied before attention and feedforward modules and after the final layer; dropout was increased in RFF and GRN to help out-of-distribution generalization.Source: §C.2 · p. 22
  • Position-wise linear layers map variable representations to u_i and v_i, and the scale parameter τ is learned in log space to ensure positivity.For models with an acyclicity constraint, diagonal predictions and their loss terms are ignored.Source: §C.2 · p. 22
  • The reported architecture and optimization specifications were selected through experiments on held-out in-distribution instances.The implementation uses Haiku in JAX.Source: §C.2 · p. 22

D Baselines

The baselines are tuned and run using domain-specific procedures, with bootstrap resampling used to estimate edge probabilities for non-Bayesian methods. Extended tests examine cross-domain generalization, including transfer between LINEAR and RFF mechanisms.

D Baselines

  • PC and IGSP parameters are selected using ground-truth graphs, a procedure unavailable in practice that favors these methods.HSIC-γ also fails to scale to d = 100, and GRN runs at d = 30 use the Gaussian CI test for IGSP because HSIC-γ exceeds 100GB.Source: D Baselines · pp. 22–23
  • Non-Bayesian edge probabilities are estimated by running each baseline on ten bootstrap datasets and counting how often each edge appears.The datasets are sampled with replacement.Source: D Baselines · p. 23
  • GRN benchmarking required numerical adjustments and omitted some failed or nonterminating runs for several baselines.LiNGAM, PC, and IGSP received small Gaussian noise; some IGSP, DCDI, and PC runs were excluded, and GraN-DAG pruning was skipped in some cases.Source: D Baselines · p. 23

E Extended Results

  • The three main AVICI models were trained for approximately four days, while test-time inference used eight CPUs and no GPU.Benchmark baselines ran on four to eight CPUs for up to 24 hours, with DCDI using a GPU.Source: E Extended Results · p. 23

E.1 AVICI generalization between LINEAR and RFF

  • With cross-function-class transfer, AVICI remains on par with most baselines, and LINEAR-trained AVICI generalizes marginally better to RFF than the reverse.On LINEAR data, RFF-trained AVICI has F1 scores of 0.19 and 0.45 for observational and interventional data; on RFF data, LINEAR-trained AVICI scores 0.27 and 0.42, respectively.Source: §E.1 · p. 23, Table 5 · p. 24

E.2 In-Distribution Benchmarking Results for d = 30

AVICI improves in the easier in-distribution setting, especially on the less-noisy SCM domains, while most baselines perform similarly to their out-of-distribution results. At d = 100, AVICI performs favorably in nonlinear RFF and challenging GRN settings, while GES/GIES are statistically indistinguishable from it on LINEAR.

E.2 In-Distribution Benchmarking Results for d = 30

  • In-distribution AVICI performance improves significantly over the out-of-distribution setting, particularly in the less-noisy SCM domains, while most baselines perform roughly the same.The datasets and their generating parameters remain unseen during training.Source: §E.2 · p. 24
  • Some GRN baselines achieve slightly better F1 scores in-distribution, likely because of a change in graph structure rather than easier simulator parameters.The authors note there is no reason to believe different generative parameters are easier for these baselines.Source: §E.2 · p. 25

E.3 Benchmarking Results for d = 100

  • At d = 100, AVICI outperforms all baselines in nonlinear RFF with and without interventions and is the only method with nontrivial GRN edge accuracy; GES/GIES are not significantly different from AVICI on LINEAR.This evaluation shifts graph structures, mechanisms or simulator parameters, noise, and problem size relative to training.Source: §E.3 · p. 25, Table 7 · p. 25

E.4 Benchmarking Results on Real-World Proteomics Data

On the proteomics dataset, benchmark results are inconclusive and vary by metric, while the reference graph and data assumptions warrant caution. In the d = 30 uncertainty analysis, AVICI is well calibrated in SCM domains and performs favorably on probabilistic metrics, including in GRN without gene-knockout data.

E.4 Benchmarking Results on Real-World Proteomics Data

  • Proteomics conclusions are limited by a challenged 17-edge reference graph, potentially unjustified causal-sufficiency and acyclicity assumptions, and possibly imperfect interventions and targets.The dataset contains continuous measurements of 11 proteins, with much of the data interventional.Source: §E.4 · p. 25
  • On proteomics, results are inconclusive: GES/GIES lead on SID, GraN-DAG on F1, and GraN-DAG, DCDI, and AVICI are most favorable on SHD.Predicted edge counts vary greatly, while almost all F1 scores fall between 0.25 and 0.30.Source: §E.4 · pp. 25–26, Table 8 · p. 26

E.5 Uncertainty quantification for d = 30

  • AVICI closely follows the calibration diagonal in SCM domains for observational and mixed data; PC, GIES, IGSP bootstraps, and DiBS are also well calibrated, while DCDI, LiNGAM, and DAG-GNN are overconfident.The baselines have worse expected calibration error because a larger proportion of AVICI predictions are well calibrated.Source: §E.5 · p. 26, Figure 6 · p. 27
  • AVICI performs favorably on probabilistic metrics across the three domains, with GES and GIES on par in LINEAR, and has nontrivial GRN accuracy even without gene-knockout data.AUROC and AUPRC evaluate the full spectrum of decision thresholds.Source: §E.5 · p. 27, Table 9 · p. 28

Related papers

Loading 2205.12934v4…