3D Causal Variational Autoencoder

• 3D CausalVAE is a deep generative framework that explicitly models causal dependencies among vector-valued latent factors in dynamic 3D environments.
• It employs a block-structured latent space with Gumbel–Softmax assignments and autoregressive priors to disentangle intervention-sensitive features from invariant information.
• Empirical results demonstrate robust recovery of causal factors, high correlation metrics, and effective sim-to-real transfer in complex spatiotemporal data.

A 3D Causal Variational Autoencoder (3D CausalVAE) is a structured representation learning framework that extends the variational autoencoder (VAE) paradigm to settings where underlying factors of variation are causally related and potentially vector-valued. Unlike classical disentanglement methods that presuppose statistical independence among latent factors, 3D CausalVAE explicitly encodes known or deduced dependencies, particularly relevant in spatiotemporal or dynamical systems with interventions. This approach enables the identification and separation of latent causes, including multidimensional attributes such as 3D object rotations and positions, and supports reasoning about interventions, counterfactuals, and dynamics. The framework is instantiated in models like CITRIS (Lippe et al., 2022), which demonstrates the recovery of causally related, multidimensional generative factors from temporal visual data under targeted interventions.

1. Structural Causal Assumptions and Generative Modeling

3D CausalVAE assumes that the data-generating process can be characterized by a set of $K$ possibly vector-valued causal factors $C_i^t \in \mathcal{D}_i^{M_i}$ evolving according to a stationary, first-order Markov Dynamic Bayesian Network (DBN) without instantaneous effects. Each causal factor at time $t$ follows

$$C_i^t = f_i(\operatorname{pa}(C_i^t), \epsilon_i^t), \quad \operatorname{pa}(C_i^t)\subseteq \{C_1^{t-1},\dots, C_K^{t-1}\},$$

where $\epsilon_i^t$ are independent noise terms and $\operatorname{pa}$ denotes the parents in the DBN.

Crucially, temporal interventions are encoded via a binary vector $I^t\in\{0,1\}^K$ at each time step. $I_i^t=1$ represents a (soft) intervention on $C_i^t$, modifying the factor's conditional transition distribution. The observational frame $X^t$ is a bijective function of all causal factors and independent observation noise $E_o^t$:

$X^t = h(C^t, E_o^t),$

where $$h: \mathcal{C}\times\mathcal{E} \rightarrow \mathcal{X}$$ is invertible and thus uniquely determines the latent configuration from the observed data (Lippe et al., 2022).

2. Latent Architecture and Inference Mechanism

The latent state $z^t\in\mathbb{R}^{M}$ is divided into $K+1$ blocks $z^t_{[i]}$, where blocks $i\ge1$ are designed to capture the minimal causal variables directly influenced by interventions on $C_i$, and block $0$ absorbs all intervention-invariant information. The prior over latents is factorizable per block, conditional on interventions:

$$p_\phi(z^{t+1} \mid z^t, I^{t+1}) = \prod_{i=0}^K p_\phi(z_{[i]}^{t+1}\mid z^t, I^{t+1}_i),$$

such that each block is only modulated by interventions on its target factor. This structure is maintained both in VAE and in hybrid architectures where a deep autoencoder (AE) provides entangled features that are then disentangled via a normalizing flow mapping.

For inference, an encoder $q_\theta(z^t|X^t)$ (or $E_{AE}$ plus $g_{flow}$) is trained together with a discrete assignment function $\psi: \{1,\dots,M\}\rightarrow\{0,1,\dots,K\}$, learned through a Gumbel–Softmax relaxation, so that each latent dimension is stably associated to a single causal block. The evidence lower bound (ELBO) for paired time steps and interventions is

$$\mathcal{L}_{ELBO} = -\mathbb{E}_{q_\theta}\left[\log p_\theta(X^{t+1}|z^{t+1})\right] + \sum_{i=0}^K \mathbb{E}_{q_\theta}\left[\mathrm{KL}\big(q_\theta(z_{[i]}^{t+1}|X^{t+1})\|\ p_\phi(z_{[i]}^{t+1}|z^t,I^{t+1}_i)\big)\right],$$

with a regularization on block 0 to incentivize collection of residual information.

An auxiliary target classifier is trained to ensure that block $i$'s latent encodes its corresponding intervention variable but not others, supporting identifiability. In the autoencoder–flow variant, a separate AE is pretrained, and the normalizing flow is trained with a similar block-structured objective.

3. Identifiability and Minimal Causal Splits

