LEADS: Learning Dynamical Systems that Generalize Across Environments
Authors: Yuan Yin
Ibrahim Ayed
Emmanuel de Bezenac
Nicolas Baskiotis
patrick gallinari
Venue: NeurIPS 2021 Poster (2021)
Keywords: differential equation
data-driven
neural network
generalization

Abstract: When modeling dynamical systems from real-world data samples, the distribution of data often changes according to the environment in which they are captured, and the dynamics of the system itself vary from one environment to another. Generalizing across environments thus challenges the conventional frameworks. The classical settings suggest either considering data as i.i.d and learning a single model to cover all situations or learning environment-specific models. Both are sub-optimal: the former disregards the discrepancies between environments leading to biased solutions, while the latter does not exploit their potential commonalities and is prone to scarcity problems. We propose LEADS, a novel framework that leverages the commonalities and discrepancies among known environments to improve model generalization. This is achieved with a tailored training formulation aiming at capturing common dynamics within a shared model while additional terms capture environment-specific dynamics. We ground our approach in theory, exhibiting a decrease in sample complexity w.r.t classical alternatives.  We show how theory and practice coincides on the simplified case of linear dynamics. Moreover, we instantiate this framework for neural networks and evaluate it experimentally on representative families of nonlinear dynamics. We show that this new setting can exploit knowledge extracted from environment-dependent data and improves generalization for both known and novel environments.

LEADS: Learning Dynamical Systems that Generalize Across Environments
Yuan Yin', Ibrahim 12 Emmanuel de Bézenac', Nicolas Baskiotis', Patrick Gallinari 'Sorbonne Université, Paris, France SIS Lab, Thales, Paris, France *Criteo AI Lab, Paris, France Lyuan. yin, ibrahim. ayedd, emmanuel de-bezenac, nicolas. baskiotis patrick. gall inari wwwrrrbonneuniieeeeeee
Abstract
When modeling dynamical systems from real-world data samples, the distribution data often changes according the environment which they are captured, and the dynamics of the system itself vary from one environment to another. Generalizing across environments thus challenges the conventional frameworks. The classical settings suggest either considering data i.idd. and learning single model to cover all situations or learning optimal: the former disregards the discrepancies between environments leading biased solutions, while the latter does not exploit their potential commonalities and is prone scarcity problems. We propose LEADS, novel framework that leverages the commonalities and discrepancies among known environments to improve model generalization. This is achieved with tailored training formulation aiming capturing common dynamics within a shared model while additional terms capture environment-o speccifc dynamics. We ground our approach theory, exhibiting a decrease in sample complexity w.r.t. classical alternatives. We show how theory and practice coincides on the simplified case of linear dynamics. Moreover, we instantiate this framework for neural networks and evaluate it experimentally on representative families of nonlinear dynamics. We show that this new setting can exploit knowledge extracted from environment-dependen data and improves generalization for both known and novel environments.
wrroomment-speifff models. Both are sub-
Introduction
Data-driven approaches offer interesting alternative and complement to physical- bhasic methods for modeling the dynamics of complex systems and are particularly promising in a wide range of settings: e.g. if the underlying dynamics inaccurate, or fails to adapt to different contexts, or if external perturbation sources and forces are not modeled. The idea of deploying machine learning (ML) to model complex dynamical systems picked momentum a few years ago, relying on recent deep learning progresses and on the development of new methods targeting the evolution of temporal and spatiotemporal systems [6. 7, 21, 30, 37].
partially known understood, if the physical model is incomplete,
is already being applied in different scientific disciplines (see e.g. [36] for recent survey) and could help accelerate scientific discovery address challenging domains such as climate [32] health [12].
However, despite promising results, current developments are limited and usually postulate an idealized setting where data is abundant and the environment does not change. the so-called "i.i.d hypothesis". practice, real-world data may be expensive or difficult to acquire. Moreover, changes in the environment may be caused by many different factors. For example, in climate modeling, there are external forces (e.g. Coriolis) which depend on the spatial location [23]; or, in health science, parameters need to be personalized for each patient for cardiac computational models [27]. More
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

