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

TL;DR

FlexCausal introduces a novel causal disentanglement framework that models complex exogenous noise and enforces structural causal relations, leading to improved disentanglement and generation fidelity in real-world data.

Contribution

The paper proposes FlexCausal, a new CDRL method using block-diagonal covariance VAE and flow-based priors to better model non-Gaussian noise and causal structures.

Findings

Outperforms existing methods on synthetic datasets.

Achieves superior causal disentanglement on real-world data.

Demonstrates high-fidelity generation with manifold-aware interventions.

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.