Stochastic Neural Simulator for Generalizing Dynamical Systems across Environments

1 Introduction

The learning of dynamical systems is a fundamental task in various scientific domains such as economics [Varian, 1981], geophysics [Samelson and Wiggins, 2006], and epidemiology [Galea et al., 2010]. However, traditional learning methods heavily rely on expert knowledge and may require extensive computational resources, primarily due to the necessity of solving a large number of complex differential equations [Fu et al., 2017; Liu et al., 2021]. A popular current approach is to utilize neural networks to construct data-driven neural simulators [Li et al., 2021; Chen et al., 2018]. These networks can mitigate the need for expert knowledge and enhance computational speed [Li et al., 2021]. They have also been successfully employed to address challenging problems in chaotic dynamics [Linot et al., 2023] and to expedite scientific discovery [Reichstein et al., 2019].

∗Corresponding authors

However, the majority of current neural simulators are constrained by the i.i.d. assumption, which fantasizes that the observed trajectories are abundant and originate from an unchanging environment [Yin et al., 2021]. In reality, data from dynamical systems are influenced by various environmental factors such as gravity, pressure, and temperature [Baradel et al., 2020; Sanchez-Gonzalez et al., 2020], leading to the inadequacy and inefficiency of existing models when applied to real-world data. While it is possible to train multiple models specific to different environments, this approach would require a substantial amount of computational resources and would fail to capture the potential commonalities in dynamics across different environments, resulting in poor predictive performance when the data from individual environments is limited or sparse [Huang et al., 2023]. Consequently, learning dynamical systems across multiple environments remains a fundamental challenge.

Currently, there are several researches exploring crossenvironmental learning models for system dynamics [Yin et al., 2021; Kirchmeyer et al., 2022]. As illustrated in Fig. 1, these models utilize a shared component and an environmentspecific component to model ordinary differential equations (ODEs) that characterize system dynamics for each environment, capturing commonalities across environments and the distinct effects of each environment's settings. In practical applications, real-world data may be costly or incomplete due to the expensive and time-consuming nature of data acquisition [Huang et al., 2020], as well as the frequent occurrence of broken sensors or damaged memory units during data collection [Tang et al., 2020]. However, these methods either require numerous parameters to be tuned in the environment-specific component, necessitating a large number of environment-specific trajectories [Yin et al., 2021], or they carry the risk of overfitting in sparse settings [Kirchmeyer et al., 2022]. Moreover, existing methods are deterministic, focusing on learning a deterministic system evolution process, even when dealing with sparse data. This is undoubtedly suboptimal.

To tackle these challenges, we propose a novel Contextinformed Neural ODE Processes (CoNDP) to learn dynamical systems from sparse observations across environments. It innovatively introduces uncertainty modeling, where environmental parameters are generated by a random control vector, retaining the advantage in few-shot learning (see Fig. 1). A context-informed encoder is proposed to produce the conditional distribution based on contextual information, i.e., p(ue |Context, ξe ). The incorporation of uncertainty enables the generation of diverse governing equations, thereby enhancing the model's generalization capabilities and supporting learning in sparse settings. Note that the parameters in the encoder are divided into two components, namely W c and W e , with W e determined by a hyper network, thereby endowing our model with greater expressive power to capture the inherent commonalities across environments and distinguishable differences among them.

The main contributions are summarized as follows:

• We propose a novel Context-informed Neural ODE Processes (CoNDP) to learn dynamical systems from sparse observations across environments. To our knowledge, this is the first work to incorporate uncertainty into modeling cross-environmental dynamical systems to enhance generalization in sparse setting.
• Extensive experiments are conducted on various complex dynamical systems in ecology, chemistry, and physics, demonstrating the effectiveness of the CoNDP with superior results compared to existing models.

2 RelatedWork and Problem Statement

2.1 RelatedWork

