cs.LGSep 29, 2026
SaveIdentifying ODEs from Unstructured Data with Causal Representation Learning
Organizations: Department of Computer Science University of Pisa · Saarland Informatics Campus Saarland University & University of Amsterdam
Abstract
We study the problem of recovering the governing ODE of a dynamical system from unstructured, high-dimensional observations such as images. Existing methods for ODE discovery typically assume direct measurements of the variables, or do not provide theoretical guarantees on the learned variables and equations. While Causal Representation Learning (CRL) methods provide guarantees on identifying variables from high-dimensional observations up to component-wise diffeomorphisms, we show that in general these variables cannot be used directly as input to equation discovery methods, which typically assume that the variables will lead to sparse equations. So we introduce SParse Equivalent Equation Discovery AutoEncoder (SPEED-AE), a framework that combines a pretrained CRL method with a component-wise autoencoder that learns transformations of variables that are amenable to sparse ODE discovery. We show that for polynomial ODEs, this additional step allows us to restrict the identifiability of each variable from polynomial to monomial diffeomorphisms. Experiments on Lotka-Volterra, Lorenz, and a two-pendulum system show that SPEED-AE improves on the disentanglement of the CRL methods and that it recovers ODEs that are closest to the ground truth, while achieving state-of-the-art forecasting performance.
Figures & tables
Figure 1: SPEED-AE pipeline: we observe high-dimensional data from an invertible mapping of the true latents . First, SPEED-AE employs a pretrained CRL model to obtain low-dimensional latents ( red block). We assume that the recovered latents correspond to the ground-truth up to a permutation and component-wise diffeomorphism . For simplicity, in the figure we assume that the permutation is the identity function. SPEED-AE learns a set of component-wise autoencoders and an ODE by combining the functions in a library , imposing sparsity of this representation ( blue block). The output latents identify the ground-truth up to permutation and monomial and, in most cases, linear transformation.
Figure 2: Comparison between SINDyAE, SINDy applied to CRL latents (using CITRIS), and SPEED-AE (with CITRIS as CRL method and SINDy as equation discovery method) on the Lotka-Volterra dataset ( , ). For each model, the left box reports scatter plots of the true versus learned latents. We then compare a test trajectory with the one predicted from only the initial condition and the discovered equation, after converting both to the true latent space, following the evaluation protocol in Section 6 . SPEED-AE is the only method that recovers a sparse equation and achieves strong long-term prediction performance.
Figure 3: Learned coefficients for Lotka-Volterra for each polynomial of the state variables (each row) and different methods (each column, where the last four are versions of SPEED-AE). The first column shows the ground truth coefficients and the black boxes highlight the non-zero coefficients.
Figure 4: Results for the Lotka-Volterra and Lorenz experiments with mean and standard deviation over seeds. For all metrics, lower is better. We limit the height of the boxes for better readability, reporting the actual scales of the outlier results inside their bars.
| Model | CorrD | Err | Err | Recovered ODE |
|---|---|---|---|---|
| SINDyAE | ||||
| CITRIS+SINDy | ||||
| SPEED-AE (C+S) |
Table 1: Results for the Pendulum experiment for the model.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
| Parameter | Lotka-Volterra | Lorenz | Pendulum |
|---|---|---|---|
| num. steps |
Table 2: Data generation parameters.
| Model | Training Time (mins) |
|---|---|
| LSTM | |
| Transformer | |
| SINDyAE | |
| MNNAE | |
| Autoencoder (CITRIS-NF) | |
| CITRIS-NF |
Table 5: Training time for each model on the Lotka-Volterra experiment.
Figure 5: Evaluation pipeline of our experiments.
Figure 6: Experimental results for the ablation on the independence from CRL method. Lotka-Volterra dataset with an additional version of SPEED-AE with an oracle CRL model that achieves perfect disentanglement and SINDy-based loss.
Figure 7: Results on the synthetic experiments with perfect disentanglement, polynomial identifiability in the CRL latents, and SPEED-AE encoder and decoders implemented as polynomials. (a) Some of the runs show strong instabilities, while others (b) remain stable but with flat dynamics. In general, the polynomial functions are hard to optimize simultaneously to the dynamics.
Figure 8: Experimental results for the Lotka-Volterra and the Lorenz experiments.
| Model | Lotka-Volterra | Lotka-Volterra | Lorenz (Stable) | Lorenz (Chaotic) |
|---|---|---|---|---|
| LSTM | 13.860 ±2.609 | 29.313 ±14.719 | 0.157 ±0.056 | 1.624 ±0.210 |
| Transformer | 16.001 ±2.076 | 25.999 ±6.919 | 0.422 ±0.076 | 1.902 ±0.210 |
| SINDyAE | 23.075 ±11.611 | 17.610 ±2.898 | 1.988 ±0.323 | 3.137 ±0.120 |
| MNNAE | 1.257 ±0.240 | 0.930 ±0.290 | 0.448 ±0.102 | (failed) |
| SPEED-AE (C+S) | 1.279 ±0.483 | 3.999 ±2.390 | 0.851 ±0.384 | 2.411 ±0.379 |
Table 6: Err: mean and standard deviation across seeds.
Figure 9: Pearson cross correlation matrices between and for all experiments and models. The rightmost matrix in each subfigure corresponds to the true correlation matrix (i.e., between and ).
Figure 10: Lotka-Volterra experiment with . Scatter plots of the found latents against the true ones . If their function is well defined, the plots well represent the maps . In the second row, we include the same plots for the CRL latents from CITRIS and DMSVAE. In most cases, SPEED-AE can go from a general diffeomorphism to a linear one.
Figure 11: Lotka-Volterra experiment woth . Scatter plots of the found latents against the true ones . If their function is well defined, the plots well represent the maps . In the second row, we include the same plots for the CRL latents from CITRIS and DMSVAE. In most cases, SPEED-AE can go from a general diffeomorphism to a linear one.
Figure 12: Identified coefficients for the Lotka-Volterra experiment with transfer of CRL models and . The true ones are in the leftmost column and highlighted by the black borders. For comparison, we include four baselines that consist of equation discovery directly applied to the CRL latents .
Figure 13: Example test trajectory of the Lotka-Volterra experiment (first dimensions of ). Ground truth (blue) against the predicted trajectory from the LSTM model (red). During rollout, the model collapses after around 500 steps.
Figure 14: Example test trajectory of the Lotka-Volterra experiment (first dimensions of ). Ground truth (blue) against the predicted trajectory from the Transformer model (red). During rollout, the model collapses after a few iterations.
Figure 15: Lorenz experiment with stable trajectories . Scatter plots of the found latents against the true ones . If their function is well defined, the plots well represent the maps . In the second row, we include the same plots for the CRL latents from CITRIS and DMSVAE. SPEED-AE models with DMSVAE as disentangler suffer more due to the CRL latents being less disentangled, although the SINDy-based loss is actually able to recover a linear map .
Figure 16: Scatter plots of the found latents against the true ones . If their function is well-defined, the plots well represent the maps . In the second row, we include the same plots for the CRL latents from CITRIS and DMSVAE. SPEED-AE models with DMSVAE as disentangler suffer more due to the CRL latents being less disentangled, although the SINDy-based loss is actually able to recover a linear map .
Figure 17: Identified coefficients for the Lorenz experiments in the stable ( , left) and chaotic case ( , right). The true ones are in the leftmost column and highlighted by the black borders. For comparison, we include four baselines that consist of equation discovery directly applied to the CRL latents .
Figure 18: Example images from the experiment with two pendulums. Pendulums are colored differently to ensure invertibility of the mixing function.
Figure 19: Pendulums experiment. Scatter plots of the found latents against the true ones . If their function is well defined, the plots well represent the maps . In the second plot, we include the same for the CRL latents from CITRIS. SPEED-AE is able to achieve linear identifiability from the general diffeomorphism of CITRIS. SINDyAE only identifies one variable, while the other likely represents a mapping between the two pendulums.