Disentangling Dynamical Systems: Causal Representation Learning Meets Local Sparse Attention

TL;DR

This paper introduces a novel method combining causal representation learning and local sparse attention to identify and disentangle system parameters from raw data, with theoretical guarantees and empirical validation.

Contribution

It provides a new identifiability theorem for disentangling system parameters without structural assumptions, using local causal structures and a sparsity-regularised transformer.

Findings

Successfully recovers disentangled representations in synthetic domains

Theoretical lower bounds on disentanglement from causal structures

Enforcing local causal structure enhances identifiability

Abstract

Parametric system identification methods estimate the parameters of explicitly defined physical systems from data. Yet, they remain constrained by the need to provide an explicit function space, typically through a predefined library of candidate functions chosen via available domain knowledge. In contrast, deep learning can demonstrably model systems of broad complexity with high fidelity, but black-box function approximation typically fails to yield explicit descriptive or disentangled representations revealing the structure of a system. We develop a novel identifiability theorem, leveraging causal representation learning, to uncover disentangled representations of system parameters without structural assumptions. We derive a graphical criterion specifying when system parameters can be uniquely disentangled from raw trajectory data, up to permutation and diffeomorphism. Crucially, our analysis demonstrates that global causal structures provide a lower bound on the disentanglement guarantees achievable when considering local state-dependent causal structures. We instantiate system parameter identification as a variational inference problem, leveraging a sparsity-regularised transformer to uncover state-dependent causal structures. We empirically validate our approach across four synthetic domains, demonstrating its ability to recover highly disentangled representations that baselines fail to recover. Corroborating our theoretical analysis, our results confirm that enforcing local causal structure is often necessary for full identifiability.