Disentanglement by means of action-induced representations

TL;DR

The paper tackles the challenge of learning disentangled representations in VAEs by introducing action-induced representations (AIR) and the stronger minimal AIR (minAIR) framework, which exploits experimental actions that couple to subsets of latent factors. It proves a disentanglement theorem showing that latent components shared across actions become disentangled, and presents VAIR, a dual-encoder VAE designed to approximate minAIR and realize action-dependent disentanglement. Across abstract, classical, and quantum physics experiments, VAIR yields interpretable latent factors (e.g., $z=(m, q/m)$ in classical dynamics and Bloch-like coordinates in quantum tomography) and consistently outperforms baselines on disentanglement metrics. By linking actions to specific degrees of freedom, the approach enables interpretable representations and paves the way for integrating active experimentation or reinforcement learning to guide discovery.

Abstract

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.

Figures (6)

• Figure 1: a) Schematic representation of the problem considered. A physical system, uniquely described by its factors of variation $\mathbf{c}$, is observed in a certain state $x$. Some elementary actions $a_i\in\mathbb{A}$ are applied to the system, leading to outcomes $y_i(\mathbf{c})$ that may only depend on a subset of $\mathbf{c}$. b) Diagrammatic overview of the theoretical representation of the problem. See main text for details. c) Top scheme represents the factors of variation $\mathbf{c}$ of a), together with the subsets $I_A$ corresponding to each action ${a_i}$. Because $c_2$ is required for both actions, it is in the intersection $I_{a_1}\cap I_{a_2}$ and by \ref{['thm:disentanglement']} will be disentangled in any minAIR (yellow color). The two lower schemes represent the same picture, here for \ref{['example:thm']} (main text), and for two different action combination sets $P_\mathbb{A} = \{\{1\}, \{2\}, \{1,2\} \}$ and $P_\mathbb{A} = \{\{1\}, \{1,2\} \}$, respectively. d) VAIR: two encoders ($E_X$ and $E_A$) compute the latent neurons' mean and variance from an input sample $x$ and action combination $A$ (or elementary action $a$), respectively. A latent sample $z$ and the $A$ are then fed into the decoder $D$ to compute the output $y_{A}$.
• Figure 2: Abstract experimenta) Factors of variation in the experiment, indicating the subsets $I_a$ associated with each of the considered actions. b) Latent neuron variances $\sigma_i$ output by $E_a$ for the two actions in the experiment, plotted as a function of training epoch. c) Latent neuron means $\mu_i$ output by $E_x$ as a function of the hidden factors $c_i$. Colors correspond to the neurons (as in panel b), while line styles indicate different factors. d) Mutual information gap (MIG) computed for two groups of factors: $c_2$, and all remaining factors excluding $c_2$. Boxplots show results from 20 independent training runs per model, each with random dataset generation and weight's initialization.
• Figure 3: Classical experimenta) Schematic representation of the experiment: we consider objects of certain mass and charge and two actions: an elastic collision with a larger object with fixed mass ($a_M$) or the activation of an electric field via a capacitor ($a_E$). We also consider the additional presence of an external magnetic field $B$ of different coupling strengths $\alpha_{ext}>0$. b) The central panel shows the learned variances by $E_A$ for the two different actions at $\alpha_{\text{ext}} = 0$. The two lower plots showcase the relationship between each latent neuron's $\mu_i$, as a function of the input trajectories' mass and charge. c) Difference between the predicted variances by $E_A$ for each action, as a function of the strength $\alpha_{\text{ext}}$ of the external magnetic field $B$. d) Same as panel b but for $\alpha_{\text{ext}} = 2$.
• Figure 4: Quantum experimenta) Schematic representation of the experiment: a quantum state $\rho$ is repeatedly measured. From the measurement outcomes, the goal is to predict the outcome of a new, given measurement of $\rho$. b) Upper plots: output variances of $E_A$ for different actions in the training set, showing that a single neuron is active for each action. Lower plots: output means of $E_X$ for the same active neurons (dashed) as a function of the tuning parameter $r$. Passive neurons (solid) are also plotted, and remain constant for all values of $r$. c) Output variances of $E_A$ for an action combination $A_N$ with $N=3$ (dark blue). Other colors show the output when each individual action $a^{(n)}$ from the combination is independently fed into $E_A$. d) Output variances of $E_A$ for action combinations $A_N$, as a function of the parameter $v^{(n)}$, restricted to neurons active for the respective $a^{(n)}$. Dashed line shows the scaling $\log(\sigma^2)\propto - v^{(n)}$
• Figure 5: Schematic representation of VAE benchmark architectures a) TC-VAE$_x$, having as input and output the observation $x$. b) VAE$_{x,a}$ having as input the observation $x$ and $a$, and predicting the output $y$. Moreover, we further consider two variants, VAE$_{D_a}$ and TC-VAE$_{D_a}$, where the action is also input to the decoder.

Theorems & Definitions (17)

• Example 1
• Definition 1
• Definition 2
• Definition 3
• Example 2
• Theorem 1
• Example 3
• Definition 4
• Theorem 2: Action-induced disentanglement, formal and detailed
• proof