Technical Summary: Nonparametric Partial Disentanglement via Mechanism Sparsity

Problem Setting

This work addresses the problem of identifiable representation learning within the framework of Causal Representation Learning (CRL). The goal is to recover high-level, semantically meaningful latent factors ($z$) from high-dimensional observations ($x$) and to uncover the causal graphical model (CGM) that governs their dynamics.

The authors consider a generative model where observations are generated by a potentially highly nonlinear mixing function (decoder) $f$ acting on latent factors $z$, corrupted by noise:
$$x_t = f(z_t) + n_t$$
The latent factors evolve over time according to a transition model $p(z_t | z_{<t}, a_{<t})$, where $a_t$ represents observed auxiliary variables (e.g., actions, environment indices, or intervention indicators). Crucially, the latent factors are assumed to be conditionally independent given their parents in a sparse causal graph $G$ (comprising $G_z$ for temporal dependencies and $G_a$ for auxiliary dependencies).

The core challenge is identifiability: without additional assumptions, the latent factors and the mixing function are unidentifiable up to arbitrary diffeomorphisms. While previous works have established identifiability under specific conditions (e.g., exponential family assumptions, known intervention targets, or strict graphical criteria), this paper seeks to relax these constraints to handle nonparametric settings and arbitrary latent causal graphs, potentially allowing for partial disentanglement.

Methodology

1. Theoretical Framework: Mechanism Sparsity Regularization

The paper introduces mechanism sparsity regularization as a new inductive bias. The central hypothesis is that the mechanisms governing the latent dynamics are sparse:

• Sparse Auxiliary Influence: Auxiliary variables $a$ affect only a few latent factors (sparse $G_a$).
• Sparse Temporal Dependencies: Latent factors depend sparsely on their past states (sparse $G_z$).

The authors formalize the relationship between the ground-truth model $\theta = (f, p, G)$ and a learned model $\hat{\theta} = (\hat{f}, \hat{p}, \hat{G})$. They define an entanglement map $v = f^{-1} \circ \hat{f}$, which maps the learned representation to the ground-truth representation.

2. Identifiability Theory

The paper develops a nonparametric identifiability theory based on the sparsity of the learned graphs $\hat{G}_a$ and $\hat{G}_z$.

• Equivalence Relations: Instead of guaranteeing identifiability up to permutation and element-wise transformation (complete disentanglement), the theory establishes identifiability up to consistency:

  • $a$-consistency ($\sim_{a}^{con}$): The entanglement map $v$ is a composition of a permutation and a $G_a$-preserving diffeomorphism.
  • $z$-consistency ($\sim_{z}^{con}$): The entanglement map $v$ is a composition of a permutation and a $G_z$- (and $G_z^\top$-) preserving diffeomorphism.
  • Partial Disentanglement: This implies that while some latent factors may remain entangled (mixed), the structure of this entanglement is strictly constrained by the sparsity pattern of the ground-truth graph.

• Sufficient Influence Assumptions: To prove these results, the authors introduce "sufficient influence" assumptions (Assumptions 6, 7, 8). These require that the Hessian (for continuous $a$) or partial differences (for discrete $a$) of the log-density of the transition model vary sufficiently to span the relevant subspaces defined by the graph structure.

• Graph Preserving Maps: A key mathematical tool is the concept of $G$-preserving maps, functions whose dependency graphs respect the sparsity constraints of $G$. The theory shows that if the learned graph is regularized to be as sparse as the ground truth, the entanglement map must be a $G$-preserving map.

• Complete Disentanglement: The paper recovers the stronger result of complete disentanglement (identifiability up to permutation and element-wise transformation) as a special case when the ground-truth graph satisfies a specific graphical criterion (Assumption 5), which generalizes previous criteria from Lachapelle et al. (2022).

3. Extension to Exponential Families

The paper connects its nonparametric results to the popular exponential family assumption. By assuming the transition model belongs to an exponential family with a one-dimensional sufficient statistic, the authors show that the "sufficient influence" assumptions can be relaxed, allowing for Markovian transitions without auxiliary variables (a limitation in the nonparametric case).