generally. data acquisition and modeling are affected by different factors such as geographical position, sensor variability, measuring cireumstances, etc. The classical paradigm either considers all the data as i.i.d. and looks for a global model, or proposes specific models for each environment. The former disregards discrepancies between the environments, thus leading to a biased solution with an averaged model which will usually perform poorly. The latter ignores the similarities between environments, thus affecting generalization performance, particularly in settings where per-environment data is limited. This is particularly problematic in dynamical settings, small changes in initial conditions lead to trajectories not covered by the training data.
In this work, we consider setting where it is explicitly assumed that the trajectories are collected from different environments. Note that in this setting, the i.i.d. hypothesis is removed twice: by considering the temporality of the data and by the existence of multiple environments. In many useful contexts the dynamies in each environment share similarities, while heing distinct which translates into changes in the data distributions. Our objective is to leverage the similarities between environments in order to improve the modeling capacity and generalization performance, while still carefully dealing with the discrepancies across environments. This brings us to consider two research questions:
ROI Does modeling the differences between environments improve generalization error w.r.t. classi- cal settings: One-For-All, where unique function is trained for all environments; and One- Per-Env., where specific function is fitted for each environment? (cf. Sec. for more details) RQ2 Is it possible to extrapolate to novel environment that has not been seen during training?
We propose LFarning Across Dynamical Systems (LEADS), novel learning methodology decom- posing the learned dynamics into shared and mnviromment-ppecij components. The learning problem is formulated such that the shared component captures the dynamics common across environments and exploits all the available data, while the environment-speritn -peeiiifi-- component only models the remain- ing dynamics, i.e. those that cannot be expressed by the former, based on environment-speccinn data. We show, under mild conditions, that the learning problem is well-posed, as the resulting decompo- sition exists and unique (Sec. 2.2). We then analyze the properties of this decomposition from sample complexity perspective. While, in general, the bounds might be too loose to be practical, a more precise study is conducted in the case of linear dynamics for which theory and practice are closer. We then instantiate this framework for more general hypothesis spaces and dynamics, lead- ing to heuristic for the control of generalization that will be validated experimentally. Overall, we show that this framework provides better generalization properties than One-Per- Env,, requiring less training data to reach the same performance level (RQ/). The shared information is also useful to extrapolate to unknown environments: the new function for this environment can be learned from very little data (RQ2). We experiment with these ideas on three representative cases (Sec. 4) where the dynamics are provided by differential equations: ODEs with the Vootara predator-prey model, and PDEs with the Gray-Scott reaction-diffusion and the more challenging incompressible Navier-Stokes equations. Experimental evidence confirms the intuition and the theoretical findings: with a similar amount of data, the approach drastically outperforms One-For-Ale and One-Per- Per-Em: Emv. settings, especially in low data regimes. Up to our knowledge, it is the first time that generalization in multiple dynamical systems is addressed from an MI. perspective!.
2 Approach
2.1 Problem setting
We consider the problem of learning models of dynamical physical processes with data acquired from a set of environments F. Throughout the paper, we will assume that the dynamics an environment e E are defined through the evolution of differential equations. This will provide in particular clear setup for the experiments and the validation. For a given problem, we consider that the dynamics of the different environments share common factors while each environments has its own specificity, resulting in a distinct model per environment. Both the general form of the differential equations and the specific terms of each environment the equation for environment e, taking its values from a bounded set A, with evolution term fe A - T.A, T.A being the tangent bundle of A. In other words, over a fixed time interval [0, T'. we have: drf di
assumed to be completely unknown. x. denotes the state of
= fe(2f)
(I)
'Code is available https: //github. com/yuan-uu LinSLEADS.
2

---

We assume that, for any fe lies in functional vector space F. In the experiments, we will consider one ODE, which case R and two PDEs, in which case is d' d'-dimensional vector field over bounded spatial domain S C TRd" The term of the data-generating dynamical system in Eq. 1 sampled from distribution for each e, i.e. fe ~ Q. From fe: we define Te, the data distribution of trajectories x. verifying Eq. 1, induced by distribution of initial states 26 Po- The data for this environment then composed of / trajectories sampled from Je. and is denoted as J. with x.!² the i-th trajectory. We will denote the full dataset by J UneR Je.
The classical empirical risk minimization (ERM) framework suggests model the data dynamics either at the global level (One-For-All, taking trajectories indiscriminately from J, or at the spe- eific enviromment level (One-Per-Bee Per-Enn.), training one model for each Te. Our aim is formulate a new learning framework with the objective of explicitly considering the existence of different environments to improve the modeling strategy w.r.t. the classical ERM settings.
2.2 LEADS framework
We decompose the dynamics into two components where I is shared across environments and Je E is specific to the environment e, so that
Ve E f De
(2)
Since we consider functional vector spaces, this additive hypothesis is not restrictive and such de- composition always exists. It is also quite natural as a sum of evolution terms can be seen as the sum of the forces acting on the system. Note that the sum of two evolution terms can lead behaviors very different from those induced by each of those terms. However, learning this decomposition from data defines an ill-posed problem: for any choice of f, there a {ge) such that Eq. 2 verified. A trivial example would be f O leading to solution where each environment is fitted separately.
Our core idea that f should capture as much of the shared dynamics is possible, while ge should focus only on the environment characteristics not captured by f. To formalize this intuition, we introduce (ge,), penalization on ge, which precise definition will depend on the considered setting. We reformulate the learning objective as the following constrained optimization problem:
dx. = (f (x)()
min M.9-
Σ ((ge) subject to V E J, Vt, cEE
(3)
dt
Minimizing Ω aims to reduce ges' complexity while correctly fitting the dynamics of each environ- ment. This argument will be made formal in the next section. Note that f will be trained on the data from all environments contrary to ges. A key question is then to determine under which conditions the minimum in Eq. is well-defined. The following proposition provides an answer (proof cf. Sup. A): Proposition 1 (Existence and Uniqueness), Assume Ω is convex, then the existence of a minimal decomposition I"; {yr JeeR E F of Eq. 3 is guaranteed. Furthermore, if S is strictly convex, this decomposition is unique.
In practice, we consider the following relaxed formulation of Eq. 3: 1 Σ
Σ (HoOo))
e.s dx
(f g.)lag dt
min f.ge E3
de
CEE
where f. are taken from hypothesis space I approximating F. À is regularization weight and the integral term constrains the learned f + ge to follow the observed dynamics. The form this objective and its effective calculation will be detailed in Sec. 4.4.
3 Improving generalization with LEADS
Defining an appropriate Ω2 is crucial for our method. In this section, show that the generalization error should decrease with the number of environments. While the bounds might be too loose for NNs, our analysis is shown adequately model the decreasing trend in the linear case, linking both our intuition and our theoretical analysis with empirical evidence. This then allows us to construct an appropriate Ω for NNs.
3.1 General case
After introducing preliminary notations and definitions, we define the hypothesis spaces associated with our multiple environment framework. Considering a first setting where all environments of interest are present at training time, prove an upper-bound of their effective size based on the
3

---

covering numbers of the approximation spaces. This allows us to quantitatively control the sample complexity of our model, depending on the number of environments m and other quantities that can be considered and optimized in practice. We then consider an extension for learning on a new and unseen environment. The bounds here are inspired by ideas initially introduced in [4]. They consider multi-task classification in vector spaces, where the task specific classifiers share a common feature extractor. Our extension considers sequences corresponding to dynamical trajectories, and a model with additive components instead of function composition in their case.
Definitions. Sample complexity theory is usually defined for supervised contexts, where for a given input z we want to predict some target y. In our setting, we want to learn trajectories (æî Josest starting from an initial condition Zy. We reformulate this problem and cast it as a standard supervised learning problem: Je being the data distribution of trajectories for environment e, as defined in Sec. 2.1, let us consider a trajectory x. ~ Je, and time T Unif ((0, T]); we define system states x=x; E A as input, and the corresponding values of derivatives y fe(x*) E T.A as the associated target. We will denote Pa the underlying distribution of (1: y). and e the associated dataset of size n.
We are searching for f.9.:At :A T.A in an approximation function space I of the original space F. Let us define 9 C 3 the effective function space from which the ges are sampled. Let ff 7+ + := {ff + g E 9} be the hypothesis space generated by function pairs (f.g), with a fixed f E F. For any h + T'A, the error on some test distribution Pe is given by ero (h) A [1((z)) y//-dP, e (xr,
and the error on the training set by Etg. (h) = (z.)ee [1(z)) v1
LEADS sample complexity. Let (E. and Cy 9) denote the capacity of 9 and 9 at certain scale e > O. Such capacity describes the approximation ability of the space. The capacity of a class of functions is defined based on covering numbers, and the precise definition is provided in Sup. B.2, Table $1. The following result is general and applies for any decomposition of the form f + ge- It states that guarantee given average test error, the minimal number of samples required is a function of both capacities and the number of environments m, and it provides step towards RQ/ (proof see Sup. B.2):
Proposition 2. Given m environments, let E1, 0 > 0,E = E Assume the number of examples n per environment satisfies
64 [
1
4 + log Eg ((((1))) (1) 3) log Ezg (+2 16'
9)
16
max
log 8
(5)
Then with probability at least 1 (over the choice of training sets (Pe)). any learner (f+g...... gm.) will satisfy t/ Σee ery, (f + 9c) S m EeED ērg (f + ge) +
f+
The contribution of F to the sample complexity decreases as m increases, while that of 9 remains the same: this is due the fact that shared functions have access to the data from all environments, which is not the case for 9e. From this finding, one infers the basis of LEADS: when learning from several environments, to control the generalization error through the decomposition fe f f+ ge- f should account for most of the complexity of fe while the complexity of ge should be controlled and minimized. We then establish an explicit link to our learning problem formulation in Eq. 3. Further in this section, we will show for linear ODEs that the optimization of ((g.) in Eq. 4 controls the capacity of the effective set by selecting geS that are as "simple" as possible.
As a corollary, we show that for a fixed total number of samples in J. the sample complexity will decrease as the number of environments increases. To see this, suppose that we have two situations corresponding to data generated respectively from m and m/b environments. The total sample complexity for each case will be respectively bounded by Oloo Cg( (é) J) + m log C.y f/. 9)) and O(b log Eg ( F) + m log C 18 I)). The latter being larger than the former, a situation with more environments presents a clear advantage. Fig. in Sec. 4 confirms this result with empirical evidence.
LEADS sample complexity for novel environments. Suppose that problem Eq. 3 has been solved for set of environments E, can we the learned model for new environment not present in the initial training set (RQ2)? Let e' be such a new environment, Per the trajectory distribution of c', generated from dynamies fer ~ Q. and P., an associated training set of size n'. The following results show that the number of required examples for reaching a given performance is much lower when training Je' with ʃ fixed this new environment than training another f' + ge' from scratch (proof see Sup. B.2).

---

Proposition 3. For all E, with 0 < 5, <
1 if the number of samples n' satisfies
64
4e( -. f 9)
16
n' 2 max
52 log
(6)
5
c2
then with probability at least 1 - 8 (over the choice of novel training set Pel), any learner f + 9e' E
f + 9 will satisfy eTP,,, (f + ge) _ Erg ,( (f+ 9e)) +
In Prop. 3 as the capacity of I no longer appears, the number of required samples now depends only on the capacity of f + 9. This sample complexity is then smaller than learning from scratch fe f + Ie' can be seen by comparing with Prop. 2 m 1.
From the previous propositions, it is clear that the environment-speclfee functions ge need to be explicitly controlled. We now introduce practical way to do that. Let w(r,)) e) be strictly increasing function w.r.t. r such that
log (e, 9) S w(),)), r sup gE9 (gg)
Minimizing Q would reduce r and then the sample complexity of our model by constraining S. Our goal is thus to construct such pair (w, 92). In the following, will first show in Sec. 3.2, how one can construct a penalization term S2 based on the covering number bound for linear approximators and linear ODEs. We show with a simple use case that the generalization error obtained in practice follows the same trend as the theoretical error bound when the number of environments varies. Inspired by this result, we then propose in Sec. 3.3 effective f to penalize the complexity of the neural networks ge-
3.2 Linear case: theoretical bounds correctly predict the trend of test error
Results in Sec. 3.1 provide general guidelines for approach. We now apply them to linear system to see how the empirical results meet the tendency predicted by theoretical bound.
Let us consider a linear ODE dæž = L, (xff) where L, : Fox is a linear transformation associated df
to the square real valued matrix F. E (R..(R). We choose hypothesis space the space of linear dx;
functions F C & (RR Rd) and instantiate a linear LEADS
dt = (Lf + (ZF), LF E F. LG. e S F. As suggested in [3], we have that (proof in Sup. B.3):
Proposition 4. for all linear maps LG, E 9, ||G|| if the input space is bounded s.t. s.t. || and the MSE loss function is bounded by c, then
log (E, S [red(2b)* / log 2dl =: w(r
w(r,e) e) is a strictly increasing function w.r.t. This indicates that we can choose ((LG)=I|Gi) F as our optimization objective in Eq. 3. The sample complexity in Eq. 5 will decrease with the size the largest possible r sup (2(La)). The optimization process will reduce ??(G.r) until minimum is reached. The maximum size of the effective hypothesis space is then bounded and decreases throughout training thanks the penalty. Then in linear case Prop. 2 becomes (proof cf. Sup. B.3):
Proposition 5. for linear maps LF E I, ||Fl],rr' LG E S. |[G|İ || and if the MSE loss function is bounded by c, given m environments and n samples per environment, with the probability 1 - 6, the generalization error upper bound is E = max (v/+ vp +4//2, v /1 where p mn
64 log 3 and q 61
[( m (1= )ed(3)))))) log 2dl for any 0 < <
In Fig. 1, we take an instance of linear ODE defined by F. QA.QT with the diag- onal A. specific to each environment Af- ter solving 3 we have at the optimum that G.
x1o-2
we optin. Enorw 22 aptim.
Bound w/o @2 sptim
Bound is optim.
3.0
Generalizatin Enror Bound
F' = F. t Eeeer Fe. Then we can take as the norm bound of 9 when Mg.e) is op- timized. Fig. 1 shows the left the test error with and without penalty and the cor- responding thcoretical bound on the right. We observe that, after applying the penalty Ω the test error is reduced well the theoretical generalization bound, as indi- cated by the arrows from the dashed line the concrete one. See Sup. B.3 for more de- tails on this experiment.
Test MSE
2.5
mmx Ml....)
tmy
2.0
trapenn
15
Number of Training Environments
Figure 1: Test error compared with corresponding the- oretical bound. The arrows indicate the changes after applying M(ge) penalty.
5