Learning dynamical systems for a single environment. Neural simulators modeling system dynamics for a single environment can be broadly categorized into two groups: discrete and continuous. Typical discrete neural simulators are constructed using recurrent neural networks such as LSTM [Hochreiter and Schmidhuber, 1997] and GRU [Cho et al.,

2014] due to the power to process sequences. These models often integrate various prior knowledge to enhance learning, such as antisymmetric relations [Chang et al., 2019], system Lagrangians [Rajchakit et al., 2021], and Lyapunov stability [Engelken et al., 2023]. However, prior knowledge requires early efforts or a broad understanding of the dynamical system, which is often impractical or unattainable. Additionally, the discrete nature of these methods makes them struggle with irregularly-sampled observations. Therefore, continuous neural simulators, such as neural ordinary differential equations (neural ODEs) [Chen et al., 2018], have garnered increasing attention. It has been found that neural ODEs can serve as powerful tools for modeling continuous time series [Weerakody et al., 2021; Kidger et al., 2020; Morrill et al., 2021] and system dynamics [Legaard et al., 2023; B¨ottcher et al., 2022; Linot et al., 2023; Gupta and Lermusiaux, 2021], especially in latent space [Rubanova et al., 2019]. ODE-RNN [Chen et al., 2018] is a representative model that combines neural ODE and recurrent neural networks, where the latter is utilized to update hidden states when observations are available. However, these methods still unrealistically assume that the observed trajectories are abundant and originate from unchanging environments, leading to ill-defined and ineffective modeling for real-world systems.

Learning dynamical systems across environments. Currently, there have been several researches on exploring the learning of dynamical systems across different environments [Yin et al., 2021; Kirchmeyer et al., 2022]. LEADS is the first attempt to address multi-environmental scenarios [Yin et al., 2021]. It utilizes a shared component and an environmentspecific component to model ODEs that characterize system dynamics in each environment, capturing both commonalities across environments and the distinct effects of each environment's settings as dxet dt = (fθc + gθe )(xet , t). Here, xet represents the system state at time t in environment e, while θc and θe respectively parameterize the commonalities and differences among environments. After the two components are trained across multiple environments, the shared component is frozen, and only the environment-specific component is tuned when adapting to new environments. It directly optimizes θe during adaptation, resulting in the need for a large number of environment-specific trajectories [Yin et al., 2021]. To alleviate the data scale issue, CoDA [Kirchmeyer et al., 2022] introduces a hyper network to produce θe . During the adaptation phase, it determines θe by tuning environmentspecific parameters ξe , where |ξe | ≪ |θe |. Consequently, CoDA learns fewer parameters than LEADS during the adaptation phase, mitigating overfitting. However, these methods are still deterministic, focusing on learning a deterministic system evolution process. This makes it difficult to generalize in sparse settings caused by expensive data acquisition [Huang et al., 2020] or malfunctioned sensors [Tang et al., 2020] in the real world. Therefore, we need to consider how to generalize to sparse settings for learning dynamical systems across environments.

2.2 Problem Statement

The objective of learning dynamical systems across environments is to construct a generalized neural simulator S. This simulator uses sparse trajectories from the past K timestamps and the current timestamp to predict future trajectories over 1 : T timestamps for any environment e. In other words, it aims to achieve the mapping: S(xet-K , · · · , xet-1 , xet0 ) → (xet1 , · · · , xetT ), where the timestamps can have non-uniform intervals and take continuous values.

We consider a dynamical system whose dynamics for environment e is governed by the following ODE:

where xet ∈ Rd is system state at time t under environment e ∈ E. The dimension of the state is denoted by d, and E signifies the set of environments. The observed data from a system under environment e consists of the sparse trajectories De = {xet-K , xet-K+1 , · · · , xet-1 , xet0 , xet1 , · · · , xetT }. We focus on learning f e based on the observed data from multiple environments, i.e., D1 ,D2 , · · · ,D|E| , to capture f e′ for any new environment e′ . This is essential for predicting the system dynamics and achieving a generalized neural simulator. Note that we are considering two tasks: when the new environment e′ belongs to the training environment set E, we refer to this as the transductive task, whereas when e′ is not in E, it is referred to as the inductive task.

3 CoNDP: Context-Informed Neural ODE Processes

