Source-linked AI summary

CausalVAE: Structured Causal Disentanglement in Variational Autoencoder

Mengyue Yang, Furui Liu, Zhitang Chen, Xinwei Shen, Jianye Hao, Jun Wang

arXiv:2004.08697v7cs.LGstat.ML

TL;DR

Semantic factors may be causally related rather than independent, so CausalVAE maps independent exogenous factors into causal representations using an SCM layer. Experiments on synthetic and real data report learned causal concepts and counterfactual interventions as expected.

  • Problem

    Many disentanglement methods assume independent factors, although concepts of interest in real-world data may be causally related.

    Source: §1 · p. 1
  • Method

    CausalVAE transforms independent exogenous factors into causal representations using a Causal Layer and Mask Layer that model a structural causal mechanism.

    Source: §1 · p. 2, §3 · p. 3
  • Results

    Experiments with synthetic and real data report learned representations of causally related concepts and expected counterfactual outputs; on CelebA, CausalVAE achieves the best alignment with concept labels.

    Source: §7 · p. 8, §6.2 · p. 8
  • Takeaways & Limitations

    The learned causal representations support interventions that produce counterfactual outputs.

    Source: §3.2 · p. 4, §7 · p. 8

    The identifiability guarantee is conditional on the theorem’s assumptions, including a measure-zero condition on {x:ϕξ(x)=0}.

    Source: §5 · pp. 5–6

Abstract

from arXiv · show

Learning disentanglement aims at finding a low dimensional representation which consists of multiple explanatory and generative factors of the observational data. The framework of variational autoencoder (VAE) is commonly used to disentangle independent factors from observations. However, in real scenarios, factors with semantics are not necessarily independent. Instead, there might be an underlying causal structure which renders these factors dependent. We thus propose a new VAE based framework named CausalVAE, which includes a Causal Layer to transform independent exogenous factors into causal endogenous ones that correspond to causally related concepts in data. We further analyze the model identifiabitily, showing that the proposed model learned from observations recovers the true one up to a certain degree by providing supervision signals (e.g. feature labels). Experiments are conducted on various datasets, including synthetic and real word benchmark CelebA. Results show that the causal representations learned by CausalVAE are semantically interpretable, and their causal relationship as a Directed Acyclic Graph (DAG) is identified with good accuracy. Furthermore, we demonstrate that the proposed CausalVAE model is able to generate counterfactual data through "do-operation" to the causal factors.

1. Introduction

Common disentanglement methods assume independent factors, but semantic concepts can be causally related. CausalVAE models those dependencies in a causal representation and supports interventions that generate counterfactual data.

1. Introduction

  • VAE-based disentanglement uses a KL-divergence regularizer to encourage latent-factor independence, with some extensions also minimizing mutual information.Examples include increasing the KL weight [3] or reducing total correlation [12] [4].Source: §1 · p. 1
  • Real-world concepts of interest may be causally related, whereas many disentanglement methods assume observations arise from independent factors.The authors argue that this calls for a framework for causal disentanglement.Source: §1 · p. 1
  • In the pendulum example, illumination-source position and pendulum angle cause shadow position and length, so manipulating shadow factors can create images without a shadow.The paper connects this intervention to the causal “do-operation” and its use for asking what-if questions.Source: §1 · pp. 1–2
  • CausalVAE combines a VAE with a Mask Layer that models latent factors through a causal DAG, and its contributions include identifiability analysis and experiments on synthetic and real-world images.The framework requires weak supervision signals, such as concept labels, while learning the causal structure rather than taking it as a prior.Source: §1 · p. 2

2. Related Works

The related work covers disentangled representation learning and causal discovery, including efforts to combine causality with disentanglement. It positions CausalVAE as using causally structured layers to implement causal disentanglement.

2. Related Works

  • The section reviews state-of-the-art disentangled representation learning, recent combinations of causality and disentanglement, and observational causal structure learning as a CausalVAE ingredient.Source: §2 · p. 2

2.1. Disentangled Representation Learning

  • Prior methods use independence regularization, including β-VAE, FactorVAE [4], and Ladder VAE, while additional labels can reduce ambiguity in learned models.The cited discussion notes that unsupervised VAE models may be unidentifiable because of rotational invariance in the standard Gaussian prior.Source: §2.1 · p. 2

