5.2.2 Identifiable Variational Autoencoders (iVAE)

Note: iVAE requires PyTorch. The IVAE estimator in pydeepcausalml.generative uses a conditional prior \(p(z \mid u)\) for identifiable representation learning.

Overview

Standard Variational Autoencoders are expressive generative models, but they suffer from a fundamental statistical weakness: their latent representations are not identifiable. In practice, many different latent spaces can produce the same data distribution, meaning the model has no way to pin down which latent structure is the “true” one. This ambiguity makes it impossible to reliably recover the underlying factors of variation, which is precisely what causal inference and disentanglement require.

Identifiable VAEs (iVAEs), introduced by Khemakhem et al. (2020), solve this problem by conditioning the latent prior on an observed auxiliary variable \(u\) — for example, a class label, a time index, or an experimental context. This single change is enough to break the symmetry that causes non-identifiability, and the authors prove rigorously that the resulting model can recover the true latent factors up to a simple coordinate-wise transformation (permutation and element-wise rescaling). In effect, iVAEs bridge deep generative modelling with nonlinear Independent Component Analysis (ICA), inheriting the theoretical guarantees of the latter.

For causal inference, identifiability matters because estimating treatment effects requires that the latent representation capture the correct confounding structure — not just a representation that happens to fit the data. iVAEs provide that assurance.

Where iVAE Fits: CEVAE, iVAE, and CausalVAE Compared

These three notebooks collectively cover the main families of deep generative models for causal inference. They share the VAE backbone but make fundamentally different assumptions about the latent space, and each is the right tool in a different setting. The table below places iVAE in that landscape before we dive into its architecture.

Model	Latent structure	What is learned	Key assumption	Best suited for
CEVAE	Independent \(z\); \(x\) are noisy proxies	Latent confounder recovered from proxy covariates	Ignorability holds in latent space	Settings where the true confounders are unobserved but leave a measurable trace in \(x\)
iVAE	Independent \(z\); prior conditioned on \(u\)	Identifiable latent factors anchored by an auxiliary variable	Sufficient variation in \(u\); injective decoder	Settings where an observed context variable (label, time, environment) can anchor the latent space
CausalVAE	DAG over \(z\); \(\varepsilon\) are independent exogenous noises	Causal graph and structural equations jointly	DAG is acyclic; weak supervision on graph structure	Settings where the causal relationships among latent factors need to be modelled and intervened on

The key conceptual distinction is what each model treats as the source of non-identifiability and how it resolves it:

CEVAE resolves confounding by inferring a latent common cause from proxy measurements. It does not attempt to make \(z\) identifiable in the ICA sense — it is concerned with removing bias, not with uniquely pinning down the latent axes.
iVAE resolves rotational non-identifiability by conditioning the prior on \(u\), giving the latent axes fixed reference points. Each component of \(z\) is independently recoverable, but \(z\) carries no explicit causal graph.
CausalVAE goes further and represents the causal structure among the latent components explicitly as a DAG, enabling do-calculus interventions directly in the latent space. This is the most expressive but also the most demanding in terms of supervision and identifiability requirements.

Reading these three notebooks in sequence — CEVAE → iVAE → CausalVAE — traces a progression from bias removal to identifiable representation to explicit causal structure.

Model Architecture

The iVAE extends the standard VAE by introducing a conditional prior \(p(z \mid u)\) in place of the fixed standard Gaussian prior \(\mathcal{N}(0, I)\). Every component of the model is described below.

Inputs

\(x\) — the observed data (e.g., a sensor reading, a patient record, an image).
\(u\) — an auxiliary variable providing context. In the diagram below this is a one-hot task identifier, but it can be any observed variable that meaningfully segments the data: class labels, experimental conditions, time periods, or treatment assignments.

Encoder \(q(z \mid x, u)\)

Unlike the standard VAE encoder which sees only \(x\), the iVAE encoder receives both \(x\) and \(u\). It passes them through shared hidden layers and outputs two parameter vectors:

\(\mu\) — the mean of the approximate posterior.
\(\log\sigma^2\) — the log-variance of the approximate posterior.

Together, these define the distribution \(q(z \mid x, u) = \mathcal{N}(\mu(x,u),\, \text{diag}(\sigma^2(x,u)))\), a Gaussian whose location and spread depend on the full observed context.

Conditional Prior \(p(z \mid u)\)