In this section, we present Context-Informed Neural ODE Processes (CoNDP) to generalize to sparse settings for learning dynamical systems across environments.

3.1 Model Overview

An overview of the CoNDP is shown in Fig. 2. It consists of four main components: an initial encoder that infers the initial latent states of the system; a context-informed encoder for learning the environment representation, where the attention weights consist of shared weights and environment-specific weights determined by environment-specific parameters and a hyper network; a process of solving initial value problems in hidden space; and a decoder for predicting future system states. Initial encoder estimates the distribution of initial latent state zt0 from initial state xet0 , i.e., p(zt0 |xet0 ). We denote the historical trajectories over the past K timestamps as Context = (xet-K , · · · , xet-1 ). The context-informed encoder is to map the Context into a conditional distribution of environmental representation ue , i.e., p(ue |Context, ξe ), which is determined by environment-specific parameters ξe and a hyper network. The environmental representation can be viewed as a random control vector of the system dynamics, which is fed into a control net to produce the environmentspecific component of the governing equation f . By sampling the initial latent state zt0 , we can solve the IVP to generate the hidden state at any time t, i.e., zt. Then, a decoder is used to provide the predictive distribution p(xe1:T |zt0 , ue ). On the contrary, all shared parameters and environment-specific parameters are learned based on the empirical error between predicted and true values.

3.2 Model Details

We now introduce the four main components of the model.

Initial encoder. Given the initial condition xet0 of the dynamical system over environment e, the initial encoder encφi generates a Gaussian distribution of the initial latent state, i.e., p(zt0 |xet0 ) = N(µz, σ2z ). The mean and variance of the distribution are obtained from

where φi is the learnable parameters that are updated during the training phase. The initial latent state zt0 determines the starting point for carrying out dynamic evolution in hidden space. Note that zt0 is a random vector, so it can generate diverse starting points in prediction. We implement encφi using a Multilayer Perception (MLP) whose output is split into two halves to represent the mean and variance respectively.

Context-informed encoder. The evolution process of dynamical systems is often influenced by external environments, and environmental characteristics are usually latent and challenging to observe. We thus design an encoder to extract environmental features from data to automatically adapt to the impact of the environment.

As in many cases, the influence of the environment features on system evolution is often manifested in delayed effects rather than instant behaviors [Glass et al., 2021; Smith, 2011; Kuang, 1993; Guglielmi et al., 2022], we employ a temporal self-attention mechanism to learn environment features from observed trajectories by summarizing system behaviors with different delays. This also offers better parallelization for accelerating training speed and alleviates the gradient issues brought by long sequences [Sankar et al., 2020]. Specifically, we use a temporal encoding [Kazemi et al., 2020] to integrate the time delay information.

Given historical trajectories over the past K timestamps under environment e as Context = (xet-K , · · · , xet-1 ), the computation of the representation of the observation at time tm involves conducting a weighted summation over each different time point tj, along with a residual connection. The attention score is determined using a transformer-based approach [Vaswani et al., 2017] and is computed as the dotproduct of observation representations derived from value, key, and query projection matrices. We calculate the representation for each historical observation xetm , where m ∈ {-1, · · · ,-K}, as follows

where σ(·) is a non-linear activation function, the temporal encoding TE(·) satisfying TE(∆t)2i = sin( ∆t 100002i/d ), and

TE(∆t)2i+1 = cos( ∆t 100002i/d ), 1 ⩽ i ⩽ d, where d is the state dimension. Wk, Wv, W c q , andW e q are the parameters in this attention layer. Here we assume the query projection matrix is different among environments by introducing W e q . This enables our model to receive different attention scores when dealing with different environments, enhancing its ability to capture the intrinsic commonalities and distinguishable differences among environments. W e q is determined by a hyper network, with its input being an environment-specific parameter ξe . This parameter can be viewed as the unobservable implicit characteristic of the environment. As it exclusively focuses on the differences among environments (the commonalities are modeled with other shared parameters in the encoder), ξe can be very low-dimensional, enabling rapid adaptation to new environments.

