Causal representation learning via mixture-based priors

In brief

• Causal Representation Learning (CRL) seeks to uncover the true underlying causal variables from high-dimensional pixels without access to interventional data.
• Existing identifiability theorems often demand unrealistic functional restrictions or supervised environment labels.
• This work partitions the latent space into causal factors (modeled via invertible diffeomorphic flows) and exogenous noise, proving identifiability using mixture-based priors.

This retrospective explainer is dated to the May 2025 ICLR Deep Generative Models Workshop. It discusses Causal Representation Learning and Inference via Mixture-Based Priors by Avinash Kori, Carles Balsells-Rodas, Ben Glocker, Yingzhen Li and Francesco Locatello.

Imagine listening to an orchestra recording. The music is a single continuous audio waveform made of millions of pressure samples. Yet your brain effortlessly separates the underlying causal score—the melody played by the violin section, which triggers the entrance of the cello—from the acoustic acoustics of the concert hall and random audience rustling.

In machine learning, Causal Representation Learning (CRL) attempts to do something similar for visual data: extracting the hidden variables that drive a scene (e.g., light source angle, pendulum elevation, object velocity) directly from raw pixels, while identifying the directed causal relationships between them.

The central hurdle is identifiability. Without strong constraints or active interventional experiments, infinitely many different mathematical transformations can reconstruct an image equally well, but only one corresponds to the true physical causes.

This paper introduces a principled framework that achieves identifiability by partitioning the latent space and modeling causal mechanisms through mixture-based priors and diffeomorphic flows.

The fundamental challenge of unsupervised CRL

Let $x \in \mathcal{X}$ be an observed high-dimensional image generated by a non-linear mixing function $g: \mathcal{Z} \to \mathcal{X}$ acting on latent variables $z \in \mathcal{Z}$:

\[x = g(z).\]

Standard non-linear Independent Component Analysis (ICA) proves that if $z$ consists of independent coordinates and $g$ is an unconstrained neural network, $z$ is fundamentally unidentifiable from observational data alone. A model can easily find an arbitrary rotation $R$ such that $\tilde{z} = R(z)$ yields the identical marginal distribution $p(x)$.

When the latent variables $z$ are causally dependent (e.g., $z_1 \to z_2$), the challenge becomes even steeper: we must not only recover the coordinates of $z$, but also discover which variables cause which.

Structured decomposition: causal factors vs exogenous noise

The paper resolves this by partitioning the latent vector $z$ into two distinct functional components:

\[z = (\mathbf{c}, \boldsymbol{\epsilon}),\]

where:

1. $\mathbf{c} = (c_1, \dots, c_d) \in \mathbb{R}^d$ represents the causal factors of interest (the endogenous variables).
2. $\boldsymbol{\epsilon} = (\epsilon_1, \dots, \epsilon_m) \in \mathbb{R}^m$ represents independent exogenous noise accounting for unexplained variability.

The causal factors are structured through a Structural Causal Model (SCM):

\[c_i = f_i\left(\mathbf{pa}(c_i), \epsilon_i\right), \quad i = 1, \dots, d,\]

where $\mathbf{pa}(c_i) \subseteq {c_1, \dots, c_{i-1}}$ in some topological causal order.

To model the non-linear mechanisms $f_i$ while maintaining exact probability density evaluation, the authors implement $f$ as a triangular diffeomorphic flow. The Jacobian matrix of this transformation is lower-triangular, reflecting the directed acyclic structure of the causal graph:

\[J_f(\mathbf{c}) = \begin{bmatrix} \frac{\partial f_1}{\partial c_1} & 0 & \dots & 0 \\ \frac{\partial f_2}{\partial c_1} & \frac{\partial f_2}{\partial c_2} & \dots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial f_d}{\partial c_1} & \frac{\partial f_d}{\partial c_2} & \dots & \frac{\partial f_d}{\partial c_d} \end{bmatrix}.\]

Mixture-based priors for identifiability

To prove that the true causal factors $\mathbf{c}$ and exogenous terms $\boldsymbol{\epsilon}$ can be uniquely identified, the framework conditions the latent variables on an observed or inferred auxiliary context variable $u$ (such as temporal time-steps or multi-environment indicators):

\[p(\mathbf{c}, \boldsymbol{\epsilon} \mid u) = \sum_{k=1}^K \pi_k(u) \, p_k(\mathbf{c}) \, p_k(\boldsymbol{\epsilon}).\]

By modeling the conditional prior as a Gaussian mixture whose component means and variances vary across contexts, the characteristic functions of the latent components become distinguishable. The paper proves that under mild full-rank conditions on the mixing coefficients $\pi_k(u)$:

1. The mapping $g$ is identified up to component-wise diffeomorphisms.
2. The causal ordering of the endogenous factors $\mathbf{c}$ is uniquely identified.
3. The exogenous noise variables $\boldsymbol{\epsilon}$ are isolated from the causal mechanisms.

This structured decomposition provides a scalable, mathematically grounded blueprint for learning causal representations without requiring costly physical interventions.

Paper and sources

• Avinash Kori, Carles Balsells-Rodas, Ben Glocker, Yingzhen Li and Francesco Locatello. Causal Representation Learning and Inference via Mixture-Based Priors. ICLR 2025 Workshop on Deep Generative Models.