# SpReME: Sparse Regression for Multi-Environment Dynamic Systems

MoonJeong Park  Youngbin Choi  Namhoon Lee  Dongwoo Kim

###### Abstract

Learning dynamical systems is a promising avenue for scientific discoveries.
However, capturing the governing dynamics in multiple environments still
remains a challenge: model-based approaches rely on the fidelity of
assumptions made for a single environment, whereas data-driven approaches
based on neural networks are often fragile on extrapolating into the future.
In this work, we develop a method of sparse regression dubbed SpReME to
discover the major dynamics that underlie multiple environments. Specifically,
SpReME shares a sparse structure of ordinary differential equation (ODE)
across different environments in common while allowing each environment to
keep the coefficients of ODE terms independently. We demonstrate that the
proposed model captures the correct dynamics from multiple environments over
four different dynamic systems with improved prediction performance.

*arg max *arg min sign Tr

##  1 Introduction

Discovering the underlying dynamics of a physical system is an extremely
challenging task in many fields of science and engineering; e.g., it took
nearly three decades for experts in fluid mechanics to uncover the true
dynamics of flow passing a cylinder (Noack et al., 2003). Recently, deep
learning-based approaches have gained great attention from the physics and
engineering communities with their promises in realistic scenarios including,
for instance, those where the physical model is unknown (Shi et al., 2015;
Wang et al., 2018), underlying dynamics are incomplete (Yin et al., 2021b) or
a form of perturbation from unknown external sources is present (Li et al.,
2021).

Many current approaches focus on learning a dynamical system from a single
environment (Long et al., 2018). The approaches with partially known physical
models especially show remarkable results in predicting future dynamics (Yin
et al., 2021b). However, learning from a single environment raises the
question of whether the learned dynamics can be generalized to the unseen
environment. For example, suppose the training dynamics contain additional
perturbation from external sources irrelevant to the governing dynamics. In
that case, the learned model is unlikely to generalize to a new environment
where we want to predict future dynamics.

To tackle this issue, some recent works propose learning to capture
commonalities across different environments based on neural networks (Yin et
al., 2021a; Kirchmeyer et al., 2022). However, these developments are mostly
data-driven approaches without utilizing any prior knowledge of the underlying
system. Hence, they are often challenged on their capability to extrapolate
for forecasting future dynamics on unseen time-steps (Xu et al., 2021).

In this work, we aim to provide a general framework that can uncover governing
dynamics from multiple environments with the help of incomplete prior
knowledge. We formulate dynamics discovery as a sparse linear regression
problem under the existence of a few major factors of governing dynamics.
Specifically, we define a linear model whose feature functions are the
candidates of the governing dynamics. We make the linear models defined for
different environments have the same sparse structure through shared binary
masks over the feature functions. The shared binary masks determine which
feature functions are used to form the linear models. To uncover the correct
sparse structure, we develop an optimization algorithm alternating updates on
the shared sparse binary masks and coefficients of features.

In experiments, we compare our model against five baseline models, including
model- and data-based approaches. Four different dynamic systems are used to
measure the performance of our model. The experimental results show that the
proposed model can accurately uncover the underlying dynamics and forecast the
evolution of systems over time. Additional analysis reveals that the proposed
model performs robust prediction under uncertain environments compared to the
other baselines.

#### Contributions

The main contributions of this paper can be summarized as follows:

  * •

We formulate a dynamic discovery in multiple environments as a sparse linear
regression problem.

  * •

We develop an optimization algorithm that can uncover the true dynamics from
multi-environmental observations.

  * •

We show that the proposed model captures true dynamics on various dynamic
systems.

##  2 Related Work

There have been proposed several model-based approaches to leverage some prior
knowledge for finding unknown dynamics (Yin et al., 2021b; Brunton et al.,
2016). Yin et al. (2021b) assume that a physical model is partially known and
tries to learn unknown parts of its dynamics using neural networks. In Brunton
et al. (2016), candidate elements of unknown dynamics are used as feature
(basis) functions to form a sparse regression problem of dynamic systems.
Under the presence of correct basis functions, the underlying dynamics are
accurately discovered. However, sufficient data from a single environment is
needed for using the above methods.

Data-driven approaches have been proposed to uncover the dynamics given the
data collected from multiple environments (Yin et al., 2021a; Kirchmeyer et
al., 2022; Norcliffe et al., 2021). Yin et al. (2021a) suggests employing two
neural network components, one for capturing common dynamics across all
environments and the other for capturing environment-specific dynamics.
Kirchmeyer et al. (2022) propose a method that utilizes a single neural
network, where the parameters are divided into two parts: common parameters
and environment-specific parameters that captures subtle differences across
different environments. This approach allows for efficient adaptation to
unseen environments, as only the environment-specific parameters need to be
re-trained. Neural ODE Processes (NDP) (Norcliffe et al., 2021) combines
Neural ODE (Chen et al., 2018) and Neural Processes (Garnelo et al., 2018).
NDP estimates the uncertainty of neural ODE while dealing with multi-
environment data. The proposed approaches for the multi-environment data are
mostly data-driven. This raises a question of their ability to forecast the
future state given the limitation of neural networks on extrapolation (Ziyin
et al., 2020; Xu et al., 2021).

Figure 1: Examples of Lorenz system. Each trajectory shows an evolution of a
three-dimensional state over time.

##  3 Methodology

In this section, we first describe a linear model for discovering dynamics for
a single environment. We then extend this to the multi-environment and
introduce our model to uncover the dynamics from multiple environments.
Finally we provide an optimization algorithm to fit the parameters of the
model given observations.