Then, we can obtain the final representation of each observation based on the following attention operations

where l = 1, 2, · · · ,L and dh is the dimension of hidden representation. We use the output of the final attention layer as the representation for each observation, denoted as h(L+1) tm .

Following the computation of representations for each observation, we aggregate the entire observed trajectory into a vector to provide conditional distribution of the environment representation, i.e., p(ue |Context, ξe ) = N(µu, σ2u), as

where ae = tanh(( 1K ∑ ˆhtm )Wa) is the average of observation representations with a non-linear transformation and ˆhtm = h(L+1) tm +TE(t), Wa is the weights, K is the number of observations for each trajectory, and encφu is implemented via the MLP with trainable parameters φu. Note that the environment representation ue is a random control vector for generating diverse governing equations of dynamical systems, further enhancing the generalization in sparse settings.

The process of solving initial value problems. After obtaining two conditional distributions, namely p(zt0 |xet0 ) and p(ue |Context, ξe ), we sample an initial latent state zt0 and an environment representation ue from these distributions. Subsequently, we can predict the latent states at any time t by solving an initial value problem in the latent space as

where θc is the shared parameters across environments and θe is environment-specific parameters generated by a control network with parameters φm, i.e., θe = conφm (ue ).

Due to the sparsity of observed data, it becomes challenging to uniquely determine a dynamical equation and ensure that the data reveals all (or at least most) of the features of the environment. However, existing methods typically only learn a deterministic dynamical process from observed data [Yin et al., 2021; Kirchmeyer et al., 2022] and are unable to capture the uncertainty and multiple solutions in the systems caused by low-resource data. Fortunately, uncertainty modeling can alleviate this issue [Norcliffe et al., 2020]. This is why we use random vectors to model initial signals and environment representations.

State decoder. Using latent states zt1 , zt2 , · · · , ztT and placing a Gaussian distribution on system states, we can obtain the predicted system states by employing a decoder as

where, j = 1, 2, · · · , T and µxtj and σ2xtj are the mean and variance of the predictive distribution of system state, i.e., p(xe1:T |zt0 , ue , t1:T ) = N(µxt1:T , σ2xt1:T ). This predictive distribution can accommodate system noise.

3.3 CoNDP as Stochastic Process

From a probabilistic perspective, our proposed CoNDP is a type of neural network-parameterized stochastic process for system states. The Kolmogorov Extension Theorem states that exchangeability and consistency conditions are sufficient to define a stochastic process [Oksendal, 2013]. We demonstrate that the CoNDP satisfies these conditions and present the following proposition.

Proposition 1 CoNDP satisfies the exchangeability and consistency conditions.

In other words, the stochastic processes we have established exist, and the CoNDP is its family of finite-dimensional distributions. As a stochastic process, it should theoretically have excellent generalization power for sparse settings. The detailed proof can be found in supplementary material 1 .

3.4 Training and Adpating

Now, we introduce the overall training procedure of the CoNDP. We randomly generate multiple environments E = {e1, e2, ...} for a dynamical system and obtain sparse trajectories through irregularly sampling observations under each environment. These trajectories are divided into two halves along the time, where the first half is [t-K, t-1] for learning environment representation, and the second half, i.e., [t0, tT ], is to predict system states. We denote the collected data as De = {xet-K , xet-K+1 , · · · , xet-1 , xet0 , xet1 , · · · , xetT } for environment e. Based on observed data from multiple environments, i.e., De1 ,De2 , ..., we jointly train all parts in an endto-end way to achieve a generalized neural simulator. Specifically, following [Norcliffe et al., 2020], we learn the model by maximizing the following variational lower bound

1 https://github.com/ljqjlu/CoNDP/blob/main/supp-material.pdf

where qzt0 ,ue = p(zt0 |xet0 )p(ue |Context, ξe ). Detailed implementations, such as neural network architectures, can be found in supplementary material.

