5.2.6 Causal Encoding Generative Modeling (CausalEGM)

Note: CausalEGM requires PyTorch. The CausalEGM estimator in pydeepcausalml.generative learns a tripartite latent decomposition for high-dimensional causal inference.

CausalEGM is a deep learning framework for nonlinear dimensionality reduction and structured causal representation learning. Its defining contribution is a causally partitioned latent space: rather than encoding all covariates into a single unstructured vector, CausalEGM forces the encoder to separate confounders (variables that affect both treatment and outcome), instruments (variables that affect treatment only), and prognostic factors (variables that affect outcome only). This tripartite decomposition — learned end-to-end from data — allows the model to identify the minimal causal subset of the high-dimensional covariate space and estimate treatment effects accurately even when the number of covariates is large relative to the sample size.

CausalEGM was introduced in Chen et al. (2023), “CausalEGM: a general causal model using encoder-decoder with generative model”, with a Bayesian extension CausalBGM appearing in 2025. It is directly applicable to genomics, electronic health records, and neuroimaging — settings where thousands of covariates are measured but only a small structured subset drives the causal mechanism.

Where CausalEGM Fits

Each notebook in this series has introduced a model with progressively richer covariate structure. CausalEGM’s novel contribution is operating natively in high-dimensional covariate spaces while maintaining explicit causal partitioning:

Model	Covariate handling	Latent structure	Continuous treatment
CEVAE	Raw \(x\) as proxy inputs	Single unstructured \(z\)	No
iVAE	Raw \(x\) + auxiliary \(u\)	Independent factors anchored by \(u\)	No
CausalVAE	Raw \(x\)	DAG over \(z\)	No
CD-VAE	Raw \(x\); MMD alignment	Unstructured \(z\)	No
DSCM	Raw \(x\); per-node modules	Full SCM in \(z\)	No
CausalEGM	Nonlinear compression of high-dim \(x\)	Tripartite \(z = [z_c \mid z_t \mid z_y]\)	Yes — full dose-response curve

Two capabilities distinguish CausalEGM from every preceding model: it operates after nonlinear compression of potentially thousands of covariates (making it suitable for genomics-scale data), and it natively handles continuous treatments by estimating the full dose-response curve \(\mathbb{E}[Y(t)]\) for any \(t \in \mathbb{R}\).

Overview: The High-Dimensional Confounding Problem

All the models covered so far (TARNet, CFRNet, CEVAE, DSCM) work well when the confounder space is low-dimensional and tractable. Real observational data — genomics, electronic health records, fMRI scans — routinely involves thousands of covariates. Conditioning directly on all of them causes the curse of dimensionality: propensity scores become unreliable in high-dimensional \(x\)-space, matching breaks down, and neural networks overfit. Even MMD-based balancing (CD-VAE) becomes statistically unreliable when comparing distributions in 1000-dimensional space with limited sample sizes.

The classical answer to this is sufficient dimension reduction (SDR): find a low-dimensional function of \(x\) that preserves all information relevant for identifying causal effects. Classical SDR methods are linear (PCA, PLS), which limits their applicability to the nonlinear relationships ubiquitous in biological and clinical data. CausalEGM generalizes SDR to the nonlinear case by parameterizing the compression with deep networks, and adds a critical structural constraint: the compressed representation must be causally organized.

Model Architecture

The Tripartite Latent Decomposition

The architectural innovation of CausalEGM is a forced partition of the latent code \(z\) into three disjoint, statistically independent components:

\[z = [\underbrace{z_c}_{\text{confounders}} \mid \underbrace{z_t}_{\text{instruments}} \mid \underbrace{z_y}_{\text{prognostic}}]\]

\(z_c \in \mathbb{R}^{d_c}\): the confounder subspace — variables that influence both treatment assignment and outcome. These are the variables that must be adjusted for to identify causal effects. Their dimension \(d_c\) is typically small (4–16) relative to the original covariate space.
\(z_t \in \mathbb{R}^{d_t}\): the instrument-like subspace — variables that shift treatment probability but have no direct effect on outcomes. Analogous to instrumental variables in econometrics; they provide identifying variation without introducing outcome-confounding.
\(z_y \in \mathbb{R}^{d_y}\): the prognostic subspace — variables that predict outcomes regardless of treatment but do not cause selection into treatment. Including them in the outcome model improves precision without introducing bias.

Independence among the three components is enforced by a discriminator-based mutual information penalty (an adversarial disentanglement loss \(\mathcal{L}_{\text{disent}}\)) that penalizes any statistical dependence between \(z_c\), \(z_t\), and \(z_y\).

