Disentanglement by means of action-induced representations

Learning interpretable representations with variational autoencoders (VAEs) is a major goal of representation learning. The main challenge lies in obtaining disentangled representations, where each latent dimension corresponds to a distinct generative factor. This difficulty is fundamentally tied to the inability to perform nonlinear independent component analysis. Here, we introduce the framework of action-induced representations (AIRs) which models representations of physical systems given experiments (or actions) that can be performed on them. We show that, in this framework, we can provably disentangle degrees of freedom w.r.t. their action dependence. We further introduce a variational AIR architecture (VAIR) that can extract AIRs and therefore achieve provable disentanglement where standard VAEs fail. Beyond state representation, VAIR also captures the action dependence of the underlying generative factors, directly linking experiments to the degrees of freedom they influence.

💡 Research Summary

The paper tackles a fundamental limitation of representation learning with variational autoencoders (VAEs): when the mapping from latent factors of variation to observations is nonlinear, disentangling those factors becomes theoretically ill‑posed because of the impossibility of nonlinear independent component analysis (ICA). To overcome this, the authors introduce the concept of action‑induced representations (AIRs), which explicitly incorporate the actions (or experiments) that are performed on a physical system together with the resulting observations.

An action‑induced dataset consists of pairs ((x, y_A(x))) where (x) is an observation from the underlying data manifold (X) and (y_A(x)) is the deterministic outcome of applying a specific action (or a combination of actions) (A) to the system. Each action typically depends only on a subset of the true latent variables (e.g., measuring gravitational force depends only on mass). By collecting data for several actions, the authors obtain a structured collection of datasets ({D_A}_{A\in\mathcal P_A}).

The AIR framework formalizes this structure. A latent space (Z) together with an encoder (\psi:X\to Z) and, for each action set (A), a projection (\pi_A:Z\to Z_A) (selecting a subset of latent dimensions) and a decoder (\phi_A:Z_A\to Y_A) must satisfy ((\phi_A\circ\pi_A\circ\psi)(x)=y_A(x)). When the representation satisfies two additional constraints—(1) (\psi) is surjective, the total latent dimensionality equals the union of all index sets (I_A), and each (\phi_A) is a continuous bijection, and (2) (Z) is an open subset of (\mathbb R^{d_Z})—the AIR is called a minimal AIR (minAIR). These conditions guarantee that no latent dimension is redundant and that each action uses exactly the dimensions it needs.

The central theoretical result (Theorem 1) shows that for any two minAIRs (one over the true latent variables (C) and one over the learned latent space (Z)), the shared latent dimensions for a given set of actions are bijectively related across the two spaces, independent of the particular action combination. Consequently, latent dimensions that are used by multiple actions become automatically disentangled from those used by only a single action. In other words, the overlap of index sets (I_A) determines which neurons encode factors that are common to several actions, and those neurons are guaranteed to be independent of the others.

To turn this theory into a practical model, the authors propose Variational AIR (VAIR), a VAE variant with two encoders:

• (E_X) processes the observation (x) and outputs only the means (\mu_i) of the latent Gaussian.
• (E_A) processes the action (or action set) (A) and outputs the variances (\sigma_i^2).

During sampling, dimensions with very small (\sigma_i^2) are effectively deterministic (active), while those with large (\sigma_i^2) are heavily noised and thus ignored (passive). This mechanism directly implements the projection (\pi_A) required by the AIR theory: the action encoder decides which latent coordinates are relevant for the current experiment. The decoder receives the sampled latent vector together with the action identifier and reconstructs the corresponding outcome (y_A). Training follows the standard ELBO with a (\beta) weight on the KL term, encouraging a polarized latent space where only the necessary dimensions deviate from the unit Gaussian prior.

The authors evaluate VAIR on three domains:

1. Synthetic benchmark – a 2‑dimensional factor dataset where each factor is observed under different action masks. VAIR achieves near‑perfect alignment between latent dimensions and ground‑truth factors, outperforming β‑VAE, FactorVAE, β‑TCVAE, and other recent disentanglement methods.
2. Classical particle – a point particle with unknown mass (m) and charge (q). Actions correspond to measuring gravitational force (depends only on (m)) and electric force (depends only on (q)). VAIR learns a latent space where one dimension encodes (m) and another encodes (q), exactly matching the theoretical minAIR.
3. Quantum tomography – simulated measurements of a two‑level quantum system under different measurement bases. Each measurement setting is treated as an action. VAIR successfully separates parameters such as Bloch‑sphere angles into distinct latent neurons, whereas standard VAEs produce entangled representations.

Across all experiments, VAIR not only yields disentangled representations but also provides an explicit mapping from actions to the latent dimensions they affect, fulfilling the promise of action‑induced disentanglement.

In conclusion, by integrating experimental actions into the generative model and enforcing the minimal AIR constraints, the paper offers a theoretically grounded and empirically validated solution to the longstanding problem of non‑linear ICA in unsupervised representation learning. The approach opens new avenues for scientific domains where interventions are natural (physics labs, robotics, causal inference) and where interpretable, factor‑wise representations are essential.