For adaptation, we assume that a small number of sparse trajectories are observed under the new environment e′ (not more than 10 in experiments). To adapt to the new environment, we freeze all the parameters of the model except the environment-specific parameter ξe . Using the same loss function as in the training procedure, we only tune the ξe in the adaptation phase. After convergence, the entire model is frozen and deployed for testing in the new environment.

4 Experiments

We test the CoNDP 2 on five complex dynamical systems in various fields to answer the following questions: 1) Can the introduction of uncertainty improve model performance, especially in the sparse setting? 2) What factors in our model have an impact on performance? 3) Has our model learned appropriate environment representations and reasonable governing laws of dynamical systems?

Dynamical systems. We consider several complex dynamical systems: Lotka-Volterra (LV) predator-prey system [Lotka, 1925] is to describe the dynamics of biological systems where two species interact. Glycolitic-Oscillator (GO) system [Ruoff et al., 2003] is to describe the dynamics of yeast glycolysis in biology. Gray-Scott (GS) [Pearson, 1993] is a reaction-diffusion system where two chemicals react while spreading over space, forming Turing patterns [Horv´ath et al., 2009]. Navier-Stokes equation (NS) is a fluid dynamics equation [Constantin and Foias¸, 1988] describing the waveform of the fluid system. The heat diffusion equation (HEAT) describes the diffusion of thermal energy [Widder, 1976]. Different parameters in dynamical systems correspond to different environments.

Building training, tuning, and testing sets. For the training set, we sample N env tr environments per system and N tra tr trajectories per environment. In addition, we sample N tra tu (up to 10) and N tra te (∼ 100) trajectories per new environment for tuning and testing, respectively. A large number of environments from each system are collected as testing environments. If the testing environments have already been encountered during training, the task of predicting trajectories becomes a transductive task. Otherwise, if the testing environments are entirely new and outside the scope of the training set, it becomes an inductive task. To evaluate the model's performance in a sparse setting, we randomly drop a fraction of environments and observed points of trajectories for the training, tuning, and testing sets. Additional details regarding the dynamical systems and experimental settings can be found in the supplementary material.

2 Our code is available at https://github.com/ljqjlu/CoNDP.

Baselines. We compare our model against several representative neural simulators for a single environment, including discrete models such as GRU [Cho et al., 2014] and LSTM [Hochreiter and Schmidhuber, 1997], as well as continuous models such as Neural ODE [Chen et al., 2018] and ODE-RNN [Chen et al., 2018]. Additionally, we compare with state-of-the-art models for learning dynamical systems across environments, namely LEADS [Yin et al., 2021] and CoDA [Kirchmeyer et al., 2022]. Note that both are unable to utilize contextual historical trajectories during the testing phase, but make predictions directly from time t0. Therefore, we fed the context into their tuning phase to ensure fairness.

4.1 Performance Evaluation under Sparse Settings

We conducted tests on five systems under different testing environments and sparsity levels. Specifically, for each type of task per system, we tested on 100 environments with the ratio of observed data ranging from 4% to 100%. From Table 1, it is evident that the models across environments significantly outperform the single environment models in all dynamical systems. This demonstrates that the former can effectively distinguish and utilize the commonalities across environments and environment-specific differences, leading to better performance in sparse settings, while the latter suffers from the i.i.d. assumption. Thanks to uncertainty modeling and the introduction of contextual information, our CoNDP achieves optimal results, reducing prediction errors by 3.15% to 8.88% compared to the second-best method. The performance on the inductive task is generally worse than that on the transductive task, which also indicates that adapting to completely new environments in sparse settings is difficult. However, our model still demonstrates generalization power for unseen environments. Fig. 3 provides a visualization of the predictive trajectories for a sampled unseen environment of the LV system under various sparsity levels. We also present predictive errors at different sparsity levels in Fig. 4. The sparse scenarios do indeed weaken model performance. Sufficient training data enables the models to learn the commonalities and differences among environments, thereby improving performance. And, the sparsity of testing data makes it difficult to capture features for new environments. Our CoNDP outperforms both LEADS and CoDA, and the gap becomes increasingly apparent as the setting becomes sparser. The results for other systems are similar and can be found in supplementary material.