Bidirectional Encoder-Decoder

CausalEGM establishes a bidirectional mapping between the high-dimensional covariate space and the structured latent space:

\[f: x \rightarrow z = [z_c, z_t, z_y] \quad \text{(encoder)}\]

\[g: z \rightarrow \hat{x} \quad \text{(decoder / generator)}\]

The bidirectionality is enforced by a reconstruction loss \(\mathcal{L}_{\text{recon}} = \|x - g(f(x))\|^2\), which ensures that the compression is lossless: the full information in \(x\) can be recovered from \(z\). The latent space is regularized toward a standard Gaussian prior via a KL term, matching the iVAE framework but with the tripartite structure imposed on top.

Treatment and Outcome Models

With the latent space partitioned, CausalEGM builds:

\[p(T \mid x) \approx p(T \mid z_c, z_t) \quad \text{(treatment model / propensity score in latent space)}\]

\[\hat{Y}(0) = H_0(z_c, z_y), \quad \hat{Y}(1) = H_1(z_c, z_y) \quad \text{(binary treatment)}\]

\[\hat{Y}(t) = H(t, z_c, z_y) \quad \text{(continuous treatment — dose-response)}\]

Crucially, the outcome heads receive only \((z_c, z_y)\) — not \(z_t\). This is what makes the instrument variables do their job: they are available for propensity estimation but blocked from the outcome model, preventing them from introducing spurious outcome variation. The treatment model similarly receives only \((z_c, z_t)\) — not \(z_y\) — ensuring prognostic factors do not confound the propensity score.

Joint Training Objective

CausalEGM minimizes a composite loss end-to-end:

\[\mathcal{L} = \lambda_{\text{recon}} \underbrace{\|x - g(f(x))\|^2}_{\text{reconstruction}} + \lambda_{\text{treat}} \underbrace{\mathcal{L}_{\text{treat}}(T, \hat{T})}_{\text{treatment prediction}} + \lambda_{\text{out}} \underbrace{\mathcal{L}_{\text{out}}(Y, \hat{Y})}_{\text{outcome prediction}} + \underbrace{\mathcal{L}_{\text{KL}} + \mathcal{L}_{\text{disent}}}_{\text{latent structure}}\]

Loss term	Purpose	Effect of increasing weight
\(\mathcal{L}_{\text{recon}}\)	Faithful compression of \(x\) into \(z\)	More information preserved; less regularization
\(\mathcal{L}_{\text{treat}}\)	Accurate propensity scores from \((z_c, z_t)\)	Stronger treatment-relevant signal in \(z_c, z_t\)
\(\mathcal{L}_{\text{out}}\)	Accurate potential outcome prediction	Stronger outcome-relevant signal in \(z_c, z_y\)
\(\mathcal{L}_{\text{disent}}\)	Independence among \(z_c\), \(z_t\), \(z_y\)	Cleaner partition; risk of under-utilising correlated features

The adversarial disentanglement loss uses a discriminator network trained to distinguish the three components; the encoder is then trained adversarially to fool the discriminator, driving the components toward statistical independence.

Why This Outperforms Prior Methods

CausalEGM’s advantage is largest when sample size is large and covariate dimension is high — precisely the settings where the curse of dimensionality hurts other methods most. Compared to models already seen in this series:

vs CEVAE: CEVAE infers a single unstructured latent \(z\) — it does not partition into confounders, instruments, and prognostic factors. The partition itself is the causal structure; without it, the downstream propensity and outcome models must operate in a mixed-signal latent space.
vs TARNet/CFRNet: These operate directly in the high-dimensional raw \(x\) space. MMD balancing is statistically unreliable in 1000-dimensional space.
vs classical SDR: Classical sufficient dimension reduction is linear. CausalEGM captures nonlinear dependencies via the deep encoder.
vs all preceding models: None of CEVAE, iVAE, CausalVAE, CD-VAE, or DSCM handle continuous treatments natively. CausalEGM directly estimates the dose-response curve \(\mathbb{E}[Y(t)]\) for any \(t \in \mathbb{R}\), making it the natural choice for pharmaceutical dose studies, environmental exposure analyses, and policy evaluation with continuous interventions.

Known Limitation: Cyclical Dependency

The bidirectional encoder-decoder introduces a structural cycle — the encoder’s output feeds back into the decoder’s input in the reconstruction loss, and both depend on each other’s gradients. This circularity technically violates the acyclicity assumption fundamental to causal DAGs and Bayesian networks. CausalEGM also uses deterministic encoder functions, which limits its ability to quantify uncertainty in the latent representation. These limitations motivated CausalBGM (2025), which replaces the deterministic encoder with a fully Bayesian posterior and restores a clean DAG structure.