This is the architectural innovation that drives identifiability. In a standard VAE the prior \(\mathcal{N}(0,I)\) is completely static — it carries no information about what the data actually is. In the iVAE, a small neural network \(g(u)\) maps the auxiliary variable to a context-specific prior mean:

\[p(z \mid u) = \mathcal{N}(g(u),\, I)\]

Because the prior now shifts as \(u\) changes, different segments of the data are anchored to different regions of the latent space. This is what prevents the latent axes from rotating freely and what ultimately enables identifiability.

In the fully general treatment (Khemakhem et al., 2020) the prior belongs to the exponential family:

\[p(z_i \mid u) = \exp\!\bigl(T_i(z_i) \cdot \lambda_i(u) - \log A(\lambda_i(u))\bigr)\]

where \(T_i(z_i)\) are sufficient statistics and \(\lambda_i(u)\) are natural parameters that vary with \(u\). The Gaussian formulation above is the special case used in most practical implementations, including the one in this notebook.

Reparameterization

To allow gradients to flow through the stochastic sampling step, the model uses the standard reparameterization trick:

\[z = \mu + \sigma \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)\]

This expresses the random sample as a deterministic function of the distributional parameters plus independent noise, making backpropagation through \(z\) straightforward.

Decoder \(p(x \mid z)\)

The decoder maps the latent code \(z\) back to data space to produce a reconstruction \(\hat{x}\). Crucially, the decoder does not see \(u\). All information needed to reconstruct \(x\) must pass through \(z\). This architectural constraint forces the encoder to compress everything about \(x\) that is not already explained by the conditional prior into the latent code.

Loss Function

The iVAE is trained by maximizing the following Evidence Lower Bound (ELBO):

\[\mathcal{L}_{\text{iVAE}} = \mathbb{E}_{q(z \mid x, u)}\!\bigl[\log p(x \mid z)\bigr] - D_{\mathrm{KL}}\!\bigl(q(z \mid x, u) \,\|\, p(z \mid u)\bigr)\]

Comparing this with the standard VAE objective:

\[\mathcal{L}_{\text{VAE}} = \mathbb{E}_{q(z \mid x)}\!\bigl[\log p(x \mid z)\bigr] - D_{\mathrm{KL}}\!\bigl(q(z \mid x) \,\|\, \mathcal{N}(0, I)\bigr)\]

the difference is entirely in the KL term. In the iVAE, the encoder posterior is compared against the context-specific prior \(p(z \mid u)\) rather than a fixed zero-mean Gaussian. This forces latent codes belonging to the same context \(u\) to cluster around the prior mean \(g(u)\), creating distinct, anchored regions in latent space for each value of \(u\).

The reconstruction term is unchanged: it measures how faithfully the decoder recovers \(x\) from \(z\).

iVAE Architecture

Why This Achieves Identifiability

Three conditions together guarantee that the latent representation is identifiable (recoverable up to permutation and element-wise scaling):

Injective mixing function. The map from latent \(z\) to data \(x\) must be one-to-one: no two distinct latent codes can produce the same observation. This prevents the model from “hiding” distinct causal factors inside the same data point.
Sufficient variation in \(u\). The auxiliary variable must genuinely shift the prior across its different values. If \(p(z \mid u)\) were identical for all \(u\), the model would collapse back to an unidentifiable standard VAE. In practice, the number of distinct values of \(u\) must be large enough relative to the latent dimensionality.
Conditional independence of latent components. Given \(u\), the components of \(z\) must be statistically independent. This is what allows each component to be recovered separately without interference from the others.

When all three conditions hold, Khemakhem et al. (2020) prove that the learned encoder recovers the true latent factors — even when the mixing function is a deep nonlinear network.

Implementation in Python

We use pydeepcausalml.generative.IVAE on synthetic multi-class data.

Check and Install Required Python Packages

The following packages are used in this notebook. Install any missing packages with the cell below:

numpy, pandas, scipy, torch, scikit-learn, matplotlib, seaborn, pydeepcausalml

Code

import importlib
import subprocess
import sys

PACKAGES = [
    "numpy", "pandas", "scipy", "torch", "scikit-learn",
    "matplotlib", "seaborn",
]

for pkg in PACKAGES:
    mod = "sklearn" if pkg == "scikit-learn" else pkg
    try:
        importlib.import_module(mod)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pkg])