2.2. Causal Discovery & Causal Disentangled Representation Learning

  • Causal discovery methods use observational data or combine it with interventions; examples include LiNGAM, NOTEARS [34], and graph search with reinforcement learning [36].Interventional approaches include identifying causal structure from intervention data [5] or invariance across domains [10] [25].Source: §2.2 · pp. 2–3
  • CausalGAN [14] supports image do-operations but requires a graph prior, whereas CausalVAE uses causally structured masking layers; Besserve et al. [1] instead model dependent factors.The authors characterize their work as the first successful implementation of causal disentanglement.Source: §2.2 · p. 3, Figure 2 · p. 3

3. Causal Disentanglement in Variational Autoencoder

CausalVAE adds an SCM-based Causal Layer to a VAE, transforming independent exogenous factors into causal representations of related concepts. Concept labels supervise the causal structure, and the model supports interventions by propagating changes through the graph.

3. Causal Disentanglement in Variational Autoencoder

  • The Causal Layer transforms independent exogenous factors into causal endogenous factors for related concepts, using a mask mechanism to propagate parental effects to children.It is the component that supports intervention or do-operation in the system.Source: §3 · p. 3

3.1. Transforming Independent Exogenous Factors into Causal Representations

  • CausalVAE represents n concepts with a DAG and uses a linear SCM, z = (I − A^T)^−1ϵ, where ϵ contains independent Gaussian exogenous factors.The resulting z is a structured causal representation; additional concept labels supervise the conditional prior and learning of A.Source: §3.1 · p. 4

3.2. Structural Causal Model Layer

  • The Mask Layer reconstructs each z_i from its parent variables via a mild nonlinear function, while learning adjacency weights and function parameters by minimizing reconstruction error.Element-wise masking removes nonparent variables from the input to the child function; the authors report that the mild nonlinearity improves stability.Source: §3.2 · p. 4
  • Intervening on z_i fixes its right-hand-side value and propagates its effect to itself and its children, while intervening on an effect does not change its cause.The directional flow follows the model's two-layer representation of causal effects.Source: §3.2 · p. 4

3.3. A Probabilistic Generative Model for CausalVAE

  • The probabilistic model treats x and u as observed, and ϵ and z as latent variables, with a decoder f and encoder h.The decoding and encoding processes include independent noise; when the noise is infinitesimal, the encoder and decoder can be regarded as deterministic.Source: §3.3 · p. 4
  • The exogenous prior is standard Gaussian, while the endogenous prior pθ(z|u) is a factorized Gaussian conditioned on additional observations u.The model uses the prior's mean and variance as sufficient statistics, with λ1(u)=u and λ2(u)≡1.Source: §3.3 · pp. 4–5

4. Learning Strategy

CausalVAE is trained by optimizing an ELBO and learning the causal graph alongside the representation. Labels constrain the graph and latent codes, while a differentiable DAG constraint enforces acyclicity.

4.1. Evidence Lower Bound of CausalVAE

  • The learning strategy jointly learns the causal representation and causal graph; it uses variational Bayes to approximate the true posterior and optimize an ELBO.The posterior is simplified using the one-to-one correspondence between ϵ and z, yielding a tractable loss under the model assumptions.Source: §4 · p. 5, §4.1 · p. 5

4.2. Learning the Causal Structure of Latent Codes

  • Additional labels u constrain the learned graph A because both A and z are unknown, supporting training and graph identifiability.The model applies a corresponding constraint to the learned latent code z as well.Source: §4.2 · p. 5
  • The adjacency matrix A is constrained to be a DAG using a continuous differentiable function that equals zero if and only if A represents a DAG.This replaces a traditional combinatorial DAG constraint in the stated training procedure.Source: §4.2 · p. 5

5. Identifiability Analysis

The identifiability analysis establishes that, under stated assumptions and additional observations, the model parameters are identifiable up to the paper's equivalence relation. The authors connect this result to recovery of the true generative model and causal representations.

