Invariance-based Learning of Latent Dynamics
Authors: Kai Lagemann
Christian Lagemann
Sach Mukherjee
Venue: ICLR 2024 poster (2024)
Keywords: Latent Dynamical Model
Invariant Decomposition
Neural Network
Spatio-Temporal Attention
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: We propose a new model class aimed at predicting dynamical trajectories from high-dimensional empirical data. This is done by combining variational autoencoders and (spatio-)temporal transformers within a  framework designed to enforce certain scientifically-motivated invariances. The models allow inference of system behavior at any continuous time and generalization well beyond the data distributions seen during training. Furthermore, the models do not require an explicit neural ODE formulation, making them efficient and highly scalable in practice. We study  behavior through simple theoretical analyses and  extensive empirical experiments. The latter investigate  the ability to predict the trajectories of complicated  systems based on finite data and show that the proposed approaches can outperform existing neural-dynamical models. We study also more general inductive bias in the context of transfer to data obtained under entirely novel system interventions. Overall, our results provide a new framework for efficiently learning complicated dynamics in a data-driven manner, with potential applications in a wide range of fields including physics, biology, and engineering.

Published as a conference paper at ICLR 2024
INVARIANCE-BASED LEARNING OF LATENT DYNAMICS
Kai Lagemann"
Statistics and Machine Learning, German Center for Neurodegenerative Diseases (DZNE), Bonn, Germany kai, wetemannnzzee....
Christian Lagemann"
Department of Mechanical Engineering, University of Washington, Seattle, USA
Sach Mukherjee
Statistics and Machine Learning, German Center for Neurodegenerative Diseases (DZNE), Bonn, Germany MRC Biostatistics Unit, University of Cambridge, Cambridge, UK
sach,
ABSTRACT
We propose a new model class aimed at predicting dynamical trajectories from high- dimensional empirical data. This is done by combining variational autoencoders and aepatio-)eeppnte transformers within framework designed to enforce certain conntffically-ootvvtte invariances. The models allow inference of system behav- ior at any continuous time and generalization well beyond the data distributions seen during training. Furthermore, the models do not require explicit neural ODE formulation, making them efficient and highly scalable in practice. We study behavior through simple theoretical analyses and extensive empirical experiments. The latter investigate the ability to predict the trajectories of complicated systems based finite data and show that the proposed approaches can outperform existing neural-dynamin models. We study also more general inductive bias in the context of transfer data obtained under entirely novel system interventions. Overall, our results provide a new framework for efficiently learning complicated dynamics in a data-driven manner, with potential applications a wide range of fields including physics, biology, and engineering.
INTRODUCTION
Dynamical models are central to our ability to understand and predict natural and engineered systems. key question in studying dynamical systems is predicting future behavior. Real-world systems often show time-varying behavior that is much too complex for straightforward statistical forecasting or extrapolation approaches. This due the fact that the temporal behavior, while potentially explained by an underlying dynamical model, can show strong, possibly abrupt changes in the ohservation/data space, precluding effective modeling via traditional curve-fitting or extrapolation. Furthermore, different instances realizations of single scientificeenineeringg system (e.g. with different initial conditions or constants) can show large differences in terms of data distributions, hence going beyond standard in-distribution assumptions of traditional data-fitting approaches.
Against this background, in recent years a wide range of sophisticated dynamical machine learning approaches have been proposed, including in particular neural ordinary differential equations Chen ct al.] and a wider class of related models (see for example Zhi et al. (2022); ; Finlay et al. (2020): ;
Duong and Atanasov (202T); Choi et al. (022);; Chen al. (2021) }) Kim al. (2021b)). Broadly speaking, these models go beyond simple curve-fitting curve-fiting/eettapllatii
schemes by leveraging suitable
inductive biases to allow learning of latent dynamical models. There has been rapid progress this
area but key challenges remain for complicated real-world systems, due to multiple factors, including
data limitations, generalization unseen settings, irregular time sampling and issues relating to long-horizon trajectories (lakovlev et al. 2023).
Thheee authors contributed equally this work.

---

Published as a conference paper at ICLR 2024
V
Embodddr
NNO
Irpul Irojectory
Temooett on rursforrn
Query
Time
Multicle Augmentatiin
Resreazntatio Lean
Latent Dynamics Loomiin
Sontio-ternporal Trenstometton
Decoding
G ey timer
Lewer Tojechalezz
Trigectory
Rocresontarie
Approgation
NN
Excoder
0.000 citeenaciioo
Latueed lobero
NotwOR
Learning of Latent Dynamics via Invariant Decomposition (LaDID), A set of high -igh-dimensional snapshots on regular or irregular time-grid serves as the input to LaDID. The trajectory is split into subpatches using Multiple Shooting Augmentation (i). The first time- points of each subpatch are used compute a subtrajectory representation: features of the selected snapshots re- re-weighted w.r.t. time and spatial location (ii), transformed to low-dimensional embedding (iii), and aggregated into one trajectory representation vi (iv). During inference, the latent dynamical model is conditioned on the representation v². Prediction is possible any continuous time by querying the latent state of any time point (v). Latent subtrajectories are sewn together by smoothness loss. Finally, the entire latent trajectory is decoded to the observation space (vi).
Figure 1: Architecture:
Motivated by these challenges, we propose a new framework, called "Latent Dynamics via Invariant Decomposition" LaDID, for learning latent dynamics from empirical data. LaDID leverages certain scientifically-mootoe motiivatee invariances to permit efficient learning and effective generalization. In numerous real-world dynamical systems, longitudinal trajectories may exhibit significant variability, e.g. due to differences in initial conditions or model constants. Each temporal trajectory, which we refer to "realization". represents a particular manifestation of the system's 's dynamics under certain conditions. A key notion underpinning LaDID is that, even when temporal trajectories from a class of scientific systems appear diverse, they can still effectively explained by an appropriate, in sense "universal", model; such model is therefore generalization, LaDID introduces factors specific each realization as inputs to its universal model. These factors are hence nullzatioonspecijs and play the role of (implicitly) encoding aspects such as the initial states of the system or specific model constants. transformer-basee based architecture used to learn all model aspects from data, including both realization-sper speciffi (RS) and palization-invarina (RI) information. At inference- time LaDID can output predictions for any continuous time. Due to the universal nature of the RI model, LaDID effectively handle substantial variations in system behavior and data distributions (e.g. due to changes in initial conditions or system constants). We empirically validate LaDID on various spatio-temporal systems with dynamics on regular and irregular time grids governed by ordinary or partial differential equations. The LaDID architecture is fast and easy train and, we show, substantially outperforms existing neural-dynamical models over range of challenging tasks. Thus, our main contributions are:
intvatiiiinnnnniaint..... To facilitate broad
We present novel framework, and associated transformer-based network, for the separation of realization-speeific information and eealizzation-invaiinnn latent dynamical systems.
We systematically study performance short- and longer-horizon prediction of wide range of complex temporal and spatio-temporal problems, comparing against a range of ofate-of-the-ane art neural-dynamical baselines.
We study the challenging case of transfer to data obtained under entirely novel system interventions via sho- learning (FSL.) approach.
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORK
Flexible neural models have been exploited to learn dynamical models, with connections drawn between deep architectures and numerical solvers for ODEs, PDEs and SDEs (Chen et al. 2018
Weinan 2017 Lu et al., 2018 Ruthotto and Haber 2020; Haber and Ruthotto 2017; Richter-Powell
et al. 2022). Algorithms rooted in neural differential equations (NODEs) have been shown to offer
benefits relative to standard recurrent neural networks (RNNs) and their variants. However, since NODEs directly relate to the problem formulation of standard ODEs, they inherit some associated limitations. Specifically, the temporal dynamies only depend on the current state but not on the history putting limit on the complexity of the trajectories that NODEs can model (Holt et al. 2022).
Improvements have been proposed that augment the latent state space to broaden the range of dynamical systems that can be learned (Dupont et al. (019)) Juhlno/uuaaova et al. (2019) suggested a combination with an autoregressive RNN updated at irregularly sampled time points. Complementary work has proposed neural controlled differential equations, mechanism to adjust the trajectory based on subsequent observation (Kidger et al. 2020 Morrill et al. (021). Massaroli et al. (2021) transferred the concept of multiple shooting to solve differential equations to the conceptual space of NODEs and [akovlev et al. (2023) extended this concept to sparse Bayesian multiple shooting, with both works evaluating latent NODEs. However, for certain types of dynamics numerical instability poses challenges for NODEs (and their variants) (Li et al. 2020). This is due to the fact that NODEs rely on numerical ODE solvers to predict the latent trajectory (forward pass) and this becomes unstable with longer time horizons (Iakovlev et al. 2023). In contrast, by exploiting RS and RI invariances our model eschews explicit neural ODEs altogether, providing an arguably simpler and faster transformer-basee scheme that can be trained in straightforward fashion, as we detail below.
The idea of leveraging invariances is a core notion in scientific modeling and is seen throughout the natural sciences at a conceptual and practical level. For instance, in the field of AI, it has been utilized in the context of video frame prediction as demonstrated by various studies (van der Wilk et al., 2018 Franceschi et al.) [2000; Kabra et al.) [202I). LaDID differs from these approaches because it uses invariances to model a kind of generalized initial condition (motivated by scientific uses of dynamical models; see below) and because it learns continuous latent trajectories (as opposed to an autoregressive model), including in the irregular time-sampling case. See Section A of the Appendix for further details related work.
3 PROBLEM STATEMENT AND CONCEPTUAL OUTLINE
We start with a high-level problem statement and outline the motivating concepts behind LaDID, deferring detailed description of the architecture itself to subsequent Sections. We focus settings in which we have (finite) observations of a system of interest at time points t E T (potentially irregularly spaced). We do not require any specifie prior information on the underlying model; rather our approach is data-driven, informed by certain very general invariances as detailed below.
For instance/realizztiin of dynamical system of interest let X, RTxCx denote high-dimensional trajectory in the observation space; T denotes the number of time steps (possibly irregular), and C, H, and W respectively the number of channels, frame height and frame width (in empirical examples we focus on image-like data inputs) of the observations. Let X=(X,))ee (eR)) denote the collection of available training trajectories; the notation emphasizes the possibility that available data spans multiple instancesrrealization Given these initial data, LaDID seeks to predict future observations x. for any continuous time and for any realization r.
From invariances to simple learning framework. We start by studying a basic, conceptual set-up that sheds light on how our assumptions lead a simple, but very general, learning framework. Intuitively, we demonstrate that while we cannot guarantee recovery of the true underlying model parameters, under mild invariance assumptions, there exists a function capable of reconstructing the true observations, even when dealing with potentially highly heterogeneous parameters and data. Importantly, we do not make any prior assumptions about the nature of these potentially complex and nonlinear functions. Instead, our learning framework simultaneously uncovers and refines these functions in a data-driven, end-to-end manner, as elaborated in Section 4
Consider an entirely general system ff in which some aspects are caalization-speciff (RS) while others are iualization-invaiian (RD) the latter model aspects are general to all instances/realization of the model, while the former not. We assume also that the RS aspects the same for all
3

---

Published as a conference paper at ICLR 2024
times; i.e. these are invariint.. Our model is not specific to physical ODE-like models, but rather generalizes invariances to permit flexible learning in a wide range of settings. To fix ideas, it is nevertheless useful to consider the specific example of a physical model class described by ODEs. For this latter setting, the model itself would be RI, while the initial conditions or system constants could be thought of as RS. As a result, the same model class can describe wide range of actual systems which share an underlying scientific basis while differing (perhaps strongly) in details.
Let f{= O,) denote the fully general model. Here, O the complete parameter set needed to specify the time-evolution; including both RS and RI parts. To make the separation clear, we write the two parts separately as x ft: f). where 8,, 9 are respectively the RS and RI parameters and indexes the realization. Suppose 0, is a candidate encoding of the RS information. We now assume that the encoding, while possibly incorrect (i.e. such that 0,, satisfies the property 3m, 30 0 0,=m( 0,0 fm), where m is a function that "corrects" 0, to give the correct RS parameter. This essentially demands that while the encoding 0, might be very different from the true value (and may even diverge from in complicated way that depends on unknown system parameters), there exists an RI transformation that recovers the true RS parameter from it, and in this sense the encoding contains all RS information. We call this the sufficient encoding assumption (SEA). In the context of dynamics involving latent variables e R9 again consider model with RS and RI parts but at the level of the latents, i.e. 2 f 8, 0). Assume the observables are given by (an unknown) function g of the hidden state Then, shown in the Appendix assumptions similar in spirit to SEA above imply existence of a mapping function that allows arbitrary queries to in principle be predicted via a straightforward learning scheme. In brief, these assumption require that we have access to an approximation و of the observation process that while possibly very far away from the true function, can nonetheless be "corrected" in specific way (see Appendix [B] for details). We emphasize that as for SEA, it is not required that we have good approximation in the usual sense of being close to the true function g. only that there exists correction of a certain form. Importantly, at no point do we actually require prior knowledge of the underlying dynamical system, its true latent variables nor any system constants or initial conditions; rather, the assumptions imply existence of mapping functions that can be learned from data, even when the underlying model is entirely unknown at the outset.
Analogy to traditional ODE formulations. To facilitate more intuitive understanding of our framework, we discuss further analogies to standard ODE problems. Typical ODE solvers comprise a formulated ODE function and some initial values (IVs) which are evolved over time using well-known integration methods, e.g. Euler, Lunge-Kutta, or DOPRI schemes. From a high-level perspective, the IVs relate to our RS representation while the ODE functions and its corresponding integration approach is our RI part. However, please note that this analogy only holds on a superficial level due to two fundamental differences: First, LaDID relies only on observations of specific system, for which underlying state variables required to apply standard ODE solvers are not known (and in fact not observed). Hence, our RS representation benefits from access to a collection of nigh-dimensional observations. Second, our RI function is continuous in time and therefore unites temporal integration and dynamics function on an abstract level. To do so, we condition our RI model on specific RS representations and only query time points from it to obtain discretized latent trajectory.
4 METHODS
Based on these initial arguments, we now put forward a specific architecture to allow learning in practice. At high level, the architecture implements the general mapping approach outlined above (further details in the Appendix), learning, in a data-driven manner, both RS and RI model components and putting these together allow prediction at any continuous time in a query realization. Implementation details and a full ELBO derivation can be found in Section C F of the Appendix.
4.1 MODEL, INFERENCE AND PREDICTION Model. The LaDID architecturd is composed of three main components: the encoder forne the invariant dynamical system feav.. and the decoder feure respectively governed by parameters Penc- Ødgn, and Ødec. The encoder is collection of three NNs: CNN processing spatial information in
"Called the sufficient mapping assumption (SMA) and realization-specife SMA (RS-SMA) for the cases in which the aatent-to-osserven mapping is itself RI or RS respectively; see Appendix for details.
A reference implementation is available under / http:/gghuuu.occoo com/k1 844k8844477777 44444777444 LabiD
4

---

