PaperlayerView exact revision on arXiv

Source-linked AI summary

Identifying ODEs from Unstructured Data with Causal Representation Learning

Alessandro Trenta, Riccardo Massidda, Davide Bacciu, Sara Magliacane

arXiv:2609.37083v1cs.LGcs.AIstat.ML

TL;DR

ODE discovery from unstructured, high-dimensional observations lacks direct state measurements, while CRL variables are not generally suited to sparse equation discovery. SPEED-AE adds component-wise transformations after pretrained CRL; across three systems, it improves disentanglement and ODE recovery while achieving state-of-the-art forecasting.

  • Problem

    For high-dimensional observations, existing equation-discovery approaches either lack guarantees that learned variables match the ground truth or assume the ODE form is known.

    Source: §1 · p. 1

  • Method

    SPEED-AE combines pretrained CRL with a component-wise autoencoder that transforms variables to support sparse ODE discovery.

    Source: §1 · p. 2, §4 · p. 5

  • Results

    Across Lotka–Volterra, Lorenz, and two-pendulum systems, SPEED-AE improves CRL disentanglement, recovers ODEs closest to ground truth, and achieves state-of-the-art forecasting.

    Source: §1 · p. 2

  • Takeaways & Limitations

    For polynomial ODEs, sparsity can restrict variable identifiability from polynomial transformations to monomial transformations.

    Source: §3 · p. 4

    The analysis addresses autonomous low-dimensional ODEs observed through an unknown invertible mixing function.

    Source: §2 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

ODE discovery from unstructured, high-dimensional observations lacks either direct state measurements or guarantees that learned variables and equations match the ground truth. SPEED-AE combines pretrained causal representation learning with component-wise transformations informed by sparse equation discovery.

1 INTRODUCTION

  • Prior work on high-dimensional observations either learns equations without guarantees that its variables match the ground truth or assumes the ODE functional form is known.Equation discovery has more commonly studied settings with direct access to state variables.Source: §1 · p. 1
  • SPEED-AE combines pretrained CRL with a component-wise autoencoder that transforms variables for sparse ODE discovery.CRL provides variables up to permutation and component-wise diffeomorphisms; SPEED-AE then builds a sparse ODE from a predefined function library.Source: §1 · p. 2, Figure 1 · p. 2
  • Across Lotka–Volterra, Lorenz, and two-pendulum experiments, SPEED-AE improves CRL disentanglement, recovers ODEs closest to ground truth, and achieves state-of-the-art forecasting.The method adds a component-wise autoencoder after pretrained CRL to make transformed variables amenable to sparse ODE discovery.Source: §1 · p. 2

2 BACKGROUND

The background defines the autonomous ODE recovery setting and reviews equation-discovery and causal-representation methods. It distinguishes sparse equation fitting on observed states from learning hidden variables and identifying them up to restricted transformations.

2 BACKGROUND

  • The target is an autonomous low-dimensional ODE observed only through high-dimensional trajectories mixed by an unknown invertible function.An autonomous system's derivative depends on its current state rather than explicitly on time.Source: §2 · p. 2
  • SINDy evaluates a library of candidate functions on measured states and seeks a sparse coefficient matrix to recover the ODE.The library commonly contains polynomial and trigonometric functions; derivatives may need to be estimated from data.Source: §2 · p. 3
  • SINDyAE extends sparse equation discovery to hidden states from high-dimensional observations by jointly learning an encoder and decoder.Related MNN approaches replace derivative fitting with a network that learns equation coefficients and simulates trajectories.Source: §2 · p. 3
  • CRL methods can identify latent variables up to permutation and component-wise diffeomorphisms, with guarantees depending on assumptions or data type.The experiments consider methods using temporal data and interventions or actions, including CITRIS and DMSVAE.Source: §2 · p. 3

3 IDENTIFIABILITY ANALYSIS FOR POLYNOMIAL ODES

The analysis shows how diffeomorphic changes of latent variables transform governing ODEs and why disentanglement or sparsity alone does not guarantee recovery of both variables and equations. Under polynomial conditions, sparsest equations restrict the variable transformations to monomials.