Under standard assumptions—smoothness, adequate support, bijectivity of $h$, faithfulness of the DBN graph, and non-deterministic interventions—it is shown that the 3D CausalVAE can identify, up to orthogonal mixing, the minimal causal variables for each $C_i$ in separate blocks, as long as all information shared due to coupled interventions or invariant subspaces is collected in the intervention-invariant block 0.

A minimal causal split of each factor $C_i^t$ is defined (nonuniquely) as

$$s_i(C_i^t) = (C^{\text{var}}_{i,t},\, C^{\text{inv}}_{i,t}),$$

where $C^{\text{var}}_{i,t}$ is the maximal entropy component that is conditionally independent of $I^t_i$ given parents. The learned representation's blocks thus recover these minimal causal subcomponents, satisfying a structural adequacy guarantee formalized in a main theorem (Lippe et al., 2022).

4. Implementation and Training Details

3D CausalVAE (in the CITRIS instantiation) uses a convolutional VAE pipeline: four strided convolutional layers (64 channels, 3×3 kernels, stride 2) interleaved with non-strided convs, mapping to a linear embedding (256 dimensions). The decoder employs bilinear upsampling, residual blocks, and a final Tanh layer for $64\times 64$ RGB outputs. The transition prior for $z_{[i]}$ employs autoregressive MADE networks (two hidden layers, SiLU activations) to predict Gaussian transitions from $z^t$ and $I^{t+1}_i$.

In the AE+Flow setting, a deep autoencoder is first trained by L₂ pixel loss; its bottleneck features are mapped via a normalizing flow (stacked coupling layers, ActNorm, and 1×1 convolutions) to the causal blocks. The Gumbel–Softmax assignment ensures “hard” membership of latent features to causal blocks.

Practical training objectives include the standard ELBO (with a weighting of block 0's KL term), flow log-likelihood under the structured prior, and cross-entropy auxiliary loss from the intervention-classification task.

5. Empirical Evaluation: Disentanglement and Generalization in 3D Data

Experiments use synthetic datasets capturing complex, temporally evolving 3D environments (e.g., Temporal-Causal3DIdent). Causal factors include $[x,y,z]$ position, $[\alpha,\beta]$ 3D object rotation angles, shape, color, and lighting properties, each with their own dynamics and targets for intervention. For example, object position is nonlinearly dependent on preceding rotations; interventions resample factors from their prior distributions with a set probability.

Evaluation metrics incorporate:

• Triplet mixing: Generating a synthetic frame from mixed latent blocks of two different frames, followed by automated measurement of the factor error via a pretrained classifier.
• Correlation structure: Quantifying the diagonal/off-diagonal $R^2$ and Spearman correlations between ground-truth factors and each latent block, including tests on out-of-distribution (OOD) data where factor independence is enforced.

Results demonstrate that 3D CausalVAE, in both VAE and AE+Flow settings, recovers all designated causal factors (diagonal $R^2 > 0.9$, off-diagonal $\approx 0$), exceeding alternatives such as SlowVAE and iVAE*, which struggle when factors are strongly temporally correlated. Factor mixing consistently produces perceptually accurate reconstructions, especially in the AE+Flow variant, which preserves multidimensional attributes with superior fidelity (Lippe et al., 2022).

6. Sim-to-Real Transfer and Robustness to Unseen Causal Configurations

A notable property is the ability to achieve direct sim-to-real generalization: a powerful autoencoder trained on (potentially mixed) real and simulated data can be kept fixed, and a normalizing flow can be trained to map its latent codes to causally disentangled blocks using a small interventional dataset. This configuration enables zero-shot transfer of disentanglement to new real data or previously unseen 3D object categories without further adaptation. Empirical results on held-out shapes maintain $R^2 > 0.93$ and triplet errors below 0.25, supporting the claim of robust generalization of the causal block structure beyond the specific intervention set used in training. This suggests an efficient path for bringing causal disentanglement frameworks developed in simulation to bear on real-world perception tasks, circumventing the need for laborious manual factorization or extensive fine-tuning.

7. Outlook and Connections to Broader Causal Representation Learning

The 3D CausalVAE paradigm, as realized in the CITRIS framework, unites several threads of contemporary research: causal inference, variational inference, deep generative models, and disentanglement. It extends scalar causal representation results to high-dimensional, vector-valued, time-varying latent factors, incorporating interventions and temporal structure explicitly. A plausible implication is stronger identifiability and interpretability in settings involving action, dynamics, and environmental variation, where purely independence-based disentanglement is inadequate. This approach suggests new directions for sim-to-real transfer, robust generalization, and machine causal reasoning for complex visual domains (Lippe et al., 2022).