Implementation in Python

We fit CausalEGM with pydeepcausalml.generative.CausalEGM on IHDP.

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())
from sklearn.linear_model import LogisticRegression

torch: 2.12.1+cu130
cuda: False

Code

set_seed(42)
run_fast = True

Global settings and hyperparameters

The run_fast = TRUE setting uses compact networks and fewer epochs for rendering. Set FALSE for the full configuration matching the published benchmark.

Load IHDP data

Build matrices and split train / val / test

We use a three-way split (70 / 15 / 15) so early-stopping validation and final evaluation are fully independent.

Feature scaling

Propensity score overlap diagnostic

Before fitting any causal model, verify that treated and control units share sufficient covariate support. CausalEGM’s propensity model operates in the low-dimensional latent space \((z_c, z_t)\), but the raw covariate overlap is a prerequisite for any method.

Fitting CausalEGM

Train the model

The dim_c, dim_t, and dim_y arguments control the dimensionality of each latent subspace. On IHDP with 25 covariates, small values (4–8 each) are sufficient; on genomics data with thousands of covariates these would typically be set to 16–64 each.

Training Diagnostics

Loss components over training

CausalEGM trains five losses simultaneously. A healthy run shows all five declining monotonically after an initial settling phase. The disentanglement loss may oscillate slightly as the discriminator and encoder reach an adversarial equilibrium — this is normal and expected.

Treatment Effect Evaluation

Potential outcomes: predicted vs ground truth

ITE scatter and distribution

ITE calibration by prediction decile

Latent Space Diagnostics

Learned propensity scores

The propensity scores produced by CausalEGM are computed in the low-dimensional latent space \((z_c, z_t)\) rather than in the raw 25-dimensional covariate space. Well-separated treated/control distributions indicate that \(z_c\) and \(z_t\) have captured the treatment assignment mechanism. Substantial overlap confirms that common support exists in the learned representation, which is required for valid causal identification.

Confounder subspace: PCA visualization

The confounder latent subspace \(z_c\) should organize observations in a way that explains treatment heterogeneity. If disentanglement has succeeded, the first two principal components of \(z_c\) should show the treatment groups mixing substantially — the confounder information is shared, not treatment-specific.

Dose-Response Estimation (Continuous Treatment)

One of CausalEGM’s unique capabilities — unavailable in any preceding model in this series — is estimating the full dose-response curve \(\mathbb{E}[Y(t)]\) for a continuous treatment \(t \in \mathbb{R}\). We demonstrate this on a synthetic dataset where the true curve is known, allowing us to evaluate estimation accuracy across the entire treatment range.

The dose-response function \(H(t, z_c, z_y)\) is fit jointly with the rest of the model; at inference time it is evaluated at any desired treatment value by passing the encoded \((z_c, z_y)\) through the continuous outcome head.

Load IHDP data and prepare (X, t, y)

Load IHDP

Code

def load_ihdp(replications: int = 2, random_state: int = 1):
    """Load IHDP semi-synthetic benchmark (CausalML format)."""
    base_url = "https://raw.githubusercontent.com/uber/causalml/master/docs/examples/data"
    cols = ["treatment", "y_factual", "y_cfactual", "mu0", "mu1"] + [f"X{i}" for i in range(25)]
    parts = []
    for i in range(1, 10):
        url = f"{base_url}/ihdp_npci_{i}.csv"
        tmp = pd.read_csv(url, header=None)
        tmp.columns = cols[: tmp.shape[1]]
        parts.append(tmp)
    df = pd.concat(parts, ignore_index=True)
    if replications > 1:
        df = pd.concat([df] * replications, ignore_index=True)
    perm = list(range(6, 25)) + list(range(6))
    xcols = [f"X{i}" for i in range(25)]
    X = df[xcols].to_numpy(dtype=float)[:, perm]
    treatment = df["treatment"].to_numpy(dtype=int)
    y = df["y_factual"].to_numpy(dtype=float)
    mu0 = df["mu0"].to_numpy(dtype=float)
    mu1 = df["mu1"].to_numpy(dtype=float)
    tau = np.where(
        treatment == 1,
        df["y_factual"] - df["y_cfactual"],
        df["y_cfactual"] - df["y_factual"],
    )
    n = len(X)
    rng = np.random.default_rng(random_state)
    val_idx = rng.choice(n, size=int(0.2 * n), replace=False)
    train_idx = np.setdiff1d(np.arange(n), val_idx)
    return df, X, treatment, y, tau, mu0, mu1, train_idx, val_idx


