LEADS: Learning Dynamical Systems that Generalize Across Environments

1 Introduction . Data-driven approaches offer an interesting alternative and complement to physical-based methods for modeling the dynamics of complex systems and are particularly promising in a wide range of settings : e.g . if the underlying dynamics are partially known or understood , if the physical model is incomplete , 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 , 9 , 7 , 21 , 30 , 2 , 37 ] . It is already being applied in different scientific disciplines ( see e.g . [ 36 ] for a recent survey ) and could help accelerate scientific discovery to address challenging domains such as climate [ 32 ] or 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 ” . In 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 as for cardiac computational models [ 27 ] . More 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . generally , data acquisition and modeling are affected by different factors such as geographical position , sensor variability , measuring circumstances , 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 , as small changes in initial conditions lead to trajectories not covered by the training data . In this work , we consider a 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 dynamics in each environment share similarities , while being 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 : RQ1 Does modeling the differences between environments improve generalization error w.r.t . classical settings : One-For-All , where a unique function is trained for all environments ; and OnePer-Env. , where a specific function is fitted for each environment ? ( cf . Sec . 4 for more details ) RQ2 Is it possible to extrapolate to a novel environment that has not been seen during training ? We propose LEarning Across Dynamical Systems ( LEADS ) , a novel learning methodology decomposing the learned dynamics into shared and environment-specific 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-specific component only models the remaining dynamics , i.e . those that can not be expressed by the former , based on environment-specific data . We show , under mild conditions , that the learning problem is well-posed , as the resulting decomposition exists and is unique ( Sec . 2.2 ) . We then analyze the properties of this decomposition from a 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 , leading to a 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 ( RQ1 ) . 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 Lotka-Volterra 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-All and One-Per-Env . 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 ML perspective1 . 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 E. Throughout the paper , we will assume that the dynamics in an environment e ∈ E are defined through the evolution of differential equations . This will provide in particular a 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 are assumed to be completely unknown . xet denotes the state of the equation for environment e , taking its values from a bounded set A , with evolution term fe : A→ TA , TA being the tangent bundle of A . In other words , over a fixed time interval [ 0 , T ] , we have : dxet dt = fe ( x e t ) ( 1 ) 1Code is available at https : //github.com/yuan-yin/LEADS . We assume that , for any e , fe lies in a functional vector space F. In the experiments , we will consider one ODE , in which case A ⊂ Rd , and two PDEs , in which case A is a d′-dimensional vector field over a bounded spatial domain S ⊂ Rd′ . The term of the data-generating dynamical system in Eq . 1 is sampled from a distribution for each e , i.e . fe ∼ Q . From fe , we define Te , the data distribution of trajectories xe· verifying Eq . 1 , induced by a distribution of initial states x e 0 ∼ P0 . The data for this environment is then composed of l trajectories sampled from Te , and is denoted as T̂e with x e , i · the i-th trajectory . We will denote the full dataset by T̂ = ⋃ e∈E T̂e . The classical empirical risk minimization ( ERM ) framework suggests to model the data dynamics either at the global level ( One-For-All ) , taking trajectories indiscriminately from T̂ , or at the specific environment level ( One-Per-Env . ) , training one model for each T̂e . Our aim is to 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 f ∈ F is shared across environments and ge ∈ F is specific to the environment e , so that ∀e ∈ E , fe = f + ge ( 2 ) Since we consider functional vector spaces , this additive hypothesis is not restrictive and such a decomposition 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 to 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 is a { ge } e∈E such that Eq . 2 is verified . A trivial example would be f = 0 leading to a solution where each environment is fitted separately . Our core idea is that f should capture as much of the shared dynamics as is possible , while ge should focus only on the environment characteristics not captured by f . To formalize this intuition , we introduce Ω ( ge ) , a penalization on ge , which precise definition will depend on the considered setting . We reformulate the learning objective as the following constrained optimization problem : min f , { ge } e∈E∈F ∑ e∈E Ω ( ge ) subject to ∀xe , i ∈ T̂ , ∀t , dxe , it dt = ( f + ge ) ( x e , i t ) ( 3 ) Minimizing Ω aims to reduce ges ’ complexity while correctly fitting the dynamics of each environment . 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 . 3 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 f ? , { g ? e } e∈E ∈ F of Eq . 3 is guaranteed . Furthermore , if Ω is strictly convex , this decomposition is unique . In practice , we consider the following relaxed formulation of Eq . 3 : min f , { ge } e∈E∈F ∑ e∈E ( 1 λ Ω ( ge ) + l∑ i=1 ∫ T 0 ∥∥∥∥dxe , itdt − ( f + ge ) ( xe , iτ ) ∥∥∥∥2dt ) ( 4 ) where f , ge are taken from a hypothesis space F̂ approximating F. λ is a regularization weight and the integral term constrains the learned f + ge to follow the observed dynamics . The form of this objective and its effective calculation will be detailed in Sec . 4.4 . 3 Improving generalization with LEADS Defining an appropriate Ω is crucial for our method . In this section , we 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 to 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 , we prove an upper-bound of their effective size based on the 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 x we want to predict some target y . In our setting , we want to learn trajectories ( xet ) 0≤t≤T starting from an initial condition x0 . We reformulate this problem and cast it as a standard supervised learning problem : Te being the data distribution of trajectories for environment e , as defined in Sec . 2.1 , let us consider a trajectory xe· ∼ Te , and time τ ∼ Unif ( [ 0 , T ] ) ; we define system states x=xeτ ∈ A as input , and the corresponding values of derivatives y = fe ( xeτ ) ∈ TA as the associated target . We will denote Pe the underlying distribution of ( x , y ) , and P̂e the associated dataset of size n. We are searching for f , ge : A→TA in an approximation function space F̂ of the original space F. Let us define Ĝ ⊆ F̂ the effective function space from which the ges are sampled . Let f + Ĝ : = { f + g : g ∈ Ĝ } be the hypothesis space generated by function pairs ( f , g ) , with a fixed f ∈ F̂ . For any h : A→ TA , the error on some test distribution Pe is given by erPe ( h ) = ∫ A×TA‖h ( x ) − y‖ 2dPe ( x , y ) and the error on the training set by êr P̂e ( h ) = 1n ∑ ( x , y ) ∈P̂e‖h ( x ) − y‖ 2 . LEADS sample complexity . Let C Ĝ ( ε , F̂ ) and C F̂ ( ε , Ĝ ) denote the capacity of F̂ and Ĝ at a certain scale ε > 0 . 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 S1 . The following result is general and applies for any decomposition of the form f + ge . It states that to guarantee a 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 a step towards RQ1 ( proof see Sup . B.2 ) : Proposition 2 . Givenm environments , let ε1 , ε2 , δ > 0 , ε = ε1 +ε2 . Assume the number of examples n per environment satisfies n ≥ max { 64 ε2 ( 1 m ( log 4 δ + log C Ĝ ( ε1 16 , F̂ ) ) + log C F̂ ( ε2 16 , Ĝ ) ) , 16 ε2 } ( 5 ) Then with probability at least 1−δ ( over the choice of training sets { P̂e } ) , any learner ( f+g1 , . . . , f+ gm ) will satisfy 1m ∑ e∈E erPe ( f + ge ) ≤ 1 m ∑ e∈E êrP̂e ( f + ge ) + ε . The contribution of F̂ to the sample complexity decreases as m increases , while that of Ĝ remains the same : this is due to the fact that shared functions f have access to the data from all environments , which is not the case for ge . From this finding , one infers the basis of LEADS : when learning from several environments , to control the generalization error through the decomposition fe = 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 Ω ( ge ) in Eq . 4 controls the capacity of the effective set Ĝ by selecting ges that are as “ simple ” as possible . As a corollary , we show that for a fixed total number of samples in T̂ , 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 O ( logC Ĝ ( ε116 , F̂ ) +m logCF̂ ( ε2 16 , Ĝ ) ) and O ( b logC Ĝ ( ε116 , F̂ ) +m logCF̂ ( ε2 16 , Ĝ ) ) . The latter being larger than the former , a situation with more environments presents a clear advantage . Fig . 4 in Sec . 4 confirms this result with empirical evidence . LEADS sample complexity for novel environments . Suppose that problem Eq . 3 has been solved for a set of environments E , can we use the learned model for a new environment not present in the initial training set ( RQ2 ) ? Let e′ be such a new environment , Pe′ the trajectory distribution of e′ , generated from dynamics fe′ ∼ Q , and P̂e′ 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 f + ge′ with f fixed on this new environment than training another f ′ + ge′ from scratch ( proof see Sup . B.2 ) . Proposition 3 . For all ε , δ with 0 < ε , δ < 1 if the number of samples n′ satisfies n′ ≥ max { 64 ε2 log 4C ( ε16 , f + Ĝ ) δ , 16 ε2 } , ( 6 ) then with probability at least 1− δ ( over the choice of novel training set P̂e′ ) , any learner f + ge′ ∈ f + Ĝ will satisfy erPe′ ( f + ge′ ) ≤ êrP̂e′ ( f + ge′ ) + ε . In Prop . 3 as the capacity of F̂ no longer appears , the number of required samples now depends only on the capacity of f + Ĝ . This sample complexity is then smaller than learning from scratch fe′ = f + ge′ as can be seen by comparing with Prop . 2 at m = 1 . From the previous propositions , it is clear that the environment-specific functions ge need to be explicitly controlled . We now introduce a practical way to do that . Let ω ( r , ε ) be a strictly increasing function w.r.t . r such that logC F̂ ( ε , Ĝ ) ≤ ω ( r , ε ) , r = supg∈Ĝ Ω ( g ) ( 7 ) Minimizing Ω would reduce r and then the sample complexity of our model by constraining Ĝ . Our goal is thus to construct such a pair ( ω , Ω ) . In the following , we will first show in Sec . 3.2 , how one can construct a penalization term Ω 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 an effective Ω to penalize the complexity of the neural networks ge .

Summary: The paper proposes an algorithm which learns to model dynamical systems consisting of multiple environments that exhibit shared and environment-specific dynamics. The key idea is to learn separate approximation functions for shared and environment-specific components. To prevent a trivial solution that circumvents the approximation of shared dynamics, the proposed objective regularizes environment-specific terms.  On the theoretical side, the paper quantifies the improvement in sample complexity for novel environments in a simplified (linear) setting for which a suitable regularization term is derived. Empirically, the method is evaluated on three datasets with nonlinear dynamics (all datasets have shared and environment-specific components) and the results support the claims about improved generalization on known and novel environments, especially when data is scarce (i.e., limited number of trajectories).