###  3.1 Linear models for dynamic systems

Consider a system of dynamics over the continuous time interval between
$[0,T]$ whose state
$\boldsymbol{x}=[x_{1},\cdots,x_{n}]^{\top}\in\mathbb{R}^{n}$ is defined in an
$n$-dimensional space. We consider the dynamics of the form

$\frac{d\boldsymbol{x}(t)}{dt}=f(\boldsymbol{x}(t)),$ (1)

where $\boldsymbol{x}(t)$ denotes the state of dynamics at time $t$. A common
interest in many science and engineering discipline is to uncover the precise
form of function $f$ that governs the observed dynamics. In many cases, the
function $f$ consists of a few terms, even if the observed dynamics seem
rather complex and chaotic (Broer & Takens, 2010). For example, a well-known
chaotic dynamic system, the Lorenz system, looks unpredictable as shown in
Figure 1 but can be formalized via the second-order terms on the three-
dimensional observations $x,y,z$ as

$\bigg{(}\frac{dx}{dt},\frac{dy}{dt},\frac{dz}{dt}\bigg{)}=\big{[}\sigma(y-x),x(\rho-z)-y,xy-\beta
z\big{]}^{\top},$ (2)

where the constants $\sigma,\rho,\beta$ instantiate the real observation.

Given a trajectory, which is a sequence of observation $\boldsymbol{x}$ over
$m$ time steps $t_{1},t_{2},...,t_{m}$, one can formulate the problem of
dynamics through the linear model defined as

$\frac{dx_{k}(t)}{dt}=\sum_{i=1}^{p}\xi_{ki}\phi_{i}(\boldsymbol{x}(t)),$ (3)

where $\phi_{i}:\mathbb{R}^{n}\rightarrow\mathbb{R}$ is a feature function,
c.f., the first and second-order terms in Equation 2, and $\xi_{ki}$ is the
coefficient of feature $\phi_{i}$, c.f, $\gamma,\rho,\beta$ in Equation 2.

Although the approach sounds plausible, it is useless if we do not know which
feature functions constitute the model. One common approach is to curate an
extensive list of feature functions and fit the model with the sparsity
constraints, such as $\ell_{1}$ regularizer, on $\xi$. In previous work, the
sparse regression approach can find the correct dynamics when the list of
feature functions contains the true candidate of the underlying dynamics
(Brunton et al., 2016; Kaheman et al., 2020; Fasel et al., 2022). When the
number of observations is insufficient, however, the model is unlikely to find
the true dynamics from the limited observations.

###  3.2 SpReME

We tackle the case where multiple trajectories are observed from various
environments, where each environment is defined through its own dynamics of
the form in Equation 1. Specifically, consider multiple environments with
different dynamics

$\frac{d\boldsymbol{x}^{(e)}(t)}{dt}=f^{(e)}(\boldsymbol{x}^{(e)}(t)),$ (4)

indexed by environment $e$. Our goal is to find the correct functional form of
the true dynamics $f^{(e)}$ with the observations from multiple environments
sharing the same functional form, i.e., $f^{(e)}$ consists of the same feature
functions $\phi_{i}$ with varying constants $\xi_{i}$ across different
environments.

One can apply the model in Equation 3 for each environment to find the true
dynamics independently. However, the resulting dynamics may not reveal the
same sparse structure across the environments, which is problematic since we
know the presence of the common structure. To make the different environments
share the same feature functions, we model the dynamics as

$\frac{d{x}^{(e)}_{k}(t)}{dt}=\sum_{i=1}^{p}(M_{ki}\xi_{ki}^{(e)})\phi_{i}(\boldsymbol{x}^{(e)}(t)),$
(5)

where $M_{ki}\in\\{0,1\\}$ is a binary mask. Through the shared binary masks,
the model enforces the same sparsity structure across different environments.

The above model can be described compactly over the $n$-dimensional state
space in matrix form

$\frac{d\boldsymbol{x}^{(e)}(t)}{dt}=(\boldsymbol{M}\circ\boldsymbol{\Xi}^{(e)})\Phi(\boldsymbol{x}^{(e)}(t)),$
(6)

where
$\Phi(\boldsymbol{x}^{(e)}(t))=[\phi_{1}(\boldsymbol{x}^{(e)}(t)),\cdots,\phi_{p}(\boldsymbol{x}^{(e)}(t))]\in\mathbb{R}^{p}$,
$\boldsymbol{\Xi}^{(e)}\in\mathbb{R}^{n\times p}$ is a collection of
$\xi^{(e)}_{ki}$, $\boldsymbol{M}\in\\{0,1\\}^{n\times p}$ is a collection of
$M_{ki}$, and $\circ$ indicates element-wise matrix multiplication.

To formulate the model in Equation 6 as a regression problem, one needs to
define a proper loss between the derivative of state
${d\boldsymbol{x}^{(e)}}/{dt}$ and predicted derivatives. However, the true
derivatives are unavailable in general. To overcome this limitation, we
utilize an ordinary differential equation (ODE) solver to use the observed
state $\boldsymbol{x}^{(e)}$ as direct supervision. The ODE solver takes
initial state $\boldsymbol{x}^{(e)}(t_{i})$, estimated functional form of
derivatives, initial time $t_{i}$, and target time $t_{j}(>t_{i})$ as inputs
to approximate state at time $t_{j}$. More precisely, the predicted state at
time $T$ is computed as ^x^(e)(t_j; t_i) &= x^(e)(t_i) + ∫_t_i^t_j dx(e)(t)dt
dt
= ODESolve(x^(e)(t_i), dx(e)(t)dt, t_i, t_j). Note that we include $t_{i}$ in
predicted state $\hat{\boldsymbol{x}}^{(e)}(t_{j};t_{i})$ to emphasize the
time step at the initial state.