def preprocess_ihdp_features(train_x, test_x, cont_cols=None):
    """Scale continuous covariates (cols 19–24 after binary-first permutation)."""
    cont_cols = cont_cols or list(range(19, 25))
    train = train_x.copy()
    test = test_x.copy()
    means = train[:, cont_cols].mean(axis=0)
    sds = train[:, cont_cols].std(axis=0)
    sds[sds == 0] = 1.0
    train[:, cont_cols] = (train[:, cont_cols] - means) / sds
    test[:, cont_cols] = (test[:, cont_cols] - means) / sds
    return train, test


def synthetic_ihdp_fallback(n=5000, p=25, random_state=42):
    rng = np.random.default_rng(random_state)
    X = rng.normal(size=(n, p))
    lin = 0.3 * X[:, 0] - 0.2 * X[:, 1] + 0.15 * X[:, 2]
    tau = 0.5 + 0.2 * np.tanh(X[:, 4])
    mu0 = lin + 0.1 * X[:, 6] ** 2
    mu1 = mu0 + tau
    ps = 1 / (1 + np.exp(-(0.2 * X[:, 0] - 0.1 * X[:, 1])))
    t = rng.binomial(1, ps)
    y = np.where(t == 1, mu1, mu0) + rng.normal(scale=1.0, size=n)
    perm = list(range(6, 25)) + list(range(6))
    return X[:, perm], t, y, mu0, mu1, tau

try:
    _, X, treatment, y, tau, mu0, mu1, _, _ = load_ihdp(replications=2 if run_fast else 10)
except Exception:
    X, treatment, y, mu0, mu1, tau = synthetic_ihdp_fallback()

n = len(X)
rng = np.random.default_rng(42)
idx = rng.permutation(n)
train_n, val_n = int(0.7 * n), int(0.85 * n)
tr_i, te_i = idx[:train_n], idx[val_n:]

X_train, t_train, y_train = X[tr_i], treatment[tr_i], y[tr_i]
X_test, mu0_te, mu1_te = X[te_i], mu0[te_i], mu1[te_i]
X_train_s, X_test_s = preprocess_ihdp_features(X_train, X_test)
print(f"Train: {len(tr_i)} | Test: {len(te_i)}")

Train: 3500 | Test: 750

Propensity score overlap diagnostic

Propensity overlap

Code

ps = LogisticRegression(max_iter=1000).fit(X_train_s, t_train).predict_proba(X_train_s)[:, 1]
plt.figure(figsize=(7, 4))
for m, c, lab in [(t_train == 0, "#185FA5", "Control"), (t_train == 1, "#993C1D", "Treated")]:
    sns.kdeplot(ps[m], fill=True, alpha=0.35, color=c, label=lab)
plt.xlabel("P(T=1 | X)")
plt.title("Propensity score overlap (train)")
plt.legend()
plt.tight_layout()
plt.show()

Fit CausalEGM

Code

from pydeepcausalml.generative import CausalEGM

egm = CausalEGM(
    dim_c=8,
    dim_t=4,
    dim_y=4,
    hidden=128,
    lambda_recon=1.0,
    lambda_treat=2.0,
    lambda_outcome=2.0,
    lambda_disent=0.5,
    epochs=60 if run_fast else 100,
    batch_size=128,
    lr=1e-3,
    random_state=42,
    verbose=True,
    log_every=10,
)
egm.fit(X_train_s, t_train, y_train)

CausalEGM(epochs=60, lr=0.001, batch_size=128, weight_decay=0.0001, early_stopping_patience=None, device=None)

Training diagnostics

Training loss components

Code

hist = egm.history_
cols = [c for c in ["loss", "recon", "treat", "outcome", "disent"] if c in hist]
if cols:
    fig, axes = plt.subplots(1, len(cols), figsize=(3 * len(cols), 3))
    if len(cols) == 1:
        axes = [axes]
    epochs = range(1, len(hist[cols[0]]) + 1)
    for ax, c in zip(axes, cols):
        ax.plot(epochs, hist[c])
        ax.set_title(c)
        ax.set_xlabel("Epoch")
    plt.suptitle("CausalEGM training losses")
    plt.tight_layout()
    plt.show()

Treatment effect evaluation

Code

ite_pred = egm.predict_ite(X_test_s)
ite_true = mu1_te - mu0_te
ate_true, ate_pred = ite_true.mean(), ite_pred.mean()
pehe = np.sqrt(np.mean((ite_pred - ite_true) ** 2))
metrics = pd.DataFrame({
    "Metric": ["True ATE", "Predicted ATE", "ATE bias", "Abs. ATE error", "sqrt(PEHE)"],
    "Value": [ate_true, ate_pred, ate_pred - ate_true, abs(ate_pred - ate_true), pehe],
})
display(metrics.round(4))

	Metric	Value