3 IDENTIFIABILITY ANALYSIS FOR POLYNOMIAL ODES

  • A smooth invertible encoder induces a diffeomorphic change of variables, under which the encoded latents follow a conjugate ODE.The transformed vector field is ˜f(˜z) = [∇h(˜z)]^-1 f(h(˜z)).Source: §3 · p. 3
  • SINDyAE's sparsity loss does not guarantee recovery of the true ODE form or ground-truth variables, even for polynomial ODEs.A toy example shows that learned variables need not correspond one-to-one with state variables, potentially making the learned equations more complicated.Source: §3 · p. 3
  • Applying equation discovery directly to CRL latents can yield an arbitrary-form ODE; the polynomial analysis assumes component-wise polynomial diffeomorphisms and non-cancellation conditions.Under these conditions, disentangled representations governed by sparse polynomial ODEs identify the true latents up to component-wise monomial transformations.Source: §3 · pp. 3–4
  • For polynomial ODEs under Assumptions A.2 and A.3, matching the minimum monomial-term count restricts each component transformation to a monomial; if q_i is constant, it is affine.The theorem states N(˜r_π(i)) ≥ N(r_i), with equality yielding the monomial transformation.Source: §3 · p. 4

4 SPARSE EQUIVALENT EQUATION DISCOVERY AUTOENCODER

SPEED-AE adds a component-wise autoencoder after a pretrained CRL method so its transformed variables support sparse ODE discovery. This retains disentanglement and, under the theorem’s assumptions for polynomial ODEs, identifies latents up to monomial transformations.

4 SPARSE EQUIVALENT EQUATION DISCOVERY AUTOENCODER

  • SPEED-AE transforms CRL variables component-wise so the new latents have sparse ODE equations while remaining disentangled with respect to the true variables.The pipeline uses a pretrained CRL encoder, then learns component-wise autoencoders together with a sparse ODE.Source: §4 · p. 5
  • Under Theorem 3.1’s assumptions, SPEED-AE identifies true latents up to simple monomial transformations without assuming a latent ODE form.The framework is agnostic to the CRL and equation-discovery methods, provided the CRL method guarantees identifiability up to permutation and component-wise diffeomorphisms.Source: §4 · p. 5
  • The framework can pair its transformed latents with either a SINDy-based or MNN-based equation-discovery loss.The SINDy variant uses a function library and coefficient matrix; the MNN variant simulates trajectories from the learned ODE.Source: §4 · p. 5
  • An x-consistency term prevents latent collapse, while an L1 penalty with sequential thresholding sparsifies the ODE coefficients.The final component-wise encoder and decoder use MLPs because polynomial versions were not flexible enough during training.Source: §4 · p. 5

5 RELATED WORK

Prior work studies equation discovery from directly observed variables, learning dynamics from hidden representations, and identifying variables from unstructured data. SPEED-AE’s related-work context is the gap between CRL’s variable-identification guarantees and sparse ODE recovery.

5 RELATED WORK

  • Sparse regression methods such as SINDy discover equations from observed variables, while MNNs and symbolic-recovery methods offer other approaches to modeling dynamics.Theoretical identifiability results for nonlinear ODEs remain limited to known functional forms or single initial conditions.Source: §5 · pp. 5–6
  • For unstructured observations, prior methods either learn latent equations without guarantees on ground-truth correspondence or assume the ODE’s functional form is known [Champion et al., 2019; Auzina et al., 2023; Linial et al., 2021].SINDyAE learns representations for sparse equations but does not provide theoretical guarantees or analysis of representation quality [Champion et al., 2019].Source: §5 · p. 6
  • CRL methods identify latent variables up to permutation and component-wise transformations, but generally do not learn ODEs or use equation sparsity to improve identifiability.The cited trajectory-specific parameter work assumes observed state variables and a single initial condition; this paper instead assumes shared ODE parameters across trajectories.Source: §5 · p. 6

6 EXPERIMENTS

SPEED-AE is evaluated on mixed high-dimensional observations of Lotka–Volterra, Lorenz, and pendulum systems using CRL and equation-discovery components. It generally improves variable disentanglement and ODE recovery over the baselines, with strong forecasting results across the reported settings.

6 EXPERIMENTS

  • The evaluation combines CITRIS or DMSVAE with SINDy or MNN, comparing these SPEED-AE variants against SINDyAE, MNNAE, and direct CRL-plus-discovery ablations.The experiments use high-dimensional observations generated by invertible mixing functions, and CRL pretraining includes interventions on each state variable with probability p = 0.01.Source: §6 · p. 6
  • Trajectory prediction uses component-wise mappings to align learned variables with ground-truth variables, then measures latent and observation-space errors.ODE recovery converts learned polynomial maps to the true-latent space and compares coefficients using CoeffSAE and CoeffSSE.Source: §6 · p. 7