5. Identifiability Analysis

  • Theorem 1 states that (f, h, C, T, λ) are ∼-identifiable under the stated assumptions, including nonzero additional observations u_i.The other listed conditions concern the characteristic function, decoder differentiability and Jacobian rank, and nonzero sufficient statistics.Source: §5 · pp. 5–6
  • The authors state that Theorem 1 places learned generative-model parameters in an identifiable family, recovering the true model up to a certain degree.The analysis attributes identifiability under supervision to the conditional prior pθ(z|u), whose sufficient statistics relate to u.Source: §5 · p. 6
  • The authors state that each z_i aligns with the additional observation for concept i and is expected to inherit the causal system's relationships.They use this alignment to characterize z as a causal representation.Source: §5 · p. 6

6. Experiments

Experiments on synthetic datasets and CelebA test whether CausalVAE learns interpretable causal representations and whether interventions behave consistently with the modeled causal system. The reported results include concept-label alignment, intervention behavior, and recovery of the causal graph.

6.1.1 Datasets:

  • The experiments use synthetic Pendulum and Flow data and two CelebA subsets, with the main text showing Pendulum and CelebA(SMILE) results.CelebA contains about 200k labeled face images; evaluation uses MIC and TIC to measure information relevance with ground-truth concept labels.Source: §6.1.1 · p. 6, §6 · p. 6

6.2. Intervention experiments

  • On Pendulum, setting a concept’s latent code to zero keeps the intervened pattern nearly fixed across images while changing shadow attributes consistently with the stated physics.The authors report similar behavior in other intervention experiments; Flow results are deferred to the appendix.Source: §6.2 · pp. 7–8, Figure 3 · p. 7
  • ConditionVAE sometimes produces no image change when manipulating effect concepts, a problem the authors associate with not explicitly modeling causal disentanglement.The experiment compares supervised and unsupervised baselines as well as ablations; CausalGAN is compared only on interventions.Source: §6.1.1 · p. 6, §6.2 · p. 8
  • CausalVAE achieves the best alignment with concept labels among compared methods on CelebA, while other methods’ factors have low correlation with the labels.The comparison uses MIC and TIC between learned representations and concept labels.Source: §6.2 · p. 8, Table 1 · p. 8
  • On CelebA, intervening on SMILE changes MOUTH OPEN, whereas intervening on MOUTH OPEN does not change SMILE.The reported intervention outcomes follow the cause-to-effect direction described for these concepts.Source: Figure 4 · p. 7, §6.2 · p. 8
  • The learned causal matrix quickly converges to the ground-truth graph as training proceeds.The authors report that this indicates CausalVAE can correctly learn relationships among factors.Source: §6.2 · p. 8, Figure 5 · p. 8

7. Conclusion

The conclusion presents CausalVAE as a model for disentangling causally related concepts and summarizes its identifiability and experimental findings.

7. Conclusion

  • CausalVAE adds an SCM layer to model the causal generation mechanism of data with causally related concepts.The authors state that identifiability holds with additional supervision and that experiments show interpretable representations and expected counterfactual interventions.Source: §7 · p. 8
  • The authors report that the model is fully identifiable given additional supervision signals.Their experiments on synthetic and real data are reported to recover causally related concepts and generate counterfactual outputs through intervention.Source: §7 · p. 8
  • Experiments on synthetic and real data show representations of causally related concepts and counterfactual outputs consistent with the understood causal system.Source: §7 · p. 8

Supplementary Material of CausalVAE: Disentangled Representation Learning

The supplementary material is attributed to the CausalVAE paper and lists its authors, affiliations, and preprint identifier.

Supplementary Material

A. Proof of Proposition 1

The proof rewrites the ELBO KL term, treats a third term as constant, and combines terms to obtain Proposition 1’s stated form.

A. Proof of Proposition 1

B.1. Proof of Theorem 1

The proof analyzes identifiability of the exogenous and inferred latent variables, extends the result to dimension reduction, and characterizes causal-graph identifiability. The graph is identifiable up to a linear equivalence class, with the Mask Layer and additional supervision supporting strict identifiability.

B.1. Proof of Theorem 1

  • The proof relates the generative and inference models to establish identifiability of the latent variables under the paper’s assumptions.It first analyzes equality of the conditional generative distributions, then relates the inference model to the generative model.Source: §B.1 · pp. 12–13