0	True ATE	0.5023
1	Predicted ATE	0.5052
2	ATE bias	0.0029
3	Abs. ATE error	0.0029
4	sqrt(PEHE)	0.1953

Potential outcomes and ITE plots

Code

# Potential outcomes via outcome head at t=0/1
with torch.no_grad():
    zc, _, zy = egm._encode(X_test_s)
    ones = torch.ones(len(zc), 1, device=egm._device)
    zeros = torch.zeros(len(zc), 1, device=egm._device)
    y1 = egm.module_.outcome_head(torch.cat([zc, zy, ones], dim=-1)).squeeze(-1).cpu().numpy() * egm.y_std_ + egm.y_mean_
    y0 = egm.module_.outcome_head(torch.cat([zc, zy, zeros], dim=-1)).squeeze(-1).cpu().numpy() * egm.y_std_ + egm.y_mean_

fig, axes = plt.subplots(1, 2, figsize=(12, 5))
axes[0].scatter(mu0_te, y0, alpha=0.45, s=10, color="#185FA5", label="Y(0)")
axes[0].scatter(mu1_te, y1, alpha=0.45, s=10, color="#993C1D", label="Y(1)")
lims = [min(mu0_te.min(), y0.min()), max(mu1_te.max(), y1.max())]
axes[0].plot(lims, lims, "k--")
axes[0].set_xlabel("Ground truth")
axes[0].set_ylabel("Predicted")
axes[0].legend()
axes[0].set_title("Potential outcomes")

axes[1].scatter(ite_true, ite_pred, alpha=0.4, s=10, color="#534AB7")
axes[1].plot([ite_true.min(), ite_true.max()], [ite_true.min(), ite_true.max()], "k--")
axes[1].set_xlabel("True ITE")
axes[1].set_ylabel("Predicted ITE")
axes[1].set_title(f"ITE scatter (PEHE={pehe:.3f})")
plt.tight_layout()
plt.show()

Latent space diagnostics

Learned propensity scores in latent space

Code

ps_latent = egm.predict_propensity(X_train_s)
plt.figure(figsize=(7, 4))
for m, c, lab in [(t_train == 0, "#185FA5", "Control"), (t_train == 1, "#993C1D", "Treated")]:
    sns.kdeplot(ps_latent[m], fill=True, alpha=0.35, color=c, label=lab)
plt.xlabel("P(T=1 | z_c, z_t)")
plt.title("Learned propensity (latent space)")
plt.legend()
plt.tight_layout()
plt.show()

Summary and Conclusions

CausalEGM addresses the high-dimensional confounding problem that limits all preceding models in this series. By learning a nonlinear compression of the covariate space and simultaneously partitioning the compressed representation into causally distinct subspaces — confounders \(z_c\), instruments \(z_t\), and prognostic factors \(z_y\) — it provides both accurate treatment effect estimation and interpretable latent structure. The tripartite decomposition is the causal structure itself, not just a post-hoc interpretation of an unstructured embedding.

Key takeaways from this notebook:

The propensity overlap check on the raw covariates is a necessary precondition, but CausalEGM’s propensity model operates in the compressed latent space \((z_c, z_t)\). Comparing the raw and latent propensity distributions reveals how much the compression changes the effective overlap.
The disentanglement loss is adversarial and may oscillate during training. If it fails to converge downward, increasing lambda_disent or reducing learning_rate_disc (to slow the discriminator) typically stabilizes it.
The confounder PCA visualization is the primary diagnostic for latent structure quality: treated and control groups should overlap in \(z_c\) if the disentanglement has succeeded, confirming that the confounder subspace captures shared variation rather than treatment-specific patterns.
The dose-response curve is CausalEGM’s most distinctive capability. On the synthetic example, the estimated curve should track the true quadratic relationship closely; large deviations suggest the outcome head needs more capacity or the lambda_outcome weight needs increasing.
CausalBGM (2025) is the natural successor when uncertainty quantification is needed, or when the cyclical encoder-decoder structure is a theoretical concern. The implementation in pydeepcausalml follows the same API.

References

Chen, B., et al. (2023). CausalEGM: a general causal model using encoder-decoder with generative model. PNAS Nexus, 2(8).
Guo, R., et al. (2025). CausalBGM: Bayesian Generative Model for Causal Inference. arXiv.
pydeepcausalml repository
CausalML documentation