Marrying Causal Representation Learning with Dynamical Systems for Science

This paper aims to connect causal representation learning with parameter identification in dynamical systems. By doing so, existing causal representation learning approaches can be used for estimating parameter in dynamical systems with identification guarantees. On the other side, it also demonstrates the applicability of causal representation learning on real-world data. Experimental evaluation on simulated and real-world climate data show the effectiveness of applying causal representation learning approaches in solving parameter identification tasks in dynamical systems.

The authors draw a connection between assumptions in the (neural) ODE and causal representation learning (more precisely, latent variable identifiability, but I will follow the authors and refer to it as CRL) literatures, by framing ODE inverse problems as latent variable problems. A particular focus is given to the case where the parametric form of the ODE is unknown, and only trajectories (i.e., time series data) are given. The learned generator and latent variables in CRL are analogous to a flexibly parametrized ODE solver and corresponding parameters, respectively. An example is given where pairs of trajectories sharing some parameter values are interpreted as same-content-different-style observations, which can be block-identified according to previous work in CRL. The authors describe a specific learning framework based on mechanistic neural networks which encode the trajectories into not only the parameters of a flexible family of ODEs, but also hyperparameters of the solver, defining a flexible generative process of observed trajectories.

This paper bridges causal representation learning (CRL) with dynamical system learning.

It introduces partially identifiable and practical models by merging methodologies from both CRL and dynamic systems.

The authors develop models capable of handling out-of-distribution classification tasks and treatment effect estimation.

Notably, they validate their approach using a wind simulator and real-world climate data, effectively demonstrating the model's potential in addressing causal questions related to climate change.

This paper proposes a theory and methodology for representation learning from dynamical systems. In particular, it proposes a model in which latent causal variables deterministically generate an observed time-series trajectory through an ODE, with the task of identifying the latent variables from observed time series. Theoretically, the authors devise identifiability conditions by drawing the causal representation learning literature, treating the ODE as the (black box) mixing function. To identify variables in practice, it is proposed to use mechanistic neural networks (MNNs) to learn the ODE mixing function. Empirical results show that the proposed method can strike a good balance between identification of the latent variables and predictive performance.

[1] Steven L Brunton, et al. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 113(15):3932–3937, 2016.

[2] Steven L Brunton, et al. Sparse identification of nonlinear dynamics with control (sindyc). IFAC-PapersOnLine, 49(18):710–715, 2016.

[3] Alan A Kaptanoglu, et al. Promoting global stability in data-driven models of quadratic nonlinear dynamics. Physical Review Fluids, 6(9):094401, 2021.

[4] Philippe Wenk, et al. Fast gaussian process based gradient matching for parameter identification in systems of nonlinear odes. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1351–1360. PMLR, 2019.

[5] Julius Von Kügelgen, et al. Self-supervised learning with data augmentations provably isolates content from style. Advances in neural information processing systems, 34: 16451–16467, 2021.

[6] Dingling Yao, et al. Multi-view causal representation learning with partial observability. In The Twelfth International Conference on Learning Representations, 2024.

[7] Klindt, David, et al. "Towards nonlinear disentanglement in natural data with temporal sparse coding." arXiv preprint arXiv:2007.10930 (2020).

[8] Yao, Weiran, et al. "Temporally disentangled representation learning." Advances in Neural Information Processing Systems 35 (2022): 26492-26503. (referred to as [a] by reviewer yeqY)

[9] Balsells-Rodas, Carles, et al. "On the Identifiability of Switching Dynamical Systems." arXiv preprint arXiv:2305.15925 (2023).

[10] Ernest Lindelöf. Sur l’application de la méthode des approximations successives aux équations différentielles ordinaires du premier ordre. Comptes rendus hebdomadaires des séances de l’Académie des sciences, 116(3):454–457, 1894.

[11] Ror Bellman and Karl Johan Åström. On structural identifiability. Mathematical biosciences, 7 (3-4):329–339, 1970.

[12] Kong, Lingjing, et al. "Identification of nonlinear latent hierarchical models." Advances in Neural Information Processing Systems 36 (2023): 2010-2032.

[13] Xie, Feng, et al. "Identification of linear non-gaussian latent hierarchical structure." International Conference on Machine Learning. PMLR, 2022.

[14] Danru Xu, et al. A sparsity principle for partially observable causal representation learning. International Conference on Machine Learning, 2024.

[15] Ahuja, Kartik, et al. "Interventional causal representation learning." International conference on machine learning. PMLR, 2023.

[16] Song, Xiangchen, et al. "Temporally disentangled representation learning under unknown nonstationarity." Advances in Neural Information Processing Systems 36 (2024).

[17] Li, Zijian, et al. "When and How: Learning Identifiable Latent States for Nonstationary Time Series Forecasting." arXiv preprint arXiv:2402.12767 (2024). (referred to as [1] by reviewer bh3V)

[18] d'Ascoli, Stéphane, et al. "Odeformer: Symbolic regression of dynamical systems with transformers.” In International Conference on Learning Representations, 2024.

[19] Locatello, Francesco, et al. "Weakly-supervised disentanglement without compromises." International conference on machine learning. PMLR, 2020.

[20] Fumero, Marco, et al. "Leveraging sparse and shared feature activations for disentangled representation learning." Advances in Neural Information Processing Systems 36 (2024).

[21] Lachapelle, Sébastien, et al. "Synergies between Disentanglement and Sparsity: Generalization and Identifiability in Multi-Task Learning." arXiv preprint arXiv:2211.14666 (2022).

[22] Adeel Pervez, et al. Mechanistic neural networks for scientific machine learning. International Conference on Machine Learning, 2024.

