---
title: "Stochastic Neural Simulator for Generalizing Dynamical Systems across Environments"
authors:
  - "Liu Jiaqi"
  - "Jiaxu Cui"
  - "Jiayi Yang"
  - "Bo Yang"
doi: "10.24963/ijcai.2024/653"
url: "https://ijcai.org/proceedings/2024/653"
pdf: "https://ijcai.org/proceedings/2024/0653.pdf"
venue: "Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence"
year: 2024
pages: "5909-5917"
track: "Main Track"
keywords:
  - "Multidisciplinary Topics and Applications: MTA: Physical sciences"
  - "Machine Learning: ML: Time series and data streams"
---

# Stochastic Neural Simulator for Generalizing Dynamical Systems across Environments

Liu Jiaqi, Jiaxu Cui, Jiayi Yang, Bo Yang

*Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence*, pages 5909-5917. DOI: [10.24963/ijcai.2024/653](https://doi.org/10.24963/ijcai.2024/653)

## Abstract

Neural simulators for modeling complex dynamical systems have been extensively studied for various real-world applications, such as weather forecasting, ocean current prediction, and computational fluid dynamics simulation. Although they have demonstrated powerful fitting and predicting, most existing models are only built to learn single-system dynamics. Several advanced researches have considered learning dynamics across environments, which can exploit the potential commonalities among the dynamics across environments and adapt to new environments. However, these methods still are prone to scarcity problems where per-environment data is sparse or limited. Therefore, we propose a novel CoNDP (Context-Informed Neural ODE Processes) to achieve learning system dynamics from sparse observations across environments. It can fully use contextual information of each environment to better capture the intrinsic commonalities across environments and distinguishable differences among environments while modeling uncertainty of system evolution, producing more accurate predictions. Intensive experiments are conducted on five complex dynamical systems in various fields. Results show that the proposed CoNDP can achieve optimal results compared with common neural simulators and state-of-the-art cross-environmental models.

## Full text

## 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

Figure 1: Comparison of state-of-the-art model principles for learning dynamical systems across environments. All models have two components, i.e., a shared component (blue) and an environmentspecific component (orange). Both LEADS [Yin et al., 2021] and CoDA [Kirchmeyer et al., 2022] are deterministic. As a comparison, our CoNDP innovatively introduces uncertainty modeling, where environmental parameters are generated by a random control vector (ue ), enhancing the model's generalization power.

![Image](https://ijcai.org/proceedings/2024/fulltext/653/image_000000_42b1a83fe380ce2e4bf64162b65b1c8a34c25bffdec9fa5426b28386f58a12cd.png)

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.

Figure 2: Overview of the CoNDP. Initial encoder estimates the distribution of initial latent state zt0 from xet0 , i.e., p(zt0 |xet0 ). The contextinformed encoder is to map the context into a conditional distribution of environment representation ue , i.e., p(ue |Context, ξe ), which is controlled by the environment-specific parameter ξe . Environment representation can be seen as a random control vector of system dynamics, which is fed into a control net to produce the environment-specific component of the governing equation f. By sampling the initial latent state, we can solve the initial value problem (IVP) to generate the hidden state at any time t, i.e., zt. Then, using a decoder to give the predictive distribution.

![Image](https://ijcai.org/proceedings/2024/fulltext/653/image_000001_09bbef43b01f9d8b2abe470ef3d921fff949f597113af6e9e3a3c037ecb45205.png)

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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.

Table 1: The average of Mean Square Error (MSE) between predictions and ground truth from various dynamical systems across all testing environments and sparsity levels for both transductive and inductive tasks. The best results are bolded.

|            | Transductive - LV (×10-5 )   | Transductive - GO (×10-4 )   | Transductive - GS (×10-4 )   | Transductive - NS (×10-4 )   | Transductive - HEAT (×10-2 )   | Inductive - LV (×10-5 )   | Inductive - GO (×10-4 )   | Inductive - GS (×10-3 )   | Inductive - NS (×10-4 )   | Inductive - HEAT (×10-2 )   |
|------------|------------------------------|------------------------------|------------------------------|------------------------------|--------------------------------|---------------------------|---------------------------|---------------------------|---------------------------|-----------------------------|
| GRU        | 7.56                         | 141.52                       | 8.41                         | 173.70                       | 367.66                         | 274.37                    | 7312.40                   | 64.62                     | 121.24                    | 274.42                      |
| LSTM       | 7.54                         | 132.58                       | 8.18                         | 167.39                       | 428.03                         | 256.79                    | 7286.87                   | 70.69                     | 118.01                    | 274.47                      |
| Neural ODE | 8.01                         | 130.43                       | 7.84                         | 169.44                       | 417.31                         | 299.41                    | 7435.18                   | 64.61                     | 118.19                    | 278.50                      |
| ODE-RNN    | 7.66                         | 138.24                       | 8.45                         | 170.11                       | 384.97                         | 282.75                    | 7951.86                   | 63.58                     | 115.20                    | 293.72                      |
| LEADS      | 3.75                         | 49.96                        | 3.34                         | 35.23                        | 113.30                         | 85.12                     | 212.31                    | 2.46                      | 33.02                     | 175.49                      |
| CoDA       | 1.69                         | 2.58                         | 1.55                         | 9.87                         | 5.07                           | 2.26                      | 5.72                      | 1.83                      | 10.70                     | 7.21                        |
| CoNDP      | 1.54                         | 2.41                         | 1.48                         | 9.33                         | 4.74                           | 2.16                      | 5.54                      | 1.73                      | 10.02 ↓                   | 6.98                        |
|            | ↓ 8.88%                      | ↓ 6.59%                      | ↓ 4.52%                      | ↓ 5.47%                      | ↓ 5.95%                        | ↓ 4.42%                   | ↓ 3.15%                   | ↓ 5.46%                   | 6.36%                     | ↓ 3.19%                     |

## 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

<!-- formula-not-decoded -->

[1 https://github.com/ljqjlu/CoNDP/blob/main/supp-material.pdf](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.](https://github.com/ljqjlu/CoNDP)

Figure 3: Comparison of predictive trajectories for a sampled unseen environment of Lotka-Volterra (LV) system under various sparsity levels. ρtrain and ρtest denote the ratio of observed data in the training and testing sets, respectively. Due to the uncertainty of our model, the solid line of our model is the mean of predictive trajectories, while the shaded area is twice the standard deviation approximating a 95% confidence interval. It can be seen that our CoNDP performs best even with few observations, and as the number of observations decreases, the uncertainty of the model increases, partially covering the truth.

![Image](https://ijcai.org/proceedings/2024/fulltext/653/image_000002_c0dfea5c669473f3e60ed91e663d3987a53a5a987c848c449563579a6f246826.png)

Figure 4: Comparison of prediction errors under different sparsity levels on Lotka-Volterra (LV) system.

![Image](https://ijcai.org/proceedings/2024/fulltext/653/image_000003_98f74039b3f47d7883a99424123c0960d156d110f5e0923b8bf4ac17897d61c1.png)

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.

Figure 6: T-SNE visualizations of the environment-specific parameter ξe from LV system (a) and the environment representations ue also from LV system (b).

![Image](https://ijcai.org/proceedings/2024/fulltext/653/image_000004_40cc3b8668b6a165135e494aa9168cbbd56541f1d039d3fa8d559dde9ad74e3d.png)

## 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.

Figure 7: The average L1 error between the coefficients of the discovered governing law and those of the true equation of LV system. The lower the coefficient error, the closer the discovered equation is to the truth.

![Image](https://ijcai.org/proceedings/2024/fulltext/653/image_000005_316c1f46bf8d1e664adca2eefa88b496af3cf1cd0ea465c6383e1d0713472f4d.png)

## 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.

## References

1. [Baradel et al., 2020] Fabien Baradel, Natalia Neverova, Julien Mille, Greg Mori, and Christian Wolf. Cophy: Counterfactual learning of physical dynamics. In Inter- national Conference on Learning Representations, 2020.
2. [Böttcher et al., 2022] Lucas Böttcher, Nino Antulov- Fantulin, and Thomas Asikis. Ai pontryagin or how artificial neural networks learn to control dynamical systems. Nature communications, 13(1):333, 2022.
3. [Brunton et al., 2016] Steven L Brunton, Joshua L Proctor, and J Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical sys- tems. Proceedings of the national academy of sciences, 113(15):3932-3937, 2016.
4. [Chang et al., 2019] Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi. AntisymmetricRNN: A dynamical system view on recurrent neural networks. In International Con- ference on Learning Representations, 2019.
5. [Chen et al., 2018] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary dif- ferential equations. Advances in neural information pro- cessing systems, 31, 2018.
6. [Cho et al., 2014] Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. In Proceedings of the 2014 Confer- ence on Empirical Methods in Natural Language Process- ing (EMNLP), page 1724. Association for Computational Linguistics, 2014.
7. [Constantin and Foias ¸, 1988] Peter Constantin and Ciprian Foias ¸. Navier-stokes equations. University of Chicago press, 1988.
8. [Engelken et al., 2023] Rainer Engelken, Fred Wolf, and Larry F Abbott. Lyapunov spectra of chaotic recurrent neural networks. Physical Review Research, 5(4):043044, 2023.
9. [Fu et al., 2017] Haohuan Fu, Conghui He, Bingwei Chen, Zekun Yin, Zhenguo Zhang, Wenqiang Zhang, Tingjian Zhang, Wei Xue, Weiguo Liu, Wanwang Yin, et al. 9- pflops nonlinear earthquake simulation on sunway taihu- light: enabling depiction of 18-hz and 8-meter scenarios. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analy- sis, pages 1-12, 2017.
10. [Galea et al., 2010] Sandro Galea, Matthew Riddle, and George A Kaplan. Causal thinking and complex system approaches in epidemiology. International journal of epi- demiology, 39(1):97-106, 2010.
11. [Garnelo et al., 2018] Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Es- lami, and Yee Whye Teh. Neural processes. In ICML Workshop on Theoretical Foundations and Applications of Deep Generative Models, 2018.
12. [Glass et al., 2021] David S Glass, Xiaofan Jin, and Ing- mar H Riedel-Kruse. Nonlinear delay differential equa- tions and their application to modeling biological network motifs. Nature communications, 12(1):1788, 2021.
13. [Guglielmi et al., 2022] Nicola Guglielmi, Elisa Iacomini, and Alex Viguerie. Delay differential equations for the spatially resolved simulation of epidemics with specific application to covid-19. Mathematical Methods in the Ap- plied Sciences, 45(8):4752-4771, 2022.
14. [Gupta and Lermusiaux, 2021] Abhinav Gupta and Pierre FJ Lermusiaux. Neural closure models for dynami- cal systems. Proceedings of the Royal Society A, 477(2252):20201004, 2021.
15. [Hochreiter and Schmidhuber, 1997] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735-1780, 1997.
16. [Horváth et al., 2009] Judit Horváth, István Szalai, and Patrick De Kepper. An experimental design method lead- ing to chemical turing patterns. Science, 324(5928):772- 775, 2009.
17. [Huang et al., 2020] Zijie Huang, Yizhou Sun, and Wei Wang. Learning continuous system dynamics from irregularly-sampled partial observations. Advances in Neural Information Processing Systems, 33:16177-16187, 2020.
18. [Huang et al., 2023] Zijie Huang, Yizhou Sun, and Wei Wang. Generalizing graph ode for learning complex sys- tem dynamics across environments. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 798-809, 2023.
19. [Kazemi et al., 2020] Seyed Mehran Kazemi, Rishab Goel, Sepehr Eghbali, Janahan Ramanan, Jaspreet Sahota, San- jay Thakur, Stella Wu, Cathal Smyth, Pascal Poupart, and Marcus Brubaker. Time2vec: Learning a vector represen- tation of time, 2020.
20. [Kidger et al., 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:6696-6707, 2020.
21. [Kirchmeyer et al., 2022] Matthieu Kirchmeyer, Yuan Yin, Jérémie Donà, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing to new physical sys- tems via context-informed dynamics model. In Interna- tional Conference on Machine Learning, pages 11283- 11301. PMLR, 2022.
22. [Kuang, 1993] Yang Kuang. Delay differential equations: with applications in population dynamics. Academic press, 1993.
23. [Legaard et al., 2023] Christian Legaard, Thomas Schranz, Gerald Schweiger, Ján Drgoňa, Basak Falay, Cláudio Gomes, Alexandros Iosifidis, Mahdi Abkar, and Peter Larsen. Constructing neural network based models for simulating dynamical systems. ACM Computing Surveys, 55(11):1-34, 2023.
24. [Li et al., 2021] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equa- tions. In International Conference on Learning Represen- tations, 2021.
25. [Linot et al., 2023] Alec J Linot, Joshua W Burby, Qi Tang, Prasanna Balaprakash, Michael D Graham, and Romit Maulik. Stabilized neural ordinary differential equations for long-time forecasting of dynamical systems. Journal of Computational Physics, 474:111838, 2023.
26. [Liu et al., 2021] Yong Liu, Xin Liu, Fang Li, Haohuan Fu, Yuling Yang, Jiawei Song, Pengpeng Zhao, Zhen Wang, Dajia Peng, Huarong Chen, et al. Closing the" quantum supremacy" gap: achieving real-time simulation of a ran- dom quantum circuit using a new sunway supercomputer. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analy- sis, pages 1-12, 2021.
27. [Lotka, 1925] Alfred James Lotka. Elements of physical bi- ology. Williams & Wilkins, 1925.
28. [Morrill et al., 2021] James Morrill, Cristopher Salvi, Patrick Kidger, and James Foster. Neural rough differ- ential equations for long time series. In International Conference on Machine Learning, pages 7829-7838.
29. PMLR, 2021.
30. [Norcliffe et al., 2020] Alexander Norcliffe, Cristian Bod- nar, Ben Day, Jacob Moss, and Pietro Liò. Neural ode processes. In International Conference on Learning Rep- resentations, 2020.
31. [Oksendal, 2013] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
32. John E Pearson. Complex patterns in a sim- ple system. Science, 261(5118):189-192, 1993.
33. [Perrone and Kao, 1975] Nicholas Perrone and Robert Kao. A general finite difference method for arbitrary meshes. Computers & Structures, 5(1):45-57, 1975.
34. [Rajchakit et al., 2021] G Rajchakit, R Sriraman, N Boon- satit, P Hammachukiattikul, Chee Peng Lim, and P Agar- wal. Exponential stability in the lagrange sense for clifford-valued recurrent neural networks with time delays. Advances in Difference Equations, 2021(1):1-21, 2021.
35. [Reichstein et al., 2019] Markus Reichstein, Gustau Camps- Valls, Bjorn Stevens, Martin Jung, Joachim Denzler, Nuno Carvalhais, and fnm Prabhat. Deep learning and process understanding for data-driven earth system science. Na- ture, 566(7743):195-204, 2019.
36. [Rubanova et al., 2019] Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud. Latent ordinary differential equa- tions for irregularly-sampled time series. Advances in neu- ral information processing systems, 32, 2019.
37. [Ruoff et al., 2003] Peter Ruoff, Melinda K Christensen, Jana Wolf, and Reinhart Heinrich. Temperature de- pendency and temperature compensation in a model of yeast glycolytic oscillations. Biophysical chemistry, 106(2):179-192, 2003.
38. [Samelson and Wiggins, 2006] Roger M Samelson and Stephen Wiggins. Lagrangian transport in geophysi- cal jets and waves: The dynamical systems approach, volume 31. Springer Science & Business Media, 2006.
39. [Sanchez-Gonzalez et al., 2020] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International conference on ma- chine learning, pages 8459-8468. PMLR, 2020.
40. [Sankar et al., 2020] Aravind Sankar, Yanhong Wu, Liang Gou, Wei Zhang, and Hao Yang. Dysat: Deep neural rep- resentation learning on dynamic graphs via self-attention networks. In Proceedings of the 13th international con- ference on web search and data mining, pages 519-527, 2020.
41. [Smith, 2011] Hal L Smith. An introduction to delay differ- ential equations with applications to the life sciences, vol- ume 57. springer New York, 2011.
42. [Tang et al., 2020] Xianfeng Tang, Huaxiu Yao, Yiwei Sun, Charu Aggarwal, Prasenjit Mitra, and Suhang Wang. Joint modeling of local and global temporal dynamics for mul- tivariate time series forecasting with missing values. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 34, pages 5956-5963, 2020.
43. Hal R Varian. Dynamical systems with ap- plications to economics. Handbook of mathematical eco- nomics, 1:93-110, 1981.
44. [Vaswani et al., 2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
45. [Weerakody et al., 2021] Philip B Weerakody, Kok Wai Wong, Guanjin Wang, and Wendell Ela. A review of irreg- ular time series data handling with gated recurrent neural networks. Neurocomputing, 441:161-178, 2021.
46. [Widder, 1976] David Vernon Widder. The heat equation, volume 67. Academic Press, 1976.
47. [Yin et al., 2021] Yuan Yin, Ibrahim Ayed, Emmanuel de Bézenac, Nicolas Baskiotis, and Patrick Gallinari. Leads: Learning dynamical systems that generalize across environments. Advances in Neural Information Processing Systems, 34:7561-7573, 2021.