---

3.3 Nonlinear case: instantiation for neural nets
The above linear case validates the ideas introduced in Prop. 2 and provides an instantiation guide and an intuition on the more complex nonlinear case. This motivates us to instantiate the general case by choosing an appropriate approximating space F and penalization function 52 from the generalization bounds for the corresponding space. Sup. B.4 of the Appendix contains additional details justifying those choices. For 3, we select the space of feed-forward neural networks with a fixed architecture. We choose the following penalty function:
(ge))
talleelii
where ||9||>== ess suplgl and I1-1 is the Lipschitz semi-norm, a is hyperparameter. This is inspired by the existing capacity bound for NNs [14] (see Sup. B.4 for details). Note that constructing tight generalization bounds for neural networks still an open research problem [26]; however, may still yield valuable intuitions and guide algorithm design. This heuristic is tested successfully on three different datasets with different architectures in the experiments (Sec.
4 Experiments
Our experiments are conducted on three families of dynamical systems described by three broad classes of differential equations. All exhibit complex and nonlinear dynamics. The first one is an ODE-driven system used for biological system modeling. The second one is a PDE-driven reaction-diffusio model, well-known in chemistry for its variety of spatiotemporal patterns. The third one is the more physically complex Navier-Stokes equation, expressing the physical laws of incompressible Newtonian fluids. To show the general validity of our framework, we will use 3 different NN architectures (MLP, ConvNet, and Fourier Neural Operator [19]). Each architecture is well-adapted to the corresponding dynamics. This also shows that the framework is valid for variety approximating functions.
4.1 Dynamics, environments, and datasets
Lotka Voltera (LV). This classical model [22] is used for describing the dynamics of interaction between predator and a prey. The dynamics follow the ODE:
du/ac = aru - Suv, de/a: = dur - -
with u, v the number of prey and predator, B, y, 5 > 0 defining how the two species interact. The system state is x = (uf, vf R.? The initial conditions un up are sampled from a uniform distribution Po. We characterize the dynamics by 0 = (*/B, 7/8) E defined by parameters 0, sampled from a uniform distribution over a parameter set O. We then sample two sets of environment parameters: one used as training environments for RQI, the other treated as novel environments. for RQ2.
θ. An environment e is then
Gray-Scott (GS). This reaction-diffusion model is famous for its complex spatiotemporal behavior given its simple equation formulation [29]. The governing PDE is:
Bu/ot Du u + F(1 u), 0w/ae = D, + uo² - (F + k)v
where the u, represent the concentrations of two chemical components in the spatial domain $ with periodic boundary conditions, the spatially discretized state at time t is x = (u, uf) )
E R2x32" Du, D, denote the diffusion coefficients respectively for u, u, and are held constant, and F, l are the reaction parameters determining the ppatio-tempora patterns of the dynamics [29]. As for the initial conditions (wo, vo) ~ Po. we consider uniform concentrations, with 3 2-by-2 squares fixed at other concentration values and positioned at uniformly sampled positions in S to trigger the reactions. An environment is defined by its parameters 0, (Fe, k.) O. We consider a set of O. parameters uniformly sampled from the environment distribution Q on e.
Navier-Stokes (NS).
We consider the Navier-Stokes PDE for incompressible flows: Ow/et
Vw + +Aw +
v
where e is the velocity field, w V x the vorticity, both v, w lie in spatial domain with periodic boundary conditions, is the viscosity and & is the constant forcing term in the domain S. The discretized state time is the vorticity x = w E R32 Note that is already contained in w. We fix 10 across the environments. We sample the initial conditions wi ~ Po as in [19]. An environment e is defined by its forcing term & E Og. We uniformly sampled a set of forcing terms from Q θ

---

One-Per-Eny
FT-NODE
LEADS
Ground bruth
One-Per-Enx
FT-NODE LEADS
GS
NS
NS
GS
NS
NS
Figure 2: Left: final states for GS and NS predicted by the two best baselines (One-Per-Ent and FT-NODE) and LEADS compared with ground truth. Different environment are arranged by row (3 in total). Right: the corresponding MAE error maps, the scale of the error map is [O, 0.6] for GS, and [0, 0.2] for NS; darker is smaller. (See Sup. D for full sequences)
One-Por-Env
FT-NODE
LEADS
Ground truth
125
1.75
100
150
1,0
0.75
125
0.8
0.8
1.0
12
0.5
Figure 3: Test predicted trajectories in phase space with two baselines (One-Per-Ent and FT-NODE) and LEADS compared with ground truth for LV for envs., one per figure from left to right. Quantity of the prey and the predator respectively on the horizontal and the vertical axis. Initial state is the rightmost end-point of the figures and it is common to all the trajectories.
Datasets. For training, we create two datasets for LV by simulating trajectories of K 20 successive points with temporal resolution A = 0.5. We use the first one as a set of training dynamics to validate the LEADS framework. We choose 10 environments and simulate 8 trajectories (thus corresponding to = 8 K data points) per environment for training. We can then easily control the number of data points and environments in experiments by taking different subsets. The second one is used to validate the improvement with LEADS while training novel environments. We simulate I trajectory (n=1-K data points) for training. We create two datasets for further validation of LEADS with GS and NS. For GS, we simulate trajectories of K = 10 steps with At 40=40. We choose parameters and simulate 1 trajectory (n=1: 1-K K data points) for training. For NS, we simulate trajectories of K- steps with At A 1. We choose foreing terms and simulate 8 trajectories (n = 8 K states) for training. For test-time evaluation, we create for each equation in each environment a test set of 32 trajectories (32- K) data points. Note that every environment dataset has the same number of trajectories and the initial conditions are fixed to equal values across the environments to ensure that the data variations only come from the dynamics themselves, i.c. for the i-th trajectory in P Ve, a zz.- LV and GS data are simulated with the DOPRIS solver in NumPy [10, 13]. NS data is simulated with the pseudo-spectral method as in [19].
4.2 Experimental settings and baselines
We validate LEADDS in two settings: in the first one all the environments in E are available once and then f and all the ges are all trained on E. In the second one, training has been performed on E' as before, and we consider novel environment e' E: the shared term f being kept fixed, the approximating function for + ge is trained on the data from e' (i.c. only ge is modified).
All environments available at once. We introduce five baselines used for comparing with LEADS: (a) One-For-All: learning on the entire dataset P over all environments with the sum of a pair of NNs f + g. with the standard ERM principle, as in [2]. Although this is equivalent to use only one function f. we use this formulation to indicate that the number of parameters is the same for this experiment and for the LEADS ones. (b) One-Per-Env: learning a specific function for each dataset Pe. For the same reason as above, we keep the sum formulation (f + g).. (c) Factored Tensor RNN or FT-RNN [33]: it modifies the recurrent neural network to integrate a one-hot environment code into each linear transformation of the network. Instead of being encoded a separate function 9e like in LEADS, the environment appears here as an extra one-hot input for the RNN linear transformations. predictors such as GRU [8] or NN [35]. (d) FT-NODE: baseline for which the same environment encoding FT-RNN incorporated in Neural ODE [7]. (e) Gradient-based Meta Learning or GBML--like method:
This can be implemented for representative SOTA (ppatio-)tempon
propose GBML-like baseline which can directly compare to our framework. It follows the
7

---

Table 1: Results for LV. GS, and NS datasets, trained on m envs, with n data points per env.
LV (m 10,n = 1.
GS (m 3,
1)
NS (m 4, n 8- K)
Method
MSE train
MSE test
MSE train
MSE test
MSE train
MSE test
One-For-All
4.57c-1
5.08+0.56 c-1
1.55e-2
1,43=0.15 c-2
5.17c-2
7.3145.29 c-2
One-Per- Eny.
2.15e-5
e.9516.96
8.48e-5
6,43 e-3
5.60e-6
1,1010.72 e-2
FI-RNN [33]
5.29e-5
6.4015.69 c-3
8,44c-6
8.19=3.09 e-3
7.40c-4
5.924:4000 c-2
FT-NODE
7.74e-5
3,40122.4
3.51e:5
3.8613.36
1.80e-4
2.96.1..9 e-2
GBML-liik
3.84c-6
5.87+5.65 e-3
1,07e-4
1.39c-4
7.37+-4.80 c-3
L.ADS no min.
3.28e-6
3.07+2.58 e-3
7.65e-5
5.53+3.43 e-3
3.20e-4
7.10+-4.24 e-3
LEADS (Ours)
5.74c-6
1.16:20,99 e-3
5.756-5
2.08 e-3
1.03c-4
5,95+3.65 e-3
principle of MAMI. [11], by training One-For-All at first which provides an initialization near to the given environments like GBML. does, then fitting it individually for each training environment. (f) LEADS no min.: ablation baseline, our proposal without the (gg=)) penalization. A comparison with the different baselines is proposed in Table 1 for the three dynamics. For concision, we provide a selection of results corresponding to training trajectory per environment for LV and GS and 8 for NS. This is the minimal training set size for each dataset. Further experimental results when varying the number of environments from to 8 are provided in Fig. and Table S3 for LV.
Learning on novel environments. We consider the following training schemes with pre-trained, fixed f: (a) Pre-trained- f-Only: only the pre-trained f is used for prediction; sanity check to ensure that f cannot predict in any novel environment without further adaptation.
(b) One-Per-Env. training from scratch on (Per) One-Per-Ent Per-Env. in the previous section. (c) Pre-trained-f-PPee we train on each dataset Pe based on pre-trained f. i.e. [ + Ge', leaving only g adjustable. We compare the test error evolution during training for 3 schemes above for comparison of convergence speed and performance. Results are given in Fig. 5.
4.3 Experimental results
All environments available at once. We show the results in Table 1. For LV systems, we confirm first that the entire dataset cannot be learned properly with a single model -ooe-For-All) when the number of environments increases. Comparing with other baselines, our method LEADS reduces the test MSE over 85% w.r.t. One-Per-Env: and over 60% w.r.t. LEADS no min., we also cut 50%-75% of error w.r.t. other baselines. Fig. shows samples of predicted trajectories in test, LEADS follows very closely the ground truth trajectory, while Mne-Per-Env under-perform in most environments. We observe the same tendency for the GS and NS systems. The error is reduced by: around 2/3 (GS) and 45% (NS) w.r.t. One-Per-Eni; over 60% (GS) and 15% (NS) w.r.t. LEADS no min.; 45-75% (GS) and 15-90% (NS) w.r.t. other baselines. In Fig. 2, the final states obtained with LEADS are qualitatively closer to the ground truth. Looking at the error maps on the right, we see that the errors are systematically reduced across all environments compared to the baselines. This shows that LEADS accumulates less errors through the integration, which suggests that LEADS alleviates overfitting.
We have also conducted larger scale ex- periment on LV (Fig. 4) to analyze the be- havior of the different training approaches as the number of environments increases. We consider three models One-For-All, One-Per-Env. and LEADS, 1, 2, 4 and 8 environments, and for each such case, we have 4 groups of curves, corresponding to 1, 2, 4 and 8 training trajectories per envi- ronment. We summarize the main observa- tions. With One-For-All (blue), the error increases as the number of environments increases: the dynamics for each environ- ment being indeed different, this introduces an increasingly large bias, and thus the data cannot be fit with one single model. The performance of One-Per-En: (in red), for which models are trained independently for each environment, is constant as expected when the number of environments changes.
trajcctoiis
10-1
One Fon-For-All One-Per-Enn LEADS (Ours)
Test (Log-scakl
10-2
10
10-4
8
2
4
Number Training Environments
Figure 4: Test enror for LV w.r.t. the number of environ- ments. We apply the models in I to environments. groups of curves correspond to models trained with 1 to 8 trajectories per env. All groups highlight the same ten- dencies: increasing One-For-All, stable One-Per-Env, and decreasing LEADS.. More results of baselines meth- ods in Sup. D.
8

---

LEADS (green) circumvents these issues and shows that the shared characteristics among the environ- ments can be leveraged so as improve generalization: is particularly effective when the number of samples per environment is small. (See Sup. D for more details on the experiments and on the results).
Learning on novel environments. We demonstrate how the pre- pre-trained dynamics can help to fit a model for novel environ- ments. We took an pre-trained by LEADS on a set of LV environments, Fig. 5 shows the evolution of the test loss during training for three systems: a f function pre-trained by LEADS on set of LV training environ- ments, a ge function trained from scratch on the new environment and LEADS that uses pre-trained and learns 9c. residue on this new environment. Pre-trained- f f-Only alone cannot predict in any novel environments. Very already surpasses the best error of the
TesMSE (Lag-saaa scald
Pre-trandd--OOll
10
One- Pere-Per-Bee Erv.
Pre-tminnd-- --luu--
10-2
2500
5000
7500
10000
12500
15000
17500
20000
Gradient step
Figure 5: Test error evolution during training on 2 novel environments for LV.
fast in the training stages, Pre- frrriieddf-Plus-Trarii model trained from scratch (indicated with dotted line). Similar results are also observed with the GS and NS datasets (cf. Sup. D, Table S5). These empirical results clearly show that the learned shared dynamics accelerates and improves the learning in novel environments.
4.4 Training and implementation details
Discussion trajectory-basse optimization. involves computing a trajectory loss (integral term in Eq. 4). However, in practice, we do not have access to the continuous trajectories at every instant but only to finite number of snapshots for the state values {*^Δ Jocks & at a temporal resolution Åt. From these discrete observed trajectories, it is still possible to recover an approximate derivative df ~ drzae using a numerical scheme A. The integral term for given sample in the objective Eq. 4 would then be estimated as
Solving the learning problem Eq. 2 in our setting,
de
Ex - (ʃ + ge) (Faa))) This is not the best solution and we have observed much better prediction performance for all models, including the baselines, when computing the error directly on the states, using an integral formulation Σ Σx= [[=(++1)4+ #(*+1)eell 112 where F((+1) is the solution given by numerical solver approximating the integral ZR + SkSt (((t)))((( + ge (x:)).))))))) )ds starting from FEAT- Comparing directly in the state space yields more accurate results for prediction as the learned network tends to correct the solver's numerical errors, first highlighted in [37].
Calculating Ω. Given finite data and time, the exact infinity norm and Lipschitz norm are both intractable. We opt for more practical forms in the experiments. For the infinity norm, we chose to minimize the empirical norm of the output vectors on known data points, this choice is motivated in Sup. C. In practice, we found out that dividing the output norm by its input norm works better: $ Ex kat) /| II,, where the #KAt are known states in the training set. For the Lipschitz norm, as suggested in [5], we optimize the sum of the spectral norms of the weight at each layer 12, 112.
We use the power iteration method in [25] for fast spectral norm approximation.
Implementation. We used 4-yyyer MLPs for LV, 4-yerr ConvNets for GS and Fourier Neural Operator (FNO) [19] for NS. For FT-RNN baseline, we adapted GRU [8] for LV and PredRNN [35] for GS and NS. We apply the Swish function [31] as the default activation function. Networks are integrated in time with RK4 (LV, GS) or Euler (NS), using the basic back-propagation through the internals of the solver. We apply an exponential Scheduled Sampling [17] with exponent of 0.99 to stabilize the training. We use the Adam optimizer [15] with the same learning rate 10-3 and (51, (2) (0.9. 0.999) across the experiments. For the hyperparamters in Eq. 8, we chose respectively X= = 5 103, 102, 10 and a = 10-3, 10-2, 10-3 for IV, GS and NS. All experiments are performed with a single NVIDIA Titan Xp GPU.
5 Related work
Recent approaches linking invariances to Dut-of-Distribution (OoD) Generalization, such as [1, 16, 34), aim at finding a single classifier that predicts well invariantly across environments with the power of extrapolating outside the known distributions. However, our dynamical systems context, the optimal regression function should be different in each environment, and modeling environment

---

bias is as important as modeling the invariant information, as both are indispensable for prediction. Thus such invariant learners are incompatible with our setting. Meta-learning methods have recently been considered for dynamical systems as in [11, 18]. Their objective is to train single model that can be quickly adapted to a novel environment with a few data-points in limited training steps. However, in general these methods do not focus on leveraging the commonalities and discrepencies in data and may suffer from overfitting at test time [24]. Multi-task learning [38] seeks for learning shared representations of inputs that exploit the domain information. Up to our knowledge current multi-task methods have not been considered for dynamical systems. [33] apply multi-task learning for interactive physical environments but do not consider the case of dynamical systems. Other approaches like [39, 28] integrate probabilistic methods into a Neural ODE, to learn a distribution of the underlying physical processes. Their focus is the uncertainty of a single system. [37] consider an additive decomposition but focus on the combination of physical and statistical components for a single process and not learning from different environments.
6 Discussions
Limitations Our framework is generic and could be used in many different contexts. On the theoretical side, the existence and uniqueness properties (Prop. 1) rely on relatively mild conditions covering a large number of situations. The complexity analysis, on the other side, is only practically relevant for simple hypothesis spaces (here linear), and then serves for developing the intuition more complex spaces (NNs here) where bounds are too loose to be informative. Another limitation is that the theory and experiments consider deterministic systems only: the experimental validation is performed on simulated deterministic data. Note however that this is the case in the vast majority of the ML literature on ODE/PDE spatio-temporal modeling [30, 20, 19, 37J. In addition, modeling complex dynamics from real world data is a problem by itself.
Conclusion We introduce LEADS, a data- data-driven framework to learn dynamics from data collected from set of distinct dynamical systems with commonalities. Experimentally validated with three families of equations, our framework can significantly improve the test performance in every envi- ronment w.r.t. classical training, especially when the number of available trajectories is limited. We further show that the dynamics extracted by LEADS can boost the learning in similar new environ- ments, which gives us a flexible framework for generalization in novel environments. More generally, we believe that this method is a promising step towards addressing the generalization problem for learning dynamical systems and has the potential to be applied to a large variety of problems.
Acknowledgements
We acknowledge financial support from the ANR AI Chairs program DLACLIM ANR- 19-CHIA 0018-01.
References
[1] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. arXivv/907.0000 [cs, stat), Mar. 2020. arXiv: 1907.02893.
Invariant Risk Minimization.
[2] I. Ayed, F. de Bézenac, A. Pajot, J. Brajard, and P. Gallinari. Learning dynamical systems from partial observations. CORR, abs/1902. 11106.1....... 11161,11222111 2019.
[3] P. L. Bartlett, D. Foster, and M. J. Telgarsky. -ormmally margin bounds for neural networks. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30, pages 6240- 6249, Curran Associates, Inc., 2017.
[4] J. Baxter. A model of inductive bias learning. J. Artif. Int. Res., 12(1):149- 198, Mar. 2000.
[5] A. Bietti, G. Mialon, D. Chen, and J. Mairal. A kernel perspective for regularizing deep neural networks. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 664 674, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
10

---

[6] S. L. Brunton, J. L. Proctor, and J. N. Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 113(15)3392-3397, 2016.
[7] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 6571 6571-6583. Curran Associates, Inc., 2018.
[8] K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau. F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 1724-1734. 1734, Doha, Qatar, Oct. 2014, Association for Computational Linguistics.
[9] E. de Bézenac, A. Pajot, and P. Gallinari. Deep learning for physical processes: Incorporating prior scientific knowledge. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. Openke- view.net, 2018.
[10] J. Dormand and P. Prince. A family of embedded runge-kutta formulae. Journal of Computa- tional and Applied Mathematics, 6(1):19 - 26, 1980.
[11] C. Finn, P. Abbeel, and S. Levine. aggnosic meta-learning for fast adaptation of deep networks. CoRR, abs/ ttss/703.033400 2017.
[12] S. Fresca, A. Manzoni, Dedè, and A. Quarteroni. Deep learning- -bsed reduced order models in cardiac electrophysiology, volume 15. 2020.
[13] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg. N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del R'10, M. Wiebe, P. Peterson, G'erard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature, 585(7825):3357-332, Sept. 2020.
[14] D. Haussler. Decision theoretic generalizations of the pac model for neural net and other learning applications. Information and Computation, 100(1):78 150, 1992.
[15] D. P. Kingma and Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[16] D. Krueger, E. Caballero, J. Jacobsen, A. Zhang, J. Binas, R. Le Priol, and A. C. Courville. Out-of disttiiibitttnnnnn generalization via risk extrapolation (rex). CoRR, abs/2003. 00688, 2020.
[17] A. Lamb, A. Goyal, Y. Zhang, Zhang, A. Courville, and Y. Bengio. Professor Foreing: New Algorithm for Training Recurrent Networks, arXiv:1610.00038 [cs, stat), Oct. 2016. arXiv: 1610.09038.
[18] S. Lee, H. Yang, and W. Seong. Identifying physical law of hamiltonian systems via meta- learning. CoRR, abs/2102.115444 2021.
[19] Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anand- kumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021.
[20] Z. Long, Y. Lu, and B. Dong. Pde-net 2.0: Learning pdes from data with A numeric-symbolic hybrid deep network. CORR, http//882.04422 2018
[21] Z. Long, Lu, X. Ma, and B. Dong. Pde-net: Learning pdes from data. In International Conference on Machine Learning, pages 3214-3222, 2018.
[22] A. J. Lotka. Flements of physical biology. Science Progress the Twentieth Century (1919- 1933), 21(82):341-343, 1926.

---

[23] G. Madec, R. Bourdallé-Badie J. Chanut, E. Clementi, A. Coward, C. Ethé, D. Iovino, D. Lea, C. Lévy, T. Lovato, N. Martin, Masson, S. Mocavero, C. Rousset, D. Storkey, M. Vancop- penolle, S. Müeller, G. Nurser, M. Bell, and G. Samson. Nemo ocean engine, Oct. 2019, Add S13 and TOP reference manuals.
[24] N. Mishra, M. Rohaninejad, X. Chen, and P. Abbeel. Meta-learing with temporal convolutions. CoRR, 0es/1707.03311 2017.
[25] T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida. Spectral normalization for generative adversarial networks. CoRR, abs/1802.009955 2018.
[26] V. Nagarajan and J. Z. Kolter. Uniform convergence may be unable to explain generalization in deep learning. H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[27] A. Neic, F. Campos, J. Prassl, S. A. Niederer, M. J. Bishop, E. J. Vigmond, and Plank. Efficient computation of electrograms and ecgs in human whole heart simulations using reaction- eikonal model. Journal of Computational Physics, 346:191 21,, 2017.
[28] A. Norcliffe, C. Bodnar, B. Day, J. Moss, and P. Liò. Neural ODE processes. In International Conference on Learning Representations, 2021.
[29] E. Pearson. Complex patterns in a simple system. Science, 261(5118); 1993.
[30] M. Raissi, P. Perdikaris, and G. E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differ- ential equations. Journal of Computational Physics, 378:686-707, 2019.
[31] Ramachandran, B. Zoph, and Q V. Le. Searching for activation functions. CoRR, abs/1710.05941, 2017.
[32] M. Reichstein, G. Camps-Vaie Valls, B. Stevens, M. Jung. J. Denzler, N. Carvalhais, and Prabhat. Deep learning and process understanding for data-driven Earth system science. Nature, 566:195- 204, 2019.
[33] Spieckermann, S. Düll, S. Udluft, A. Hentschel, and Runkler. Exploiting similarity in system identification tasks with recurrent neural networks. Neurocomputing, 169:343 - 349, 2015. Learning for Visual Semantic Understanding in Big Data ESANN 2014 Industrial Data Processing and Analysis.
[34] D. Teney, E. Abbasnejad, and A. van den Hengel. Unshuffling Data for Improved Generalization. 2020.
[35] Y. Wang, M. Long, J. Wang. Z. Gao, and P. S. Yu. Predrnn: Recurrent neural networks for predictive learning using spatiotemporal Istms. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garntt, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[36] J. D. Willard, X. Jia, S. Xu, M. Steinbach, and Kumar. Integrating physics- modeling with machine learning: A survey. volume 1, pages 1-34, 2020.
[37] Y. Yin, V. Le Guen, J. Dona, E. de Bezenac, I. Ayed, N. Thome, and P. Gallinari. Augment- ing physical models with deep networks for complex dynamics forecasting. In International Conference on Learning Representations, 2021.
[38] Y. Zhang and Q Yang. A survey on multi-task learning. CoRR, abs/1707.000 707.081 14,1 2017.
[39] Çagatay Yıldiz, M. Heinonen, and H. Lähdesmäki. Vde: Deep generative second order odes with bayesian neural networks, 2019.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In Sec. 6.
(c) Did you discuss any potential negative societal impacts of your work? [No]
The only relevant societal impact is around the computational cost, while we use very limited computation power (maximum with single GPU).
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions all theoretical results? [Yes] In Sec. 2.1 and Sec. 3.
(b) Did you include complete proofs of all theoretical results? [Yes]
In the supplemental material.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [Yes]
We will provide the code the supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] In Sec.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] In Sec. 4.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Sec. 4.
4. If you are using existing assets code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] In Sec. 4.
(b) Did you mention the license of the assets? [Yes] In the supplemental material.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "qth090c1bzp", "reviewer": "Reviewer_hYVh", "summary": "This proposes a simple but effective method to learn a dynamics prediction model that can generalize different environments. The authors also provide theoretical and experimental results to demonstrate the superiority of the proposed method over baselines.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Advantages:\n   1. Learning a dynamics prediction model that can generalize across different environments is important and practical for many real-world applications, and the authors provide a simple and intuitive method to solve it with theoretical guarantees.\n   2. The experimental results shown in the main paper and Appendix are satisfactory.\n\nDisadvantage:\n  1. I notice that the proposed method can only be applied to deterministic dynamics prediction, which limits its potential application scenarios.\n  2. The hyperparameters \\lambda and \\alpha varies significantly across different environments. Can the authors provide the results about the sensitivity of hyperparameters?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See the section \"Main Review\" ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "LDWQgUw5jFT", "reviewer": "Reviewer_35iL", "summary": "This paper presents an approach called LEADS to learn *decomposed* dynamics models of system instances (\"environments\") from a family of systems, e.g., Lotka-Volterra systems with randomly sampled parameters. The main idea is to additively decompose the dynamics into a shared component $f$ which does not depend on the specific environment and an environment-specific component $g_e$. It is shown both theoretically and experimentally that it is crucial to limit the complexity of $g_e$. This allows to adapt to new, unseen environments by leveraging \"knowledge\" absorbed in $f$ and additionaly using only a few interactions with the new environment to learn a new $g_{e'}$. The performed experiments qualitatively and quantitatively confirm the hypothesis that the proposed algorithm LEADS is more accurate than training a single model for all environments and even a single model for each environment.\n\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**\n\nThe paper suggests an additive decomposition of environment dynamics into a *shared* and *environment-specific* part. To my understanding, the main methodological novelty is that *both* parts are learned while a minimization objective on the environment-specific part is imposed. With regard to this, some related work should be discussed more precisely in my opinion:\n\n  * Another instance of dynamics models which are decomposed into a *shared* part and an *environment-specific* part are \"context-aware dynamics models\" or \"context-conditional dynamics models\", e.g. by (Lee et al., 2020). I think the paper would benefit from mentioning these approaches.\n\n  * An additive decomposition of a gradient field of an ODE has been described in (Yin et al. 2021), who also highlight the need for a minimization objective on one of the components (Yin et al., 2021, eq. 2) and show existence and uniqueness properties (Yin et al., 2021, Prop. 1+2). I think the paper of Yin et al. (2021) should be, due to this overlap, discussed in more detail in the paper under review.\n\n  * I think the statement in l. 75/76 \"Up to our knowledge, this is the first time it is addressed from an ML viewpoint\" needs clarification. To me it is not immediately clear what is meant by \"this\". Has the problem/question already been addressed from a different, non-ML viewpoint?\n\n\n**Quality**\n\n* The chosen approach is well-motivated by mathematical derivations, e.g. regarding limiting the complexity of the environment-specific dynamics components. \n\n* All experiments are performed on a set of ODE/PDE dynamics (Lotka-Volterra, Gray-Scott and Navier-Stokes), showing that the reported effects are valid for multiple systems.\n\n* Reasonable ablation and comparison experiments are performed highlighting the importance of the minimization objective.\n\n\n**Clarity**\n\n* The paper is well-structured and clearly written. The relevant research questions are clearly outlined.\n\n* The figures nicely illustrate the results obtained using the method and the main statements of the paper. Some minor comments:\n  * Fig. 2: I would appreciate a colorbar for the MSE error maps to at least estimate the quantitative difference in the predictions between LEADS and its ablations / baseline methods.\n  * Fig. 4: Coordinate axes ('lines') at the bottom and on the right (including 'ticks') would help to immediately see that the y coordinate of the curves is not determined by \"Training trajectories per env.\", which I first falsely assumed to be a second \"axis\" on the left.\n  * For consistency, I suggest to introduce coordinate axes also in Fig. 1, 3, 5.\n\n* The derivations and experiments for the 'simplified' linear case help in understanding the mechanics of the method before considering the more-complicated NN models.\n\n\n**Significance and conclusion**:\nIn my opinion the proposed idea of decomposing dynamics into a \"general\" component and an \"environment-specific\" component is very interesting and could, for example, help the community to build more sample-efficient model-based RL algorithms which are able to quickly adapt to changing environments. Although the idea of additive decomposition seems simple, it requires significant theoretical effort to concretize this idea, especially regarding the required complexity regularization. This paper provides a strong theoretical formulation and experimental evaluation of the proposed method. Noting that I haven't checked the mathematical derivations in detail, I vote for acceptance.\n\n\n**References:**\n\n (Lee et al., 2020): \n Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, Jinwoo Shin  \n \"Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning\"  \n ICML2020, https://arxiv.org/abs/2005.06800  \n \n (Yin et al., 2021):  \n Yuan Yin, Vincent Le Guen, Jérémie Donà, Emmanuel de Bézenac, Ibrahim Ayed, Nicolas Thome, Patrick Gallinari  \n \"Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting\"  \n ICLR2021, https://arxiv.org/abs/2010.04456", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "* The authors clearly mention the limitations of their work.\n\n* Societal impacts are not explicitly addressed. In my opinion, this work does not exhibit extraordinary societal impacts in addition to those induced by basic reasearch on dynamics modelling methods in general.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "D_ISc5X0xe3", "reviewer": "Reviewer_juRp", "summary": "This paper proposes a new framework (LEADS) for learning dynamical systems (ODEs, PDEs) with the aim of generalising across related, but slightly different environments. It uses an additive two-component model, comprising a shared model and an environment-specific model. The paper presents some generic sample complexity results, then looks at the special case of a linear ODE. The paper claims this motivated a choice for an objective in the nonlinear NN case. This is then experimentally investigated in the small data regime, with just a few example trajectories, and just a few different environments. The paper claims the proposed framework significantly outperforms other baselines and ablations.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**\n\nI’m not an expert in this area, but the theoretical propositions seem to be generic, from their reference [4], with some mild adaptation to their setting - this doesn’t seem that original in my opinion. Further, they don’t properly reference this with respect to [4], stating only that the theory was inspired by [4]. More detail is needed to properly assess the originality of this section. I would have also liked to see more detail in the related work section. To the best of my knowledge, the proposed method seems to be original. \n\n**Quality**\n\nI couldn’t follow the proof of Prop 2 in supplemental, I don’t understand why the factor of 4/delta appears in c_g term, but not in c_f term. I haven’t tried to verify the proof of other statements. I didn’t understand why the linear case motivated taking equation (7) for the NN case. I would like a scale for the MSE error maps in fig 2 to assess quality. Otherwise, qualitative results in figs2,3 seem good quality. Table 1: Needs more detail on how significance window was calculated. I felt for LV and NS that LEADS wasn’t significantly better than the ablation LEADS no min. For GS, I felt LEADS wasn’t significantly better than FT-NODE (note typo in table 1, FT-ODE -> FT-NODE). To complement fig 5, I would like to see results on GS and NS. I would have also liked to see performance differences in the regime of larger number of environments and trajectories. \n\n**Clarity**\n\nThe prose in most of the paper was acceptable, but I found the math difficult to read, especially prop 2-4. Are these propositions really necessary for the main paper? They didn’t help me as a non-expert reader to understand what was going on, or why this was important. All figures would benefit from being larger. \n\n**Significance**\n\nAs noted in the limitations section, this only investigated simulated, deterministic data. I found this surprising because when considering generalization, this is typically done in the setting where the data contains some noise, and one tries not to overfit to that noise.  \nMoreover, considering this paper is using neural networks, which are typically better in the larger data regime, this makes me question the significance of the results, compared to the baselines (especially those that use neural networks). \nAs also noted in limitations section, the complexity analysis is only relevant in simple hypothesis spaces, such as the linear case. But already this wasn’t enough for the simple NN experiments in this paper. \n\n----\n\nEdit:\n\nFollowing discussion period, I have now increased my score from 4 to 7, based on author's better explanation of equation (7), and their other proposed changes.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Some technical limitations were discussed in the work. No societal impact was discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "55dYWDQvTg5", "reviewer": "Reviewer_27Ns", "summary": "The paper proposes an algorithm which learns to model dynamical systems\nconsisting of multiple environments that exhibit shared and\nenvironment-specific dynamics. The key idea is to learn separate approximation\nfunctions for shared and environment-specific components. To prevent a\ntrivial solution that circumvents the approximation of shared dynamics, the\nproposed objective regularizes environment-specific terms.\n\nOn the theoretical side, the paper quantifies the improvement in sample\ncomplexity for novel environments in a simplified (linear) setting for which a\nsuitable regularization term is derived. Empirically, the method is evaluated\non three datasets with nonlinear dynamics (all datasets have shared and\nenvironment-specific components) and the results support the claims about\nimproved generalization on known and novel environments, especially when data\nis scarce (i.e., limited number of trajectories).\n", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Even though the paper is outside my domain of expertise, it was mostly easy\nto follow, due to the clear and unambiguous writing style and notation, as well\nas the excellent explanations throughout the paper. The authors are transparent\nabout the limitations of the proposed approach. The distinction from previous\nwork in related fields appears somewhat superficial (see Comments and Questions). \n\nThe presented empirical results show consistent and significant improvements\ncompared to all baselines, supporting the claims about improved generalization\non known and novel environments. However, I think that certain ablations (see\nComments and Questions) could provide a more complete picture of the results\nand the suitability of the employed heuristics (e.g., the regularization term\nfor neural networks). Due to my lack of domain knowledge, I cannot evaluate the\noriginality and significance of this work in the context of dynamical systems.\n\n\n**Comments and Questions:**\n\n- Could you give an example for a real-world dynamical system that exhibits the\n  properties assumed in this work (i.e., shared and\n  environment-specific mechanisms across a bunch of environments)? Is there a\n  particular system that you imagine in practice or are the three presented\n  datasets the closest examples that you have for such a dynamical system?\n\n- In the related work section, the distinction from work in related fields\n  seems somewhat hand-wavy. With respect to meta-learning approaches, even if\n  they do not focus on generalization on training environments, that does not\n  imply that they do not generalize well on training environments. Even if\n  meta-learning approaches or ODE2VAE (Yildiz et al., 2019) focus on a single\n  model or a single dynamical system, it could still be interesting to include\n  the methods as a benchmark.  For example, does their test error also increase\n  with the number of training environments (similar to *One-For-All* in Figure\n  4)?\n\n- Did you consider running an ablation for the hyperparameter lambda? It could\n  be instructive to see the generalization performance for different\n  values of lambda. \n\n- From Appendix B.4 and Appendix C, I understand that the regularizer from\n  Equation (7) is approximated via heuristics that introduce additional\n  hyperparameters. For example, the choice of $p$-norm (Table S4) seems to have\n  a non-negiglible effect on the results of the LV experiment. In light of this\n  variablility, did you consider to include the results for different $p$-norms\n  in Table 1 directly?\n\n- For Figure 4, do you also have results for FT-RNN and FT-ODE? It would be\n  interesting to see how robust the methods perform with respect to the number\n  of training environments.\n\n- Lines 103f: could you point out the source or theorem for why the additive\n  hypothesis is not restrictive given functional vector spaces?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": " ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: From the SAC: This is an instance where a good rebuttal has helped. The original decision on this paper was a reject, but the SAC felt that your rebuttal was very thorough and detailed. Hence I am recommending moving this paper to an accept. Please do make sure in the next version of the paper to take all reviewer comments into account as well as to merge into the paper the additional discussion in your rebuttal (perhaps, due to space considerations, into the appendices). 
Decision: Accept (Poster)