Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FlexCausal: Flexible Causal Disentanglement via Structural Flow Priors and Manifold-Aware Interventions

Yutao Jin, Yuang Tao, Junyong Zhai

TL;DR

FlexCausal addresses the limitations of Gaussian priors and diagonal posteriors in causal disentangled representation learning by integrating a Block-Diagonal VAE with a Flow-based Exogenous Prior, supervised alignment, and a Counterfactual Consistency Loss. It introduces a manifold-aware directional intervention to maintain realism in counterfactuals and utilizes an Additive Noise Model with non-Gaussian exogenous noises to improve identifiability. The model achieves superior identifiability and distributional fidelity across synthetic and real-world datasets, with ablations underscoring the importance of the flow prior, consistency loss, and block-diagonal structure. Overall, FlexCausal advances robust, on-manifold causal generation and reliable counterfactual reasoning in complex environments, enabling more faithful and controllable causal synthesis.

Abstract

Causal Disentangled Representation Learning(CDRL) aims to learn and disentangle low dimensional representations and their underlying causal structure from observations. However, existing disentanglement methods rely on a standard mean-field approximation with a diagonal posterior covariance, which decorrelates all latent dimensions. Additionally, these methods often assume isotropic Gaussian priors for exogenous noise, failing to capture the complex, non-Gaussian statistical properties prevalent in real-world causal factors. Therefore, we propose FlexCausal, a novel CDRL framework based on a block-diagonal covariance VAE. FlexCausal utilizes a Factorized Flow-based Prior to realistically model the complex densities of exogenous noise, effectively decoupling the learning of causal mechanisms from distributional statistics. By integrating supervised alignment objectives with counterfactual consistency constraints, our framework ensures a precise structural correspondence between the learned latent subspaces and the ground-truth causal relations. Finally, we introduce a manifold-aware relative intervention strategy to ensure high-fidelity generation. Experimental results on both synthetic and real-world datasets demonstrate that FlexCausal significantly outperforms other methods.

FlexCausal: Flexible Causal Disentanglement via Structural Flow Priors and Manifold-Aware Interventions

TL;DR

FlexCausal addresses the limitations of Gaussian priors and diagonal posteriors in causal disentangled representation learning by integrating a Block-Diagonal VAE with a Flow-based Exogenous Prior, supervised alignment, and a Counterfactual Consistency Loss. It introduces a manifold-aware directional intervention to maintain realism in counterfactuals and utilizes an Additive Noise Model with non-Gaussian exogenous noises to improve identifiability. The model achieves superior identifiability and distributional fidelity across synthetic and real-world datasets, with ablations underscoring the importance of the flow prior, consistency loss, and block-diagonal structure. Overall, FlexCausal advances robust, on-manifold causal generation and reliable counterfactual reasoning in complex environments, enabling more faithful and controllable causal synthesis.

Abstract

Causal Disentangled Representation Learning(CDRL) aims to learn and disentangle low dimensional representations and their underlying causal structure from observations. However, existing disentanglement methods rely on a standard mean-field approximation with a diagonal posterior covariance, which decorrelates all latent dimensions. Additionally, these methods often assume isotropic Gaussian priors for exogenous noise, failing to capture the complex, non-Gaussian statistical properties prevalent in real-world causal factors. Therefore, we propose FlexCausal, a novel CDRL framework based on a block-diagonal covariance VAE. FlexCausal utilizes a Factorized Flow-based Prior to realistically model the complex densities of exogenous noise, effectively decoupling the learning of causal mechanisms from distributional statistics. By integrating supervised alignment objectives with counterfactual consistency constraints, our framework ensures a precise structural correspondence between the learned latent subspaces and the ground-truth causal relations. Finally, we introduce a manifold-aware relative intervention strategy to ensure high-fidelity generation. Experimental results on both synthetic and real-world datasets demonstrate that FlexCausal significantly outperforms other methods.
Paper Structure (35 sections, 1 theorem, 35 equations, 9 figures, 4 tables, 1 algorithm)

This paper contains 35 sections, 1 theorem, 35 equations, 9 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Disentangled Representation Learning and Causal Representation Learning
  4. Causal Disentangled Representation Learning
  5. Deep Generative Model and Counterfactual Editing
  6. Primaries
  7. Methodology
  8. Structure-Aware Representation Learning via Block-Diagonal VAE
  9. Flow-based Exogenous Prior
  10. Consistency Loss
  11. Optimization Objective
  12. Manifold-Aware Directional Intervention
  13. Experiment
  14. Experimental Setup
  15. Datasets.
  16. ...and 20 more sections

Key Result

Proposition 3.6

Let $\mathcal{G}$ be a Directed Acyclic Graph (DAG) representing the causal structure. Under Definition anm, the mapping $\mathcal{T}: \mathcal{Z} \to \mathcal{N}$ defined by $n_k = z_k - f_k(\text{PA}(z_k))$ is a volume-preserving bijective transformation. Specifically, the absolute determinant of This property implies that the log-likelihood in the causal latent space equates to the log-likelih

Figures (9)

  • Figure 1: CRDL aims to disentangle latent causal factors from observation and enable controllable causal generation. Our FlexCausal effectively disentangles causal representations. Interventions on parent nodes will influence their child nodes, while interventions on child nodes leave the parent nodes invariant.
  • Figure 2: An overview of FlexCausal. (a) Overall Framework: The model encodes input images $x$ into latent variables $z$, which are then reconstructed via the decoder. (b) Structural Causal Mechanism: The latent space is structured as a causal graph defined by the adjacency matrix $A$. For each node $z_k$, its value is determined by its parents $\text{PA}(z_k)$ through a nonlinear structural function $f_{k}$. (c) Flow Prior: To model complex priors, we utilize Normalizing Flows ($T_{\psi, k}$) to transform the simple base distribution into the flexible posterior distribution of the residuals. (d) Counterfactual Consistency Constraint: This explicitly penalizes violations of structural equations under intervention, thereby ensuring the reliability of generated counterfactuals.
  • Figure 3: Visualization of the Learned Latent Space on the Filter Dataset. (a) Latent Manifold Topology: The t-SNE visualization reveals the manifold structure of latent space. (b) GT-Latent Alignment: The scatter plots demonstrate a strong linear correlation between latent projections and ground truth labels (with $R^2 > 0.99$ for almost every concept), confirming precise semantic alignment. (c) Posterior Density: The learned latent distribution closely matches the ground truth distribution of the underlying factors, verifying the expressive capability of our Flow-based prior to model complex non-Gaussian densities.
  • Figure 4: Filter Dataset
  • Figure 5: Resuls of CausalVAE on CelebA-Smile, CelebA-Age, Pendulum, Filter
  • ...and 4 more figures

Theorems & Definitions (5)

  • Definition 3.2
  • Definition 3.4
  • Proposition 3.6
  • Definition 4.1: Causal Partition
  • proof