Modality-Specific Causal VAEs

• The paper introduces a modality-specific causal VAE framework that leverages block-wise latent disentanglement and structural sparsity to recover fine-grained biomedical mechanisms.
• The model employs modality-specific encoder-decoder architectures and a learnable adjacency matrix to capture causal dependencies across different data types.
• Empirical results demonstrate near-perfect factor recovery and superior performance against baselines in both simulated and real-world biomedical experiments.

A modality-specific causal variational autoencoder (VAE) is a generative modeling framework designed for multimodal datasets, prevalent in biomedical domains and human phenotype research. This architecture seeks to identify interpretable, component-wise causal factors for each data modality, using nonparametric latent distributions and formal identifiability guarantees. Unlike earlier models dependent on restrictive parametric forms or yielding only coarse identification, the modality-specific causal VAE integrates structural sparsity constraints and block-wise latent disentanglement to recover fine-grained biomedical mechanisms with interpretability and identifiability essential for scientific investigation (Sun et al., 2024).

1. Generative Model and Variational Factorization

The modality-specific causal VAE assumes $M$ data modalities indexed by $m$, each with a block of causal latent factors $z^{(m)} \in \mathbb{R}^{d^{(m)}}$ and exogenous “style” variables $\eta^{(m)} \in \mathbb{R}^{d_\eta^{(m)}}$. The generative model is defined:

$$p(x, \{z^{(m)}, \eta^{(m)}\}_{m=1}^M) = \left[ \prod_{m=1}^M p(x^{(m)} \mid z^{(m)}, \eta^{(m)})\,p(\eta^{(m)}) \right]\,p(z | \epsilon)\,p(\epsilon)$$

where $z^{(m)}$ is generated via a modality-specific structural function $g_z^{(m)}$ and noise $\epsilon^{(m)}$:

$$z^{(m)} = g_z^{(m)} (\mathrm{Pa}(z^{(m)}), \epsilon^{(m)}),\quad x^{(m)} = g_x^{(m)} (z^{(m)}, \eta^{(m)})$$

with noise $$\epsilon = (\epsilon_1,\dots,\epsilon_{d_z}) \sim \mathcal{N}(0, I)$$. In practice:

$$p(x^{(m)} \mid z^{(m)}, \eta^{(m)}) = \mathcal{N}\left(x^{(m)}; \mathrm{Dec}^{(m)}(z^{(m)}, \eta^{(m)}), \sigma^2 I \right),\quad p(\eta^{(m)}) = \mathcal{N}(0, I)$$

The variational posterior adopts an amortized modality-wise factorization:

$$q\left(\{z^{(m)}, \eta^{(m)}\}_{m=1}^M \mid x^{(1)}, \ldots, x^{(M)}\right) = \prod_{m=1}^M q\left(z^{(m)}, \eta^{(m)} \mid x^{(m)}\right)$$

with Gaussian factors for each latent and exogenous variable.

2. Neural Architecture and Latent Integration

For each modality, the encoder $\mathrm{En}^{(m)}$ is parameterized by a dedicated neural network—typically a multilayer perceptron for tabular, a convolutional net for images, or RNN/1D-CNN for time series. The encoder outputs means and log-variances for both $z^{(m)}$ and $\eta^{(m)}$. Decoders $\mathrm{Dec}^{(m)}$ mirror this structure, receiving the concatenated latents $(z^{(m)}, \eta^{(m)})$ to reconstruct $x^{(m)}$.

No global shared latent block is introduced; block identification is enforced per modality, and subsequent causal alignment across modalities is managed via the downstream graph-structured latent flow. This ensures that cross-modalities interactions stem from the learned causal structure rather than shared nuisance factors.

3. Structural Sparsity and Graph-Structured Causality

Causal dependencies among latent factors are encoded with a learnable adjacency matrix $A \in \mathbb{R}^{d_z \times d_z}$, where $d_z = \sum_m d^{(m)}$. Each $(i, j)$ entry of $A$ represents a directed edge indicating whether $z_j$ causally influences $z_i$.

Within the normalizing-flow parameterization of $g_z$, each latent coordinate $\epsilon_i$ is obtained by combining its parent latents $\mathrm{Pa}(z_i)$ via a flow block masked by $A_{i, \mathrm{Pa}(i)}$. Structural sparsity is imposed by adding an $\ell_1$ penalty to the adjacency matrix:

$$\mathcal{L}_\mathrm{Sp} = \lambda_\mathrm{sp} \|A\|_1 = \lambda_\mathrm{sp} \sum_{i, j} |A_{i, j}|$$