With the help of the ODE solver, our objective can be formalized as the
following optimization problem: _M, ΞL(M, Ξ) = ∑_e ∑_t ℓ(x^(e)(t), ^x^(e)(t)),
subject to a certain sparsity constraint on the mask matrix $\boldsymbol{M}$
with a loss function $\ell$.

###  3.3 Optimization algorithm

Algorithm 1 Optimization algoritnm

Input: dynamics state
$\\{[\boldsymbol{x}^{(e)}(t_{1}),\cdots,\boldsymbol{x}^{(e)}(t_{m})]\\}_{e=1}^{E}$
Output: mask $\boldsymbol{M}$, coefficient
$\\{\boldsymbol{\Xi}^{(e)}\\}_{e=1}^{E},\forall e$

1: Initialize $\boldsymbol{\Xi}^{(e)}$ via sparse regression in each
environment

2: Initialize $\tilde{\boldsymbol{M}}$ $\triangleright$ Equation 7

3: for $\tau=1,\cdots\mathcal{T}$ do

4: Update $\boldsymbol{M}$ from $\tilde{\boldsymbol{M}}$ $\triangleright$
Section 3.3

5:
$\boldsymbol{\Xi}^{(e)}\leftarrow\boldsymbol{\Xi}^{(e)}-\alpha_{\Xi}\nabla\mathcal{L}_{\Xi},\forall
e$ $\triangleright$ Section 3.3

6: Prune $\boldsymbol{\Xi}^{(e)}$ smaller than threshold $\eta_{\Xi}$

7:
$\tilde{\boldsymbol{M}}\leftarrow\tilde{\boldsymbol{M}}-\alpha_{M}\nabla\mathcal{L}_{\tilde{\boldsymbol{M}}}$
$\triangleright$ Section 3.3

8: Prune $\tilde{\boldsymbol{M}}$ smaller than threshold $\eta_{M}$

9: end for

10: Update $\boldsymbol{M}$ from $\tilde{\boldsymbol{M}}$ $\triangleright$
Section 3.3

We provide an optimization algorithm alternating optimization for
$\boldsymbol{M}$ and $\boldsymbol{\Xi}$ to solve Section 3.2.

#### Coefficients $\Xi$ update

Given mask $\boldsymbol{M}$, we formulate the regression problem as a
minimization of the following loss for coefficient matrices
$\\{\boldsymbol{\Xi}^{(e)}\\}_{e=1}^{E}$:

L_Ξ(Ξ) = ∑_e ∑_i ℓ_δ( x^(e)(t_i+η),^x^(e)(t_i+η; t_i^’)) where
$t_{i}^{{}^{\prime}}$ expresses $t_{\eta\lceil\frac{i}{\eta}\rceil}$, and
$\ell_{\delta}$ is the Huber loss (Huber, 1964) with $\delta=1$.
$\hat{\boldsymbol{x}}^{(e)}(t_{i+\eta};t_{i}^{{}^{\prime}})$ is the predicted
state at time $t_{i+\eta}$ obtained from an ODE solver using
$(\boldsymbol{M}\circ\boldsymbol{\Xi}^{(e)})\Phi(\boldsymbol{x}^{(e)}(t))$ as
a differential equation and $\boldsymbol{x}^{(e)}(t_{i}^{{}^{\prime}})$ as an
initial state, i.e.,
$\text{ODESolve}(\boldsymbol{x}^{(e)}(t_{i}),(\boldsymbol{M}\circ\boldsymbol{\Xi}^{(e)})\Phi(\boldsymbol{x}^{(e)}(t)),t_{i}^{{}^{\prime}},t_{i+\eta})$.
We compute the gradient by backpropagating the ODE solver directly and use the
gradient descent to update $\boldsymbol{\Xi}^{(e)}$ with step size
$\alpha_{\Xi}$.

To compute the loss, we use $\eta$-step prediction for each state, i.e.,
$\hat{\boldsymbol{x}}^{(e)}(t_{i+\eta};t_{i}^{{}^{\prime}})$ is only predicted
from $\boldsymbol{x}^{(e)}(t_{i}^{{}^{\prime}})$. If $\eta$ is one, we perform
a one-step prediction for the state where we predict the immediate next time
step given a current time step via ODE solver. If $\eta$ is large, we predict
the behavior of the state in a long horizon.

#### Mask $\boldsymbol{M}$ update