4. Estimation Procedure

To implement these ideas, the authors propose a learning algorithm based on Variational Autoencoders (VAEs):

• Model: A VAE learns the decoder $\hat{f}$ and the transition model $\hat{p}$.
• Sparsity Constraint: Instead of using a sparsity penalty (regularization), the method employs a sparsity constraint. The adjacency matrices $\hat{G}_a$ and $\hat{G}_z$ are treated as parameters of Bernoulli distributions, and the optimization is constrained such that the expected number of edges does not exceed a budget $\beta$ (ideally set to the ground-truth edge count).
• Optimization: The constrained optimization is solved using a Lagrangian formulation with dual restarts.

Key Contributions

1. Novel Principle: Introduction of mechanism sparsity regularization as a principle for disentanglement, applicable when latent factors depend sparsely on auxiliary variables or past states.
2. Nonparametric Identifiability: Development of a fully nonparametric identifiability theory that does not require the exponential family assumption, allowing for arbitrary latent causal graphs (provided no instantaneous causal links exist between latents at the same time step).
3. Partial Disentanglement Theory: Formalization of partial disentanglement via new equivalence relations ($a$-consistency and $z$-consistency). The theory precisely predicts the structure of the "entanglement graph" (which variables remain mixed) based on the ground-truth graph structure.
4. Generalization of Interventions: Demonstration that the framework handles multi-node interventions with unknown targets, a setting where previous methods were often limited to single-node interventions or known targets.
5. New Evaluation Metric: Introduction of $R_{con}$, a metric to quantify how close a learned representation is to being $a$-consistent or $z$-consistent, bridging the gap between complete disentanglement (MCC) and general representation quality ($R$).
6. Empirical Validation: Implementation of a sparsity-constrained VAE and validation on synthetic datasets, showing that enforcing sparsity improves both graph recovery (SHD) and disentanglement (MCC, $R_{con}$).

Results

• Theoretical: Theorems 1, 2, and 3 establish that under sufficient influence and sparsity regularization, the learned model is identifiable up to $a$-consistency, $z$-consistency, or $(a,z)$-consistency. Proposition 7 shows that if the graphical criterion holds, this reduces to complete disentanglement.
• Experimental:
  • On datasets satisfying the graphical criterion (e.g., diagonal or specific non-diagonal graphs), the sparsity-constrained VAE achieves near-perfect Mean Correlation Coefficient (MCC) and low Structural Hamming Distance (SHD), outperforming baselines like TCVAE, iVAE, PCL, and SlowVAE.
  • On datasets violating the graphical criterion (allowing only partial disentanglement), the method achieves high $R_{con}$ scores, confirming that the learned representations respect the predicted entanglement structure, whereas baselines fail to achieve partial disentanglement.
  • Experiments demonstrate that the method works for both continuous auxiliary variables (actions) and discrete interventions (including multi-target interventions).
  • The method is robust to heteroscedastic variance (non-exponential family cases) in the nonparametric setting.

Significance and Claims

The paper claims to provide a significantly extended version of Lachapelle et al. (2022). Its primary significance lies in:

• Relaxing Assumptions: Moving away from the restrictive exponential family assumption to a general nonparametric setting.
• Broadening Scope: Allowing for arbitrary latent causal graphs rather than requiring specific structures for identifiability.
• Formalizing Partial Disentanglement: Providing a rigorous mathematical framework for understanding and predicting when and how latent factors will remain entangled, rather than treating partial disentanglement as a failure mode.
• Handling Unknown Targets: Offering a theoretical justification for using sparse interventions with unknown targets to achieve disentanglement.

The authors position this work as a theoretical foundation for future empirical investigations in realistic scenarios, such as single-cell biology and video analysis, where mechanisms are naturally sparse but complex. They explicitly state that their work does not claim to solve disentanglement in all settings but provides sufficient conditions under which it is possible, even in the presence of non-linear mixing and complex causal structures.