Causal Representation Learning

Causal Representation Learning

Causal Representation Learning (CRL) seeks to learn latent representations that correspond to the underlying causal variables generating observed data. Unlike standard representation learning, which optimizes for task performance, CRL aims for representations that are causally disentangled — each latent dimension corresponds to an independent causal factor, and the causal graph among these factors is recovered. Such representations generalize robustly across distribution shifts and enable sample-efficient adaptation to new interventional regimes.

Motivation: Why Causality in Representations?

Standard deep learning models learn statistical patterns that break under distribution shift. A classifier trained to distinguish cows from horses fails when cows appear on beaches (a spurious correlation: cow ↔ green pasture). If the model instead learned the causal factor “animal body shape” independently of “background scene,” it would generalize correctly.

Formally, we observe data $x \sim p(x)$ generated by latent causal variables $z = (z_1, \ldots, z_n)$ via a mixing function $x = f(z)$. The goal is to recover $z$ and the causal graph $\mathcal{G}$ over $z$ from $x$ alone.

Identifiability Theory

A core question in CRL is identifiability: under what conditions can we uniquely recover the causal variables from observations?

Linear ICA (Fully Identifiable)

When the mixing function $f$ is linear and the independent components are non-Gaussian, ICA uniquely identifies the sources (Comon, 1994). This is the gold standard of identifiability.

Nonlinear ICA (Not Identifiable Without Auxiliary Information)

Nonlinear ICA is generally not identifiable from i.i.d. data — many different source distributions and mixing functions can produce the same observed distribution. Hyvärinen & Morioka (2017) showed that auxiliary information (time or segment labels) restores identifiability:

$$\log p(x \mid u) = \sum_{i} \log p_i(s_i \mid u) - \log |\det J_f(x)|$$

where $u$ is an auxiliary variable (e.g., time index, domain label) and $J_f$ is the Jacobian of the mixing function.

Weak Supervision for CRL Identifiability

Recent theoretical results (Khemakhem et al., 2020; Von Kügelgen et al., 2021) show that:

• iVAE (Identifiable VAE): using auxiliary labels $u$ makes nonlinear ICA identifiable
• Interventional data: access to paired data under different interventions recovers causal variables
• Sparse changes assumption: if each domain shift changes only a sparse subset of causal factors, linear identifiability follows

iVAE: Identifiable Variational Autoencoder

import torch
import torch.nn as nn
import torch.nn.functional as F


class iVAE(nn.Module):
    """
    Identifiable VAE (Khemakhem et al., 2020).
    Recovers independent components from observed data with auxiliary labels.
    """

    def __init__(self, obs_dim: int, latent_dim: int, aux_dim: int, hidden: int = 256):
        super().__init__()
        # Encoder: q(z | x, u)
        self.encoder = nn.Sequential(
            nn.Linear(obs_dim + aux_dim, hidden), nn.LeakyReLU(),
            nn.Linear(hidden, hidden), nn.LeakyReLU(),
        )
        self.enc_mu = nn.Linear(hidden, latent_dim)
        self.enc_logvar = nn.Linear(hidden, latent_dim)

        # Decoder: p(x | z)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden), nn.LeakyReLU(),
            nn.Linear(hidden, hidden), nn.LeakyReLU(),
            nn.Linear(hidden, obs_dim),
        )

        # Prior network: p(z | u) — factorized exponential family
        self.prior_mu = nn.Linear(aux_dim, latent_dim)
        self.prior_logvar = nn.Linear(aux_dim, latent_dim)

    def encode(self, x: torch.Tensor, u: torch.Tensor):
        xu = torch.cat([x, u], dim=-1)
        h = self.encoder(xu)
        return self.enc_mu(h), self.enc_logvar(h)

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x: torch.Tensor, u: torch.Tensor):
        mu_q, logvar_q = self.encode(x, u)
        z = self.reparameterize(mu_q, logvar_q)
        x_recon = self.decoder(z)

        # Prior conditioned on auxiliary variable
        mu_p = self.prior_mu(u)
        logvar_p = self.prior_logvar(u)

        # ELBO = reconstruction + KL(q(z|x,u) || p(z|u))
        recon_loss = F.mse_loss(x_recon, x, reduction="sum")
        kl = -0.5 * torch.sum(
            1 + logvar_q - logvar_p - ((mu_q - mu_p) ** 2 + logvar_q.exp()) / logvar_p.exp()
        )
        return recon_loss + kl, x_recon, z