Given the coefficient of features $\\{\boldsymbol{\Xi}^{(e)}\\}_{e=1}^{E}$,
optimizing $\boldsymbol{M}$ requires enumerating all possible configurations
of binary masks, which is intractable in general. We relax the binary mask to
continuous matrix $\tilde{\boldsymbol{M}}\in\mathbb{R}^{p\times n}$ to
approximate the discrete mask. The continuous relaxation allows us to use a
continuous optimization method. With the relaxation, we minimize the following
objective to update $\tilde{\boldsymbol{M}}$: L_M(~M) =& ∑_e ∑_i
(x^(e)(t_i+η)-^x^(e)(t_i+η; t_i^’))^2
\+ λ(1+s)^τ∑_eσ(~M)_1, where $\sigma$ is a sigmoid function making the relaxed
mask ranges between zero and one, and $\lambda$, $s$, and $\tau$ control the
importance of the regularization term. To obtain the predicted state at
$i+\eta$-th step
${\hat{\boldsymbol{x}}^{(e)}(t_{i+\eta};t_{i}^{{}^{\prime}})}$, we use the
integrated state from the initial state $t_{i}^{{}^{\prime}}$ to $t_{i+\eta}$
using an ODE solver, i.e.,
$\text{ODESolve}(\boldsymbol{x}^{(e)}(t_{i}),(\sigma(\tilde{\boldsymbol{M}})\circ\boldsymbol{\Xi}^{(e)})\Phi(\boldsymbol{x}^{(e)}(t)),t_{i}^{{}^{\prime}},t_{i+\eta})$
as done in the coefficient $\boldsymbol{\Xi}$ update. We introduce a scheduler
$(1+s)^{\tau}$ for the regularizer weight, where $\tau$ indicates an iteration
step. When $s>0$, the weight of the regularizer increases as the training step
increases. We find that increasing weight helps to explore the sparse
structure more in practice. The relaxed mask is updated with the gradient
descent, where the gradient is computed through the ODE solver via back-
propagation.