6.1 LOTKA-VOLTERRA AND LORENZ WITH POLYNOMIAL LIBRARIES

  • With perfect CRL disentanglement, SPEED-AE has negligible correlation-discrepancy and ODE-recovery errors; polynomial autoencoders were insufficiently flexible or stable in training.The final setup therefore uses MLP encoders and decoders.Source: §6 · pp. 7–8
  • SPEED-AE recovers Lotka–Volterra coefficients closest to ground truth, especially with CITRIS; direct CRL-plus-discovery combinations and baselines produce less sparse equations.The coefficient figure displays polynomials by row and methods by column, with ground truth first and nonzero coefficients boxed.Source: Figure 3 · p. 7, §6.1 · p. 9
  • Across Lotka–Volterra and stable Lorenz, SPEED-AE outperforms both baselines on correlation and ODE recovery; under chaotic Lorenz, it remains unmatched in ODE recovery.MNNAE’s strong forecasting is countered by poor correlation and ODE recovery, and it fails in more than one seed for chaotic Lorenz.Source: Figure 4 · p. 8, §6.1 · p. 9

6.2 PENDULUMS AND NON-POLYNOMIAL LIBRARIES

  • On pendulum images, SPEED-AE reduces forecasting errors by a factor of 3–4 relative to SINDyAE, whose learned variables fail to identify and disentangle both pendulums.SPEED-AE recovers nearly identical ODEs for the two pendulums, close to their shared ground-truth dynamics; CITRIS+SINDy forecasts accurately but yields less simple equations.Source: §6.2 · p. 9

7 CONCLUSION

SPEED-AE combines a transformation of disentangled variables with sparse equation discovery, because neither disentanglement nor sparsity alone suffices to recover structural variables and their governing equations. For polynomial ODEs, the appendix shows why sparsity alone does not ensure disentangled variables.

7 CONCLUSION

  • SPEED-AE learns a componentwise transformation of disentangled variables, informed by the sparsity of the learned equations, to recover structural variables and governing equations from high-dimensional data.Its results depend on each pipeline step; the theoretical analysis is limited by non-cancellation assumptions and polynomial diffeomorphisms.Source: §7 · pp. 9–10

A.1 PRELIMINARIES

  • Under a change of variables z = h(˜z), the conjugate ODE is ˙˜z = [∇h(˜z)]−1f(h(˜z)).The transformed trajectory starts from the inverse-mapped initial condition.Source: §A.1 · p. 14

A.2 SINDYAE LOSS CANNOT GUARANTEE DISENTANGLEMENT FOR POLYNOMIAL ODES

  • Even for polynomial ODEs and polynomial changes of variables, the SINDy loss cannot guarantee that learned variables are disentangled.The appendix gives polynomial ODEs related by a diffeomorphism that are equivalent under the SINDyAE loss because they have the same number of terms.Source: §A.2 · pp. 14–15
  • A non-disentangled representation can therefore have learned equations more complicated than the ground-truth equations.The text presents this as a general possibility, not an inevitable outcome.Source: §A.2 · p. 15

A.3 IDENTIFIABILITY OF LATENT VARIABLES AND ODE

For polynomial ODEs related by componentwise polynomial diffeomorphisms, the appendix relates equation sparsity to the form of those transformations. Its conclusions rely on assumptions that rule out cancellations during polynomial expansions.

A.3 IDENTIFIABILITY OF LATENT VARIABLES AND ODE

  • The term-count argument assumes polynomial changes of variables and no cancellations in polynomial powers or transformed remainders.Assumptions A.2 and A.3 specify these non-cancellation conditions.Source: §A.3 · p. 15
  • If the transformed remainder has the same number of terms as the original remainder, each componentwise polynomial diffeomorphism must be monomial; when q_i is constant, it must be affine.Under the stated assumptions, the transformed remainder cannot have fewer terms than the original.Source: §A.3 · pp. 16–17
  • When q_i is constant, the admissible componentwise transformation is affine.The text derives this from the condition that h′_i divides q_i(h_i).Source: §A.3 · p. 18
  • For a q_i with a single root of multiplicity k, the appendix gives admissible transformations of the form h_i(˜z_i) = a_i(˜z_i − b_i)^k + c_i.The stated condition requires h′ to have a single root with arbitrary multiplicity.Source: §A.3 · p. 18
  • For more general q_i, the authors report no satisfactory property for h_i and give a counterexample to a general characterization.The analysis assumes h is polynomial; polynomiality does not follow merely from both ODEs being polynomial.Source: §A.3 · pp. 18–19