This constraint incentivizes parsimonious cross-modal relationships, which is empirically natural for biomedical systems displaying sparse inter-modality causality.

4. Objective Formulation and Independence Constraints

The total objective aggregates reconstruction, independence, and sparsity terms:

$$\mathcal{L} = \underbrace{\sum_{m=1}^M \mathbb{E}_q[\|x^{(m)} - \mathrm{Dec}^{(m)}(z^{(m)}, \eta^{(m)})\|_2^2]}_{\mathcal{L}_\mathrm{Recon}} + \underbrace{\alpha_\mathrm{Ind}\,\mathrm{KL}[q(\gamma)\,\|\,\mathcal{N}(0,I)]}_{\mathcal{L}_\mathrm{Ind}} + \underbrace{\alpha_\mathrm{Sp}\,\,\|A\|_1}_{\mathcal{L}_\mathrm{Sp}}$$

where $$\gamma = (\eta^{(1)}, \dots, \eta^{(M)}, \epsilon_1, \dots, \epsilon_{d_z})$$ collects all nuisance-type latents and the KL constraint enforces their independence. Hyperparameters $\alpha_\mathrm{Ind}$ and $\alpha_\mathrm{Sp}$ balance the independence and sparsity penalties, respectively.

5. Identifiability Guarantees and Theoretical Results

The framework provides formal identifiability results under mild nonparametric smoothness and sparsity conditions. Theorem 4.1 ("Subspace Identifiability") states that, given smooth invertibility of modality-specific mixing maps $g_x^{(m)}$ and local linear independence of the Jacobian, each block $z^{(m)}$ is recoverable up to a smooth invertible transformation:

$\hat z^{(m)} = h^{(m)}(z^{(m)})$

Under further cross-modal sparsity assumptions (Theorem 4.2, "Component-wise Identifiability"), each scalar latent is identified up to permutation and a one-dimensional invertible map:

$\hat z^{(m)}_{\pi(i)} = h_i(z^{(m)}_i)$

Component-wise identifiability is achieved by ensuring that modality-specific nuisance variables $\eta^{(m)}$ are disentangled, exploiting their lack of cross-modal causal influence, and by penalizing extra inter-modal edges through the $\ell_1$ loss on $A$. This drives each estimated latent toward correspondence with a true independent source.

6. Optimization and Training Workflow

Training leverages the Adam optimizer (learning rate $10^{-3}$, batch size $256$). All terms in the loss are differentiable; $\ell_1$ regularization on $A$ is applied via subgradient methods. The normalizing flows deliver closed-form log-determinant Jacobians, incorporated into the model’s likelihood for gradient-based optimization.

End-to-end joint training encompasses the encoders $\{\mathrm{En}^{(m)}\}$, decoders $\{\mathrm{Dec}^{(m)}\}$, normalizing-flow parameters for $g_z$, and the adjacency mask $A$. Post-convergence, thresholding $A$ yields a binary representation of the latent causal graph.

7. Empirical Performance and Biomedical Relevance

Numerical simulations employed up to four modalities (dimension $15$–$20$ per modality), latent dimension $2$–$3$, and sparse causal graphs. The modality-specific causal VAE consistently achieved near-perfect factor recovery (Mean Correlation Coefficient $>0.98$, structural Hamming distance $=0$). Competing methods (BetaVAE, single-modality CausalVAE, multimodal contrastive learning) failed to recover independent sources or recovered only latent subspaces.

An ablation confirmed the theoretical prediction: factor recovery sharply improves as inter-modal sparsity increases. In the "Variant MNIST" experiment, the architecture identified cause and effect variables in paired modalities with high fidelity, outperforming baselines:

Method	$R^2$	MCC
MCL	0.48±0.01	0.82±0.02
BetaVAE	0.22±0.00	0.03±0.00
CausalVAE	0.02±0.01	0.14±0.01
Ours	0.89±0.05	0.87±0.02

For a large-scale human phenotype dataset (fundus imaging, sleep time series, tabular measures), the pipeline reconstructed latent causal skeletons consistent with established biomedical findings—for example, Sleep-latent 1→Oxygen saturation and fundus-latent→hand-grip strength—validating the real-world reliability and interpretability of causal discoveries (Sun et al., 2024).

A plausible implication is that modality-specific causal VAEs can provide fine-grained mechanistic insights in multimodal biomedical studies, where standard VAEs or contrastive approaches fail to achieve component-wise identifiability under realistic sparsity and nonparametric conditions.