try:
    import pydeepcausalml  # noqa: F401
except ImportError:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "-q",
         "git+https://github.com/zia207/PyDeepCausalML.git"]
    )

import pydeepcausalml
print("pydeepcausalml", pydeepcausalml.__version__, "ready.")

pydeepcausalml 0.2.0 ready.

Verify imports

Code

import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from pydeepcausalml import set_seed

print("torch:", torch.__version__)
print("cuda:", torch.cuda.is_available())

torch: 2.12.1+cu130
cuda: False

Code

set_seed(42)

Fitting the iVAE on Synthetic Data

Generating Synthetic Data

We construct a synthetic dataset with three class-specific Gaussian clusters in 13-dimensional input space. The class label plays the role of the auxiliary variable \(u\): each class defines a distinct prior mean in the latent space. This setup is deliberately analogous to a multi-environment experiment, where each “environment” or “task” anchors a different region of the latent space — exactly the structure that iVAE identifiability theory exploits.

The data are standardized before training so that each input feature contributes on the same scale to the reconstruction loss.

Model Construction and Training

We instantiate the iVAE with a 4-dimensional latent space — substantially lower than the 13-dimensional input — and 128-unit hidden layers. The latent dimensionality is deliberately small so that the model must find a compact, structured representation rather than simply memorizing the inputs.

Training uses the Adam optimizer with mild weight decay. Gradient clipping is applied to stabilize the early epochs when the KL term can otherwise dominate and prevent the encoder from learning useful representations.

Training and Validation Loss

The plot below shows the negative ELBO over training. A well-behaved run sees both curves decline quickly in the early epochs (reconstruction improving) and then converge smoothly (the KL term settling as the posterior anchors to the conditional prior). A persistent gap between training and validation loss would suggest overfitting to the small synthetic dataset; approximate parity indicates that the model has learned a generalizable latent structure.

Reconstruction Quality

After training, we pass the test set through the full encoder–decoder pipeline and compare the reconstructions \(\hat{x}\) with the original inputs \(x\). The scatter plot on the left overlays original and reconstructed points in the space of the first two input features; the parity plot on the right plots original against reconstructed values directly — perfect reconstruction would align all points on the diagonal.

Latent Space Structure

A key diagnostic for an identifiable model is whether the learned latent space organizes observations by their true generating class. The left panel below plots the first two latent dimensions on the test set, colored by class \(u\). Clear separation indicates that the conditional prior has successfully anchored each class to a distinct region. The right panel shows the same class coloring in the original 13-dimensional input space (projected to the first two features), providing a reference for how much structure already existed in the raw data.

Latent–Feature Correlation Heatmap

The heatmap below reports Pearson correlations between each of the four learned latent dimensions (\(z_1\)–\(z_4\)) and the first five input features on the test set. Strong correlations (positive or negative) indicate that a latent dimension has specialized to capture variation in a particular input feature. In an ideally identifiable model, each latent dimension should align predominantly with one or a small number of independent sources of variation rather than diffusely mixing information from all features.

Load data and prepare synthetic multi-class data

Synthetic multi-class data

Code