B SPEED-AE IMPLEMENTATION DETAILS

SPEED-AE adds componentwise encoders and decoders to a pretrained CRL representation and fits sparse ODEs in the transformed variables. The appendix also describes its loss variants, comparison baselines, and the CRL methods used.

B SPEED-AE IMPLEMENTATION DETAILS

  • SPEED-AE uses componentwise encoders to transform CRL variables and decoders to reconstruct them, while learning ODE coefficients for a candidate-function library.Its SINDy-based loss uses push-forwarded derivatives and an L1 sparsity term with sequential thresholding every 100 epochs.Source: B SPEED-AE IMPLEMENTATION DETAILS · p. 19
  • The SINDy-based SPEED-AE variant fits derivatives, whereas the MNN-based variant predicts trajectories from initial conditions without requiring derivatives.The MNN-based training data use fixed-length trajectory windows, with almost every data point considered as an initial condition.Source: B SPEED-AE IMPLEMENTATION DETAILS · p. 19

C BASELINES

  • The appendix compares SPEED-AE with SINDyAE and MNNAE, as well as autoregressive LSTM and Transformer forecasting models.SINDyAE combines an autoencoder with a SINDy function library; MNNAE combines an autoencoder with a mechanistic network that predicts ODE parameters and trajectories.Source: C BASELINES · pp. 19–21

D.1 CITRIS

  • CITRIS-NF uses temporally intervened sequences, assumes each true variable can be intervened on at any time, and applies an invertible normalizing flow after a pretrained image autoencoder.CITRIS requires the intervention target at each time step.Source: §D.1 · p. 21

D.2 DMSVAE

  • DMSVAE exploits sparsity between observed intervention targets and state variables because sparse transition graphs among state variables are described as uncommon in most dynamical systems.Its model learns a causal graph through masks on transition-model inputs and applies an L0 sparsity loss to that graph.Source: §D.2 · pp. 21–22

E DETAILS ON THE EXPERIMENTAL SETUP

The experiments generate standard and interventional trajectories, pretrain CRL representations, and train SPEED-AE and baseline models with validation-based selection. The full pipeline takes approximately the same or less training time than SINDyAE and MNNAE on the reported Lotka–Volterra setup.

E DETAILS ON THE EXPERIMENTAL SETUP

  • Standard trajectories are generated from the ground-truth ODE, while interventional trajectories randomly intervene on a variable with probability 0.01 at each step.Standard trajectories are generated within the bounds of the interventional data, then both are mapped to high-dimensional observations.Source: E DETAILS ON THE EXPERIMENTAL SETUP · p. 22
  • CITRIS or DMSVAE is selected using validation R2 between assigned CRL latents and true variables, averaged across variables.The best hyperparameter combinations are reported in Table 3.Source: E DETAILS ON THE EXPERIMENTAL SETUP · p. 22, Table 3 · p. 23
  • SPEED-AE trains on CRL-encoded trajectories after reordering latents by the learned assignment, using finite differences for SINDy loss or 50-step batches for MNN loss.Validation forecasting selects hyperparameters and the best model; trajectories are evolved with a Runge-Kutta 45 solver.Source: E DETAILS ON THE EXPERIMENTAL SETUP · p. 23
  • SINDyAE and MNNAE use comparable MLP architectures, with their respective training and validation procedures described for the baselines.The final evaluation uses test trajectories.Source: E DETAILS ON THE EXPERIMENTAL SETUP · pp. 23–24

E.1 HARDWARE AND COMPUTATIONAL COSTS

  • The full SPEED-AE pipeline, including disentanglement, requires approximately the same or less time than SINDyAE and MNNAE on Lotka–Volterra.Training times are reported for that experiment as a reference.Source: §E.1 · p. 24

F EVALUATION PIPELINE

The evaluation converts models’ outputs to a common reference space so their learned variables and equations can be compared with the ground truth. It assesses latent correlations, ODE recovery, and long-term prediction from initial conditions.