Once $\tilde{\boldsymbol{M}}$ is optimized, we prune the mask based on the
predefined threshold $\kappa_{M}$ by setting the corresponding entries to
negative infinity to binarize the continuous mask. We also prune the entries
if the corresponding coefficients are zero across all environments since these
masks are not updated during optimization. After pruning, binary mask
$\boldsymbol{M}$ can be obtained by thresholding: M_ij = {0, &if σ(~M_ij) = 0
,
1, otherwise, which is again used to update the coefficient. In practice, we
find that further pruning the coefficients based on the predefined threshold
$\kappa_{\Xi}$ improves the sparsity after the coefficient optimization step.

#### Initialization

The use of an ODE solver makes the entire optimization difficult since
randomly initialized coefficients may result in a significant loss leading to
unstable optimization. Therefore, it is necessary to begin with a meaningful
initialization in order to alleviate this challenge.

To initialize the coefficient and mask matrix, we use the sparse regression
method of Brunton et al. (2016), which can be seen as a special case of our
model for a single environment without a mask matrix. We can directly
initialize the coefficients $\boldsymbol{\Xi}^{(e)}$ from the results of
sparse regression for each environment with the lowest validation loss. To
initialize the relaxed mask matrix $\tilde{\boldsymbol{M}}$, we use statistics
of the coefficients as follows:

$\tilde{M}_{ij}=\sigma^{-1}\left(\alpha\frac{\sum_{e}\mathds{1}_{\Xi_{ij}^{(e)}\neq
0}}{\\#\text{ofenvironments}}\right),$ (7)

where $\alpha$, fixed to $0.7$, is introduced to alleviate the gradient
vanishing probelm.

The overall optimization algorithm is described in Algorithm 1. Note that the
optimization steps for the masking matrix are similar to the straight through
estimator (Bengio et al., 2013). Our algorithm alternates the optimization
over the continuous mask and the optimization over the coefficients through
the quantization.

#### Validation

At train time, the ground truth of the mask is not available. Given the
assumption that the learned model extrapolates well on future dynamics, we
measure the extrapolation performance as a validation criterion. We split the
given data
$\\{[\boldsymbol{x}^{(e)}(t_{1}),\cdots,\boldsymbol{x}^{(e)}(t_{m})]\\}_{e=1}^{E}$
into data for timestamp $1$ to $v$ and $v+1$ to $m$.We use the data point from
$1$ to $v$ for training and $v+1$ to $m$ for validation. The mean-squared
error of prediction integrated from time $v+1$ to $m$ without regularization
term, i.e., Section 3.3 with $\lambda=0$, is used as a validation loss.

#### Adaptation

When the trained model encounters a new environment $e^{*}$, the model needs
to adapt to the new environment with partial observations
$[\boldsymbol{x}^{(e^{*})}(t_{1}),\cdots,\boldsymbol{x}^{(e^{*})}(t_{m})]$
available. To adapt to the new environment, we only optimize coefficient
$\boldsymbol{\Xi}^{(e^{*})}$ since the structure of dynamics is captured in
optimized mask $\boldsymbol{M}$. We optimize randomly initialized
$\boldsymbol{\Xi}^{(e^{*})}$ using gradient descent with the loss in Section
3.3 until convergence.

##  4 Experimental Setting

We conduct experiments on four systems of ODE: one linear ordinary
differential equation (linear model) and three complex non-linear differential
equations (Lotka-Volterra, damped pendulum, and Lorenz models).

| train | adaptation | test
---|---|---|---
horizon | dt | # env. | # traj. | horizon | dt | # traj. | horizon | dt | # traj.
Linear | 4 | 0.05 | 9 | 8 | 4 | 0.05 | 1 | 10 | 0.025 | 16
Lorenz | 4 | 0.05 | 9 | 12 | 4 | 0.05 | 1 | 10 | 0.025 | 16
Lotka Volterra | 10 | 0.30 | 9 | 4 | 10 | 0.30 | 1 | 25 | 0.150 | 32
Damped Pendulum | 4 | 0.20 | 9 | 8 | 4 | 0.20 | 1 | 10 | 0.100 | 32
Table 1: Description of parameters used to generate trajectories for each ODE
system. dt , # traj. # env. indicate time intervals, the number of
trajectories, and the number of environments, respectively. The training set
contains # of trajectories for each environment. We use shorter time intervals
for the test set than those for training and adaptation to measure the
interpolation performance.

| Parameter | Mean | Variance
---|---|---|---
Linear | $\\{\alpha,\beta,\gamma,\delta,\omega\\}$ | $\\{-0.1,2,-2,-0.1,-0.3\\}$ | 0.01
Lorenz | $\\{\sigma,\rho,\beta\\}$ | $\\{10,8/3,28\\}$ | 0.02
DP | $\\{\alpha,\omega_{0}\\}$ | $\\{0.5,0.98\\}$ | 0.01

Table 2: Means and variances of ODE parameters for each system. We use a
normal distribution to sample ODE parameters. DP refers to the damped
pendulum. For Lotka-Volterra, we use the same setting used in Kirchmeyer et
al. (2022).

###  4.1 Dynamics

We describe the details of the four systems of ODE used to measure the
performance of dynamic discovery in experiments.

Linear A linear model is defined via the following ODE:

$\frac{d\boldsymbol{x}}{dt}=\big{[}\alpha x_{0}+\beta x_{1},\gamma
x_{0}+\delta x_{1},\omega x_{2}\big{]}^{\top}$

where $\alpha,\beta,\gamma,\delta,\omega$ control the dynamics of the system.

Lorenz A non-linear model representing atmospheric convection with the
following dynamics:

$\bigg{(}\frac{dx}{dt},\frac{dy}{dt},\frac{dz}{dt}\bigg{)}=\big{[}\sigma(y-x),x(\rho-z)-y,xy-\beta
z\big{]}^{\top}$

where $\sigma,\rho,\beta>0$ define the dynamics. This is a chaotic system
where small changes in the initial state produce a drastic difference in
relatively short time (Sparrow, 1982).

Lotka-Volterra A system for the dynamics of biological interaction between
predator and prey can be modeled via (Bacaër, 2011):

$\bigg{(}\frac{du}{dt},\frac{dv}{dt}\bigg{)}=\big{[}\alpha u-\beta uv,\delta
uv-\gamma v]^{\top}$

where $u,v$ refer to the number of prey and predator, respectively, and
$\alpha,\beta,\gamma,\delta$ control the dynamics of the population.

Damped Pendulum A model for the dynamics of a damped pendulum (Quiroga &
Ospina-Henao, 2017) is driven by ODE:

$\frac{d\theta^{2}}{dt^{2}}+\alpha\frac{d\theta}{dt}+\omega_{0}^{2}\sin(\theta)=0,$

where damping factor $\alpha$ and initial angular velocity $\omega_{0}$
defines damping pattern of angle $\theta$. The state $\boldsymbol{x}$ is
defined as $[\theta,d\theta/dt]^{\top}$.

###  4.2 Dataset Generation

Multi-Environment Generation To mimic multiple environments, we use various
configurations of ODE parameters for each dynamics. Each set of parameters is
sampled from the normal distribution except with Lotka-Volterra, where
parameter sets defined in prior works (Kirchmeyer et al., 2022) are used. We
describe the detailed statistics used to generate the environment in Table 2.

Trajectory Generation To generate trajectories given environment and initial
condition, we use LSODA ODE solver (Hindmarsh & Petzold, Sep 2005) to compute
the state transition over a finite time horizon with a fixed time interval.
For a given environment, we sample various initial states. For linear and
Lorenz models, the initial state for each trajectory is sampled from the
standard normal distribution, i.e., $\mathcal{N}(0,1)$. For Lotka-Volterra,
each dimension of the initial state is sampled from the unit-mean and unit-
variance normal distribution, i.e., $\mathcal{N}(1,1)$. For the damped
pendulum, the same method used in Yin et al. (2021b) is used for sampling the
initial states.

Table 1 describes the detailed statistics for the train, adaptation, and test
data sets, including the value of time horizon, time interval, and the number
of trajectories for each environment. The adaptation and test data share the
same environment, i.e., the parameters of ODEs are the same, but the
trajectories are generated independently with different initial states. For
validation of effective adaptation with few data, we use a single trajectory
for the adaptation dataset.

###  4.3 Prediction Tasks

We measure the prediction performance from two different perspectives: in-
domain and out-of-domain predictions. In-domain prediction measures the
performance of the trained model in the same environment used for training. In
this setting, the additional adaptation to the test set is not performed.
Instead, the environment used for training is used to measure the prediction
performance. Out-of-domain prediction measures the performance of the trained
model in a different environment from the training. In this setting, we use
the adaptation set to optimize the feature coefficient $\boldsymbol{\Xi}$ and
then use the test set to measure the prediction performance.

For both in-domain and out-of-domain predictions, we predict the state of the
trajectories within a training/adaptation time horizon (interpolation) and
beyond the training/adaptation time horizon (extrapolation). For the within-
time horizon state prediction, we set the interval of test data two times
shorter than the training set. By doing so, we can predict the middle points
between training states. For the beyond-time horizon prediction, we use two
and a half times longer horizon than the training set.

| In-domain | Out-of-domain
---|---|---
| Linear | Lorenz | LV | DP | Linear | Lorenz | LV | DP
LEADS | 0.4726 | 4.9767 | 0.0756 | 0.0016 | 0.0666 | 41.1628 | 0.0412 | 0.0002
CoDA | 0.0546 | 26.2325 | 0.2064 | 0.0413 | 0.0012 | 60.2433 | 0.1486 | 0.0012
SINDy | 0.0360 | 70.6965 | - | 2.76e-05 | - | - | - | -
Intersection | N/A | N/A | N/A | N/A | 0.0269 | 287.0103 | 0.2559 | 2.7516e-6
Union | N/A | N/A | N/A | N/A | - | 70.1099 | - | -
SpReME | 3.89e-11 | 12.4632 | 2.90e-06 | 2.85e-10 | 9.37e-13 | 0.1831 | 7.35e-7 | 3.76e-11
Table 3: Test MSE in adaptation and in-domain state prediction. LV and DP refer to Lotka Volterra and Damped Pendulum, respectively. SINDy is trained with an adaptation dataset. Bold denotes the best result, and underline denotes the second best result for each dataset. Numerically unstable results are denoted with a hyphen (-). N/A means we did not report the performances since the results are the same as those of SINDy. | Linear | Lorenz | LV | DP
---|---|---|---|---
Precision | 1 | 0.7 | 1 | 1
Recall | 1 | 1 | 1 | 1
$\\#$ candidates ($p$) | 56 | 56 | 21 | 24
Table 4: The precision and recall of the trained mask obtained from the lowest
validation loss. LV and DP refer to Lotka Volterra and Damped Pendulum,
respectively.

###  4.4 Baseline Models

We compare SpReME with two families of baselines. The first family is based on
SINDy (Brunton et al., 2016). We use the original SINDy and its two variants,
SINDy-Intersection and SINDy-Union, for handling multi-environment data. The
second family is based on neural network approaches, including LEADS (Yin et
al., 2021a) and CoDA (Kirchmeyer et al., 2022). LEADS and CoDA are proposed to
deal with multi-environment data. For all baseline models, we search hyper-
parameters via the validation set. In what follows, we describe the details of
each baseline model.

SINDy. SINDy is proposed to uncover the dynamics given a single trajectory
from a single environment. To evaluate the performance of SINDy in our multi-
environment setting, we train a SINDy in each environment without considering
multiple environments. SINDy can use multiple trajectories from a single
environment since they only differ in their initial states.

SINDy-Intersection. SINDy-Intersection is a variant of SINDy for multi-
environments. We train a SINDy for each training trajectory independently and
then obtain a mask from the learned feature coefficients across all training
environments. The entries of the binary mask are set to one if the
corresponding coefficients across all environments are non-zero: ~M_ij = {1,
&if Ξ^(e)_ij ≠0 , ∀e
0, otherwise, where $\boldsymbol{\Xi}^{(e)}$ is the coefficient matrix of
SINDy for environment $e$. For out-of-domain prediction, the feature
coefficients are trained with our adaptation method.

| Linear | Lorenz | LV | DP
---|---|---|---|---
LEADS | 0.0861 | 44.5342 | 0.0668 | 0.0002
CoDA | 0.0013 | 56.1133 | 0.2262 | 0.0012
SINDy | - | - | - | -
Intersection | 0.0345 | 288.3880 | 0.3328 | 1.5800e-06
Union | - | 67.6532 | - | -
SpReME | 1.3040 e-12 | 0.2917 | 7.6972 e-7 | 4.1123 e-11

Table 5: The results of extrapolation in out-of-domain state prediction. LV
and DP refer to Lotka Volterra and Damped Pendulum, respectively.

SINDy-Union. SINDy-Union is the second variant for the multi-environments,
where we set the entry mask to one if any of the coefficients is not zero from
the learned SINDy. To be specific, the entries of masks are obtained via ~M_ij
= {1, &if Ξ^(e)_ij ≠0 , ∃e
0, otherwise, where $\boldsymbol{\Xi}^{(e)}$ is the coefficient matrix of
SINDy for environment $e$. The same procedure used for SINDy-Intersection is
taken to measure the performance on the test set.

LEADS. LEADS learns dynamics from multi-environments through the combination
of the first neural network capturing common dynamics across all environments
and the second neural network capturing the environment-specific dynamics.
Specifically, LEADS models the dynamics via

$\frac{d\boldsymbol{x}^{(e)}(t)}{dt}=(f+g^{(e)})(\boldsymbol{x}^{(e)}(t)),$
(8)

where $f$ is a shared neural network across all environments, and $g_{e}$ is
an environment-specific neural network. For adaptation in unobserved
environment $e^{*}$, it only trains the new environment-specific neural
network $g^{(e^{*})}$.

CoDA. CoDA is a meta-learning approach for multi-environments. The method
employs a single neural network, denoted as $f$, whose parameters consist of
two parts: a set of shared parameters that are common across all environments
and a small number of environment-specific parameters:

$\frac{d\boldsymbol{x}^{(e)}(t)}{dt}=f_{\theta+W\phi^{(e)}}(\boldsymbol{x}^{(e)}(t)),$
(9)

where $\theta\in\mathbb{R}^{q}$ are shared parameters,
$\phi^{(e)}\in\mathbb{R}^{a}$ are environment-specific parameters, and
$W\in\mathbb{R}^{q\times a}$ is a matrix for parameter combination. It adapts
to the unobserved environment by training environment-specific parameters
while keeping the shared parameter fixed.

Figure 2: Comparison between the adaptation state prediction of SpReME, LEADS
and CoDA for Lorenz system. The solid black lines are ground truth, the dotted
orange lines are interpolation results, and the dashed green lines are
extrapolation results. With the Lorenz model, we provide enlarged views of the
extrapolated region (red boxed), changing dramatically from the training
trajectory.

###  4.5 Implementation of SpReME

We set $v$ as $0.8m$ for validataion data. To update coefficient
$\boldsymbol{\Xi}$, we set $\eta$ as $1$, and set $\eta$ as $m$ for mask
optimization. For a successful application of SpReME, it is important to
select the candidate feature functions carefully. To construct candidate
$\Phi(\boldsymbol{x}^{(e)})$, we use the constant function and polynomial
functions with degree five as candidate functions for all dynamics following
the settings from (Brunton et al., 2016). For the system of the damped
pendulum, we additionally use trigonometric feature function $\sin(x_{1}(t))$,
$\sin(x_{2}(t))$ and $\sin(x_{1}(t)+x_{2}(t))$ as candidates. In all
experiments, we use Adam optimizer (Kingma & Ba, 2014). The Runge-Kutta 4th
order (RK4) solver (Akinsola, 2023) is used to perform ODESolve. Note that RK4
solver is used for all baseline models, and the same candidates $\Phi$ are
used for SINDy-based models.

##  5 Experimental Results

In this section, we provide the experimental results of various prediction
tasks. We report the mean-squared error (MSE) as a metric for prediction
tasks.

###  5.1 Main Results

#### In-domain prediction

The result in Table 3 shows the in-domain prediction performance of prediction
models on the four dynamics. SpReME performs better than the other baselines
except for the Lorenz dataset. The Lorenz dataset is the only case where
SpReME fail to capture the true mask structure. We provide additional details
on the accuracy of the binary mask in Table 4. Note that we did not report the
performances of SINDy-Intersection and SINDy-Union since the results are the
same as those of SINDy.

#### Out-of-domain prediction

The results of out-of-domain prediction tasks are shown in Table 3. SpReME
outperforms the other baseline models over all four dynamics for the out-of-
domain prediction. Unlike the results of in-domain prediction, SpReME
outperform the other models on the Lorenz dataset in this case. Figure 2
visualizes the result. Note that SINDy cannot make a stable prediction in this
setting since, for the out-of-domain task, SINDy can only use the adaptation
set for optimization leading to unstable coefficients of features. The results
only focus on the extrapolation is also provided in Table 5. For all models,
the prediction on extrapolation is generally worse than the overall error
reported in Table 3. Still, SpReME outperforms the baselines in this region.

###  5.2 Analysis

#### Environmental uncertainty

To understand how different models cope with environmental uncertainty in
optimization, we vary the configuration of the environment by changing the
variance of normal distribution, which is used to sample the parameter of ODE.
We keep the mean of the normal distribution the same while changing the
variance from $0.01$ to $0.5$. Figure 3 shows the prediction performance on
the out-of-domain setting with varying variance in the linear and damped
pendulum systems. SpReME shows stable results even with increasing
uncertainty, whereas the other models show increasing error with increasing
uncertainty in general.

#### Prediction horizon in optimization

When optimizing the model parameters, we make a prediction over $\eta$-time
steps via an ODE solver and back-propagate over the multiple time steps. We
measure the performances in terms of the mean-squared error and computation
time with varying $\eta$ and report the results in Figure 4. As the results
show, a longer prediction horizon may result in a worse performance. We
conjecture that the accumulated errors in the ODE solver while integrating
over long horizon results in unreliable performance. The computation time is
measured until the model parameter is converged. As the results suggest, a
proper prediction horizon can achieve prediction performance and efficient
computation together.

---
|

Figure 3: The results of state prediction with increasing uncertainty in
training environments. $x$-axis and $y$-axis represent the variance between
environments and the MSE performance, respectively. The plots show out-of-
domain prediction results on the linear and damped pendulum systems,
respectively.

---
|

Figure 4: The result of adaptation performance of models (blue) and time
(orange) to learn the mask when the predicted time in mask optimization is
varied. $x$-axis represents the predicted time in mask optimization.
Experiments are conducted under linear and damped pendulum systems,
respectively.

##  6 Conclusion

We propose SpReME for discovering dynamics from multiple environments data
with incomplete prior knowledge. Unlike prior work, SpReME use incomplete
prior knowledge while dealing with multiple environments. Through the masking
scheme, we extract commonalities between different environments. Our model
adapts to new environments based on the mask. The experimental results show
that SpReME outperforms LEADS and CoDA in most cases.

#### Acknowledgement

This work was partly supported by Institute of Information & communications
Technology Planning & Evaluation (IITP) grant funded by the Korea government
(MSIT) (No.2019-0-01906, Artificial Intelligence Graduate School Program
(POSTECH) and No.2022-0-00959, (part2) Few-Shot learning of Causal Inference
in Vision and Language for Decision Making) and National Research Foundation
of Korea (NRF) grant funded by the Korea government (MSIT)
(No.2022R1F1A1064569).

## References

  * Akinsola (2023) Akinsola, V.  Numerical methods: Euler and runge-kutta.  In Shah, D. K. and Ali, D. A. (eds.), _Qualitative and Computational Aspects of Dynamical Systems_ , chapter 2. IntechOpen, Rijeka, 2023\.  doi: 10.5772/intechopen.108533.  URL https://doi.org/10.5772/intechopen.108533.
  * Bacaër (2011) Bacaër, N.  _Lotka, Volterra and the predator–prey system (1920–1926)_ , pp. 71–76.  Springer London, London, 2011.  ISBN 978-0-85729-115-8.  doi: 10.1007/978-0-85729-115-8˙13.  URL https://doi.org/10.1007/978-0-85729-115-8_13.
  * Bengio et al. (2013) Bengio, Y., Léonard, N., and Courville, A. C.  Estimating or propagating gradients through stochastic neurons for conditional computation.  _CoRR_ , abs/1308.3432, 2013.  URL http://dblp.uni-trier.de/db/journals/corr/corr1308.html#BengioLC13.
  * Broer & Takens (2010) Broer, H. and Takens, F.  _Dynamical Systems and Chaos_.  Applied Mathematical Sciences. Springer New York, 2010.  ISBN 9781441968708.  URL https://books.google.co.kr/books?id=yaov2qvj5YQC.
  * Brunton et al. (2016) Brunton, S. L., Proctor, J. L., and Kutz, J. N.  Discovering governing equations from data by sparse identification of nonlinear dynamical systems.  _PNAS_ , 113(15):3932–3937, 2016.  doi: 10.1073/pnas.1517384113.  URL https://www.pnas.org/doi/abs/10.1073/pnas.1517384113.
  * Chen et al. (2018) Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K.  Neural ordinary differential equations.  In _NeurIPS_ , 2018.  URL https://proceedings.neurips.cc/paper/2018/file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf.
  * Fasel et al. (2022) Fasel, U., Kutz, J., Brunton, B., and Brunton, S.  Ensemble-sindy: Robust sparse model discovery in the low-data, high-noise limit, with active learning and control.  _Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences_ , 478, 04 2022.  doi: 10.1098/rspa.2021.0904.
  * Garnelo et al. (2018) Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F., Rezende, D. J., Eslami, S. M. A., and Teh, Y. W.  Neural processes.  _CoRR_ , abs/1807.01622, 2018.  URL http://arxiv.org/abs/1807.01622.
  * Hindmarsh & Petzold (Sep 2005) Hindmarsh, A. C. and Petzold, L. R.  Lsoda, ordinary differential equation solver for stiff or non-stiff system, Sep 2005.  URL http://inis.iaea.org/search/search.aspx?orig_q=RN:41086668.
  * Huber (1964) Huber, P. J.  Robust estimation of a location parameter.  _The Annals of Mathematical Statistics_ , 35(1):73–101, 1964.  ISSN 00034851.  URL http://www.jstor.org/stable/2238020.
  * Kaheman et al. (2020) Kaheman, K., Kutz, J. N., and Brunton, S. L.  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), oct 2020.  doi: 10.1098/rspa.2020.0279.  URL https://doi.org/10.1098%2Frspa.2020.0279.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kirchmeyer et al. (2022) Kirchmeyer, M., Yin, Y., Dona, J., Baskiotis, N., Rakotomamonjy, A., and Gallinari, P.  Generalizing to new physical systems via context-informed dynamics model.  In _ICML_ , 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/kirchmeyer22a.html.
  * Li et al. (2021) Li, S. H. Q., Ratliff, L. J., and Açikmese, B.  Disturbance decoupling for gradient-based multi-agent learning with quadratic costs.  _L-CSS_ , 5:223–228, 2021.
  * Long et al. (2018) Long, Z., Lu, Y., Ma, X., and Dong, B.  PDE-net: Learning PDEs from data.  In _ICML_ , 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/long18a.html.
  * Noack et al. (2003) Noack, B. R., Afanasiev, K., MORZYŃSKI, M., Tadmor, G., and Thiele, F.  A hierarchy of low-dimensional models for the transient and post-transient cylinder wake.  _JFM_ , 497:335–363, 2003.
  * Norcliffe et al. (2021) Norcliffe, A., Bodnar, C., Day, B., Moss, J., and Liò, P.  Neural ODE processes.  In _ICLR_ , 2021.  URL https://openreview.net/forum?id=27acGyyI1BY.
  * Quiroga & Ospina-Henao (2017) Quiroga, G. D. and Ospina-Henao, P. A.  Dynamics of damped oscillations: physical pendulum.  _Eur. J. Phys._ , 38(6):065005, oct 2017.  doi: 10.1088/1361-6404/aa8961.  URL https://dx.doi.org/10.1088/1361-6404/aa8961.
  * Shi et al. (2015) Shi, X., Chen, Z., Wang, H., Yeung, D.-Y., Wong, W.-k., and Woo, W.-c.  Convolutional lstm network: A machine learning approach for precipitation nowcasting.  In _NeurIPS_ , 2015.
  * Sparrow (1982) Sparrow, C.  _Introduction and Simple Properties_ , pp. 1–12.  Springer New York, New York, NY, 1982.  ISBN 978-1-4612-5767-7.  doi: 10.1007/978-1-4612-5767-7˙1.  URL https://doi.org/10.1007/978-1-4612-5767-7_1.
  * Wang et al. (2018) Wang, Y., Gao, Z., Long, M., Wang, J., and Yu, P. S.  PredRNN++: Towards a resolution of the deep-in-time dilemma in spatiotemporal predictive learning.  In _ICML_ , 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/wang18b.html.
  * Xu et al. (2021) Xu, K., Zhang, M., Li, J., Du, S. S., Kawarabayashi, K.-i., and Jegelka, S.  How neural networks extrapolate: from feedforward to graph neural networks.  _ICLR_ , 2021.  URL https://par.nsf.gov/biblio/10277315.
  * Yin et al. (2021a) Yin, Y., Ayed, I., de Bézenac, E., Baskiotis, N., and Gallinari, P.  LEADS: learning dynamical systems that generalize across environments.  _CoRR_ , abs/2106.04546, 2021a.  URL https://arxiv.org/abs/2106.04546.
  * Yin et al. (2021b) Yin, Y., GUEN, V. L., DONA, J., de Bezenac, E., Ayed, I., THOME, N., and patrick gallinari.  Augmenting physical models with deep networks for complex dynamics forecasting.  In _ICLR_ , 2021b.  URL https://openreview.net/forum?id=kmG8vRXTFv.
  * Ziyin et al. (2020) Ziyin, L., Hartwig, T., and Ueda, M.  Neural networks fail to learn periodic functions and how to fix it.  In _NeurIPS_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/file/1160453108d3e537255e9f7b931f4e90-Paper.pdf.