Figure 5: Ablation study and sensitivity analysis for CoNDP.

4.2 Ablation Study and Sensitivity Analysis

To analyze the rationality behind our design, we conducted two variants of the CoNDP. One variant involved our model without the environment-specific attention weight W e q , resulting in the removal of the hyper-network and making the attention weights shared across all environments. The other variant was our model without both the environment-specific component and the ODE process in latent space, directly using neural networks to calculate solutions while bypassing the integration process. These two variants can be viewed as modified versions of Neural Processes [Garnelo et al., 2018] and Neural ODE Processes [Norcliffe et al., 2020], respectively. We tested them on the LV system. From Fig. 5(a), it is evident that our CoNDP outperforms the variants under various levels of sparsity, demonstrating the effectiveness of the design of the context-informed encoder and the ODE component. We also tested the effect of different sizes of context for our model and provided the prediction error for the inductive task on the LV system in Fig. 5(b). Although intuitively, a longer length of context would better reflect the features of the trajectories in the environments, we found that sparsity also affects model performance, with large and irregularlysampled intervals making it difficult to obtain obvious features. Empirically, we observed that model performance was already satisfactory when the context size was around 30. Thus, the size of the context was set to 30 in our experiments.

4.3 Environment Representations

To confirm whether the context-informed encoder can acquire reasonable environment representations and better distinguish the inherent commonalities and differences among environments, we visualize the optimized environment-specific parameter ξe and the learned environment representation ue for a variety of new testing environments by two-dimensional projection. In terms of the optimized ξe , we observe that they are distinguishable under different environments and are close to each other for similar environments, as shown in Fig. 6(a). For the learned ue , we conducted tests on the GrayScott equations and only allowed the reaction rate of one chemical to vary to form new environments. From Fig. 6(b), we see that they are clustered within three reaction rates, indicating that our model can effectively distinguish different environments from the observed trajectories.

4.4 Discovering Governing Law of Systems

To evaluate the capability of the CoNDP in capturing the fundamental governing law of the system, we employed a popular sparse symbolic regression, namely SINDy [Brunton et al., 2016], to discover the governing law from the predicted trajectories in unseen environments. However, the relatively large and irregularly-sampled intervals in our settings posed a significant challenge for symbolic regression techniques when using the finite difference method [Perrone and Kao, 1975] to approximate temporal derivatives. We, thus, generated densely-sampled observations from the predictive trajectories to enable the discovery of the governing law, which also serves to validate our model's capacity for temporal hyper-resolution sampling. As depicted in Fig. 7, it becomes difficult to discover the governing law of the system when fewer observations are provided. Our CoNDP demonstrates the ability to discover more accurate equations compared to others, particularly in sparse settings. The discovered equations under different sparsity levels can be found in supplementary material. This reveals that the adaptive attention and probabilistic modeling offer greater flexibility in capturing the essential governing law from limited observations.

5 Conclusion

In this study, we propose a novel stochastic neural simulator (CoNDP) designed to learn dynamical systems across environments, thereby overcoming the unrealistic i.i.d. assumption on environments. Comprehensive experiments have demonstrated that it outperforms several representative models designed for a single environment, as well as state-ofthe-art models focused on learning across environments. Our analysis reveals that our model excels in distinguishing the inherent commonalities and differences among environments, and the introduction of uncertainty enhances generalization in sparse settings. Nevertheless, challenges persist in effectively modeling counterfactual physical systems and integrating system control for decision-making support. Additionally, the influence of different environments also exists in the network dynamics of the real world. Addressing how to extend the model to high-dimensional network dynamics modeling is an important topic for future research.

Acknowledgments

This work was supported by the National Science and Technology Major Project under Grant No. 2021ZD0112500; the National Natural Science Foundation of China under Grant Nos. U22A2098, 62172185, 62206105 and 62202200; the Jilin Province Youth Science and Technology Talent Support Project under Grant No. QT202225; the Key Science and Technology Development Plan of Jilin Province under Grant No. 20240302078GX.