Published as a conference paper at ICLR 2024
the observation space, transformer utilizing temporal attention and learnable mapping function. Since we want to predict future observations based on a few observations, we only use the first K datapoints in time and process these in a shared convolutional encoder (green trapezoid in Figure [T (ii)). We employ a shallow CNN that compresses the input to 1/16 of the initial input size using four ReLU activated and batch-normalized convolutional layers. The resulting tensors are then flattened and mapped linearly to a single vector. Next, we use a transformer on the K output vectors of the convolutional encoder, applying temporal attention reweigh vectors. We tested two approaches (Bulat et al. 2021; lakovlev et al. 2023) with comparable performance which are discussed in the Appendix more detail. For each of the k E K time aware representations Pl we sample a latent embedding using the reparameterization trick, i.c. gmmb (f (px The final trajectory representation sb" is the output of an aggregation over all K tokens. In our implementation, we choose a simple yet effective mean- mean-aggregation which can be changed based on the task hand. The second important part of our proposed framework is the dynamical model fōaun. We utilized a three layer MLP which can also be interchanged by other functions. To obtain a latent trajectory, we condition the latent dynamical model and roll-out the latent trajectory x based on the queried time points ig represented through a time encoding which we choose
our end-to-end learned trajectory representation eh
set of different sine and cosine waves with different wave length. Finally, we map all data points of our latent trajectory back to the original observation space. Our decoder module Ses.. is kept very simple consisting of four deconvolutional layers.
The key novelty of our approach lies in the unique structure of the latent space mimicking the interplay of realization-speccfie information and realization-invariant dynamical model similar to the frame of differential equations. However, we can significantly reduce computational costs as we are not forced to solve explicitly any differential equation since we rely on an effective end-to-end learning scheme.
Generative model, inference and optimization. We now turn the descriptive technical context of method to a probabilistic model. Our graphical model (see Figure the Appendix) consists of trainable parameters $ = hene U Øayn U Ødee. a random variable " which additionally acts as global random variable at the level of latent states st, and observations 2t,,- The index bg refers to specific queried time point within a trajectory. The joint distribution is given by
p(z, 2, v²) = p(z, (1)p)())))))))) p(ala)p((ll"))p(0)
(I)
Our graphical model assumes these independencies: (i) The dataset contains i..d trajec-
tories of varying length. (ii) The observation of trajectory Ji, at time ty is conditionally independent of xi,-, at time ty-1, given latent states zi, and trajectory representation t²:
12,, 16*)) L. p( 1 1=, vt"). Analyzing data with this graphical model involves computing posterior distributions of hidden variables given observations
p(r, 2, )²r) p(z, 4 (x) J p(x for )p( (ddder
To effectively process long-horizon time series data, we apply a variant of multiple shooting. However, since our model does not rely on an explicit ODE formulation, we are not concerned with turning an initial value problem into a boundary value problem (Massaroli et al. (022). Instead, we incorporate a Bayesian continuity prior (Hegde et al. [022 lakovlev et al. 2023) to extend the multiple-shootin framework from deterministic neural ODEs to probabilistic context. Our approach dissects each realization o².r into a series of N overlapping subtrajectories and independently condenses each patch into latent representation. Within this Bayesian multiple shooting framework, the smoothness prior connects the patches via
N II (( (T-1)
p(=[w")
to form a cohesive global trajectory. We leverage the independence of trajectory representations in subpatches i.e. p(=])1²) IL p(2)117f). For the continuity prior, we follow Hegde et al. (2022) and
place Gaussian prior on the error between consecutive subtrajectories, i.e. A ~ N (0,0 aa) entailing exact overlapping if A + O. This yields our continuity prior
p(zn Zn 1, 1)-1) N(( I=m" vi (1-)),
5

---

Published as a conference paper at ICLR 2024
Multiple Shooting
Multi- Multi-length Training
Increase training trajectory lengths
Neg. Log Like lhood jonmenation
Phase
Phase
Phase N
...
Variational Representation Loss [Representation spoct)
Use firt detapoints
Trajectory Average
N(0.1)
Evaluation
RMSE contly last sampless
Testing
Variational Smoothness
NN
Time Bagaa
(Erbokling (1)
Prediction entire Trajectory
Figure 2: Training scheme with losses. and test/evaluation procedure. Top left: Multiple Shooting Multi-length Training. An input trajectory split into subpatches. Subtrajectory length increased in multiple phases. Bottom left: Testing: only the first few points are used to roll-out the latent trajectory and transformed to the observational space. Evaluation: Last 7 samples of the predicted trajectory are used to compute the evaluation metrics, the average of the summed normalized mean squared error and a time history diagram showing the error evolution. Right: Loss consisting of three parts: negative log likelihood loss penalize reconstruction errors representation loss to define gradient field between representations, smoothness loss to penalize jumps between latent subpatches.
where the time index -L refers to the last time point of subpatch. The prior trajectory representation set to a Gaussian, i.e. p(x)") ~ (O, 1). With the priors introduced above, we get the following generative model (we drop the subpatch index for improved readability):
p((gmb/>) (x) (f (xxx, oo (I)))
(5)
p(b" [x) fagg (1gmi))
(6)
p( w", x) feav (vb"", tq)
(7)
p<x]2) -N(Ssa (2), ddcc)
(8)
For inference, we use Gaussian approximations and dee 10-2, We then seek minimize the KL divergence KL[g(z, v² [P((,, s [r)], essentially equivalent to maximizing the ELBO in eq. A full derivation of this ELBO can be found in Section Cf the Appendix.
max E (a
In (n..)
(i) likelihood
(ii) representation prior
(ii)) smoothness prior
EXPERIMENTAL SET-UP
Experiments are structured into four different series that shed light on the performance of LaDID, We provide a short overview of the experimental set-up in the following. Further details can be found in Section in the Appendix.
Datasets. We consider a wide range of physical systems ranging from relatively simple ODE- based datasets complex turbulence-driven fluid flows. Specifically, we consider high-dimensional observations (p=16, 384) from: a nonlinear swinging pendulum; swinging double pendulum; realistic simulations of the two-dimensional wave equation; lambda-omega reaction- -dfffion system; the -vo-dimensional incompressible Navier-Stokes equations; and the fluid flow around blunt body solved via the latticed Boltzmann equations. This extensive range sheds light on performance on complex datasets relevant to real-world use-cases, including models frequently used in the literature on dynamical modeling. Regular and irregular time grids are included. We study also the challenging problem of making predictions in completely novel setting obtained by intervention on the system. This is similar in spirit to experiments seen in causal Al and in this case involves generating small datasets intervened dynamical systems (either via modifying the underlying systems, for example by changing the gravitational constant or the mass pendulum, or via augmenting the realization-specific observation, e.g. by changing the length of pendulum the

---

Published as a conference paper at ICLR 2024
location of a simulated cylinder) and fine-tuning pre- pre-trained model on fraction of the data in the target setting. We direct the interested reader to Appendix[]fov further details.
Training. Training is carried out in multi-phase schedule w.I.t. the multiple shooting loss in eq. 9 In the different phases, we split the input trajectory into overlapping patches and start learning by predicting one step ahead. We double the number of prediction steps per patch every 3000 epochs meaning that learning is done on longer patches with decreased number of patches per trajectory (where trajectory length is not divisible by number of steps, we omit the last patch and scale the loss accordingly). In the final phase, training is carried out on the entire trajectory. All network architectures are implemented in the open source framework Py PyTorch (Paszke et al. 2019). Further training details and hyperparameters can be found in Appendix F
Testing. We test the trained models on entirely unseen trajectories. During testing, the first k=10 trajectory points are provided to the trained model. Based on these samples, an RS representation eh" is computed and used to roll out the trajectory to the time points of interest. Finally, predictions and ground truth observations are compared using the evaluation metries below.
Evaluation metries. We consider mean squared error (MSE) over trajectories: inference runs over 27 steps with MSE computed over the last T timesteps, allowing assessment for relatively distant times (relative to the reconstruction MSE). We set T = 60 for all experiments, with MSE normalized w.r.t. average (true) intensity, as recommended in Zhong et al. (2021); ); [Botev et al. (202T). Additionally, we provide time history diagrams plotting root mean square error (RMSE) against normalized time (mapping the interval [T. 27' to the unit interval). Metrics are averaged across all test trajectories and five runs, with mean and 75% inter-quantile ranges (IQR) reported. Subsampled predictions and pixelwise T.2 error of one (randomly chosen) trajectory is shown for visual inspection; we acknowledge that these cannot always be representative and should be considered alongside formal metries. See Figure for intuition on the train/test procedure and metrics.
Comparisons. We compare our approach to recent models from the literature, including ODE-RNN (Rubanova et al. 2019). NDP (Norcliffe et al. 2021), ODE2VAE (Yildiz et al., 2019). and MSVI siaeie--immuattedde
(lakovlev et al.) 2023). In common with LaDID these models feature structures and seek to learn low- low-dimensional latent dynamics. ODE2VAE simulates latent trajectories in straightforward fashion using BNN to model the underlying dynamics. In contrast, ODE-RNN, NDP and MSVI leverage neural ODE solvers to integrate latent states forward in time. Further details regarding these baselines can be found in Section G of the Appendix.
RESULTS
First, we examined performance synthetic data for which the training and test data come from the same dynamical system. This body of experiments test whether the model can learn to map from finite, empirical dataset to an effective latent dynamical model. Second, we examine few- shot generalization to data obtained from systems subject to nontrivial intervention (and in that sense strongly out-of-distribution interventions, i.e. interventions upon the mass or length of the pendulum, changes to the Reynolds number, or variations the camera view on the observed system, and apply the learned inductive bias to new and unseen interventional regimes in a few-shot learning setting. This tests the hypothesis that the inductive bias of our learned latent dynamical models can be useful proxy for dynamical systems exposed to a number of interventions.
In particular, we train our model
set of trajectories under
6.1 BENCHMARK COMPARISONS TO STATE-OF-THE--O MODELS FOR ODE AND PDE PROBLEMS
We begin by investigating whether LaDID can learn latent dynamical models in the conventional case in which the training and test data come from the same system. We evaluate the performance of ODE-RNN, ODE2VAF, NODEP, MSVI and LaDID on the data described in Section and Section || of the Appendix with increasing order of difficulty, starting with the non-linear mechanical swing systems with underlying ODEs, before moving to non-linear cases based PDEs (reaction-diffusii system, 2D wave equation, von Kármán vortex street at the transition from laminar to turbulent flows, and Naiver-Stoke equations). Due limited space, we only present results for subset of performed experiments but refer the interested reader to Appendix [o a detailed presentation of all results.
Applications to ODE-based systems. For visual inspection and intuition, Figure provides predicted observations á. of a few time points of one test trajectory of the single pendulum dataset for all tested algorithms, followed by the ground truth trajectory and the pixelwise Lggerror. In addition, Figure
7

---

Published as a conference paper at ICLR 2024
presents the normalized MSE over entire trajectories averaged across the entire test dataset and the evolution the RMSE over time for the second half of the predicted observations averaged over all test trajectories (see Section 5 is provided in the Appendix. Across all ODE-based datasets LaDID achieves the lowest normalized MSE. The time history diagram (see Figure K. in the Appendix) reveals gains using LaDID for long-horizon predictions relative all other algorithms tested.
This can also seen by visual inspection in Fig- as for other approaches the predicted states at later time points deviate from the ground truth trajectory substantially while LaDID's predic- tions essentially follow the ground truth. Con- sidering only the baselines, one can observe that MSVI (a sophisticated, recently proposed ap- proach), predicts accurately within short-term horizon but nonetheless fails long-hone hong-horizon predictions. The results for the challenging double pendulum test case can be found in the Appendix.
normalized MSE
30
ODE-RNN
ODEZVAE
MSVI
10-
Ours
models
Figure 3: Test errors Normalized MSE
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODE2VAE
MSVI
Ours
truth
0.00 0.15
Figure 4: Left: predicted test trajectory at various timesteps 6, right: corresponding pixelwise error.
Applications to PDE-based processes. We additionally evaluated all baselines and our proposed
method on PDE-based processes. Due to space restrictions, focus analysis the flow evolution characterized by the Navier-Stokes equation in the two dimensional case, which is of great
importance in many engineering tasks, e.g. the analysis of internal air flow in a combustion en- al., 2022). drag reduction concepts in the transportation and energy sector Lagemann al.) 2023a);
normalize Mst
1.0
2.19 1.56)
gine (Lagemann
0.8
NDP
DDE2VAE
(Gowree et 2018
Ours
Mäteling et al. 2023). and many more. Re- sults Figure show that LaDID clearly out- performs all considered comparators. The nor- malized MSE is the lowest and the averaged
0.2
0.0
models
Figure Test error Normalized MSE
RMSE is also the lowest at any time. This is echoed in the other experiments whose results are presented in detail in Section K in the Appendix.
Overall, these results support the notion that LaDID achieves good performance for challenging ODE and PDE based systems. We direct the interested reader Section K for the complete collection of experimental results supporting this statement. In this context, Table H.2 of the Appendix highlights the massively reduced computational ressources required during training and inference since LaDID eschews an explicit neural ODE formulation (including the costly ODE solvers), making it efficient and highly scalable in practice.
Performance on regular and irregular time grids. Here, we study the performance of LaDID on regular and irregular time grids and compare it to other neural-dynamical models (which are able deal with irregular time series data). As shown Figure L. and Figure L.2 in the Appendix, the proposed LaDID performs very similarly on both types of the time grids relative to both ODE- based benchmark examples and challenging PDE-based real-world systems, outperforming existing methods demonstrating strong and robust performance on irregularly sampled data.
Effects of relevant network modules. LaDID leverages three key features: a reconstruction embedding, seatio-temporal attention module and specifically designed loss heuristic leam temporal dynamics from empirical data. We investigated the importance of these modules (results appear in the Section of the Appendix). First, compared LaDID with ablated counterparts, e.g.
8

---

Published as a conference paper at ICLR 2024
pure reconstruction loss and loss combinations either using the representation smoothness loss. Overall, the proposed loss heuristic appears to stabilize training and yields the lowest MSE and IQR values. Second, we compared LaDID to counterparts trained on ablated attention modules. Empirical results underline the utility of the applied spatio-temporal attention. Finally, Table N.3 further shows the usefulness of the representation-speciine encoding. This representation encoding be thought of learning-enhanced initial value stabilizing the temporal evolution of latent trajectory dynamics. Moreover we study the effect of restricted training trajectories on the performance of LaDID in Section K.7 of the Appendix to better understand efficiency under limited data.
6.2 GENERALIZING TO NOVEL SYSTEMS VIA FEW-SHOT LEARNING
Here, we assess LaDID's s ability to generalize to novel system obtained by nontrivial intervention on the system coefficients themselves (e.g., mass, length, Reynolds number). Such changes can induce large changes to data distributions and can be viewed through a causal lens (see also Appendix [) In particular, we train dynamical model on a set of interventions and fine-tune it to new intervention regimes with only few samples, finally evaluating performance on entirely unseen dataset. We compare the performance of our prior-based few-shot learning model with a model trained solely on the fine-tuning dataset bssrrach--traines model). In our first experiment, we use the single pendulum dataset and test the transferability hypothesis on fine-tuning datasets varying sizes. The results show that the prior-based model outperforms the scratch- and achieves comparable performance to the model trained on the full dataset with a fine-tuning dataset size of 32%. At fine-tuning dataset size of 8%, LaDID produces partially erroneous but still usable predictions, which are only slightly worse than the predictions of an advanced NODE based model, MSVI, trained the full dataset. Further results including robustness to input noise and color shifts in the observation space appear in Section Min in the Appendix.
model at all fine-tuning dataset sizes,
Second, we investigate the effect of interventions ferability to new observation settings on the von Kármán vortex street dataset. We re-simulate different cylinder locations (shifted the cylinder to the left, right, up, and down) and evaluate the performance under different fine-tuning dataset sizes. The results show that the prior- based model consistently outperforms the -rruchhttraieed model and produces accurate and usable predictions under new observation conditions with a fine-tuning dataset size of as little as 8%. These findings support our hypothesis that LaDID is capable of extracting general dynamical models from training data. Additional transfer learning experiments are detailed in the Section P of the Appendix, studying the model's performance when jointly trained on a dataset encompassing Reynolds numbers of Re [100, 250, 500] and subsequently applied for zero-shot predictions on unseen He numbers.
the observation process by testing the trans-
normalized MSE
25
dataset scratch
0.20
full dataset
scratch
20
prior based
0.15
prlor based
0.10
10
5-
0.05
0.00
4%
8% percentage of seen data samples
16%
32%
4%
8%
16%
32%
percentage seen samples
(a) single pendulum
(b) vortex street
Figure 6: Test errors for a set of transfer learning experiments.
7 CONCLUSIONS
In this paper, we presented a novel approach called LaDID aimed at end-to-end learning of la- tent dynamical models. LaDID uses a novel transformer-basee architecture that leverages certain coentiifcalllymooivva invariances to allow separation a universal dynamics module and encoded realization specific information. We demonstated strong performance several new and challenging test cases and well-known benchmarks. Additionally, we showed that LaDID can generalize to systems under nontrivial intervention (when trained on the un-intervened system) using few-shot learning. Currently, while LaDID accommodates irregular time sampling, data acquired from ir- regular spatial grids will need further work. A future research direction is explore graph-based methodologies to address this specific challenge.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work was partly supported by the German Federal Ministry of Education and Research (BMBF) project "LODE", the UK Medical Research Council (MC-UU-00002/117 and the National Institute for Health Research (Cambridge Biomedical Research Centre at the Cambridge University Hospitals NHS Foundation Trust). The work of CL funded by the Deutsche Forschungsgemeinscaatt within the Walter Benjamin fellowship LA 5508/1-1. We gratefully acknowledge the Gauss Centre for Supercomputing V. for supporting this project via computing time the GCS Supercomputers.
REFERENCES
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32, 2019.
Roberto Benzi, Sauro Succi, and Massimo Vergassola. The lattice Boltzmann equation: theory and applications. Physics Reports, 222(3):145-197, 1992.
Prabhu Lal Bhatnagar, Eugene P Gross, and Max Krook. A model for collision processes in gases. I. Small amplitude processes in charged and neutral one-component systems. Physical review, 94(3): 511, 1954.
Hans Georg Bock and Karl-Josef Plitt. A multiple shooting algorithm for direct solution of optimal control problems. IFAC Proceedings Volumes, 17(2):1033-1688, 1984.
Aleksandar Botev, Andrew Jaegle, Peter Wirnsberger, Daniel Hennes, and Irina Higgins. Which priors matter? Benchmarking models for learning latent dynamics. Advances in Neural Information Processing Systems, 34, 2021.
Manuel Brenner, Florian Hess, Jonas M Mikhaeil, Leonard F Bereska, Zahra Monfared, Po-Chen Kuo, and Daniel Durstewitz. Tractable dendritic RNNs for reconstructing nonlinear dynamical systems. In International Conference on Machine Learning, pages 2292-2320. PMLR, 2022.
Steven L. Brunton, Joshua Proctor, and Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 113(15):3932-3937 2016
Adrian Bulat, Juan Manuel Perez Rua, Swathikiran Sudhakaran, Brais Martinez, and Georgios Tzimiropoulos. Space-time mixing attention for video transformer. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 19594-19607. Curran Associates, Inc., 2021.
Ling Cai, Krzysztof Janowicz, Gengehen Mai, Bo Yan, and Rui Zhu. Traffic transformer: Capturing the continuity and periodicity of time series for traffic forecasting. Transactions in GIS, 24(3):
736-755, 2020.
Kathleen Champion, Bethany Lusch, Nathan Kutz, and Steven L Brunton. Data-driven discovery of coordinates and governing equations. Proceedings of the National Academy of Sciences, 116(5)) 22445-22451, 2019.
Ricky T. Q. Chen, Brandon Amos, and Maximilian Nickel. Learning Neural Event Functions for
Ordinary Differential Equations. International Conference Learning Representations, 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances Neural Information Processing Systems. 31, 2018.
Matthew Choi, Daniel Flam-Shepherd, Thi Ha Kyaw, and Alán Aspuru-Guzik. Learning quantum dynamics with latent neural ordinary differential equations. Physical Review A, 105(4):04243,, 2022.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks, arXiv preprint XxXi:2003.046600 2020.
10

---

Published as a conference paper at ICLR 2024
Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. GRU-ODE BRU--SD--Baee: Continuous modeling of sporadically-obsevee time series. Advances in neural information processing systems, 32, 2019.
Thai Duong and Nikolay Atanasov. Hamiltonian- based Neural ODE Networks on the SE(3) Manifold For Dynamics Learning and Control. In Proceedings of Robotics: Science and Systems, 2021.
Emilien Dupont, Arnad Doucet, and Yee Whye Teh. Augmented neural odes. Advances in Neural Information Processing Systems, 32, 2019.
Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world jacobian and kinetic regularization. In International conference on machine learning, pages 3154-3164. PMLR, 2020.
Mare Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying hamiltonian and lagrangian neural networks via explicit constraints. Advances in Neural Information Processing Systems, 33, 2020.
lean-Yves Franceschi, Edouard Delasalles, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. Stochastic latent residual video prediction. In International Conference on Machine Learning, pages 2233-3246. PMLR, 2020.
Zhihan Gao, Xingjian Shi, Hao Wang, Yi Zhu, Yuyang Bernie Wang, Mu Li, and Dit-Yan Yeung. Earthformer: Exploring space-time transformers for earth system forecasting. Advances in Neural Information Processing Systems, 35:25390-25403, 2022.
Arnab Ghosh, Harkirat Behl, Emilien Dupont, Philip Torr, and Vinay Namboodiri. Steer: Simple temporal regularization for neural ode. Advances in Neural Information Processing Systems, 33: 14831-14845 2020.
Erwin R Gowree, Chetan Jagadeesh, Edward Talboys, Christian Lagemann, and Christoph Brücker. Vortices enable the complex aerobatics of peregrine falcons. Communications biology, 1(1):27, 2018.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse Problems, 34 (T):014004 2017.
Dieter Hänel. Molekulare Gasdynamik: Einführung die kinetische Theorie der Gase und Lattice- Maltamann-Methoddae Springer- Verlag, 2006.
Pashupati Hegde, Çağatay Yildız, Harri Lähdesmäki, Samuel Kaski, and Markus Heinonen. Varia- tional multiple shooting for bayesian odes with gaussian processes. In James Cussens and Kun Zhang, editors, Proceedings of the Thirty-Eighth Conference on Uncertainty Artificial Intelli- gence, volume 180 of Proceedings of Machine Learning Research, pages 790-799. PMLR, 01-05 Aug 2022.
Markus Heinonen, Cagatay Yildiz, Henrik Mannerström, Jukka Intosalmi, and Harri Lähdesmäki. Learning unknown ODE models with Gaussian processes. In International Conference on Machine Learning, pages 1959-1968. PMLR, 2018.
Florian Hess, Zahra Monfared, Manuel Brenner, and Daniel Durstewitz. Generalized teacher forcing for learning chaotic dynamics. In International Conference on Machine Learning. PMLR, 2023.
Samuel Holt, Zhaozhi Qian, and Mihaela van der Schaar. Neural Laplace: Learning diverse classes of differential equations in the Laplace domain. In International Conference on Machine Learning. pages 881 1-832. PMLR, 2022.
Valerii lakovlev, Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Latent neural ODEs with sparse bayesian multiple shooting. In The Eleventh International Conference on Learning Representations, 2023.
11

---

Published as a conference paper at ICLR 2024
Junteng Jia and Austin R Benson. Neural jump stochastic differential equations. Advances in Neural Information Processing Systems, 32, 2019.
Rishabh Kabra, Daniel Zoran, Goker Erdogan, Loic Matthey, Antonia Creswell, Matt Botvinick, Alexander Lerchner, and Chris Burgess. Simone: View-invariant, temporally-abstrr abstracted object representations via unsupervised video decomposition. Advances in Neural Information Processing Systems, 34:20146-20155 2021.
Kadierdan Kaheman, J. Nathan Kutz, and Steven L. Brunton. SINDy-PI: a robust algorithm for parallel implicit sparse identification of nonlinear dynamics. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 476(2242), 2020.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time series. Advances in Neural Information Processing Systems, 33, 2020.
Suyong Kim, Weiqi Ji, Sili Deng, Yingbo Ma, and Christopher Rackauckas. Stiff neural ordinary differential equations. Chaos: An Interdisciplinary Journal of Nonlinear Science, 31(9), 2021a.
Timothy D Kim, Thomas Luo, Jonathan W Pillow, and Carlos Brody. Inferring latent dynamics underlying neural population activity via neural differential equations. In International Conference on Machine Learning, pages 5551-5561. PMLR, 2021b.
Diederik P Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational bayes.
arXiv preprint
Georgia Koppe, Hazem Toutounji, Peter Kirsch, Stefanie Lis, and Daniel Durstewitz. Identifying nonlinear dynamical systems via generative recurrent neural networks with applications to IMRI. PLoS computational biology, 1508)cc 2019.
Christian Lagemann, Kai Lagemann, Sach Mukherjee, and Wolfgang Schroeder. Generalization of deep recurrent optical flow estimation for particle-image velocimetry data. Measurement Science and Technology, 2022.
Esther Lagemann, Marian Albers, Christian Lagemann, and Wolfgang Schröder. Impact of reynolds number on the drag reduction mechanism of spanwise travelling surface waves. Flow, Turbulence and Combustion, pages 1-14, 2023a.
Kai Lagemann, Christian Lagemann, Bernd Taschler, and Sach Mukherjee. Deep learning of causal structures in high dimensions under data limitations. Nature Machine Intelligence, 5(11):1306- 1316, 2023b.
Longyuan Li. Jian Yao, Li Wenliang, Tong He, Tianjun Xiao, Junchi Yan. David Wipf, and Zheng Zhang. Grin: Generative relation and intention network for multi-agent trajectory prediction. Advances in Neural Information Processing Systems, 34:27 107-177777718, 2021.
Shiyang Li. Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xin Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. Advances neural information processing systems, 32, 2019.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In The Eighth International Conference on Learning Representations, 2020.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint wrXiv:2207.05209 2022.
Yuxuan Liang, Yutong Xia, Songyu Ke, Yiwei Wang, Qingsong Wen, Junbo Zhang, Yu Zheng, and Roger Zimmermann. Airformer: Predicting nationwide air quality in china with transformers. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 14329-14337, 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 171.0017 2017.
12

---

Published as a conference paper at ICLR 2024
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint wrXiv:1910.03330
2019.
Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International Conference on Machine Learning, pages 3276-3226 -285.. PMLR, 2018.
Michael Lutter, Christian Ritter, and Jan Peters. Deep lagrangian networks: Using physics as model prior for deep learning. In International Conference on Learning Representations, 2019.
Stefano Massaroli, Michael Poli, Sho Sonoda, Taiji Suzuki, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Differentiable multiple shooting layers. Advances in Neural Information Processing Systems, 16532-16566444 2021.
Esther Mäteling, Marian Albers, and Wolfgang Schröder. How spanwise travelling transversal surface waves change the hear-wall flow. Journal of Fluid Mechanics, 957:A30, 2023.
James Morrill, Cristopher Salvi, Patrick Kidger, and James Foster. Neural rough differential equations for long time series. International Conference on Machine Learning, pages 7829-7838. PMLR, 2021.
Alexander Noreliffe, Cristian Bodnar, Ben Day, Jacob Moss, and Pietro Liò. Neural {ode) processes. In International Conference on Learning Representations, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach De Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019.
Jaideep Pathak, Brian Hunt, Michelle Girvan, Zhixin Lu, and Edward Ott. Model-free prediction of large spatiotemporally chaotic systems from data: A reservoir computing approach. Physical review letters, 120(2):02400,, 2018.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms. The MIT Press, 2017.
Michael Poli, Stefano Massaroli, Junyoung Park, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Graph neural ordinary differential equations. arXiv preprint arXiv:1911.07532. 2019.
Yue-Hong Qian, Dominique d' d'Humières. and Pierre Lallemand. Lattice BGK models for Navier- Stokes equation. Europhysics letters, 17(6))479, 1992.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Jack Eichter-Powell. Yaron Lipman, and Ricky T. Q. Chen. Neural Conservation Laws: A Divergence- Free Perspective. In Advances in Neural Information Processing Systems, volume 35, 2022.
Yulia Rubanova, Ricky T. Q. Chen, and David K Duvenaud. Latent ordinary differential equations for eregularly-samplle time series. In Advances in Neural Information Processing Systems, volume 32, 2019.
Lars Ruthotto and Eldad Haber. Deep neural networks motivated by partial differential equations. Journal of Mathematical Imaging and Vision, 62:352-364, 2020.
Dominik Schmidt, Georgia Koppe, Zahra Monfared, Max Beutelspacher, and Daniel Durstewitz. Identifying nonlinear dynamical systems with multiple time scales and long-range dependencies. In International Conference on Learning Representations. PMLR, 2021.
Bernhard Schoelkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. In International Conference on Machine Learning. PMLR, 2012.
13

---

Published as a conference paper at ICLR 2024
B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchbrenner, A. Goyal, and Y. Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612-63) 2021.
Shreshth Tuli, Giuliano Casale, and Nicholas R Jennings. TranAD: Deep Transformer Networks for Anomaly Detection in Multivariate Time Series Data. Proceedings of VLDB, 1506):2011-221, 2022.
Evren Mert Turan and Johannes Jäschke. Multiple shooting for training neural differential equations on time series. IEEE Control Systems Letters, 6: 6:1897-1902, 2021.
Mark der Wilk, Matthias Bauer, ST John, and James Hensman. Learning invariances using the marginal likelihood. Advances in Neural Information Processing Systems, 31, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Pantelis R Vlachas, Wonmin Byeon, Zhong Y Wan, Themistoklis P Sapsis, and Petros Koumout- sakos. Data-driven forecasting of high-dimensional chaotic systems with long short-term memory networks. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, (7412213)201170444, 2018.
Pantelis Vlachas, Jaideep Pathak, Brian R Hunt, Themistoklis Sapsis, Michelle Girvan, Edward Ott, and Petros Koumoutsakos. Backpropagation algorithms and reservoir computing in recurrent neural networks for the forecasting of complex spatiotemporal dynamics. Neural Networks, 126: 191-217, 2020.
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. In International Conference on Learning Representations. PMLR, 2021.
Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, S Yu Philip, and Ming- sheng Long. Predrnn: recurrent neural network for spatiotemporal predictive learning. IEEE Transactions Pattern Analysis and Machine Intelligence, 45(2):2208-2233 2022.
Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
Jiehui Xu, Haixu Wu, Jianmin Wang, and Mingsheng Long. Anomaly transformer: Time series anomaly detection with association discrepancy. In International Conference on Learning Repre- sentations, 2022.
Mingxing Xu, Wenrui Dai, Chunmiao Liu, Xing Gao, Weiyao Lin, Guo-Jun Qi, and Hongkai Xiong. Spatial- -emporll transformer networks for traffic flow forecasting. arXiv preprint arXiv:2001.02908 2020.
Chao-Han Huck Yang, Yun-Yur Tsai, and Pin-Yu Chen. Voice2series: Reprogramming acoustic models for time series classification. In International conference on machine learning, pages 11808-11819, PMLR, 2021.
Cagatay Yildiz, Markus Heinonen, and Harri Lahdesmaki. ODE Deep generative second order ODEs with Bayesian neural networks. In Advances in Neural Information Processing Systems, volume 32, 2019.
Yuan Yin, Matthieu Kirchmeyer, Jean- Continuous PDE Dynamics Forecasting with Implicit Neural Representations. In International Conference on Learning Representations, 2023.
Jean-Yves Yves Franceschi, Alain Rakotomamonjy, and Patrick Gallinari.
Cunjun Yu, Xiao Ma, Jiawei Ren, Haiyu Zhao, and Shuai Yi. Spatio-temporal graph transformer networks for pedestrian trajectory prediction. In Computer Vision-ECCC 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XII 16, pages 507-523, Springer, 2020.
14

---

Published as a conference paper at ICLR 2024
Wang Zhang, Tsui- Wei Weng, Subhro Das, Alexandre Megretski, Luca Daniel, and Lam M Nguyen. Concernet: A contrastive learning based framework for automated conservation law discovery and trustworthy dynamical system prediction. arXiv preprint 2rX.v:2302.05573, 2023.
Weiming Zhi, Tin Lai, Lionel Ott, Edwin V Bonilla, and Fabio Ramos. Learning efficient and robust ordinary differential equations via invertible neural networks. In International Conference on Machine Learning, pages 27060-27074. PMLR, 2022.
Yaofeng Desmond Zhong, Biswadip Dey. and Amit Chakraborty. Symplectic ODE-Net: Learning Hamiltonian Dynamics with Control. In International Conference on Learning Representations, 2020.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Benchmarking energy-conservie neural networks for learning dynamies from data. In Proceedings of the 3rd Conference Learning for Dynamics and Control, pages 1218-1229, PMLR, 2021.
Tian Zhou, Ziqing Ma, Qingsong Wen, Xue Wang, Liang Sun, and Rong Jin. Fedformer: Frequency enhanced decomposed transformer for long-term series forecasting. International Conference on Machine Learning, pages 27268 272866 PMLR, 2022.
15

---

Published as a conference paper at ICLR 2024
A EXTENDED RELATED WORK
We begin by discussing the wide variety of methods proposed to learn dynamical models from data. Our focus is on ML-based approaches.
Operator learning and approximating dynamics. Brunton et al. (2016); Kaheman et al. (2020)
access Koopman theory to learn coefficients of a finite set of non-linear functions to approximate
the observed dynamics. Other approaches (Crammer al.] 2020; Lutter et al., 2019 Greydanus
itt al., 2019 Finzi et al. [020 Zhong et al. |020) Bai et al. (019)) induce Hamiltoman and La- grangian priors into neural networks exploiting the reformulations of Newton's 's equations of motion in energy-conservane
dynamics. This attractive inductive bias allows for accurate predictions over a significantly wider time-horizon but involves strong assumptions that may not hold for many real-world systems. Purely data-driven ML approaches have attracted much recent attention and have shown strong performance in many areas. For instance, Heinonen et al. (2018) learns a data- driven approximation of dynamical system via Gaussian Processes. Recent works focus the field of operator learning that seeks to learn mappings between
functions. Two prominent approaches are DeepONet (Lu et al. 2019) and Fourier Neural Operator (FNO; (Li et al. 2020)). DeepONet which is based on physics-informed neural networks (Raissi et al. 20T9) can query any coordinate in the domain for value of the output function. However, the input function must be observed on a set of predefined locations, requiring the same observation grid for all observations, for training and testing. FNO is an instance of neural operators, a family of approaches that integrate kernels over the spatial domain. Since this operation can be expensive, FNO addresses the problem by employing the fast Fourier transform (FFT) to transform the inputs into the spectral domain. As a consequence it cannot be used with irregular grids. Li et al. (2022) introduce an FNO extension to handle more flexible geometries, but it is tailored for design problems. To summarize, despite promising results for several applications, current operator approaches still face generalization limitations, and more importantly require access to the system's 's variable space to work sufficiently. As result, the performance of these approaches drops dramatically if applied to longitudinal system representations and not the system itself.
Incorporating strong physical priors in dynamics modeling. In real-world dynamical systems it is often infeasible to have reliable analytical model of the underlying processes. In such cases, a more general approach is learn and capture the latent dynamics of the data using an architecture that incorporates an appropriate inductive bias. Hamiltonian and Lagrangian mechanics offer distinct mathematical reformulations of Newton's 's equations of motion, specifically for mergy-conservativv dynamics. The conservation of energy in these systems enables predictions of the system's 's state over significantly longer time horizons, both forward and backward, compared to the training data. This property makes them attractive biases to incorporate into deep neural networks. Despite the different coordinate frames they employ (Hamiltonian using position g and momentum p. and Lagrangian using position and velocity (). both formalisms describe the same underlying dynamics, allowing for seamless translation between the two without any loss of generality. One advantage of these models is that they only need to infer the Hamiltonian or the Lagrangian, without the additional burden of learning a state representation. This also simplifies evaluation, as requires calculating the distance between the ground truth states and the states predicted by the model along the trajectory. To tackle this challenge, several approaches have been proposed that augment physics-inspired models with encoder/decoder modules Cranmer al. (2020): Lutter al. (0119; Greydanus al. (2019); Finzi iet al.) (2020); Zhong et al. (2020); Bai et al. (201)); Wang et al. (2021); Zhang et al. (0023). These modules facilitate the inference of low-dimensional states from high-dimensional pixel observations bridging the gap between visual input and the underlying dynamics.
Neural ODEs. In recent line of research, the ability of NNs to approximate arbitrary functions has been exploited for unknown dynamical models, with connections drawn between deep architectures and numerical solvers for ODEs, PDEs and SDEs (Chen et al. 2018; Weinan 2017 Lu et al.) 2018, Ruthotto and Haber 2020 Haber and Ruthotto 2017 7 Richter-Powell et al. 2022). This has given rise to new algorithms rooted in neural differential equations that have been shown to provide improved modeling of a variety of dynamical systems relative to standard recurrent neural networks (RNNs) and their variants. However, Neural ODEs inherit some key limitations of ODEs. Specifically, the temporal dynamics only depends on the current state but not on the history. This puts theoretical limit on the complexity of the trajectories that ODEs can model, and leads to practical
16

---

Published as a conference paper at ICLR 2024
consequences (e.g. ODE trajectories cannot intersect). Some existing works mitigate this issue by introducing latent variables (Rubanova et al. 2019) or higher order terms (Yildiz et al. (201)). Other attempts include combining recurrent neural nets with neural ODE dynamics (De Brouwer et al. 2019), where latent trajectories are updated upon observations, as well as upon Hamiltonian (Zhong et al., 2020), Lagrangian (Lutter al. 20T9), second-order (Yildiz et al. 2019), graph neural network based dynamics (Poli et al. 2019). 2019]. Dupont et al. (2119) put forward a simple augmentation to the state space broaden the range of dynamical systems that can be learned; Rearned et al. (2020) propose neural controlled differential equations; a mechanism to adjust the trajectory based on subsequent observation denoting the natural equivalent to recurrent neural networks, the work of Morrill et al. (202T) incorporate concepts of rough theory to extend applicability to trajectories spanning thousands of steps, and and Benson (2019) explore extensions to account for interrupting stochastic events. However, they still operate in the ODE framework and cannot model broader classes of differential equations. Another limitation of Neural ODE and extensions is that they struggle to model certain types of dynamics due to numerical instability. This is because Neural ODE relies on a numerical ODE solver to predict the trajectory (forward pass) and to compute the network gradients (backward pass). Two common scenarios where standard numerical ODE solvers fail are (1) systems with piecewise external forcing or abrupt changes (i.e. discontinuities) and (2) stiff ODEs (Ghosh al., 2020), both are common in engineering and biological systems. Some existing works address this limitation by using more powerful numerical solvers. Specifically, when modelling stiff systems, Neural ODE requires special treatment its computation: either using a very small step size or a specialized IVP numerical solver (Kim et al. 2021a). Both lead to a substantial increase in computation cost.
RNNs for modeling spatio-temporal dynamic systems. RNNs, with their capability of capturing temporal dependencies in time sequences, have found extensive use in dynamic system modeling. A notable example is the data-driven forecasting method for high-dimensional chaotic systems using LSTM networks, as proposed by Vlachas et al. (018)). These networks can be used to infer high-dimensional dynamical systems in a reduced order space and have proven to be effective nonlinear approximators of their attractor. In subsequent study|Vlachas et al. (200)) studied backpropagation algorithms and explored the merit of Reservoir Computing (RC) (Pathak et al. (2018) and Backpropagation through time (BPTT) in RNNs for predicting complex spatio-temporal dynamics. The study revealed that RC outperforms BPTT in terms of predictive performance and capturing long-term statisties when the full state dynamics are available for training, but in the case of reduced order data, large scale RC models appear to be more unstable and tend to diverge faster compared to the BPTT algorithms. On the other hand, BPTT methods demonstrated superior forecasting performance in reduced order systems and better captured dynamics in these challenging conditions. In a different line of research, Koppe et al. (2019) proposed a state space model (SSM) leveraging generative piccewise-lineat recurrent neural networks (PLRNN) to infer dynamics from neuroimaging data (IMRI). They outline the superior interpretability of the PLRNN model with respect to latent embeddings and showed that it is amenable to systematic dynamical systems analysis while allowing a straightforward transformation into an equivalent continuous time dynamical system. In follow-up study, the same group introduced state space model based dynamically interpretable and mathematically tractable RLRNN augmented by a linear spline basis expansion (Brenner et al. (2022)). This approach retains all the theoretically appealing properties of the simple PLRNN while boosting its performance and accuracy when dealing with arbitrary nonlinear dynamical systems in comparatively low dimensions. That is, they empirically showed that the dendritically expanded PLRNN achieves better reconstructions with fewer parameters and dimensions on various dynamical systems benchmarks. In a similar vein, the same group of authors proposed a simple modification of the teacher forcing which allows for faithful reconstructions on noisy real-world data and yields an even more improved interpretability HHes et al.] (2023). In complementary work, /Shhmidt et al. (202T) targeted a simpler regularization scheme for vanilla RNNs based on ReLU activation which aims to solve long-range dependency problems. Precisely, the proposed approach aims to express slow time scales while retaining a simple mathematical structure which makes their dynamical properties partly analytically accessible. Both proposed theorems demonstrate strong link between the dynamics of regularized RNNs and their gradients, showcasing the effectiveness of their methodology across various benchmarks. Recently, Wang et al. (2022) presented PredRNN, ofte-offthe-art recurrent network for predictive learning of spatiotemporal sequences that allows to model systematically visual appearances and temporal
17

---

Published as a conference paper at ICLR 2024
dynamics of spatiotemporal observations. Based on the overall goal to generate future images by learning from the historical context, the authors proposed that visual dynamics have modular structures that can be learned with compositional subsystems using convolutional gated networks. In PredRNN, a pair of memory cells are explicitly decoupled to form unified representations of the complex environment and operate in nearly independent transition manners. Based novel zigzag memory flow, this network enables the learning of visual dynamics at different levels of RNNs.
Multiple shooting approaches. Multiple shooting breaks up the time grid of an observed trajectory into some number of segments, and using single shooting method to solve for each segment. The idea behind this multiple shooting method stems from the observation that long integration of dynamics can be counterproducttive for transcribing continuous optimal control problems, and this problem can be tackled by limiting the integration over arbitrarily short time intervals (Bock and Pliti 1984). Hence, multiple shooting solves the ODE separately on each interval. In the context of neural ODE modeling, Turan and Jäschke (202I) proposed to combine multiple shooting with existing approaches drastically reducing the complexity of the loss landspape. That is, for long range time horizon optimal control oscillation data driven problems, multiple shooting can be used to train the neural networks which helps to escape "flattened out" fitting for oscillation data. Concurrently, Massaroli et al. (202T) showed that multiple shooting can also be combined with neural differential equation and introduced differential multiple shooting layers. [akkvllv [aakvveel et al. (2023) extended this concept to sparse Bayesian multiple shooting, with both works evaluating latent NODEs.
Attention-bann based time-series modeling, The innovation of attention-based layers in deep learning (Vaswani et al. 2017) has been central in recent years in natural language processing, computer vision, and speech processing. Transformers have shown the ability to model long-range dependencies and interactions in sequential data and thus are relevant for time series modeling. Many variants of transformers have been proposed to address special challenges in time series modeling and have been successfully applied to various time series tasks, such as forecasting (Li et al., 2019 /hou et al., 2022), anomaly detection (Xu et al. 2022 Tuli et al. (022), and classification (Yang et al. 2021). In spatio-temporal forecasting, both temporal and spatiotemporal dependencies are taken into account in time series models for accurate forecasting. For instance in challenging computer vision tasks, such as traffic analysis and trajectory forecasting, Cai et al. (2020); Xu et al. (2020) propose an autoencoding structure that uses self- ceff-attention module to capture temporal-tempie
dependencies and graph neural network module to model spatial dependencies. In Yu et al. (2020) an seeention-based graph convolution mechanism introduced that is able learn complicated temporal-spatial attention pattern to trajectory prediction. Targeting a different research topic, Earthformer (Gao et al.) 2022) proposes a cuboid attention for efficient space-time modeling, which decomposes the data into cuboids and applies cuboid-level self-attention in parallel. It is shown that this framework achieves good performance in weather and climate forecasting. Recently, Liang al. (2023) proposed a novel spatial self-attention module and causal temporal self- module to efficintlyy capture spatial correlations and temporal dependencies, respectively. Furthermore, it enhances transformers with latent variables to capture data uncertainty. Multivariate and spatio-temporal time series are becoming increasingly common in applications, calling for additional techniques to handle high dimensionality, especially the ability to capture the underlying relationships among dimensions. Recently, several studies have demonstrated that the combination of GNN and attentions could bring significant performance improvements in multimodal forecasting (Li et al. 2021). enabling better understanding of the spatio-temporal dynamics and latent causality.
Our approach is inspired by this body of work in that we also use neural networks to learn latent dynamical models. However, two key differences are as follows. First, our models are specifically designed to separate input trajectories into realization-specii leverage particular kind of scientifically motivated metttiiiallly-motvatteded inductive bias from the outset. Second, exploiting these invariances allows us to eschew explicit neural ODEs altogether, providing an arguably simpler, transformer- based scheme that can be trained in a straightforward fashion, but that, as we show, achieves excellent performance on unseen data from complex dynamical systems and that can even be leveraged for few-shot learning to generalize to nontrivial system interventions.
and iaalization-invarian and hence,
18

---

Published as a conference paper at ICLR 2024
B DETAILS OF RISK-BASED FRAMING
In the main text we introduced the sufficient encoding assumption or SEA. Here we show how under SEA, the learning task can be framed in particularly simple way. Note that in this section, consider simple, abstract version of the full problem, with the aim of understanding how the problem can be usefully placed in learning framework. The notation in this section follows from Section 3.2 in the main text and is otherwise self-contained and intended to be readable. We emphasize that the development below does not map one-to-one with the architecture in Section 3.3 in the main text, but rather explains, at a high level, why such a scheme is at all possible in this context. The architecture in the main text can be regarded as one way to instantiate the general ideas below.
To fix ideas, we first consider the simple case with no latent space and then consider the latent dynamics case. Consider supervised training of a NN, with training (input, output) pairs of the form
Or) input
3
output
where, in Section 3.2 in the main text, 0, is candidate RS encoding, is time point and z! the observation at time t in realization r. The goal is to learn a mapping that takes the input (the RS encoding and query time point) and yields the correct system output, for any test pair (t', r²). Consider the prediction a{{) = f(t: m(0,; A, 0m), 9). and note that we can write the RHS as h(t,, () where @ = (0, 0m)) is RI parameter. This emphasizes the fact that the RHS is in fact function (here h, itself involving combination of f and m) of only the inputs (t. f.), and therefore potentially learnable from training pairs. Note that the parameters of h are entirely RI and hence the only RS information is carried by the encoding 0,. It is easy to see that under SEA this construction provides the correct output by writing the RHS as: f(t; (0,; 8 0))),)
ft: f,,, 8) 2;5", i.e. the correct output. Thus, combining encoding 0, and function h allows prediction of the time evolution of any realisation. In other words, even if the RS encoding is distant from the true RS parameter, under SEA there exists a RI function that can correct it, and the second NN aims to learn a function h which combines these RI elements to provide the desired mapping.
Next, we briefly summarise arguments that extend the analysis to the case of latent dynamics. As in the main text, consider dynamics at the level of latent variables with RS and RI parts but at the level of the latents, i.c. 2 ft; 0,,, 0). Recall that we assume the observables are given by (an unknown) function g of the hidden state 2.
R4 and again consider model
We first consider the case in which the latent-to-observed mapping is RI, and then the more general case of a RS mapping.
Case I: RI mapping. Assume the observable is given as x.' = g=<;;; 0g), where g RR - TR is the (true) observation process and fg is an RI parameter. Further, assume that we have an estimate 0, of the RS parameter that satisfies the sufficient encoding assumption (SEA). In a similar fashion, assume we have an estimate 6g. which may be incorrect (in the sense of fg * fg)) but that satisfies:
3my 30m, 0, my mg(0g (0g:0. m,,))
(10)
That is, 0g admits an RI correction. As above, the correction is unknown and may potentially depend on true, RS parameters. Note also that subject to the existence of correction the estimate (and implied mapping) may be potentially arbitrarily incorrect. In analogy to SEA, we call this the sufficient mapping assumption (SMA).
Now, consider training of a NN, with training (input, output) pairs of the form
(tt,8,) (t.)
(11)
We want to understand whether supervised learning of a model to predict output for arbitrary queries (t, r) is possible. This is not obvious, since we now have training data only at the level of the observables, but the actual dynamics operate at the level of latents. Consider the following function hSMA:
hsMa(l, 0, 0) = (t: 0, 0m); my (0g; 0, 0m,,)
(12)
19

---

Published as a conference paper at ICLR 2024
where G (0, mm m,g)) is RI parameter.
Under SEA and SMA it is easy to see that hSMA provides the correct output, since:
hsMalt, 0, θ)
- ((:f(t:m(Be 0m): (g(gg; g:: 8, 0m,))
= ((f 0,, 0);
g(2)f
Case II: RS mapping. Suppose now the mapping is RS, with the model specification as above but with the observation step =7 = g(z;; 0"), where og is an RS parameter. This means that the latent-observable relationship is itself non-constant and instead varies between realizations.
Assume we have candidate estimate 0g which may be incorrect in the sense of 0g * 0g but that satisfies:
3mg 30m, 0 my (8 m,.,))
(13)
In analogy to SMA, we call this the realization-specific sufficient mapping assumption or RS-SMA. Now to create training sets, we extend the formulation require input triples, as:
{11, 0,, 07 2t (t.r)
(14)
Note that 0r in the input triples may be incorrect, but only needs to satisfy RS-SMA. As in Case I, we have training data only at the level of observables (not latents) but want to understand whether supervised learning of model to predict output for arbitrary queries (f,r) is possible. Consider the function hRS-SMA
θ) = gff (( 0, m,), 0): mg ( 0, 0m,)) hRS- 0r, 0)
(15)
where = (0, Om, 0m,) ) is an RI parameter. In a similar manner to Case I, it is easy to see that hRS SMA A provides the correct output under SEA and RS-SMA, since:
hns-sma 0. i θ) g (t: m00 0, 0): my my(0); 0 m,)))
(l 0r, 0):03
g 0g)
x!
The simple arguments in the foregoing, based on an abstract version of the problem of interest, are intended explain why (under the assumptions stated) is possible address the learning problem in relatively straightforward manner. The key idea is that subject to the sufficient encoding/mapping assumptions there exists function (conceptually a combination of various elements as above) that permits input-output mapping required. Note that for successful prediction, it is not required to actually explicitly identify the elements the combination (and nor do we attempt to do so). The architecture propose in the main text represents one way learn such function.
C
MODEL, APPROXIMATE POSTERIOR AND ELBO
Here, we provide details about our model, the approximate posterior and the ELBO.
Joint distribution. Our graphical model consists of parameters + &ene U Uddn U thtces random variable 1²' which additionally acts as global random variable at the level of latent states 2,, and observations Iz,, The index ty refers to specific queried time point within trajectory. The joint distribution is given by
p(x, z, 2.) 2(1,221²))
(p(o)
(16)
with
(17)
20

---

Published as a conference paper at ICLR 2024
Figure A1: Graphical model: The graph consists of three random variables, the realization-specifcc trajectory representation v²", the latent states and the observations x. Fixed parameters & are represented by the grey node. Here, the plate notation D) indicates the independence between trajectories within the dataset, while the plate notation L corresponds to the number of queried time points of the roll-out and denotes the conditional independence of predicted states given the trajectory representation " and the RI parameters 4, i.e. P(zt, w², ④) P(2t, W8, D); i 10); j.
p(=l0") II
(18)
p<>|z) Nifore (2), Odec).
(19)
Approximate posterior. Our graphical model assumes the following independencies: (i) The dataset contains i.i.d. trajectories of varying length. (ii) The observation of trajectory x;, at time by is conditionally independent of of, at time f--1, given latent states z and trajectory represen- tation " p(2t, It,, v²") L p(x (2 v²)) Analyzing data with this graphical model involves computing posterior distributions hidden variables given observations
p(5, 1p))
(l)))")))= = [p(c|))p(2/1T yo dzde
(20)
ELBO. To effectively manage long-horizon time series data, adopt variation of the multiple shooting approach. However, our model differs from those based on explicit ODE formulations, thus avoiding the need transform initial value problems into boundary value problems (Massaroli et al. 202T). Instead, we incorporate a Bayesian continuity prior (Hegde et al. 2022, lakovlev et al. framework from deterministic neural ODEs to a probabilistic context. Our approach dissects each realization Z into a series of N overlapping subtrajectories and independently condenses each patch into a latent representation. Within this Bayesian multiple shooting framework, the smoothness prior connects the patches via
2023) to extend the multiple-shootee
p(=[v)) II ((z2) 10..p//nnnnn-1//// 16 oh-1))
(21)
to form cohesive global trajectory. We leverage the independence of trajectory representations in subpatches i.c. p(2,(1/") 1 ((2))15). For the continuity prior, we follow Hegde et al. (2022) and
place Gaussian prior on the error between consecutive subtrajectories, i.e. A ~ NV a entailing exact overlapping if A + This yields our continuity prior
plzn|zn 1, W vh-1) ([ "
Ui (h-1),
(22)
where the time index -t refers to the last time point of subpatch. The prior trajectory representation set to Gaussian, i.e. (()))) ~ (N 1).
21

---

Published as a conference paper at ICLR 2024
ELBO:
In p(z, 2, 20)) ((2, 17)
In ln((x)
dzdy
(23)
(22 (5)
/ (2))))))) In g(2, 15r)
2,
d2dp6
(24)
/ (22,))) pl p(z)2))(=11))) )((""))(()))))) dzdy?
(25)
((2)((1)))
/ ((2)6) p(alz)p(())) I p(zn (ilnn/llpppp (zn 2n- v
d2dib"
(26)
((2)g(1())
N / ((2,)") II p(zn 2 dzzdy + 2(2, In p(1%)) ((v%)
dody"
(27)
n=1
I (12, plz1 46)
Im p(zn 10--)))1) d2dyy
(28)
((2n)
This equivilant to maximizing eq. (.c,:=/)((z,1vii)
9 of the manuscript with short term notation pa(#n) =
N
N
max E (2a
In p(r (nntrnlz))((nnnl (/))
Σ
Σ
E Et (L(g(z))llp(a 1nn v())
n=1
n=2
(29)
D ELBO COMPUTATION AND FORECASTING
To find the approximate posterior which minimizes the Kullback-Leibler divergence KL|g( viv) |p(2 [x
(30)
we maximize the ELBO which for our model given via
max (g
Σ KL(g(e
Σ
(a (L(g(zn)llpznnlal 1.)
(i) likelihood
(ii) representation prior
(ii) smoothness prior
(31)
The ELBO is computed as follows:
1. Encode - to obtain P
TA
2. Sample latent embeddings via zemb ~ (OFF
3. Compute latent trajectory 21.N
4. Compute ELBO via data likelihood estimation, and KL divergence of representation and smoothness prior.
Since approach leverage amortized variational inference, the reparameterization trick is applied when sampling l.mb Expectations of the data likelihood are computed using Monte Carlo integration with one sample. KL divergence is computed in closed form.
Forecasting Inference. Similar to training, also in forecasting a RS vector is derived for specific test trajectory aggregating its first K observations as detailed above. This latent representation is then evolved time based the learned dynamics model. Note that in contrast to NODE based
22

---

Published as a conference paper at ICLR 2024
approaches, here we can query arbitrary points in time which in fact allows to directly infer only the latent trajectory state of the query time steps. No evaluations on intermediate grid points are required saving substantial computational resources. Finally, inferred latent trajectory points are mapped back to the observation using the trained decoder.
E MODEL ARCHITECTURE
Encoder. The encoder is a collection of three NNs. First, features from the input observations -f- are extracted using a convolutional neural network (CNN) parameterized by Pene and shared across all representations and patches. Specifically, the CNN encoder has the following architecture: three convolution layers (5x5 kernel, stride 2. padding 2) with batch norm and ReLU activations, one convolution layer (2x2 kernel, stride 2) with batch norm and ReLU activation. The channels of the respective CNN layers are doubled throughout. Finally, the downsampled image features are flattened and linearly projected to the output dimension. Hence, our encoder transforms
the sequence of input observations to a sequence of feature vectors, 2, z(ene)-r -eme))==fe =fe (x)_).))
Then, we compute the trajectory representation and the latent embedding as follows. Each input patch split into two disjoint sets by time. The first & E K data points nn kct (n=(c(enc) to compute a trajectory specific representation distribution vo~"eon (1=
(ene).") are used N (ur, σ,)) and
Mr, fe (it ())e))") In cases of irregularly sampled trajectories, we use a time threshold 7 to define the representation set, (== (())))))) (t<r)) We model fox transformer network
with temporal attention. In other words, we consider the sequence feature vectors ((nc)., 2 as an time-ordered sequence of tokens and transform each token according to the temporal distance to the other tokens. We compared two approaches of temporal attention (lakovlev et al. 2023 Bulat et al., 2021) which performed roughly similar. First, temporal reweighting is performed as introduced in lakovlev et al. )023): Ct3y = /t// Ek= Cas with G5 = exp( Wkh + In(e) - t where ( denotes the dot product, K. WQ. and Wv represent the weight matrices for the query. key, value as in regular attention. and 5 are constants. Hence, the larger the distance It t, grows, the stronger the time-aware attention reduced (lakovlev et al. 2023). The parameter determines the distance threshold beyond which the scaling of regular attention occurs by at least c. Moreover, parameter p governs the shape of the scaling curve. This methods works best for most of the dynamical systems. Second, we tested temporal attention approach as defined in (Bulat et al.
2021 This time aware attention is given by CTA(()) -r= softmarl WWops.Keee! Wr vi Fimally, the trajectory representation sb" is obtained applying mean-aggregation of the temporally transformed representation tokens.
Dynamies model. With initial density given by the encoder networks ge (211= 1t), the density for all queried latent points (on continuous time grid) can be predicted by 24, ~N ~N( of, with 1I, =f ( (ty, 2,, )² Note that this approach allows for latent state predictions at any time since the learned dynamics module fe.... is continuous in time and our variational model utilizes encoders only for obtaining the initial latent distribution. We also make use of the reparameterization trick tackle uncertanties both, the latent states and in the trajectory representations (Kingmal and Welling 2013). In our implementation soa... consists of three linear layers, with the first two Tollowed by ReL.U nonlinearity.
The decoder maps the latent trajectory points back to the observational space. Hence, our implementation is fairly simple and comprises a set of transposed convolutional layers. In particular, it first projects latent trajectory points linearly followed by four transposed convolution layers (2x2 kernel, stride 2) with batch norm and ReLU nonlinearities. Finally, convolutional layer (5x5 kernel, padding 2) with sigmoid function computes our output distribution. The channel dimension of the four transposed convolution layers is halved subsequently from layer to layer.
Decoder.
23

---

Published as a conference paper at ICLR 2024
F TRAINING DETAILS AND HYPERPARAMETERS
Our implementation uses the Py framework (Paszke et al. 2019). All modules are initialized from scratch using random weights. During training, Adam W-Ottimizer (Loshchilov and Hutter 2017) is applied starting at an initial learning rate Ep = 0.0003. An exponential learning rate scheduler is applied showing the best results in the current study. Every network is trained for 30 000 epochs. At initialization, we start training at subpatch length of which is doubled every 3000 epochs. After the CNN encoder, attention layers are stacked each using 4 attention heads. A relative Sin/Cos embedding is used as position encoding followed by a linear layer. The input resolution of the observational image data is 128 x 128 px. All computations are run on a single GPU node equipped with one Nvidia A 1100 (40 GB) and global batch size of 16 is used. A full training run on the single pendulum, the double pendulum, the wave equation and the Navier-Stoker equation dataset requires approx. 14 h. A full training run on the reaction-diffusio system and the von Kármán vortex street requires approx. 8 h.
Table F.1: Training hyperparameters
Hyperparameter LR schedule
Value
Exp. decay
Initial LR
3e-4
Weight Decay
0.01
Global batch size
16
Parallel GPUs
Input resolution
128
128 px
Number of input timesteps
10
Initial subpatch length
Number of epochs per subpatch length
3000
Latent dimension
32
attention mechanism
spatio-temporal
Number of attention blocks
Number of attention heads
Position Encoding
relative Sin/Cos encoding
G DETAILS ON COMPARISON BASELINES
All baseline methods are similar our proposed approach in the sense that they encode longitudinal observations in a latent space, simulate low-dimensional latent trajectory and decode such to obtain future observations. Similar to our approach, NDP (Norcliffe al.] 202T) encodes observations into two latent variables (an "initial state" and "global control of an ODE") and aggregates latent representations a global representation via averaging. Dynamies are modeled using MLPs or convolutions and integrated over time using neural ODFs. The decoder outputs a Bernoulli distribution from which the prediction is sampled.
From a high-level perspective, MSVI (lakovlev et al. 2023) works similar, but leverages slightly modified encoder. Here, a transformer module added while dynamics function and decoder are Bayesian MLPs and CNNs whose parameterisation is assumed to be Gaussian. Similar to work, amortized multiple shooting is employed leveraging smoothness prior. Hence, training is performed based on sophisticated loss incorporating data, continuity, dynamics, and decoder priors.
Likewise, ODE2VAL (Yildiz et al., 2019) represents a variational inference framework based on Bayesian Neural Networks. Similar to our work, it encodes a number of observations into latent initial state, which are explicitly shaped by physics-mi motivated mhoivcc-motivated prior, c.g., the latent space is separated into a velocity and position part. Subsequently, these high-order dynamics are approximated by a BNN and are evolved over time. The decoder is similar to MSVI where both BNN priors are assumed to be Gaussian.
Similar to the other baselines, ODE-RNN (Rubanova et al. 2019 represents family of time series models whose hidden state dynamics are specified by Neural ODEs. Hence, the trained ODE-RNN models rely on latent ODEs and can naturally handle arbitrary time gaps between observations. In this context, note that ODE-RNN is not applicable to spatio-temporal system dynamics by default. As result, the same CNN encoding part is used as in our approach while the dynamics function and
24

---

Published as a conference paper at ICLR 2024
decoder similar to the ones used in MSVI. In this way, we aim provide useful comparison. All models are trained and tested according to default parameters and code provided the original papers. Please note that for comparisons these baseline methods, time grid datasets since ODE2VAE's encoder exclusively applicable to evenly spaced grids. For irregularly sampled observations, we report comparisons MSVI.
forced rely regular
H INFERENCE AND TRAINING TIMES
Table H.2: Comparison of training and inference time as well as the number of trainable model parameters for all methods applied to the single pendulum test case. Note that for MSVI (block size 1 / 8), no inference times can be given since inference requires roll-out across the full trajectory. All tests are performed NVIDIA A100 40 Gb with AMD EPYC 7742 processor.
forward
forward pass (inference)
trainable
backward pass
parameters
[ms]]
[ms]
M params]
ODE-RNN
851.95
351.68
10.51
NDP
388.69
163.15
1.13
ODE2VAE
571.77
68.04
3.04
MSVI (block size 60)
1192.86
48.37
1.51
MSVI (block size 8)
285.79
(N/A)
1.51
MSVI (block size 1)
60.08
(N/A)
1.51
LaDID (ours)
48.11
15.09
1.35
I DATASET DETAILS
Here, we provide details about the datasets used in this work, their underlying mathematical for- mulations and their implementation details. Partially, some of the datasets we selected are used in literature to demonstrate the effectiveness of neural based temporal modeling approaches, e.g., a swinging pendulum reaction- -ifftiioniifuiiin system. However, we also consider more unknown test cases which represent complex real-world applications such as the chaotic double-pendulum or fluid flow applications driven by complex set of PDEs, i.e., the Navier-Stokes equations.
1.1 SWINGING PENDULUM
For the first dataset we consider synthetic videos of a nonlinear pendulum simulated in two spatial dimensions. Typically, a nonlinear swinging pendulum is governed by the following second order differential equation:
d= dt
sin
(32)
with = denoting the angle of the pendulum. Overall, simulated 500 trajectories with different initial conditions. For each trajectory,
da is sampled initial angle z and its angular velocity at
uniformly from ~ 2(0, 2T) and di seconds. The training, validation and test dataset is split into 400, 50 and 50 trajectories, respectively. The swinging pendulum is rendered in black/white image space over 128 pixels for each spatial dimension. Hence, each observation a high-dimensional image representation (16384 dimensions flatted 128 x 128 px² image) of an instantaneous state of the second-order ODE.
de U(=m/2 n /2). All trajectories are simulated for = 3
1.2 SWINGING DOUBLE PENDULUM
To increase the complexity of the second dataset, we selected the kinematics of a nonlinear double pendulum motion. The pendulums are treated as two point masses with the upper pendulum being denoted by the subscript "I" " and the lower one by subscript "2". The kinematics of this nearly chaotic
25

---

Published as a conference paper at ICLR 2024
system is governed by the following set of ordinary differential equations:
2221
g(2ma1 m22 sin 21 mag sin 222) sin((2) - 22)mm2( dey de
L2 + dar " cos(2) - -22))
or
dia
L1 (2m1 + m2 m2 cos(221 222))
(33)
22 dta
2sin
22)( dt La [mm m2) + g(mi m2) cos =1 + di
Lgme cos(21 (21 22))
L2(2m² + my m2 cos(22: 222))
(34)
with my denoting the mass and the length of each pendulum respectively, and g the gravitational constant. Again, we simulated 500 trajectories split in sets of 400, 50 and 50 samples for training, validation and testing. The initial condition for (21, 22) and ( (da) dzz at dt ) are uniformly sampled in the
range (((0, 2T) and (- 222). The double pendulum is rendered in RGB color space over 128 pixels for each spatial dimension with the first pendulum colored in red and the second one in green. Hence, each observation a high-dimensional image representation (16384 x dimensions flatted 128 x 128 px² RGB image) of an instantancous double pendulum state.
1.3 REACTION-DIFFUSSE EQUATION
Many real-world applications interest originate from dynamics governed by partial differential equations with more complex interactions between spatial and temporal dynamics. One such set of PDES we selected as test case based on a lambda-omega reaction-diffsiio system which is described by the following equations:
du
du du
(1 (u² + +)) B(u + v2 + di
+
(35)
dt
dr
dyj
dy
dv
d²y
(uu² r02)u (1 (u² + ))v v
(36)
di
dx
dy²
with d2) = 0.1 denoting diffusion constants and 8 1. This set of equations generates spiral wave formation which can be approximated by two oscillating spiral modes. The system is simulated from a single initial condition from = 0 to 10 At 0.05 time intervals for total number of 10 000 samples. The initial conditions is defined as
u(r, y, 0) = tanh (Vz² + 82 cos
(x + iy) - v y22
(37)
u(z, y. 0) = tanh (V2² + v2 sin
(x + by) v x2 - y2
(38)
The simulation is performed over spatial domain of (z E [-10, 10] and y E [-10, 10) grid with 128 points in each spatial dimension. We split this simulation into trajectories 50 consecutive samples resulting in 200 in dependant realisations. We use 160 randomly sampled trajectories for training, 20 trajectories for validation and the remaining 20 trajectories for testing. Source code of the simulation can be found in Champion al. (2019).
1.4 TWO-DIMENSIONAL WAVE EQUATION
A classical example of a hyperbolic PDE is the two-dimensional wave equation describing the temporal and spatial propagation of waves such as sound or water waves. Wave equations are important for variety of fields including acoustics, electromagnetics and fluid dynamics. In two dimensions, the wave equation can be described as follows:
8u
277u,
(39)
0+2
with ? denoting the Laplacian operator in R² and is constant speed of the wave propagation. The initial displacement up is a Gaussian function
(x b)2
Ug = aeep
(40)
2r
26

---

Published as a conference paper at ICLR 2024
where the amplitude of the peak displacement a, the location of the peak displacement b and the standard deviation r are uniformly sampled from a ~ (2, 4), b ~ (1(111)) and r ~ (0.255, 0.30), respectively. Similar to Yin et al. (002)), the inital velocity gradient 2u is set to zero. The wave simulations are performed over spatial domain of (a: E [-1.1] and y E [-1, on grid with 128 points in each spatial dimension. Overall, 500 independent trajectories (individual initial conditions) computed which are split in 400 randomly sampled trajectories for training, 50 trajectories for validation and the remaining 50 trajectories for testing.
1.5 NAVIER-STOKES EQUATIONS
To ultimately test the performance of our model on complex real-world data, we simulated fluid flows governed by complex set of partial differential equations called Navier-Stokes equations. Overall, two flow cases of different nature are considered, e.g., the temporal evolution of generic initial vorticity fields and the flow around an obstacle characterized by the formations of dominant vortex patterns also known as the von Kármán vortex street.
Due to the characteristics of the selected flow fields, consider the incompressible two-dimensional Wavier-Stokes equations given by
Ou ot
(u V)u- uv Vv. p
(41)
Here, denotes the velocity in two dimensions, and p are the time and pressure, and v is the kinematic viscosity. For the generic test case, we solve this set of PDEs in its vorticity form and chose initial conditions
described in Li et al. (2020). Simulations are performed over spatial domain of (I E [-1,1] and E [-1, on grid with 128 points in each spatial dimension. Overall, 500 independent trajectories (individual initial vorticity fields) are computed which are split in 400 randomly sampled trajectories for training, 50 trajectories for validation and the remaining 50 trajectories for testing.
1.6 FLOW AROUND A BLUNT BODY
The second fluid flow case mimics an engineering inspired applications and captures the flow around blunt cylinder body, also known von Kármán vortex street. von Kármán vortices manifest in repeating pattern of swirling vortices caused by the unsteady flow separation around blunt bodies and occur when the inertial forces in a flow are significantly greater than the viscous forces. A large dynamic viscosity of a fluid suppresses vortices, whereas a higher density, velocity, and larger size of the flowed object provide for more dynamies and a less ordered flow pattern. If the factors that increase the inertial forces are put in relation to the viscosity, dimensionless measure- the Reynolds number- is obtained that can be used to characterize a flow regime. the Reynolds number is larger than Re > 80, the two vortices in the wake of the body become unstable until they finally detach periodically. The detached vortices remain stable for a while until they slowly dissociate again the flow due to friction, and finally disappear. The incompressible vortex street is simulated using an open-source Lattice- Boriiccan solver due to computational efficiency. The governing equation is the Boltzmann equation with the simplified right-hand side (RHS) Bhatnagar- Greaaar-rroo-rrora Krook (BGK) collision term (Bhatnagar et al. 1954
0f 0t
0f + C$ Oak
-
(42)
These particle probability density functions (PPDFs) f= (,1) describe the probability to find fluid partiele around location T with particle velocity Ç at time (enzi et al. 1992). The hend side (LHS) describes the evolution of fluid particles in space and time, while the RHS describes the collision of particles. The collision process is governed by the relaxation parameter 1/r with the relaxation time 7 to reach the Maxwellian equilibrium state fou. The discretized form of equation 42 yield the Pattice-BGk equation
(F. t( 1
fa(E +
(
(43)
27

---

Published as a conference paper at ICLR 2024
The standard D2Q9 discretization scheme with nine PPDFs (Qian al. 1992 is applied. The equilibrium PPDF is given by
k = = (1 (1+ SK c²
+
(44)
2c
2c
where the quantities wk are weighting factors for the D2Q9 scheme given by 4/9 for it E 0, 1/9 for k E ......4, 4, and 1/36 for & E 5.... 99 and ū is the fluid velocity. C2 denotes the speed of sound. The makroscopic variables can be obtained from the moments of the PPDFs. Within the continuum limit, i.c. for small Knudsen numbers, the Navier-Stoker equations can directly be derived from the Boltzmann equation and the BGK model (Hänel 2006). We simulated three different Reynolds numbers Re = 100, 250, 500 for 425 000 iterations with a mesh size of 128 point in vertical and 256 points in horizontal direction. We skipped the first 25 000 iterations to ensure a developed flow field and extracted velocity snapshot every 100 iterations. The simulation is performed over a spatial domain of (x E [-20, 20] and y E [-10, 10). We split this simulation into trajectories of 50 consecutive samples resulting in 200 in dependant realisations. We use 160 randomly sam- pled trajectories for training. 20 trajectories for validation and the remaining 20 trajectories for testing.
J DETAILED DESCRIPTION OF EXPERIMENTS
We structured our experiments into four different series that shed light on the performance of LaDID in different situations with unique contextual information and challenges.
1. ODE based systems: Here, we investigate the performance for simulated kinematic systems. In these experiments, the right-hand side of the ODE remains constant, so the underlying dynamical system is the same for training and test cases. However, every trajectory starts from a unique initial conditions that lead to unique trajectories in the observational space (image space). The set of initial conditions of the training data and the test data is entirely disjoint. We tested two kinematic systems:
(a) Single Pendulum: We simulated 400 training trajectories at an image resolution of 128 x 128 pixels and tested on 50 distinct test trajectories. The dataset generation is described in Subsection [.I The results are presented in Figure 3 and K.T
(b) Double Pendulum: We simulated 400 training trajectories at an image resolution of 128 x 128 pixels and tested on 50 distinct test trajectories. The dataset generation is described in Subsection 1.2 The results are presented in Figure
2. PDE based systems: The series of experiments focuses on the performance of LaDID in PDE based systems. Similar, to the experimental setup of the ODE experiments, the dynamical system described by the right-hand side of the PDE remains constant throughout this set of experiments. Every trajectory corresponds to the temporal evolution under the dynamical system given a unique initial condition. Here, we focus on four different PDE based dynamical systems:
(a) Reaction-Diffusson Equation: The dataset contains 160 training trajectories and 20 distinct test trajectories at image resolution of 128 128 pixels. The generation of the datasets is given in Subsection I.3 of the Appendix and the results are shown in Figure K.3
(b) Wave Equation: The dataset comprises 400 training trajectories and 50 unique test trajectories, all at an image resolution of 128 x 128 pixels. Details on how the datasets were generated are provided Subsection I.4 of the Appendix, and the corresponding results are depicted in Figure K.4]
(c) Navier-Stokes Equations: The dataset comprises 400 training trajectories and 50 unique test trajectories, each with an image resolution of 128 x 128 pixels. Detailed information on the dataset generation is provided in Subsection I.5 of the Appendix, with corresponding results displayed in Figure K.5
(d) Flow around a blunt body: We generated 160 training trajectories with an image resolution of 128 x 128 pixels and conducted testing on 20 unique test trajectories. The details of dataset generation can be found in Subsection [.6] The corresponding results
illustrated in Figure? 3 and K.
28

---

Published as a conference paper at ICLR 2024
3. Regular Irregular time grid: Here, we delve into the impact of time grid regularity on performance. We systematically compare the model's predictions when trained on datasets featuring regular time grid against those with irregular time grid. The objective is to study how the model adapts to variations in the temporal structure of the input data and assess its generalization capabilities across different time grid configurations. These experiments aim to provide insights into robustness and flexibility in handling diverse time representations, shedding light on suitability for real-world scenarios with irregularly sampled temporal data. We study two test cases:
(a) Single Pendulum
(b) Reaction-Diffusion Equation
The regular case is identical to the experiments for the ODE/DDE systems described above. For cases with irregular time grid, we uniformly sample the time points along the time axis and compute the corresponding observations in the image space. The description of the datasets is detailed in the Subsection I.1 and I.3] respectively. The results for the single pendulum test is shown in Figure I.. and in Figure 1.2 for the reaction-diffusion equation.
4. Transfer Learning:
Few-shot learning and zero-shot learning under interventions upon the dynamical system In this set of experiments, we focus on evaluating transferability to new dynamical systems, emphasizing scenarios where the target systems share similarities with the training distribution but have not been encountered during the training phase. We test this in two different ways. First, we adopt few-shot learning paradigm, wherein pretrained model fine-tuned using a limited number of trajectories from the new system. This set-up allows us to assess the model's adaptability to novel dynamical systems and to examine its capacity generalize knowledge from the training set to previously unseen systems. To study transferability, we used the single pendulum dataset as described in Subsection [.1 of the Appendix and fine-tuned the model on dataset that comprises Nhew [16, 32, 64, 128] fine-tuning trajectories (which corresponds to p [4%, 8%, 16%, 32%] of the training dataset size) obtained through unique interventions on the mass and length of the pendulum. Additionally, all initial conditions of the trajectories are unique. We conducted testing on 50 unique test trajectories that were subject to new interventions and never seen during fine-tuning. Note that interventions here yield dynamical systems that are qualitatively different from the others and that sense unique. In other words, the right-hand side of the underlying dynamical system is different for all interventions. The results are presented in Section 6.2 and Figure 6al of the main text and Section in the Appendix.
Furthermore, we extend our investigation to encompass zero-shot learning, wherein our approach is initially trained on broader class of dynamical systems. Subsequently, assess the model's ability to predict entirely new dynamical systems without additional training on specific instances. These experiments aim elucidate ability to leverage learned knowledge to make accurate predictions for diverse dynamical systems. We chose the flow around blunt cylinder as our test case and generated a training dataset that comprises 120 trajectories for three different Reynold's number Re [100. 250, 500). We test zero-shot transferability for three new Reynold's numbers Renew = [175, 300, 400). The results are presented in Subsection O and Figures P.I and P.2 of the Appendix.
Interventions upon the measurement process of the dynamical system In this series of experiments, we explore the robustness and adaptability of our approach under interventions that mimic variations in the measurement process of the dynamical system. Specifically, we examine the impact of additive Gaussian noise and salt pepper noise, aiming to understand performance under observational variability. Additionally, we systematically alter channels the RGB space to investigate the model's response to changes in color representations, providing insights into its ability to handle variations in the visual input domain. These effects are studied using the single- seggl-pendulum test case. The results are presented in Subsection M.1 for perturbed input images using Gaussian additive and salt & pepper noise. Subsection M.2 shows the effects of alteration in the RGB channels.
29

---

Published as a conference paper at ICLR 2024
Furthermore, extend analysis to include experiments involving few-shot learning, where the model is presented with trajectories captured from entirely new camera angles and positions. This design enables to assess the model's capacity adapt and generalize to novel observational conditions with minimal exposure to the new data. These experiments collectively provide information and environmental changes, This is relevant for real-world applications with varying observational conditions. In this particular experiment, we selected the flow around blunt cylinder our test case. Utilizing a pretrained model with a fixed camera position, we proceeded to fine-tune it using dataset comprising Nneww [6. 12, 24, 48) distinctive trajectories captured from various camera positions. This corresponds to subsets representing p = [4%, 8%, 16%, 32%] the training dataset size, providing systematic exploration of model adaptation to diverse camera perspectives. Testing is carried on an entirely unseen dataset with n = 20 test trajectories under new camera positions. The results are presented in Section 6.2 and Figure (b of the main manuscript.
resilience diverse sources of measurement noise
K ADDITIONAL RESULTS
K.1
SWINGING SINGLE PENDULUM
0.155
NDP
NDP
normalized MSE
MSVI
MSVI
20
RSME
0.075
0.000
models
0.5
0.8
normalized time
(a) Normalized MSE
(b) RMSE distribution time
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODE2VAE
MSVI
Ours
truth
(c) Example trajectory
Figure K.1: Test errors and exemplary test trajectories of different models for the single pendulum test case.
30

---

Published as a conference paper at ICLR 2024
K.2
SWINGING DOUBLE PENDULUM
0.14
0.7
0.12
0.6
normalized MSE
0.10
0.5
0.4
RSME
0.08
0.3
0.06
ODE-ANN
QDE-ANN
0.2
0.04
NDP
QDE2VAE
ODE2VAE
MSVI
0.02
Curs
0.0
models
0.0
0.2
0.4
0.6
0.8
1.0
normalized time
(a) Normalized MSE
(b) RMSE distribution over time
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODF7VAF
MSVI
Ours
truth
0.00
0.05
0.10
0.15
0.20
0.25
0.30
0.35
0.40
(c) Example trajectory
Figure K.2: Test errors and exemplary test trajectories different models for the double pendulum test case.
31

---

Published as a conference paper at ICLR 2024
K.3
LAMBDA-OMEGA REACTION-DIFFUSION SYSTEM
4./1
0.64
ODE RNN
ODE-RNN
0.20
0.08
ODE2VAE
ODEZVAE
normalized MSE
MEVI
MSVI
0.15
0.06
RSME
0:10
0.04
0.35
0.00
0.30
models
0.0
0.2
0.4
0.6
0.8
1.0
normalized time
(a) Normalized MSE
(b) RMSE distribution over time
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODE2VAE
MSVI
Ours
truth
0.00
0.01
0.02
0.03
0.04
0.05
(c) Example trajectory
Figure K.3: Test errors and exemplary test trajectories of different models for the lambda-omega system.
32

---

Published as a conference paper at ICLR 2024
K.4
TWO-DIMENSIONAL WAVE EQUATION
0.14
2.60
NDP
0.12
NDP
12
ODE2VAE
CDE2VAE
normalized MSE
MSVI
0.10
MSVI
OUrS
Curs
RSME
0.08
0.06
0.0
models
norma lized time
(a) Normalized MSE
(b) RMSE distribution time
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODE2VAE
MSVI
Ours
truth
0.00
0.05
0.10
0.15
0.20
0.25
(c) Example trajectory
Figure K.4: Test errors and exemplary test trajectories of different models for the wave equation test case.
33

---

Published as a conference paper at ICLR 2024
K.5
DAVIER-STOKES EQUATIONS
2:19K 1.50)
0.175
NDP
NDP
0.0
ODE2VAE
0.150
ODEZVAE
normalized MSE
Ours
MSVI
0.125
Curs
MSVI
0.6-
RSME
0.2
0.0
models
norma lized time
(a) Normalized MSE
(b) RMSE distribution time
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODE2VAE
MSVI
Ours
truth
0.00
0.05
0.10
0.15
0.20
0.25
(c) Example trajectory
Figure K.5: Test errors and exemplary test trajectories of different models for the Navier-Stokes equations test case.
34

---

Published as a conference paper at ICLR 2024
K.6
LATTICE BOLTZMANN EQUATIONS
0.61
0.23
ODE ANN
ND?
0.08
ODEZVAE
ODE2VAE
normalized MSE
MEVI
0.06
0.06
R5ME
0.04
0.04
0.02
models
normalized time
(a) Normalized MSE
(b) RMSE distribution over time
Predictions
Pixelwise L2-error
ODE-RNN
NDP
ODF2VAF
MSVI
Ours
truth
0.00
0.01
0.02
0.03
0.04
0.05
(c) Example trajectory
Figure K.6: Test errors and exemplary test trajectories of different models for the von Kármán vortex street case.
35

---

Published as a conference paper at ICLR 2024
K.7 RESTRICTED ACCESS TRAINING TRAJECTORIES
In this experiment, we explore the impact of limiting the number of trajectories available during training- Specifically, we assess the performance of LaDID when utilizing only 4%, 8%, 16%, and 32% of the single pendulum dataset, which originally consists of 400 training trajectories. Please Section [.] for details on the data generation. In each experiment, our model trained from scratch with restricted number of trajectories. Comparisons are drawn against baseline algorithms which are trained on the full dataset. The results indicate that even with just 8% of the data, the model achieves reasonable predictions. Further, 16%, LaDID surpasses the performance of all baseline algorithms.
25
LGD D- DS
normalized MSE
WEUT-TUTDE
20-
ODERVRE
15
0
4%
8%
16%
32%
percentage of data samples
Figure K.7: Comparison of prediction performance of LaDID as the number of training trajectories varies. The model's accuracy is evaluated using different percentages of the standard dataset, ranging from 4% to 32%. Full dataset performances are represented by horizontal lines for LaDID and all baselines.
K.8 INVESTIGATING TEALIZATION-SPECIIA AND REALIZATION INVVRIIAN INFORMATION
LaDID is rooted in the idea of RS/RI decomposition, as reflected in both model architecture and the theory. In this section we provide empirical study investigate the RS/RI aspect. In particular, to better understand model intermediates, we carried out an experiment tracking the learned trajectory representation J" over set of pendulum configurations (single pendulum) for which we know the true initial conditions be possible to recover the latter from the former. To test this, we use a trained LaDID model and compute the RS trajectory representations s² for every initial condition of the training dataset. In the next step, we learn simple mapping function that maps the trajectory representations vsr" to the initial conditions i.e. g R - R² (g is chosen to be a shallow 2 layer MLP with approximately 500 parameters). Training and testing are performed on two mutually disjoint sets trajectory representations. That is, for training original training dataset using the initially trained LaDID encoder in a 1-step multiple shooting setting. This means that split observed trajectory into multiple subtrajectories and compute for each of these subtrajectories an independent πy. For testing, we compute similar trajectory representations whast. but this time from the dataset ensuring that these are entirely disjoint from the training dataset since neither the LaDID encoder nor the learned mapping have seen these samples before. At testing, with mapping function g and LaDID model fixed, we then predict the ground truth initial conditions, i.c. ग. g(b²). We present the results as scatter plots in Figure K.8, the good correspondence supports the notion that vi indeed captures the RS information. In addition, to understand whether the computed trajectory representation of any point of trajectory holds information on the corresponding ground truth initial condition (as in ODE/PDEs), we compute for each time step of trajectory its corresponding independent trajectory representation and predict the n.. 's as outlined before. The results are presented in Figs. K.8c & K.&d] The x-axis denotes the time, i.e. the time point from which would start forecasting the dynamical behaviour, while the y-axis presents the predicted and ground truth initial conditions, i.e. angle O and Ô and angular
of the underlying ODE system. According to the theory, it should
computed the trajectory representations vhrain from the
36

---

Published as a conference paper at ICLR 2024
velocity and Ô respectively (note that markers shown in the figures are separate experiments with independent trajectory representatioos/iiitaal conditions).
360
150
predicteo anguiar velocity (0egiS)
100
predicted [deg)
120
-50
-120
-100
-150
360
-150
360
-240
-120
0
120
240
360
angle [deg]
angular (degts)
(a) Distribution of the predicted angles +
(b) Distribution of the predicted angular velocity θ
150
100
aegular velocity (dega)
200
andle e [doe]
50
-100
-50
-200
d
300
0.0
0.4
0.8
0.0
0.2
0.4
0.6
0.8
1.0
nommllized time
normalizod time
(c) Temporal distribution the angle trajecto- (d) Temporal distribution of the angular velocity θ for ries
trajectories
Figure K.8: Comparison of the recovered (predicted nonlinear mapping function) and true initial conditions of the swinging pendulum dataset: (a) Predicted vs true angles ④ of the entire test dataset, (b) Predicted vs true angular & of the entire test dataset. (c) Temporal distribution of the angles θ for five different trajectories. Note that the each marker corresponds to an initial trajectory representation al (d) Temporal distribution the angular velocity é for five different trajectories.
37

---

Published as a conference paper at ICLR 2024
L REGULAR AND IRREGULAR TIME GRIDS
ours irregular
0,20
regular MSVI requiar
irreqular
20
regular
ours irreglarr
MSVI- regular
MSVI irregu ar
normllizdd MSE
15-
RSME
0.05
0.00
0
models
0.0
0.2
0.4
0.6
0.8
1.0
normalized time
(a) Normalized MSE
(b) RMSE distribution over
Figure I.1: Test errors and exemplary test trajectories of different models for regular and irregular time grids the single pendulum test case.
ireeullrr
0.04
reguiar
0.010
MSVI Imreullrr
regular
6.014
ours Irregllar irregult
regular
MSE
0.012
0.010
RSMF
0.008
0.006
0.002
0.000
models
0.0
0.2
0.4
0.6
0.8
1.0
normalized lime
(a) Normalized MSE
(b) RMSE distribution over time
Figure L.2:
Test errors and exemplary test trajectories of different models for regular and irregular time grids of the lambda-omega reaction-diffusion system.
M ROBUSTNESS TO INPUT NOISE AND APPEARANCE MODIFICATION
M. ROBUSTNESS TO INPUT NOISE
In the following experiments, we study the impact of zero-mean Gaussian noise and salt-and- pepper noise on performance. Specifically, we vary the standard deviation for Gaussian noise (a [0.1, 0.2, 0.3, 0.4)) and the density for salt-and-pepper noise (p [10%, 35%, 50%, 70%) to systematically assess robustness. To provide visual representation, we show example trajectories that exhibit the effects of these noisy inputs.
Next, we evaluate the summed MSE over trajectorics, averaged across the test dataset. Additionally, we introduce the RMSE distribution plotted over time detailed in Section [ of the main text, offering insights into the model's 's response to noise over the entire trajectory duration. These results demonstrate robustness against the tested sources of noise for LaDID. We minimal impact for small noise levels, with an expected increase in error values as noise intensity increases.
38

---

Published as a conference paper at ICLR 2024
0.0
density 0%
0.1
density 10%
0= 0.2
density 35%
0= 0.3
density 50%
0= 0.4
density 70%
(a) Gaussian noise
(b) Salt& Salt&&Pepper noise
Figure M.1: Visualization of test dataset augmentations to test the robustness input noise.
d 0.0
02
0.4
0.14
normalized MSC
0.12
RSME
0.04
0.02
0.00
nolse density
normalzedd time
(a) Normalized MSE
(b) RMSE distribution over time
Predictions
Pixelwise L2-error
0= 0.1
0.2
8
0.3
8
3
g= 0.4 truth
0.40
(c) Example trajectory
Figure M.2: Test errors and exemplary test trajectories of the single pendulum test (gray-scale) corrupted by Gaussian Noise with various noise levels.
39

---

Published as a conference paper at ICLR 2024
denaity
danalty
dwisty- -U%
deruity 25%
danucty /CS
density- 20%
donsity -00
dwaaty- 100%
denuty 30%
16
0.175
normalized MSF
14
0.125
10-
0.100
0.075
0.025
0.000
nolse density
0.2
norma ized time
(a) Normalized MSE
(b) RMSE distribution over time
Fredictions
Pixelwise LZ-error
density = 0%
density 10%
density 35%
density = 50%
density = 70%
truth
0.00
0.05
0.10
0.15
0.20
0.25
0.30
0.35
0.40
(c) Example trajectory
Figure M.3: Test errors and exemplary test trajectories of the single pendulum test case (gray-scale) corrupted by Salt&Pepper noise with various noise levels.
40

---

Published as a conference paper at ICLR 2024
M.2 ROBUSTNESS MODIFICATION THE INPUT APPEARANCE
We explore the impact of random scaling and flipping applied to the RGB channels of input images. This investigation aims assess how the model responds such spatial transformations. To provide a visual understanding, we start by showing example trajectories, allowing an examination of the perturbed inputs resulting from these transformations.
Following the visual inspection, we present quantitative analysis, presenting the summed MSE over trajectories and averaged across the test dataset. Additionally, we introduce the RMSE distribution plotted over time, providing insights into the model's performance over the trajectory duration. The results indicate only slightly higher enror values under the influence of these spatial transformations.
(a) clean dataset
(b) flipped RGB channels
(c) scaled RGB channels
Figure M.4: Visualization of test dataset augmentations to test the robustness to modifications of the input appearance.
0.08
ReA--fip
0.06
RGB-fip
0.07
normalized MSE
0.05
0.05
0.04
0.03
0.03
0.02
0.02
0.00
0.00
augmentations
0.0
1.0
normallzed time
Normalized MSE
RMSE distribution time
Predictions
Pixelwise L2-error
clean
RGB-flip
RGB-scale
0.00
0.05
0.10
0.15
0.20
0.25
0.30
035
0.10
(c) Example trajectory
Figure M.5: Test errors and exemplary test trajectories of different RGB single pendulum test cases transformed by color augmentations.
41

---

Published as a conference paper at ICLR 2024
N
ABLATIONS EFFECTS OF RELEVANT NETWORK MODULES
Table N.3: Errors for ablated LaDID models trained on the single pendulum test case.
loss
attention mechanism
representation
heuristics
encoding
ablation
mean
TQR
ablation
mean
IQR
ablation
mean
IQR
reconstruction
2.66
1.04
no attention
7.79
6.84
w./o. encoding
2.83
2.02
reconstruction & representation reconstruction & smoothness
2.17
0.99
spatial attention
2.81
1.47
encoding
2.02
0.88
2.04
0.92
temporal attention
2.41
0.92
spatio-temporal attention
full loss
2.02
0.88
2.02
0.88
O GENERALIZATION UNDER INTERVENTIONS
We evaluated our model's ability to generalize with respect aspects of the data generating distribu- tion that are often unknown. Consider vectorized differential equation of the form
d dt x = f(x,t)
(45)
with general initial conditions <+(6) = xy. Using an appropriate discretization scheme allows us to foretell the future of the system based on its past states and we can directly read off the causal structure for such a system. This underlying causal structure remains the same when changing coefficients on the right hand side which is summarized by the principle of independent causal mechanisms (Peters)
ct al. 2017, Schoelkopf et al., 2012; Schölkopf et al., 2021 Lagemann et al. 2023b):
Independent Causal Mechanism: Let X be the outcome variable, and let M1, M2. M, be distinct mechanisms or factors that independently contribute to X. Then, we can write:
X fM1, M2, Mn.)
(46)
where f' represents the functional relationship between the mechanisms and the outcome. In other words, the outcome variable determined by the independent contributions of each mechanism M,, and these mechanisms operate independently of each other and do not inform each other.
For example, consider a system of ODEs that describes the dynamics of a population of predator and prey animals:
dr
dt or bry
(47)
dy -cy + dry
(48)
dt
where æ represents the population of prey, y represents the population of predators, and a, b, c, d are parameters that describe the interactions between the two populations. In this system, x and represent independent causal mechanisms that contribute to the dynamics of the population. The equations describe how the population of prey and predators changes over time as a result of their interactions. By solving the system of ODEs, we can study how changes in the parameters affect the long-term behavior of the system, and how interventions can be used to control the population dynamics. Mathematically, an intervention typically represented as a modification of the equations that describe the system, by setting the value of one more variables to a fixed value or function. This modification represents the assumption that the variable(s) being intervened upon is no longer subject to external influences, and its value is determined by the intervention. In formal terms, an intervention upon a set nodes in the causal structure of a system (X, : E I} means any manipulation of the system that alters its state behaviour, including changes in the initial conditions, modifying the parameters, or adding or removing variables or equations. When observing a latent process, which moves the deterministic setup of an ODE to probabilistic case, this means that the
42

---

Published as a conference paper at ICLR 2024
conditional distribution when observing the state of node given its parents Pa(X;) is replaced by new, predefined distribution. Thus, the joint probability distribution as system under an intervention changes to
((X) = II P((() II (a()))(X)))
(49)
EEI
where ((X))Pa(Xi)) indicates the conditional distribution of node X; in its general form.
43

---

Published as a conference paper at ICLR 2024
P ADDITIONAL TRANSFER LEARNING RESULTS
In this set of experiments, our focus turns to the robustness and generalization capabilities in the face of novel fluid flow scenarios. Initially trained on a dataset comprising trajectories of fluid flows around a blunt object at Reynolds numbers Re [100, 250, 500J, we consider zero-shot predictions on qualitatively different flows with Reynolds numbers Renew [175, 300, 400). Without additional training on these new systems, we assess adaptability and accuracy in capturing the dynamics of unseen scenarios.
To evaluate performance, we present the summed MSE averaged across the test set, providing an overall measure of prediction accuracy. Additionally, we introduce the RMSE distribution over time, offering insights into the temporal dynamics of prediction errors.
Furthermore, we include comparative analysis between the ground truth and predicted trajectories for each Reynolds number, offering visual representation of the model's predictive capabilities in these novel conditions. As anticipated. we observe an increase in error for systems with unseen Reynolds numbers. However, in these challenging scenarios, LaDID shows reasonable predictions in the observation space. These results underscore LaDID's adaptability and potential for generalization.
0.008
Ra-100
normalized MSE
Re-250
0.006
R2-400
0.004
0.002
0.000
Reynolds number Re
(a) Normalized MSE
0.011
Re-100
Rc-175
Re-250
0.010
Re-300
Re=400
0.009
Ra-500
RSME
0.008
0.007
0.006
0.005
0.0
0.2
0.4
0.6
0.8
1.0
normalized time
(b) RMSE distribution over
Figure P.1: Transfer Learning: Test enrors of various vortex street test cases characterized by changing Reynolds number.
44

---

Published as a conference paper at ICLR 2024
Predictions
Pixelwise L2-error
prediction truth
0.00
0.01
0.02
0.03
0.04
0.05
(a) Re = 100
Predictions
Pixelwise L2 12-error
prediction truth
0.00
0.01
0.02
0.03
0.04
0.05
(b) Re 175
Predictions
Pixelwise L2-error
prediction truth
0.00
0.01
0.02
0.03
0.04
0.05
(c) Re 250
Predictions
Pixelwise L2-error
prediction truth
0.00
0.01
0.02
0.03
0.04
0.05
(d) Re 300
Predictions
Pixelwise L2-error
prediction truth
(e) Re 400
Predictions
Pixelwise L2-error
prediction truth
0.90
(f) Re = 500
Figure P.2: Transfer Learning: Visualization of test trajectories of different vortex street test cases characterized by a changing Reynolds number.
Q LICENSES BASELINES
MSVI (Iakovlev et al. 2023): MIT license
ODE2VAE Yildiz et all
2019): The license status is unclear.
NDP (Nocllifee et al. 2021)
MIT license
ODE-RNN (Rubanova et al. 2019) MIT license
45


Reviews:
[{"review_id": "ApDmLhti3d", "reviewer": "Reviewer_apHo", "summary": "The manuscript introduces a novel deep trajectory inference model, referred to as LaDID. This model designed to predict dynamical trajectories by disentangling invariant and variant dynamical factors. LaDID employs a convolutional autoencoder unit for dimension reduction and a transformer unit for learning temporal structure, eliminating the need for ODE formulations and enabling efficient long trajectory prediction without added computational costs.\n\nIn the study, the authors evaluated LaDID's performance using six simulated datasets. They conducted a comprehensive comparison of LaDID with four ODE-based models, namely ODE-RNN, ODE2VAE, NDP, and MSVI. The results demonstrated LaDID's superiority in all experiments, including challenging few-shot learning tasks.\n\nI appreciate the problem tackled by the author and find the concept of invariant-latent representation learning in dynamical systems intriguing. The paper appears to have a solid metric in place. Nevertheless, there are certain crucial aspects that need to be improved. I am open to potentially revising my assessment following the author's rebuttal.", "questions": "- If I understand correctly, the paper addresses only the initial state as the specific (variant) trajectory factor, right? What about other sources of variability, like measurement noise?\n\n- In section 4.1, Model, the main text says *\"To obtain a latent trajectory,  ... we choose as a set of different sine and cosine waves with different wave length.\"* Could you explain it? \n\n- Could you please explain the choice of using MSE/RMSE for evaluation? It seems that capturing the overall structure of dynamics (direction, shape, etc.) might be more crucial than the absolute amount of misprediction.\n\n- Can LaDID effectively handle time-variant dynamical systems? If yes, how?\n\n- In the graphical model (Figure A1), $\\psi^r$ is out of the box \"$N$\". Then, how is $\\psi^r$ a function of \"$n$\" in Eq. 3 and Eq. 9?\n\n- The text mentions *\" ... mean-aggregation, which can be changed based on the task at hand.\"* Could you elaborate on the role and task-dependent nature of mean-aggregation?\n\n- In Eq. 9, within the representation prior term, none of the variables appear to be a function of \"$n$.\" Can you clarify this discrepancy?\n\n- What does \"$D$\" represent in Figure A1?\n\n- The authors stated that *\"our models are specifically designed to exploit certain invariances that are important in classical scientific models.\"* Could you provide more detail on these \"certain invariances\"?\n\n- Beyond classical scientific examples, can LaDID effectively learn developmental trajectories involving bifurcations and other complex phenomena?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The manuscript tackles an intriguing and challenging problem related to trajectory inference in dynamical systems. The proposed approach involves effectively decomposing the latent factors into variant and invariant components, which is fairly innovative.\n\n- The authors conducted a comprehensive performance evaluation across six complex dynamical systems. \n\n- The proposed model consistently outperforms its related counterparts across a wide range of tasks. \n\n- The proposed model showcases capabilities in handling long trajectories, even when provided with limited training samples. This adaptability to data constraints is a significant strength, as it opens up opportunities for more efficient and practical application in dynamical system modeling.", "weaknesses": "**Major comment:**\n\nMy prominent concern in the paper pertains to the insufficient explanation for why their model outperforms others. The main contribution of the paper lacks clear description and justification. As pointed out by the authors, both LaDID and MSVI (Iakovlev et al., 2023) share several similarities, such as using CNN encoders/decoders and the same transformer. The primary distinction between these models lies in replacing neural ODE with MLPs to learn $f_{\\phi_{dyn}}$ and introducing a smoothness loss to discourage abrupt transitions between latent subpatches. It remains unclear how these modifications enhance the process of learning latent dynamics in comparison to ODE-based methods.\n\n**Minor Points:**\n\n- Figure 1 could benefit from a more comprehensive overview of their framework. Clarifying how each process in the top row corresponds to each step in the bottom row would enhance understanding.\n- The paper's writing can be improved, and some content from the Appendix may be integrated into the main text.\n- The descriptions of the experiments are somewhat confusing and could benefit from increased clarity and explanation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fiYbIbBICg", "reviewer": "Reviewer_1fGG", "summary": "The authors propose a general framework for predicting dynamical trajectories based from data. The model architecture specificially aims at extracting realization-invariant and realization-specific information from different observations of underlying dynamical processes, iterated e.g. from different initial states or with different model parameters. To this end they employ a sophisticated encoder architecture, based on CNNs, temporal self-attention and a multi-shooting augmentation to obtain a condensed representation that can be used as a conditioning for a dynamics forecast. The learned latent dynamics model can be queried at irregular times, and the predicted dynamics decoded via a CNN-based decoder model and compared with the data. The model shows SOTA performance on 6 complicated and high-dimensional observations of dynamical processes, such as fluid flows. The author also investigate transfer learning between different parameter settings of the same underlying systems via few-shot learning.", "questions": "Model parameters:\n\nHow parameter intense is your approach? Would it struggle to infer meaningful dynamics from only single or few trajectories, compared to approaches with a stronger inductive bias, such as SINDy (Brunton et al., 2015), or experimental settings, where there are often only single or few observed trajectories. The models in Table H2 have between 1 and 10 million parameters, and are trained on observations with 400 initial states each used for training, which hints at a data and parameter intese regime given some of the relatively simple underlying dynamical processes. How does your model behave for sparser data? Figure 6 indicates that for comparable performance on the pendulum data, comparable performance requires hundreds of observed trajectories?\n\nRelated Work: \n\nI understand that it's hard to provide an exhaustive overview over related methods, and I appreciate the extended related work. I however want to point out that there are many RNN based approaches specifically designed for inferring complex dynamics and leveraging inferred models for long-term time series forecasting (Reservoir Computing, see Pathak et al. 2018), LSTMs (Vlachas et al., 2018) or RNNs (see e.g. Hess et al., ICML 2023), and could be mentioned here as well.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The presentation, including the figures and writing, is high quality. I also appreciate the detailed appendix. While I have seen many of the components employed in similar settings, such as sequential variational autoencoders or Latent-ODEs, the proposed architecture looks like an interesting new combination of all of them. The idea of transfer learning and generalization for dynamics models is an important field of study. I like the approach of using an expressive encoder to inform the dynamics model about the specifics of the observed time series by encoding it in the initial state, and the fact that the obtained model is continuous in time.", "weaknesses": "General problem setting:\n\nWhat is meant by invariant? To my understanding the term “invariant” in this paper mostly refers to the flow of the underlying system (i.e. the equations given in Appendix 1), while the realization-specific components are different initial conditions from which different observations start. This is also the situation primarily studied in the experimental section. The authors claim to show \"generalization well beyond the data distributions seen during training\".\nHowever, the setting they mostly investigate is just the standard setting in dynamical systems reconstruction, where underlying dynamical processes/flow fields are inferred from multiple or even single observations. I don't think this constitutes generalization beyond the data distribution in the usual sense.\n\nMore interesting to my mind is the question of transfer learning e.g. between different dynamical regimes (as studied out in 6.2.) with different parameter settings of the underlying system, or even between different but in some way related systems. Relating this to a more formal language common in dynamical systems theory might be helpful, since for dynamical systems the concept of invariants, such as topological invariants, fixed points, periodic orbits, Lyapunov spectra, invariant measures etc. is well studied. While the abstract mentions the motivation to “ enforce certain scientifically-motivated invariances” these invariances are not really made explicit in the related work or in the analysis of the obtained results and remained somewhat elusive to me.\nFurther, while you mention that your approach “is in principle not specific to physical ODE-like models”, these are the only systems you actually study in the experiments, which naturally suggests are framing in these terms.\n\nTransfer learning setting (6.2.)\n\nYour framework is aimed at transfer learning from “potentially highly heterogeneous parameters and data”, which as mentioned seems the more novel application of your framework. Most of the empirical section however focuses on multiple instantionations of the same underlying system. The results from 6.2., where transfer learning between different model parameters is learned, however indicate that the transfer effect is relatively weka, with advantages over learning from scratch remaining within the error bars, and comparative performance requiring a significant portion of the new training set (if I understood it correctly corresponding to hundreds of tractories from the altered system).\n\nInterventions/nonstationarities in parameters for dynamical systems can also be viewed through the lense of bifurcations. Increasing the Reynolds number can e.g. induce turbulence, as you rightly point out. From your description it is not altogether clear what changes in the dynamics your parameter changes in this section induce, and whether they lead to qualitative differences.\nSince to my understanding this transfer is the primary motivation for introducing your framework, I found this section to be a little too short and and not in-depth enough.\nI also wasn't 100% sure if you indeed jointly trained on different Reynolds numbers for the flow around a blunt body (Re = 100, 250, 500), which would be an interesting extension of your study in 6.2.\n\nComparisons and Experimental Evaluation\n \nAll datasets studied in the comparison methods are based on solutions to synthetic ODE systems, and include several PDEs. The claim that you provided evaluation on real-world datasets I find a little misleading, since none of the benchmarks are based on experimental/real world data, while of course being inspired by real systems.\n\nSelecting comparisons is always a debatable point, but since all benchmark systems chosen are relatively simple/polynomial form, other classes of comparisons could be valuable: approaches by Brunton et al. based on SINDy and its variants are also frequently studies similar fluid flows. Neural Fourier Operators, which you also mention in the related work have been particularly successful for PDEs, so they might be a better choice here.\n \nAlgorithms like ODE-RNN are tailored towards irregularly sampled, but not necessarily spatially extended data. ODE-RNN in general has very poor performance everywhere in the paper (see all figures including ODE-RNN, e.g. Fig 4 or Fig. J6 b), where they seem to have learned nothing meaningful at all, which indicates that there either is a bug in the training, or the method is not well suited as a comparison.\n\nOverall, since I think the general setting studied here is important and the presentation is good, if my concerns can be adressed I am also willing to adjust my final score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gn23mRTC4C", "reviewer": "Reviewer_8oY3", "summary": "The authors consider an encoder decoder model for predicting the trajectory of time series from high dimensional observations. A key difference of the model is that the authors do not impose a particular type of known structure on the dynamics such that the dynamics must follow this structure. The authors then introduce a probabilistic multiple shooting method such that the method can be trained. The authors illustrate the proposed method on a number of different datasets to validate its performance. The main question in the performance evaluation is how well the method works on unseen conditions, i.e. changing the parameters of the true underlying dynamical system.", "questions": "Is there anything in the loss that guarantees that the factorization in terms of the RI and RS components? If not how can one impose this on the latent representation?\n\nAssumptions on the latent space being Gaussian seem a bit restrictive, is there a reason for making this assumption? Specifically am wondering how this affects the motivation that very little structure should be imposed on the latent distribution. Will the smoothness and Gaussian terms will force the learned dynamics to resemble something that satisfies an ODE?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The premise of the paper is very well founded. In general, most time series approaches with latent dynamics consider some model from the applied mathematics literature (e.g. SDEs, ODEs) and then parameterize it according to a neural network. However, the authors consider a more general framework where they split the dynamics according to realization-specific components and invariant components. This induces a new algebra that should somehow be given by the parameterization of the model. The empirical results are also very good and suggest that the model outperforms the baseline methods.", "weaknesses": "While the motivation of the paper is nice, there exist a number of limitations in the implementation and in the model itself. For one, there is no guarantee that the model will converge to something that has these separable parameter sets within the motivation of the work. The loss function seems rather standard in the sense that it is the standard loss when one parameterizes a latent variable time series model. It would have been nice to show under what assumptions does the model decompose into the disjoint parameter sets -- the invariant ones and the instance-specific ones. \n\nWhile this is a small weakness, the paper seems to claim a bit too much (though I may have also misunderstood). For example, the authors mention they make no assumptions on the structure of the model in the sense of an ODE type of factorization, but still the fact that the first $K$ steps of the time series are used for conditioning does induce a particular structure. In some sense, the ODE interpretation is amenable since one can use the known properties of ODEs to probe the model.\n\nA major goal is to generalize the output on a new domains where some invariant properties are maintained, but the experiments only demonstrated this in the case of changing parameters. However, maybe it would be helpful to see the behavior under shifts in the observation distribution. For example, could the distribution of the observation space change in a trivial way (e.g. different color map for the wave equation example) and could the distribution still be predicted?\n\nThere are also some properties that the authors impose in their loss (e.g. smoothness of the latent space) but that assumption seems to violate some of the motivating ideas behind the paper. Smoothness almost implies a differential equation in the latent space.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vo8TKSjyX0", "reviewer": "Reviewer_Fm8i", "summary": "The paper introduces a novel framework called \"Latent Dynamics via Invariant Decomposition\" (LaDID) for predicting dynamical trajectories from high-dimensional empirical data. LaDID merges variational autoencoders with spatio-temporal attention, emphasizing scientifically-motivated invariances, allowing the model to predict system behavior at any continuous time and generalize beyond seen training data. Using a transformer-based architecture, it distinguishes between system-specific and universal dynamics, showing efficiency and scalability in experiments. The method is validated on various spatio-temporal systems, LaDID outperforms current neural-dynamical models.", "questions": "1. In the performance comparison figure 4, is the X axis the time axis? The results surprised me because the prediction from the first 3 models is bad starting from the beginning of the prediction, which means the error mostly comes from decoder reconstruction rather than dynamics prediction. Especially for ODE2VAE, VAE should be able to reconstruct the pendulum image fairly well.\n\n2. In Figure 5, the MSE from the proposed method is 10-100 times smaller than the prior works. This is another surprise as such a difference usually happens between white box model and black box model. Could there be bugs in your code, or is there any unfair comparison? I’m not trying to be mean but it does not look reasonable for me unless formal claims/evidence is provided.\n\n3. In the latent dynamics model (part v in figure 1), how does the query time t work? Does this model need to solve it iteratively like ODE-like solver to time t or the neural network directly takes t as input and outputs the latent state at time t.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper uniquely combines variational autoencoders with spatio-temporal attention\n\n2. empirical tests on various spatio-temporal systems demonstrate that LaDID not only has practical utility but also massively outperforms existing neural-dynamical models", "weaknesses": "1.The definition of a system with realization-specific and realization-invariant properties is not clear. In chapter 3, I understand the authors try to explain its difference to the ode like systems and want to extend the ODE-like system setting to a more general settings with RS and RI properties, but I cannot find a specific example which is included in the RS-RI setting but excluded in the ODE setting. Can the author provide a useful example?\n\nIt seems like all the experiments in the paper can be treated as ODE-like dynamical system, so I was wondering what is the major difference here?\n\n2.There is a missing literature review regarding invariant learning and its application in data-driven dynamical system learning, there are a few branches of related work worth mentioning and they are very close to the topic in this paper:\nPhysics/Conservation invariant:\n\n\"Hamiltonian neural networks.\" Advances in neural information processing systems 32 (2019).\n\nLagrangian Neural Networks. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations.\n\nConCerNet: A Contrastive Learning Based Framework for Automated Conservation Law Discovery and Trustworthy Dynamical System Prediction. International Conference on Machine Learning 2023.\n\nSymmetry invariant: Incorporating Symmetry into Deep Dynamics Models for Improved Generalization, ICLR 2021\n\nActually the Hamiltonian Neural network paper provides an image-based pendulum example which is similar to the experiment in this paper. \n\n3.Regarding the novelty of the model structure, I feel like it still belongs to the encoder-latent dynamics-decoder model class like many other prior works. Despite the author claimed the model decomposes the realization-specific (RS) and realization-invariant (RI) information, there is no evidence proving the claim both experimentally and theoretically.\n\n4.Honestly I’m a bit skeptical regarding the massive improvement of the proposed method over the existing work, especially in figure 4 and 5 (see my questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)