[22] Feuerriegel, Stefan, et al. "Causal machine learning for predicting treatment outcomes." Nature Medicine 30.4 (2024): 958-968.

We are extremely grateful to the reviewers and AC for their time and valuable feedback. We very much appreciate that they found the problem we are tackling is “very interesting and potentially impactful”(bh3V), “of great importance for e.g. science”(3NaZ), “incredibly important”(aFgN), and “intriguing and emerging”(yeqY). We are happy to see that our idea of applying CRL techniques to dynamical systems is “interesting and novel” (3NaZ), “original and potentially of great significance to both communities”(aFgN).

Novelty and contribution (3NaZ, yeqY)

• As discussed in lines 73-78, this paper’s main contribution is providing clear parameter-identifiability statements for dynamical systems, whereas numerous previous works on ODE discovery [1, 2, 3, 4] refrained from doing so by explicitly stating that it is unknown which settings yield identifiability (lines 165-166).
• To the best of our knowledge, our paper is the first to target real-world scientific applications and successfully demonstrate identification results on raw measurements, unlike previous CRL work which relies on synthetic or heavily pre-processed data (e.g. manually rendered images[5, 6] or extracted avatar skeleton [9]).

Validity of the technical assumptions (3NaZ, yeqY)

• Assumptions 2.1 and 2.2 are standard and necessary for parameter identification in dynamical system [10, 11]. Many ODEs satisfy these assumptions, including Lotka-Volterra ODE, Van der Pol oscillator, and chaotic systems such as the Lorenz attractor and the Rössler attractor.
• Further, we justify in Table 1 that CRL assumptions (3.1-3.3) align with standard assumptions 2.1 and 2.2, establishing the theoretical ground for cross-pollination between these two fields.
• Nevertheless, we acknowledge the limitation of the determinism assumption, as discussed in Section 6.

Identifiability with ODE with specific properties (3NaZ, aFgN, yeqY)

• Corollary 3.1 provides full identifiability for stationary parameters in ODEs with known functional forms, including any parametric form (linear, polynomial, even nonlinear). This is supported by various empirical studies from prior literature on equation discovery [1, 2, 3, 4] and attached additional experimental results.
• For ODEs with partially known properties, such as a sparse linear combination $f_{\theta}(x, t) = \sum\_{i = 1}^{m} \theta_i \phi_i(x)$ of various basis functions $\phi_i$ within a comprehensive dictionary (e.g., SINDy-like scenarios [1, 2, 3]), Corollary 3.1 ensures parameter identifiability under a sparsity constraint. Notably, the ODE relates to the parameters $\theta$ linearly, as detailed in the main paper (lines 154-157) and the appendix (lines 653-663).
• Many existing CRL works assume certain properties on the generating process, such as sparsity [14] or specific functional class [15], which can be directly imported into our framework by replacing the multiview approach.

Why not model dynamics in the latent space (bh3V,3NaZ,yeqY)

• We model dynamics directly in the observational space because the underlying parameters $\theta$ directly determine the whole trajectory: $(x_{t_1}, \dots, x_{t_N}) = F(\theta)$.
• Almost all CRL works assume the latent variables directly influence the observation, except for a few considering hierarchical latent models [12, 13], which we will thoroughly discuss in the updated manuscript.
• As discussed in Section 6, we agree assuming direct observation of the states is limiting, and relaxing this assumption is an interesting future direction. However, since this paper is the first to formally connect CRL-identifiability with inverse problems in dynamical systems, we believe that exploring this perspective is beyond the scope of the current study.

Additional experiments (see attached pdf)

• We provide additional experiment results on specific dynamical systems with known functional forms, including eight highly complex systems from ODEBench[18] and the Cart-Pole system inspired by [8].
  • For each system, we sample 100 tuples of parameters $\theta$ within a valid range to preserve properties like chaos. Each problem is solved exactly as described in Corollary 3.1.
  • The root-mean-square deviation (RMSE) is averaged across the parameter dimensions and sample size (100), reported via mean and std.
  • We observe highly accurate point estimates for all stationary system parameters $\theta$, validating Corollary 3.1 across various experimental settings. Due to space constraints, we cannot enclose results for all 63 systems from ODEBbench[18]. Complete results and implementation details will be provided in the revised manuscript.
• Comparing ATE estimates from non-identified and identified representations for SST-V2:
  • The attached Figure 1 illustrates the estimated ATE from the non-identified representation lacks a discernible pattern, while the identified one exhibits a noisy yet clear increasing trend, indicating the global warming effect.
  • This is because the non-identified representation failed to isolate the covariates $\theta$, leading to biased treatment effect estimates. To estimate treatment effects, the covariates (i.e., the latitude-related parameters we identify) mustn’t be influenced by the treatment (i.e., the climate zones). Otherwise, they become confounders, leading to incorrect estimates [22].
  • We apologize for not explaining this clearly in the current draft. In the final version, we will add the additional graphical model and the non-identified baseline (attached Figure 1) together with a thorough explanation about treatment effect estimation with covariates and why CRL-identifiability (up to bijection) is necessary to avoid confounding.

We would like to express our sincere gratitude to all the reviewers for their constructive feedback. The reviews and discussions have significantly enhanced the quality of our paper, and we will incorporate all suggested changes in the updated manuscript.

We also extend our thanks to the area chair and the reviewers for their time and effort. We greatly appreciate your engagement and will acknowledge your anonymous contribution in the revised paper.

This paper proposes a connection between causal representation learning and dynamic system, ODE, to establish the identifiability of the dynamics system. Reviewers recognise the contribution of the paper, although the technical novelty is limited. During the rebuttal period, the author's responses addresses many concerns from the reviewers, so it is recommended to include these in the revised version.