n_samples, input_dim, n_aux = 400, 13, 3
rng = np.random.default_rng(42)
n_per = [n_samples // n_aux] * n_aux
n_per[-1] += n_samples - sum(n_per)

X_list, y_list = [], []
for k in range(n_aux):
    mean_k = rng.normal(size=input_dim) * 2 + k * 3
    cov = np.eye(input_dim) * 0.5 + rng.uniform(0, 0.1, (input_dim, input_dim))
    cov = (cov + cov.T) / 2
    X_k = rng.multivariate_normal(mean_k, cov, size=n_per[k])
    X_list.append(X_k)
    y_list.append(np.full(n_per[k], k))

X_raw = np.vstack(X_list)
y = np.concatenate(y_list)
X = (X_raw - X_raw.mean(0)) / (X_raw.std(0) + 1e-8)

train_idx, test_idx = [], []
for cls in np.unique(y):
    idx = np.where(y == cls)[0]
    n_test = int(0.2 * len(idx))
    test_samp = rng.choice(idx, size=n_test, replace=False)
    train_idx.extend(np.setdiff1d(idx, test_samp))
    test_idx.extend(test_samp)

x_train = X[train_idx]
x_test = X[test_idx]
u_train = pd.get_dummies(y[train_idx], dtype=float).to_numpy()
u_test = pd.get_dummies(y[test_idx], dtype=float).to_numpy()
print(f"Synthetic: {X.shape[0]} samples, {input_dim} features, {n_aux} classes")

Synthetic: 400 samples, 13 features, 3 classes

Fit iVAE

Model Construction and Training

Code

from pydeepcausalml.generative import IVAE

ivae = IVAE(
    latent_dim=4,
    hidden=128,
    beta_kl=1.0,
    epochs=60,
    batch_size=32,
    lr=1e-3,
    random_state=42,
    verbose=True,
    log_every=15,
)
ivae.fit(x_train, u_train)
z_test = ivae.transform(x_test)
print("Test latent shape:", z_test.shape)

Test latent shape: (78, 4)

Training and Validation Loss

Training loss

Code

hist = ivae.history_
if hist.get("loss"):
    plt.figure(figsize=(7, 4))
    plt.plot(range(1, len(hist["loss"]) + 1), hist["loss"])
    plt.xlabel("Epoch")
    plt.ylabel("Negative ELBO")
    plt.title("iVAE training loss")
    plt.tight_layout()
    plt.show()

Latent space structure

Code

plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
for k, label in enumerate(np.unique(y[test_idx])):
    mask = y[test_idx] == label
    plt.scatter(z_test[mask, 0], z_test[mask, 1], s=12, alpha=0.7, label=f"Class {label}")
plt.xlabel("z1")
plt.ylabel("z2")
plt.title("Learned latent space")
plt.legend(fontsize=8)

plt.subplot(1, 2, 2)
for k, label in enumerate(np.unique(y[test_idx])):
    mask = y[test_idx] == label
    plt.scatter(x_test[mask, 0], x_test[mask, 1], s=12, alpha=0.7, label=f"Class {label}")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.title("Input space (first 2 features)")
plt.tight_layout()
plt.show()

Latent–feature correlation heatmap

Code

n_show = min(5, input_dim)
corr = np.corrcoef(z_test.T, x_test[:, :n_show].T)[: z_test.shape[1], z_test.shape[1] :]
plt.figure(figsize=(7, 4))
sns.heatmap(
    corr,
    annot=True,
    fmt=".2f",
    cmap="RdBu_r",
    center=0,
    xticklabels=[f"F{j+1}" for j in range(n_show)],
    yticklabels=[f"z{j+1}" for j in range(z_test.shape[1])],
)
plt.title("Pearson correlation: latent dims vs input features")
plt.tight_layout()
plt.show()

Summary and Conclusions

iVAEs extend standard VAEs by replacing the uninformative fixed prior with a conditional prior \(p(z \mid u)\) that depends on an observed auxiliary variable. This single change has profound theoretical consequences: under mild regularity conditions, the learned latent representation is provably identifiable up to permutation and element-wise scaling, meaning the model recovers the true underlying factors of variation rather than an arbitrary rotation of them. For causal inference, this property is essential — estimating average or individual treatment effects requires that the model capture the correct confounding structure, not merely a structure that fits the data.

The synthetic experiment above illustrates the key behaviors to look for in practice:

The negative ELBO should decline smoothly for both training and validation, with the two curves remaining close together.
Reconstructed points should overlap tightly with the originals in input space.
The latent space should organize observations by their auxiliary class, reflecting the distinct prior anchors learned for each value of \(u\).
Each latent dimension should show concentrated correlations with a small subset of input features rather than diffuse correlations across all of them.

Future directions for this framework include applying iVAEs to real-world observational datasets where the auxiliary variable is a treatment indicator or an experimental environment, integrating the identifiable latent representation with downstream causal estimators (e.g., meta-learners or double machine learning), and exploring extensions that relax the requirement for a fully observed \(u\) by learning the auxiliary structure jointly with the generative model.

References

Khemakhem, I., Kingma, D. P., Monti, R. P., & Hyvärinen, A. (2020). Variational Autoencoders and Nonlinear ICA: A Unifying Framework. Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics (AISTATS).
Xie, F., et al. (2024). Causal Effect Estimation using Identifiable Variational AutoEncoder with Latent Confounders and Post-Treatment Variables. arXiv:2408.07219.
Ilkhem Dridi. iVAE: Reference implementation. GitHub.
PyDeepCausalML: Python package for deep-learning-based causal inference