F EVALUATION PIPELINE

  • Models use different latent spaces, so evaluation converts their representations and recovered ODEs to the true latent space, which is withheld until evaluation.SINDyAE and MNNAE require a variable assignment based on the best Mean Correlation Coefficient across permutations.Source: F EVALUATION PIPELINE · pp. 24, 25, 27
  • Correlation Discrepancy compares true and cross-correlation matrices by averaging their absolute elementwise error.The matrices can be calculated on training or test trajectories.Source: F EVALUATION PIPELINE · p. 25
  • Long-term prediction is evaluated by encoding an initial condition, rolling out the learned dynamics, and mapping predicted latents to true latents with an oracle MLP.The time-averaged L2 distance or another metric can then be calculated between true and converted latents.Source: F EVALUATION PIPELINE · p. 25
  • ODE discovery evaluation fits component-wise polynomial maps of degree 1 to 3 between model and true latents, stopping when the score exceeds 0.99.The fitted map transforms the discovered ODE to the true latent space; implicit terms are penalized because they indicate nonlinear rather than linear maps.Source: F EVALUATION PIPELINE · pp. 25, 27

G ABLATIONS

The ablations find that SPEED-AE performs best when its input latents are well disentangled, while polynomial autoencoders proved difficult to train. Additional results report strong equation-recovery comparisons and forecasting limits across tasks.

G ABLATIONS

  • With oracle perfect disentanglement, both SINDy-loss and MNN-loss SPEED-AE achieve the lowest errors across every metric.The MNN-based version performs slightly worse; the authors conclude performance is nearly perfect when disentanglement is good enough.Source: G ABLATIONS · p. 28
  • Polynomial encoders and decoders caused solution blow-ups or flat dynamics, and training remained ineffective even with simple oracle latent transformations.The authors attribute the difficulty to optimizing the autoencoder and dynamics simultaneously and use MLPs instead.Source: G ABLATIONS · p. 28, Figure 7 · p. 28

H ADDITIONAL EXPERIMENTAL RESULTS

  • SINDyAE and MNNAE are out of scale relative to SPEED-AE combinations in many Correlation Discrepancy and recovered ODE coefficient error plots.The uncut results are shown in Figure 8.Source: H ADDITIONAL EXPERIMENTAL RESULTS · p. 29, Figure 8 · p. 29
  • On 5000-step Lotka–Volterra trajectories, LSTM and Transformer predictions collapse before 1000 steps and have errors almost 20 times higher than SPEED-AE.MNNAE generally forecasts well but fails on chaotic Lorenz, and its performance comes without latent identifiability.Source: H ADDITIONAL EXPERIMENTAL RESULTS · p. 29, Table 6 · p. 31

H.1 LOTKA-VOLTERRA EXPERIMENT

Across the three systems, SPEED-AE generally improves latent identifiability over CRL representations and recovers sparse, ground-truth-like equations in Lotka–Volterra. The Lorenz and pendulum results also show how performance depends on the CRL latents and system structure.

H.1 LOTKA-VOLTERRA EXPERIMENT

  • SPEED-AE improves on CRL latent identifiability in almost every Lotka–Volterra case, often reducing a general diffeomorphism to a linear mapping.The scatter plots compare found and true latents, with CRL latents shown for CITRIS and DMSVAE.Source: §H.1 · p. 31, Figure 11 · p. 33
  • For transferred CRL representations, SPEED-AE recovers the sparsest and most similar Lotka–Volterra equations, while direct equation discovery on CRL latents sacrifices sparsity and interpretability for forecasting.The identified coefficients are displayed in Figure 12.Source: §H.1 · p. 31, Figure 12 · p. 34

H.2 LORENZ EXPERIMENT

  • In Lorenz experiments, SPEED-AE improves cross-correlation and identifiability over CRL latents, while DMSVAE performs significantly worse than CITRIS.SINDyAE struggles to disentangle stable Lorenz latents, whereas MNNAE forecasts strongly but learns variables that do not easily relate to the true ones.Source: §H.2 · p. 37, Figure 15 · p. 37, Figure 16 · p. 38

H.3 PENDULUMS EXPERIMENT

  • For the two pendulums, SPEED-AE achieves linear identifiability from CITRIS latents, while SINDyAE identifies only one pendulum correctly across all 10 seeds.The authors suggest SINDyAE's other variable may map between the pendulums because their dynamics are equal apart from a constant phase.Source: §H.3 · p. 40, Figure 19 · p. 41
Loading 2609.37083v1…