SlowVAE: Temporal Consistency for CRL

SlowVAE (Klindt et al., 2021) leverages the slowness principle — causal factors change slowly over time — to learn identifiable representations from video:

$$\mathcal{L}{\text{slow}} = \mathbb{E}{t}\left[\sum_i \left(\frac{z_i^{(t)} - z_i^{(t-1)}}{\sigma_i}\right)^2\right]$$

The prior factorizes across dimensions with learned temporal autocorrelations, encouraging each latent dimension to represent a slowly-varying causal factor.

class SlowVAEPrior(nn.Module):
    def __init__(self, latent_dim: int):
        super().__init__()
        # Learnable log-rate parameters controlling temporal slowness per dimension
        self.log_rates = nn.Parameter(torch.zeros(latent_dim))

    def log_prob(self, z_t: torch.Tensor, z_prev: torch.Tensor) -> torch.Tensor:
        rates = self.log_rates.exp()
        # Laplace prior on differences
        diff = (z_t - z_prev).abs()
        return -(rates * diff).sum(dim=-1) + self.log_rates.sum()

Causal VAE: Encoding the Causal Graph

CausalVAE (Yang et al., 2021) parameterizes the causal relationships among latent variables explicitly using a structural causal model (SCM):

$$z = Az + \epsilon, \quad \epsilon \sim p(\epsilon)$$

where $A$ is a learned DAG adjacency matrix over latent variables:

class CausalLayer(nn.Module):
    """
    Linear SCM layer: z = (I - A)^{-1} eps
    A is constrained to be a DAG via NOTEARS-style continuous constraint.
    """

    def __init__(self, latent_dim: int):
        super().__init__()
        self.A = nn.Parameter(torch.zeros(latent_dim, latent_dim))

    def dag_constraint(self) -> torch.Tensor:
        """NOTEARS constraint: tr(e^{A ∘ A}) - d == 0 for DAGs."""
        d = self.A.shape[0]
        AoA = self.A * self.A
        return torch.trace(torch.matrix_exp(AoA)) - d

    def forward(self, eps: torch.Tensor) -> torch.Tensor:
        A_masked = self.A * (1 - torch.eye(self.A.shape[0], device=self.A.device))
        I_minus_A_inv = torch.inverse(torch.eye(self.A.shape[0], device=self.A.device) - A_masked)
        return eps @ I_minus_A_inv.T

Connection to OOD Generalization

CRL representations satisfy the Invariant Causal Prediction (ICP) property: the relationship between causal parents $PA_Y$ and the target $Y$ is invariant across environments. This connects to:

• IRM (Invariant Risk Minimization): learns a feature map $\phi$ such that the optimal linear classifier on top of $\phi$ is the same across all training environments
• DomainBed: benchmark evaluating OOD generalization on 7 datasets, where causal methods often outperform ERM

def irm_penalty(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    """IRM gradient penalty — penalizes environment-specific optimal classifiers."""
    scale = torch.ones(1, requires_grad=True).to(logits.device)
    loss = F.cross_entropy(logits * scale, labels)
    grad = torch.autograd.grad(loss, scale, create_graph=True)[0]
    return grad ** 2

Evaluation: Disentanglement Metrics

Summary

Causal Representation Learning bridges representation learning and causal inference, seeking latent spaces whose dimensions correspond to independent causal mechanisms. Identifiability theory has matured significantly — weak supervision via auxiliary labels, paired observations under interventions, or temporal structure can restore unique recoverability in nonlinear settings. As models face increasingly varied deployment environments, causally structured representations offer a principled route to robustness that purely statistical approaches cannot provide.