B.2. Extension of Definition 1

  • When observation dimension d exceeds latent dimension n, the first n components of v each depend on a different z_j, preserving the theorem’s conclusion.Source: §B.2 · p. 13

B.3. Identifiability of Causal Graph

  • The causal graph is identifiable up to a linear equivalence class; the nonlinear Mask Layer and non-Gaussian supervision support strict identifiability.Source: §B.3 · p. 13

C. Implementation Details

The experiments report better representation alignment than all baselines and an almost-correct learned CelebA graph after five epochs. The section also describes dataset generation and preparation, training settings, and graph pretraining.

C. Implementation Details

  • Training and inference use one NVIDIA Tesla P40 GPU; the latent representation is extended to a matrix with one k-dimensional subspace per concept.Source: C. Implementation Details · p. 13

C.1.1 Sythetic Simulator

  • Synthetic Pendulum scenes vary pendulum and light angles, with projection laws calculating shadows; water scenes vary ball size, original water level, and hole location, with ball size and level determining height and parabola law generating flow.The Pendulum set contains about 7k images: 6k for training and 1k for inference. The flow set has 8k samples: 6k for training and 2k for testing.Source: §C.1.1 · p. 14

C.1.2 Data Preprocess of CelebA

  • CelebA’s 20K face images are split 85% for training and 15% for testing, then cropped to facial features and resized to 128 × 128 pixels.Source: §C.1.2 · p. 14

C.2.1 Synthetic

  • Synthetic models train for 80 epochs; CausalVAE uses α = 0.3, (β, γ) = (1, 1), and a condition prior N(u, |u|).Latent subdimensions k are 4 for synthetic experiments and 32 for CelebA; the network design follows Khemakhem et al.’s strategy to satisfy Theorem 1’s assumption (ii).Source: C. Implementation Details · p. 13, §C.2.1 · p. 14

C.2.2 CelebA

  • CelebA experiments use AGE, GENDER, and BEARD labels; CausalVAE uses N(u, I), and its graph pretraining takes 10 epochs within an 80-epoch schedule.The VAE-based methods sample latent codes from a conditional Gaussian and rescale representation variance by 0.1.Source: §C.2.2 · pp. 14–15

C.3. The Pretrain Step for Causal Graph Learning

  • Graph pretraining optimizes A from labels u with an augmented Lagrangian while gradients through the rest of the model are truncated.Source: §C.3 · p. 15

D.1. The Property of Learned Representation

  • CausalVAE’s learned representations outperform all supervised and unsupervised baselines in alignment with the conditional prior.Source: §D.1 · p. 15

D.2. The Learned Graph

  • The CelebA (BEARD) graph converges after five epochs from an initialization with all entries of A set to 0.5, yielding an almost-correct graph.Source: §D.2 · p. 15

D.3. Intervention Results

Interventions on Flow and CelebA show CausalVAE changing cause descendants without the reverse effects reported for comparison models. The section also specifies how an intervention is applied to a learned representation.

D.3. Intervention Results

  • To intervene, the method encodes an observed image, fixes the selected concept’s z_i and g_i(A_i ◦ z), then decodes the modified z.Source: §D.3 · p. 15
  • CausalVAE interventions on Flow change child concepts when BALL SIZE or HOLE is controlled, while controlling WATER HEIGHT does not change BALL SIZE; ConditionVAE also changes causes under effect interventions.The comparison indicates that supervision alone does not guarantee causal disentanglement.Source: §D.3 · pp. 15–16, Figure 3 · p. 16
  • On CelebA (BEARD), CausalVAE interventions change BEARD with GENDER and BALD with AGE, while intervening on BEARD leaves GENDER unchanged; CausalGAN’s intervention performance is reported as worse.Source: §D.3 · p. 16, Figure 11 · p. 19, Figure 12 · p. 20
  • On CelebA (SMILE), CausalVAE changes EYES with GENDER and MOUTH OPEN with SMILE, while intervening on MOUTH OPEN does not change SMILE; CausalGAN’s SMILE control can alter GENDER.Source: §D.3 · p. 16, Figure 13 · p. 21, Figure 14 · p. 22
Loading 2004.08697v7…