Technical Summary: Causal Representation Learning from Network Data (GraCE-VAE)

Problem Statement

Current causal representation learning methods often assume that observed variables are unstructured feature vectors, ignoring the rich relational context inherent in many scientific domains (e.g., protein-protein interactions, pathway-gene memberships). While existing graph-based inference methods can utilize network structures, they typically assume the causal graph is known or only consider observable interventions on known nodes. Conversely, causal disentanglement methods that learn latent Structural Causal Models (SCMs) from observational and interventional data often fail to leverage known interaction networks among measured entities. This creates a gap in learning latent causal structures and intervention effects while exploiting observed context networks.

The paper addresses the challenge of discovering latent causal factors and their interdependencies from network data enriched with intervention information, specifically in settings where variables are associated with a known interaction network (e.g., biological pathways).

Methodology: GraCE-VAE

The authors propose GraCE-VAE (Graph-Aware Causal Effects Variational AutoEncoder), a framework that integrates Graph Neural Networks (GNNs) with deep generative modeling to discover latent causal graphs from heterogeneous network data.

Architecture

GraCE-VAE employs an encoder-decoder split designed to separate the use of structural context from the generative causal model:

1. Graph-Aware Encoder:

   • Utilizes a GNN (specifically instantiated as GraphSAGE) to process a heterogeneous context graph containing primary observed entities ($X$, e.g., genes) and auxiliary entities ($H$, e.g., pathways).
   • The encoder conditions on the observed attributes of $X$, the auxiliary view $H$, and the adjacency structures ($A_{XX}, A_{XH}, A_{HH}$).
   • It produces structure-aware embeddings that parameterize the posterior distribution over exogenous noise variables ($Z$), effectively infusing network context into the amortized inference process.
   • Crucially, the context graph is used only for inference; it does not enter the decoder.

2. Causal Decoder (Latent SCM):

   • The decoder instantiates a Structural Causal Model (SCM) over latent variables $U$.
   • It learns a Directed Acyclic Graph (DAG) representing the causal structure among latent factors.
   • Intervention Modeling: Interventions are modeled via a two-stage process. An intervention encoder maps intervention indicators to a soft selection over latent variables (identifying the target) and a proposed mechanism modification (e.g., an additive shift). The decoder then applies this modification to the selected latent node's causal mechanism while leaving others unchanged.
   • A mixing function $f_\theta$ maps the latent causal variables $U$ back to the observed space $X$.

Learning Objective

The model is trained using a weighted sum of three components:

1. ELBO Term: Reconstructs observational and interventional data distributions.
2. Distribution-Alignment Term: Uses Maximum Mean Discrepancy (MMD) to align the model-generated counterfactual distributions with the empirical interventional distributions.
3. Regularization: Includes sparsity penalties on the learned DAG adjacency matrix and graph-based regularizers.

Identifiability

The paper establishes that GraCE-VAE inherits the identifiability guarantees of prior causal discrepancy VAEs (specifically CMVAE [34]). Under the assumptions of:

• Linear interventional faithfulness.
• Availability of both observational and interventional data.
• Samples being i.i.d. within each intervention regime.
• Each latent variable being subject to at least one intervention.
• An auxiliary-view assumption where the pathway information $H$ aids inference but does not provide additional generative information for $X$ once $U$ is known.

The latent causal DAG and intervention targets are identifiable up to the standard equivalence class (permutations of equivalent factors).

Key Contributions

1. Framework for Structured Context: Proposes a variational graph autoencoder that jointly learns a latent causal graph and intervention effects, guided by prior knowledge from heterogeneous networks (measured entities and group nodes). Unlike methods assuming fully observed variables or given causal graphs, GraCE-VAE discovers latent structure.
2. Theoretical Clarification: Clarifies that because the generative model class matches prior causal discrepancy VAEs and the graph-aware encoder only modifies amortized inference, GraCE-VAE inherits existing identifiability guarantees without requiring new theorems for a different decoder family.
3. Empirical Validation: Demonstrates through experiments on three CRISPR perturbation datasets (norman, replogle-small, replogle-large) that leveraging structured biological context improves the prediction of interventional outcomes, particularly for unseen perturbation combinations. Ablation studies confirm that both the graph-based encoder and the causal decoder are essential for performance.

Experimental Results

Experiments were conducted on three datasets involving genetic perturbations:

• Single Interventions: GraCE-VAE achieved strong $R^2$ performance, outperforming baselines like CMVAE (which ignores structure) and SENA. It was competitive with or superior to CMVAE-multihot and VGAE (an ablation without the causal decoder).
• Unseen Double Interventions: The model was trained on single interventions and evaluated on unseen double interventions. GraCE-VAE significantly outperformed CMVAE and SENA in predicting these complex combinations, achieving an average $R^2$ of 0.7845 compared to CMVAE's 0.7382. This highlights the synergy between the GNN's encoding of structured context and the causal decoder's ability to generalize.
• Ablation Studies:
  • GNN Design: Different GNN architectures (GCN, GAT, GraphSAGE) were tested; GraphSAGE performed best. Increasing layer depth did not yield further improvements.
  • Context Levels: Incorporating richer edge information (gene-gene, pathway-gene, pathway-pathway) generally improved single-intervention prediction, though gains plateaued or slightly reversed for double interventions, suggesting a trade-off between context richness and overfitting.
  • Corrupted Graphs: Replacing 50% of biological edges with random edges consistently degraded performance, confirming that meaningful biological topology drives the gains.
  • Priors: Adding embeddings from single-cell foundation models (scFM) or alternative intervention encoders did not consistently outperform the default GraCE-VAE configuration, suggesting the gains come from the specific integration of graph structure and causal modeling rather than just adding more priors.

Significance and Claims

The paper claims that GraCE-VAE offers a new approach to uncovering latent causal factors in complex network systems by unifying relational information and intervention-based causal inference. Its primary significance lies in:

• Bridging a Gap: It addresses the lack of methods that learn latent causal structures while exploiting observed context networks among measured entities.
• Generalization: It demonstrates superior ability to extrapolate to unseen intervention combinations (e.g., double perturbations) by leveraging structured biological context, a critical capability for scientific discovery.
• Theoretical Soundness: It maintains rigorous identifiability guarantees under standard assumptions, proving that incorporating graph structure into the encoder does not compromise the theoretical recoverability of the latent causal model.

The authors note limitations, including reliance on interventional data coverage, the use of simplified intervention functional forms (additive shifts), and the specific auxiliary-view assumption required for identifiability. They conclude that while the learned DAGs provide biologically plausible hypotheses (e.g., aligning with known cell-cycle or differentiation pathways), they should be viewed as hypothesis spaces for counterfactual simulation rather than definitive causal validations without further external validation.