Source-linked AI summary
Identifiability Guarantees for Causal Disentanglement from Soft Interventions
Jiaqi Zhang, Chandler Squires, Kristjan Greenewald, Akash Srivastava, Karthikeyan Shanmugam, Caroline Uhler
TL;DR
Causal disentanglement seeks latent causal variables, but observational data alone generally cannot identify them, and biological interventions may be unpaired. This paper develops an autoencoding variational Bayes approach and proves identifiability under soft interventions for unobserved variables, up to an equivalence class. The results support predicting unseen intervention combinations in the infinite-data limit, with an experimental demonstration on genetic data.
Problem
Causal disentanglement must recover latent variables and their causal structure even though observed-variable methods can be limiting and observational data alone generally cannot identify latent causal representations.
Source: §1 · p. 1Method
Results
The latent causal model is identifiable up to an equivalence class and supports predicting unseen intervention combinations in the infinite-data limit.
Source: §1 · p. 2Takeaways & Limitations
The framework was applied to genetic data, where identification and generalization to multi-node interventions were demonstrated experimentally.
Source: §7 · p. 11Full identification may be unnecessary or inefficient for particular downstream tasks, motivating task-specific identifiability criteria.
Source: §7.1 · p. 12
Abstract
from arXiv · showhide
Causal disentanglement aims to uncover a representation of data using latent variables that are interrelated through a causal model. Such a representation is identifiable if the latent model that explains the data is unique. In this paper, we focus on the scenario where unpaired observational and interventional data are available, with each intervention changing the mechanism of a latent variable. When the causal variables are fully observed, statistically consistent algorithms have been developed to identify the causal model under faithfulness assumptions. We here show that identifiability can still be achieved with unobserved causal variables, given a generalized notion of faithfulness. Our results guarantee that we can recover the latent causal model up to an equivalence class and predict the effect of unseen combinations of interventions, in the limit of infinite data. We implement our causal disentanglement framework by developing an autoencoding variational Bayes algorithm and apply it to the problem of predicting combinatorial perturbation effects in genomics.
1 Introduction
Causal disentanglement learns latent variables and their causal structure, addressing the limits of learning on observed pixels or high-dimensional measurements. The paper studies unpaired interventions and establishes identifiability for unobserved causal variables, while situating this approach among prior work.
1 Introduction
- Observed-variable causal learning can be limiting for perceptual and very-high-dimensional data, where individual measurements may lack useful causal meaning or need aggregation.For example, genes in the same pathway may function together and should be represented as one causal variable.Source: §1 · p. 1
- Causal disentanglement represents observations through a small set of latent causal variables, such as object-related pixel abstractions or gene groups corresponding to pathways.The observations are mapped from latent variables through a mixing function.Source: §1 · p. 2
- Observational data alone generally cannot recover latent causal representations, and the number of possible causal DAGs grows super-exponentially with dimension.Prior work discusses this unidentifiability under an empty-graph assumption.Source: §1 · p. 2
- Biological assays can yield unpaired interventional data because measuring a cell often destroys it, leaving one measurement per intervention condition.CRISPR experiments can perturb single or multiple genes.Source: §1 · p. 2
- The paper establishes soft-intervention identifiability for unobserved latent variables under generalized faithfulness, up to an equivalence class that supports predicting unseen intervention combinations with infinite data.It proposes an autoencoding variational Bayes approach and applies it to genomic perturbation prediction.Source: §1 · p. 2
1.1 Related Work
- Related work often assumes latent independence or known distributions, whereas this work uses a causal DAG and intervention data without those assumptions.Prior approaches to causal disentanglement have also used paired counterfactuals or restricted mixing maps; this work considers unpaired data and soft interventions [Ahuja et al., 2022b] [Liu et al., 2022].Source: §1.1 · pp. 2–3
- Causal structure learning from observed variables uses search and interventional scores, while interventional data is generally necessary for identifiability without further assumptions on the relating functions.Methods include exact search, greedy search, and scores incorporating interventions [Kuipers and Moffa, 2022].Source: §1.1 · p. 2
2 Problem Setup
The setup treats observations as mixtures of unobserved latent variables whose joint distribution follows an unknown DAG. It uses unpaired observational and interventional samples, with interventions changing single latent-variable mechanisms.
2 Problem Setup
- Observed variables X are generated by an unknown deterministic, potentially nonlinear mixing function f of latent variables U, whose distribution factorizes according to an unknown DAG G.The DAG has p latent nodes, and the induced distribution over X is the observed-data distribution.Source: §2 · p. 3
- A single-node intervention changes the target variable’s conditional mechanism while leaving other mechanisms unchanged; soft interventions may preserve dependence on the target’s parents.Hard interventions instead replace the conditional distribution and remove that parent dependence; a shift intervention is an example of a soft intervention [Rothenh¨ausler et al., 2015].Source: §2 · p. 3
- The data comprise unpaired observational and interventional samples, and the latent dimension, DAG, and intervention targets are unknown.The setup assumes at least one intervention per latent node, without knowing which interventions share a target; in linear SCMs, one per node is necessary in the worst case.Source: §2 · pp. 3–4
3 Equivalence Class for Causal Disentanglement
Because multiple latent models can generate the same observations, causal disentanglement is identified only up to an equivalence class. The paper accounts for latent-variable permutations and affine transformations that can be offset by changes to the mixing function.
3 Equivalence Class for Causal Disentanglement
- The latent model can only be identified up to the class of models that generate the same observed samples; the mixing function and mechanisms can be inferred once the variables, graph, and targets are identified.The equivalence class is defined relative to the observational and interventional samples.Source: §3 · p. 4
- Reindexing latent variables by a permutation preserves the observed data when the mixing function, DAG, and intervention targets are reindexed accordingly.This applies to both observational and interventional distributions.Source: §3 · p. 4
- Rescaling and shifting each latent variable also preserves the observations when the mixing function is adjusted to compensate.The scale constants must be nonzero.Source: §3 · p. 4
- CD-equivalence comprises latent-variable permutations and nonzero affine transformations, with the DAG and intervention targets permuted consistently.For the graph and targets alone, equivalence is defined by permutation.Source: §3 · p. 4
4 Identifiability Results
The section establishes identifiability guarantees for causal disentanglement from soft interventions, showing how interventions refine what observational data alone can identify.
4.1 Preliminaries
- With observational data, Assumption 1 identifies latent dimension and variables only up to an invertible linear transformation, leaving the DAG unidentified.The paper notes that any causal DAG can explain the observational data.Source: §4.1 · p. 5
4.2 Identifying ancestral relations
- Under Assumptions 1 and 2 for the intervention sets, the method identifies the transitive closure of the permuted DAG together with the permuted intervention targets, for some shared permutation.Source: §4.2 · p. 6
4.3 Identifying direct edges
- Under Assumptions 1–3, the causal model and interventions are identifiable up to their CD-equivalence class; for polytrees, Assumption 3 follows from Assumption 2 under mild regularity conditions.If the DAG is the sparsest within its transitive closure, Assumptions 1 and 2 suffice, as stated in the discussion of polytrees.Source: §4.3 · p. 7
4.4 Further remarks
- For genome-wide knock-down data, the assumptions successively support identifying cell-program count, genes sharing programs, and regulatory relationships between programs.The paper links these steps to Assumptions 1, 2, and 3, respectively.Source: §4.4 · p. 7
- The recovered model can predict effects of unseen combinations of interventions with distinct targets, but general soft interventions still prevent element-wise recovery of U.The paper describes the prediction guarantee for combinatorial interventions while noting that adding one latent variable to another can yield a valid explanation.Source: §4.4 · p. 7
5 Discrepancy-based VAE Formulation
The framework fits a latent causal model with an autoencoding variational Bayes objective, then compares virtual counterfactual samples with unpaired interventional data. Its loss lower-bounds the paired-data log-likelihood and is consistent in the infinite-data limit under suitable conditions.
5 Discrepancy-based VAE Formulation
- The framework selects a gradient-based optimization approach from available DAG-search strategies, and the loss uses a discrepancy measure to compare generated and real intervention samples.The paper uses MMD rather than adversarial discrepancy measures to avoid their saddle-point optimization challenges.Source: §5 · pp. 8–9
- Autoencoding variational Bayes replaces the log-likelihood objective with a variational lower bound and parameterizes latent causal mechanisms using independent exogenous noise.The encoder approximates qϕ(Z | X), while the decoder models observations through a deep structural causal model and a mixing function.Source: §5 · p. 8
- The decoder keeps the mixing function fixed across interventions while replacing observational mechanisms with intervention-specific counterparts.An intervention encoder estimates each intervention target and shift, which supports generating virtual counterfactual samples for observational inputs.Source: §5 · pp. 8–9
- The loss is consistent in the infinite-data limit under suitable conditions, enabling consistent prediction of multi-node intervention effects.It also lower-bounds the paired-data log-likelihood that would be optimized if counterfactual pairs were available.Source: §5 · p. 9
6 Experiments
The experiments evaluate the model on Perturb-seq data, training with observational and single-gene intervention samples while reserving double-gene interventions for testing. The model matches single-intervention distributions and reports strong R2 values, while some unseen intervention pairs yield different subpopulations.
6 Experiments
- The biological dataset contains 8,907 unperturbed cells and 99,590 perturbed cells targeting one or two of 105 genes.Each cell is represented by expression measurements for 5,000 highly variable genes.Source: §6 · pp. 9–10
- The test setup reserves double-node interventions entirely and holds out 96 cells from each single-node intervention with more than 800 cells.Training uses unperturbed cells and single-node intervention data; the latent dimension is set to 105.Source: §6 · p. 10
- 0.99 average R2 over all genes and 0.95 over most differentially expressed genes were obtained across the 105 single-node interventions.The model also discovers intervention-distribution subpopulations, including for KLF1.Source: §6 · p. 10
- 0.98 average R2 over all genes and 0.88 over most differentially expressed genes were obtained for 112 double-node interventions unseen during training.Some individual pairs recover different subpopulations, and the authors note that these pairs could involve non-additive target-gene effects.Source: Figure 6 · p. 10
- The learned DAG over seven groups of target genes includes DUSP9→MAPK1 and DUSP9→ETS2, edges reported as validated in.MAPK1 and ETS2 distributions overlap in the UMAP visualization, and the model maps them to the same latent node.Source: §6 · p. 11
7 Conclusion
The paper establishes identifiability results for causal disentanglement from single-node interventions and presents an autoencoding variational Bayes framework. Its conclusion and discussion describe the scope, practical limitations, and relation to contemporaneous work.
7 Conclusion
- The paper derives identifiability results from single-node interventions and presents an autoencoding variational Bayes framework, with experiments on genetic data demonstrating latent-structure identification and multi-node generalization.Source: §7 · p. 11
7.1 Limitations and Future Work
- The theory assumes access to an intervention on every latent node, which provides partial knowledge of latent variables; intervening on all variables may be impossible, and single-node interventions can be unrealistic.The authors cite multi-targeting drugs as an example and describe the current setting as a foundation toward general causal representation learning.Source: §7.1 · p. 11
- Identifiability is analyzed primarily in the infinite-data regime, leaving sample complexity for costly interventional data as an open direction.Source: §7.1 · p. 11
- The proposed DAG-learning alternatives inherit combinatorial optimization, so performance may require large samples and be sensitive to initialization.The authors identify optimization quality and robustness as areas for improvement.Source: §7.1 · p. 12
7.2 Discussion of Contemporaneous Works.
- The discussion distinguishes concurrent guarantees by their assumptions, including linear mixing, paired interventions, restrictions on latent graphs, or a known latent graph.The paper contrasts its setting with those of Varici et al., Buchholz et al., von K¨ugelgen et al., Jiang and Aragam, and Liang et al.Source: §7.2 · p. 12
A.1 Remarks on Assumption 1
- For linear injective mixing, a non-empty interior of the latent support suffices for the stated functional-class assumption; the appendix proves full row rank in this case.Source: §A.1 · p. 18
A.2 Proof of Lemma 1
This appendix part proves identifiability of the causal graph up to transitive closure from soft-intervention data under generalized faithfulness assumptions, using increasingly informative representations and sparsest topological representations. It also establishes observational limits on identifying latent variables and the DAG, and relates intervention targets and representation coefficients to topological order and descendants.
A.2 Proof of Lemma 1
- Under Assumption 1, observational data identify the dimension p and latent variables U only up to a nonsingular linear transformation and shift.The proof’s minimum-dimension argument shows any alternative representation has the same dimension and is related by an affine transformation; such transformations also generate the same observational data.Source: §A.2 · pp. 18–19
- With observational data alone, any DAG can explain the data, so the underlying graph is not identifiable up to a nontrivial equivalence class.The example uses an empty ground-truth graph and shows that any DAG can explain its observational distribution.Source: §A.2 · p. 19
B.1 Faithfulness Assumptions
- Assumptions 2 and 3 encode influentiality and adjacency faithfulness for linearly mixed causal variables, while Lemma 3 gives sufficient conditions for Assumption 2 using finite moments.The paper states that these conditions differ from prior faithfulness conditions because only a linear mixture of causal variables is observed; on a polytree, Assumption 2 plus a support regularity condition implies Assumption 3.Source: §B.1 · pp. 22–23
B.2 Summary of representations
- The appendix organizes a sequence of representations increasingly related to U, summarized in Table 1, and under Assumption 1 takes X = UΛ + b without loss of generality.Lemma 1 and Remark 1 justify this affine form.Source: Table 1 · p. 23, §B.2 · p. 23
B.3 Proof of Theorem 1
- For K ≥ p, the proof iteratively identifies sink nodes; under Assumptions 1 and 2, a topological representation exists and is topologically consistent with U.Among transformed topological representations, choosing one whose associated graph has the fewest edges yields TS(Gτ).Source: §B.3 · p. 24
B.3.1 Topologically ordered representations
- A topological representation is a transformed representation of X paired with an assignment of intervention targets, defined by marginal-distribution changes consistent with those targets.The representation is written as ˜U = XΠ for a nonsingular matrix Π.Source: §B.3 · p. 24, §B.3.1 · p. 24
B.3 Proof of Theorem 1
- The results establish identification of G up to its transitive closure.The proof constructs a topological representation and selects a transformed representation whose associated graph has the fewest edges; this yields TS(Gτ).Source: §B.3 · p. 24
B.3.1 Topologically ordered representations
- Under Assumption 2, a topological representation of X exists and can be recovered by reordering the nodes of G.The lemma establishes existence; the preceding passage identifies reordering as the recovery operation.Source: §B.3.1 · p. 24
- For a topological representation, ˜Γτ(j),i = 0 when i < j, its diagonal entries are nonzero, and the interventions target Uτ(1), …, Uτ(p) in order.Thus, the representation’s linear transformation respects a topological ordering.Source: §B.3.1 · pp. 24–25
B.3.2 Sparsest topological representation
- If τ(j) is a descendant of τ(i), some intervention targeting τ(k), with i ≤ k < j and τ(k) a descendant of τ(i), changes the marginal distribution of ˜Uj.Source: §B.3.2 · p. 25
- The sparsest topological representation yields a graph equal to the transitive closure of the underlying graph in topological order.The associated graph is chosen to have the fewest edges over invertible upper-triangular transformations.Source: §B.3.2 · p. 27
- In a sparsest topological representation, ˜Γτ(j),l = 0 whenever τ(l) is not a descendant of τ(j).This strengthens the consistency property established for general topological representations.Source: §B.3.2 · p. 27
B.3.3 Proof of Theorem 1
Under Assumptions 1 and 2, the proof recovers the causal graph up to transitive closure and permutation, and the intervention targets up to the same permutation.
B.3.3 Proof of Theorem 1
- Theorem 1 identifies the graph up to transitive closure and permutation, and intervention targets up to that same permutation.The representation is constructed in a topological ordering of the graph.Source: §B.3.3 · p. 28
- The proof identifies additional intervention targets by finding the earliest topological coordinate whose marginal distribution changes under an intervention.For coordinates before the intervention target, the marginal remains unchanged; at the target coordinate, it changes.Source: §B.3.3 · p. 28
B.4 Proof of Theorem 2
The proof establishes recovery of the latent graph and intervention targets up to CD-equivalence, and shows that identified models support prediction under combinatorial interventions. The appendix then describes a scalable discrepancy-based VAE for learning these models.
B.4 Proof of Theorem 2
- Under Assumptions 1, 2, and 3, the latent graph and intervention targets are identifiable up to their CD-equivalence class.The proof uses sparsity to distinguish the graph from alternatives sharing its transitive closure.Source: §B.4 · pp. 28, 30
B.5 Proof of Theorem 3
- The identified model generates the same observational distribution as the ground-truth SCM under any combinatorial intervention drawn from the available intervention set.The proof matches the conditional distributions node by node and then the resulting joint distribution.Source: §B.5 · pp. 30–31
C Details on Discrepancy-based VAE
- A discrepancy-based VAE is proposed because the identifiability proofs do not yield an algorithmically efficient learning approach.The framework uses reparameterization and gradient-based optimization.Source: C Details on Discrepancy-based VAE · p. 31
C.1 Maximum Mean Discrepancy
- The section introduces maximum mean discrepancy and notes that its empirical squared estimate is unbiased.The supplied passage defines the measure and its empirical counterpart.Source: §C.1 · pp. 31–32
C.2 Discrepancy VAE Details
- The VAE uses one decoder for observational and interventional observations, while encoding a shared posterior over Z and transforming it to latent variables.The theoretical results assume noiseless measurements; consistency under noisy measurements is left for future work.Source: §C.2 · p. 32
D Lower Bound to Paired Log-Likelihood
The appendix shows that the unpaired-data loss lower bounds a paired-data objective, then establishes recovery of the CD-equivalence class under idealized VAE conditions. It also notes a practical challenge in enforcing the sparsity conditions used by the theory.
D Lower Bound to Paired Log-Likelihood
- For a Gaussian kernel with sufficiently large width, the proposed loss lower bounds the paired-data ELBO and, consequently, the paired-data log-likelihood.The bound assumes a fixed variance and uses a constant depending on the variance and data dimension.Source: D Lower Bound to Paired Log-Likelihood · p. 33
E.1 CD-Equivalence Class
- With a true VAE prior and Dirac decoder and encoder distributions, minimizing the loss recovers the ground-truth CD-equivalence class.The result assumes Assumptions 1, 2, and 3 and that the intervention targets cover the latent variables.Source: §E.1 · pp. 34–35
- In practice, gradient-based optimization may not find a DAG with both the sparsest transitive closure and the sparsest DAG realizing it.The authors instead encourage sparser DAGs by penalizing the adjacency matrix.Source: §E.1 · p. 35
E.2 Consistency for Multi-Node Interventions
For additive-noise SCMs with shift interventions, the paper shows how to encode combinations of single-node interventions as multi-node interventions. This permits generating samples for those combinations from the learned model in the infinite-data limit.
E.2 Consistency for Multi-Node Interventions
- In the infinite-data limit, encoded shift interventions combine by adding their encoded vectors, allowing generation for corresponding multi-node interventions.The result covers additive Gaussian noise models in which interventions shift the means of exogenous noises; overlapping targets add their shift values.Source: §E.2 · pp. 35–36
F Discrepancy-based VAE Implementation Details
The implementation combines a VAE with a learned intervention encoder and discrepancy-based training. Biological experiments use a compact neural-network setup and a 100-epoch training schedule.
VAE Parameterization
- The encoder and decoder are neural networks, while a fully connected network encodes interventions with a temperature-annealed softmax that approaches one-hot assignments at test time.The exogenous variables use the reparameterization trick and a standard isotropic normal prior.Source: F Discrepancy-based VAE Implementation… · p. 36
Loss Functions
- Training uses Gaussian-kernel MMD discrepancies with dyadically spaced widths, alongside scheduled discrepancy and KL coefficients.The discrepancy coefficient begins after 5 epochs and the KL coefficient after 10; both rise to their maximum by halfway through training.Source: F Discrepancy-based VAE Implementation… · p. 36
Optimization
- The implementation uses Adam with PyTorch default parameters and a learning rate of 0.001.The structural equations are linear with shift interventions; nonlinear mixing from latent U to observed X is said to limit the practical loss of expressivity.Source: F Discrepancy-based VAE Implementation… · p. 36
Biological Data
- Biological-data models use two-layer fully connected encoders with leaky ReLU and 128 hidden units, train for 100 epochs, and take under 45 minutes on one GPU.The intervention encoder also uses 128 hidden units; batch size is 32.Source: F Discrepancy-based VAE Implementation… · p. 36
G Extended Results on Biological Dataset
The extended biological results assess single- and double-gene interventions and display the learned structure between gene programs. Single-node samples closely match observed distributions, while double-node predictions vary by intervention pair.
G.1 Single-node interventions
- Generated samples closely match actual distributions for the 11 illustrated single-target interventions, with empirical MMD reported across all 105 single interventions.Figure 11 records MMD in different batches.Source: Figure 9 · p. 37, §G.1 · p. 37, Figure 11 · p. 38
G.3 Structure Learning
- The learned structure visualization shows regulatory relationships between gene programs alongside descriptions of each program.The figure also lists the full set of genes in each program.Source: §G.3 · p. 38
G.2 Double-node interventions
- Double-node predictions match actual samples for some pairs but not others, including accurate predictions for ETS2+CEBPE and SET+IRF1 but not ETS2+DUSP9.The extended results show 11 random pairs and report MMD, RMSE, and R2 summaries across all 112 interventions; formal benchmarking against prior studies is left for future work.Source: §G.2 · p. 38, Figure 13 · p. 39
H Extended Experiments
- The extended experiments introduce ablations of architecture components on biological data and a simple simulation study of different tasks.Further simulation details and results are summarized in the following section.Source: H Extended Experiments · p. 39
H.1 Ablation Studies
The ablations compare the final model with versions lacking discrepancy loss or the causal layer. A five-node simulation examines intervention-target sensitivity and zero-shot double-node extrapolation.
H.1 Ablation Studies
- The ablations compare the final architecture against models without discrepancy loss and a model without the causal layer, using the same data split, schedule, and learning rate.The discrepancy-free alternatives use conditional-VAE-style learning and differ in whether inference samples exogenous variables or encodes observational samples.Source: §H.1 · pp. 39–40
- The authors state that the encoder, decoder, DSCM, and intervention encoder are needed to learn distributions and the latent causal graph in this setting.The biological ablation table reports testing metrics and standard errors over 14 single-node and 112 double-node interventions.Source: §H.1 · p. 40, Table 3 · p. 40
- The final model works best in general, while removing discrepancy loss performs much worse through mode collapse; removing the causal layer slightly improves double-node MMD.The authors suggest this may reflect better capture of non-additive double interventions without an imposed structure; MMD is considered more meaningful than R2 for distribution generation.Source: §H.1 · p. 40, Table 3 · p. 40
H.2 Simulation
- The simulation uses a five-node graph with 2,048 samples per interventional dataset, mapped into a 10-dimensional observation space for visualizing zero-shot double-node samples.Table 4 reports MMD for held-out single- and double-node interventions, training MMD, and structural Hamming distance.Source: §H.2 · p. 41, Table 4 · p. 41
- In simulation, single-node generations are similar across target transposition distances, but zero-distance targets yield marginally better single-node results and a larger benefit for double-node extrapolation.The authors report sensitivity of learned targets to initialization, attributing it to DAG combinatorics and small sample sizes; zero-distance training results may help with model selection.Source: §H.2 · p. 41, Table 4 · p. 41