MetaPhysiCa: Improving OOD Robustness in Physics-informed Machine Learning
Authors: S Chandra Mouli
Muhammad Alam
Bruno Ribeiro
Venue: ICLR 2024 spotlight (2024)
Keywords: physics-informed machine learning
OOD robustness
meta learning
causal structure discovery
Primary area: neurosymbolic & hybrid AI systems (physics-informed, logic & formal reasoning, etc.)

Abstract: A fundamental challenge in physics-informed machine learning (PIML) is the design of robust PIML methods for out-of-distribution (OOD) forecasting tasks. These OOD tasks require learning-to-learn from observations of the same (ODE) dynamical system with different unknown ODE parameters, and demand accurate forecasts even under out-of-support initial conditions and out-of-support ODE parameters. In this work we propose to improve the OOD robustness of PIML via a meta-learning procedure for causal structure discovery. Using three different OOD tasks, we empirically observe that the proposed approach significantly outperforms existing state-of-the-art PIML and deep learning methods (with $2\times$ to $28\times$ lower OOD errors).

Published as a conference paper at ICLR 2024
METAPHYSICA: IMPROVING OOD ROBUSTNESS IN PHYSICS-INFORMER MACHINE LEARNING
S Chandra Mouli
Muhammad Ashraful Alam Department of Electrical and Computer Engineering, Purdue University, Indiana, US
Department of Computer Science Purdue University, Indiana, US
Bruno Ribeiro
Department of Computer Science Purdue University, Indiana, US
ABSTRACT
A fundamental challenge in physics-infommn machine learning (PIML) is the design of robust PIML methods for onttof-distribution (OOD) forecasting tasks. These OOD tasks require learning-to-leam from observations of the same (ODE) dynamical system with different unknown ODE parameters, and demand accurate forecasts even under out-of-support initial conditions and out-of-support ODE parameters. In this work propose to improve the OOD robustness of PIML via meta-learning procedure for causal structure discovery. Using three different OOD tasks, we empirically observe that the proposed approach significantly outperforms existing taate-of-theaare art PIML and deep learning methods (with 2 to 28 lower OOD errors).
INTRODUCTION
inysice-informed machine learning (PIML) (e.g., (Willard et al., 2020; Xingjian al., 2015; Lusch et al., 2018; Yeo & Melnyk, 2019; Raissi et al., 2018; Kochkov et al., 2021)) seeks to combine the strengths physics and machine learning models and has positively impacted fields as diverse as biological sciences (Yazdani al., 2020), climate science (Faghmous & Kumar, 2014), turbulence modeling (Ling et al., 2016; Wang et al., 2020a), among others. PIML achieves substantial success in tasks where the test data comes from the same distribution the training data (in findiitrribttinn tasks).
This work expands on existing PIML approaches to forecast a parametric dynamical system (ODE), (OOD) scenarios. In tasks, OOD robustness is tied interven-
with focus tions over the initial system state and the unknown ODE parameters (illustrated in Figure ((a,)), not arbitrary interventions the system evolves from the initial state (see (Rubenstein al., 2016) for the effect of arbitrary interventions in physics models). In this setting, we observe that existing state- of-the-art PIML models perform significantly worse OOD than in-distribution, with some scenarios being challenging even for PIML methods designed with OOD robustness in mind (Wang et al., 2021b; Kirchmeyer al., 2022). This is because the standard ML part of PIML still learns spurious associations and performs poorly in our OOD setting. While there are PIML methods (Brunton et al., 2016; Martius & Lampert, 2016) robust to OOD initial conditions, they are not robust to shifts in ODE parameters as they are transductive, i.e., they do not combine knowledge from diverse training trajectories with different ODE parameters.
We then propose approach for more robust dynamical system forecasting that leverages causal structure discovery (c.g., Zheng et al. (2018)) combined with invariant risk minimization (Arjovsky al., 2019; Krueger et al., 2021) learn the underlying ODE structure, and uses meta-learning to learn from different trajectories of the same dynamical system. More precisely, our contributions are:
1. We show that state-of-the of-te-of-the-at PIML and deep learning methods fail in test examples with OOD initial conditions and/or OOD system parameters. Prior work (Wang et al., 2021a) showed that deep learning-only methods fail in OOD tasks, and argued physics models and PIML. methods

---

Published as a conference paper at ICLR 2024
Predict future states of the same dynamical
(a)
Observations from unknown dynamical system:
(b)
system with OOD initial states & OOD parameters
dx
(22;
W
dt
W
W
W (
I nearly vertical
acute
Observe
Predict
Time
Figure 1: Dynamical system OOD problem definition. (a) Training data consists of multiple observations from the same dynamical system with different initial conditions and parameters. (b) At test, are given observations till f, (red solid) and the goal is to predict the future observations till Ir (gray dashed). The initial conditions and the ODE parameters can be OOD in test.
would succeed, including proposed solution (Wang al., 2021b). Here show that PIML methods also fail (or perform poorly) OOD, including the solution in Wang et al. (2021b).
Meta-learning framework for learning from diverse training trajectories: By meta learning we mean the definition in (Thrun & Pratt, 1998, Chapter 1.2) adapted to our setting. Given (a) a family of MM tasks or experiments observing the same dynamical system with different parameters, (b)
for each task i, and (c) performance measure (prediction error) for each task; our algorithm will meta learn such that performance on each task improves with more observations per task and with the number of tasks. Our architecture uses parameters shared across all tasks to allow transfer of knowledge between multiple tasks.
training observations or trajectory x .
3. Learning the ODE via causal structure discovery: We define family of structural causal models, each identifying a different ODE, and perform a search to find the underlying dynamical system (assumed to be in the family). We propose a causal structure discovery approach via continuous optimization with E eerreuularization and an invariant risk minimization-type penalty (Krueger et al., 2021), both of which are empirically shown to be necessary. We further show in Theorem 1 that the proposed approach learns the true causal structure under certain identifiability conditions.
The proposed method is then empirically validated using three commonly-used simulated physics tasks (with measurement noise): Damped pendulum systems (Yin et al., 2021), predator-prey sys- tems (Wang et al., 2021a), and epidemic modeling (Wang et al., 2021a).
2 DYNAMICAL SYSTEM FORECASTING OOD SETTING
We formally describe the dynamical system forecasting problem under Definition 1 (Dynamical system forecasting problem). The dynamical system is described as an ordinary differential equation (ODE) as follows:
scenarios.
dx: di
(wr;; W"),
where x Id is the hidden state of the system at time t, eb is unknown deterministic function with unknown ODE parameters W We define noisy observations of the dynamical system X₁ ~ N(x,, I where o is the noise variance. It, denotes the initial state at time to-
Training data (Figure la): In training, we are given a set of M tasks (or trajectories) from the
dynamical system in Equation (1), T()) x()) 1o..... .....- (p0) (1.....M), of lengths T(1) +
respectively, with
(i) Dig
(arta)
w(O (WW')"
where for each task i, (5) is the initial state at time to, (())= are the (hidden) task-specific ODE parameters, and (to,-...
1 ( are regularly-space discrete time steps.
OOD test data (Figure 1b); At test, we are given (r + 1)-lenghh initial observations Tite)
X(ee) X!!t), from the same dynamical system Equation (1) with an arbitrary initial state sto
and ODE parameters Wite possibly outside their respective training supports, supp( (xt.))
and supp (PtP(("WW)))
Forecasting goal is to predict the future observations X(to
((())
(t0) Lytin
till some time Tite)
2

---

Published as a conference paper at ICLR 2024
(a) Standard Neural Networks
(b) Transductive PIML
(c) Indiuctive PIML
Poor OOD approximation standard activations (ReLU, Tanh)
Assumes trajectories training and test have the same W"
network terms (red) are not robust 0OD
Pendulum mation: NeuralODE
Input: Inttaal obsensations
predictions of #o ground truth di
Training data
Noural network
Wr
predicts physics
parametors
Traniductivv PIML Model
Neural network corrects the physics model
Physica model
Time
Figure 2: Shows OOD challenges of existing methods. (a) Standard Neural networks (c.g., (Chen et al., 2018)), when trained the motion of damped pendulum, predict accurately in the training domain (green shaded), but predict a linear function outside the training domain. (b) Transductive PIML methods (e.g., (Raissi et al., 201 Brunton et al., 2016)) assume all training/test trajectories have the same dynamical system parameters W which is not the general setting. (c) Inductive PIML methods (e.g.. (Yin et al., 2021; Mehta et al., 2021)) use a neural network to predict parameters of a known physics model and another neural network residual correction. These two neural networks face the same robustness issues discussed in (a).
In summary, we are given training trajectories from that may have (a) different initial conditions, and (b) different unknown ODE parameters. We observe a test trajectory from time = to.... i, and we wish to forecast its future after time f,. The test trajectory can have an OOD initial condition and OOD ODE parameters.
Illustrative example. Figure shows an example of an out-of-distribution task for forecasting the motion of a pendulum with friction. The state x, = [0,,we]] E R describes the angle made by the pendulum with the vertical and the corresponding angular velocity at time t. The true (unknown) function va describing this dynamical system is given by vi( we]; W*) ) = [W -ax sin (0.) -=00---w) with W (a", p²) denoting the ODE parameters relating to the pendulum's 's period and the damping coefficient. They are daseddepn,,n c.g., different tasks can have different damping coefficients.
1. In training, observe M noisy trajectories (tasks) of motion over time t = 0,0.1, 10 from experiments where a pendulum dropped from different acute angles g( π /2 in materials with different damping coefficients p(*)*
2. In test, the pendulum is dropped from nearly vertical angles, π - 0.1 < 0 < π in a material
with much higher damping coefficient p The test trajectory is observed over smaller time window t S and the goal is to predict future pendulum states till time 10.
An ML or PIML model for this forecasting task must be able to learn from diverse trajectories of the same dynamical system, and be robust to shifts in distribution between training and test.
3 RELATED WORK & THEIR LIMITATIONS
Next we describe different classes of existing approaches that are commonly used for the dynamical system forecasting and their inherent challenges autrof-distributio.
3.1 NEURAL NETWORK METHODS
Deep learning's 's ability to model complex phenomena has allowed to make great strides in many physics applications (Lusch et al., 2018; Yeo & Melnyk, 2019; Kochkov et al., 2021; Dang et al., 2022; Brandstetter et al., 2022b). However, standard deep learning methods are known to learn spurious correlations and tend to fail when the test distribution of the inputs are different from that in training (Wang et al., 2021a; Geirhos et al., 2020). Figure 3b depicts the out-of--istribun failure of several deep learning methods from NeuralODE (Chen et al., 2018) to more complex meta learning approaches (Wang et al., 2021b) in our damped pendulum example (more details of the experiment is in Section 5). While Ad (Wang et al., 2021b) and CoDA (Kirchmeyer et al., 2022) use meta-learning adapt new ODE parameters, they are not robust to OOD initial conditions.
In standard deep learning tasks, Xu et al. (2021) show that an MLP's failure to extrapolate to out-of- distribution can be traced to absence of algorithmic alignment, which is appropriate combination of basis and activation functions within the architecture for the task. For example, the outputs of an MLP with ReLU activations will be linear far from the training domain even when trained to predict a sine/quadratic function. For dynamical system forecasting, our Figure 2a depicts the results of a similar experiment for standard sequence model (NeuralODF)): the model can approximate the
3

---

Published as a conference paper at ICLR 2024
2044
Trat
(a) In-distribution predictions of 02
(b) Predictions of 0, under OOD X
Figure 3: (Damped pendulum.) Predict pendulum motion (a) in-distribution, when dropped from acute angles and (b) OOD w.r.t initial conditions and parameters, when a different pendulum is dropped from nearly vertical angles. Figures show ground truth curve (blue stars) in- and out- of-distribution along with predictions from different models. Existing methods perform well in-distribution, but perform poorly OOD even in this simple example.
target sine function in the training domain (green region) but predicts a linear function far outside the training domain. This means that we need algorithmic alignment (i.e., to include appropriate basis functions) in order to make accurate forecasts in OOD tasks. In Appendix D.6 we show that even our proposed method without algorithmic aligment cannot properly extrapolate OOD (but it still does better than baselines).
3.2
MACHINE LEARNING (PIML)
To alleviate the challenges of standard neural networks, several PIML methods have been proposed (e.g., (Willard et al., 2020; Wang et al., 2020a; Faghmous & Kumar, 2014; Daw et al., 2017)) that utilize physics-based domain knowledge for better predictions. The physics-based knowledge vary across methods, for example, (a) a dictionary of basis functions (e.g.. sin, cos, di (Schmidt & Lipson, 2009; Brunton al., 2016; Martius & Lampert, 2016; Raissi, 2018: Cranmer et al., 2020a) related to the task, (b) a completely specified physics model (Raissi et al., 2017a; Raissi, 2018; Jiang et al., 2019) or with missing terms (Yin et al., 2021), and (c) different domain -omain-specific physical constraints such as energy conservation (Greydanus et al., 2019; Cranmer et al., 2020b) or symmetries (Wang al., 2020b; Finzi et al., 2021; Brandstetter et al., 2022a). While these PIML methods improve upon standard neural networks, damped pendulum results in Figure 3b show that they are generally not designed for OOD forecasting tasks. To study reasons for this failure, we categorize these methods into inductive and transductive based on their assumptions over ground-truth ODE parameters W*,
Transductive PIML methods. Transductive inference focuses on predicting missing parts from the observed data. In PIML, transductive inference methods focus forecasting tasks with the same dynamical system parameter, and treat a test task with a different parameter as an unrelated task. For instance, SINDy (Brunton et al., 2016), EQL (Martius & Lampert, 2016), and related methods (Raissi, 2018; Chen, 2021), learn the ODE based on a dictionary of basis functions for a specific ground-trun -ruth ODE parameter W These transductive methods (Figure 2b) do not transfer knowledge Icarnt in training to predict test examples with a different W This forces these methods to forecast simply based on the initial observations of the test task alone, often leading to poor performance. Another class of transductive methods (Raissi et al., 20;7;;;; Yu et al., 2022) assume that the ODE parameters W' remain constant across all training and test tasks, and regularize neural networks to respect given physics model. They have also been shown to be challenging to train for harder differential equations (Krishnapriyan et al., 2021) return trivial solutions (Leiteritz & Pflüger, 2021). Causal PINNs (Wang et al., 2022) ensure that predictions at any time less than are accurately resolved before predictions at time i, but do not allow for causal interventions to initial states and unknown parameters of the dynamical system. These methods will perform poorly if different training tasks have different ODE parameters.
Inductive PIML methods. Taking the opposite approach, inductive inference focuses on learning rules from the training data that can be applied to unseen test examples. Inductive methods dominate PIML approaches but are fragile OOD, since the learned rules are learned within the scope of the training data and are not guarantee to work outside the training data scope. For example, APHYNITY (Yin et al., 2021) and NDS (Mchta et al., 2021) are such inductive methods that augment neural network to known incomplete physics model where the parameters of the physies model are predicted inductively using a recurrent network. As illustrated in Figure 2c, these methods are able to learn from training tasks with different ODE parameters (+)), However, the recurrent network in APHYNITY fails OOD and often returns incorrect physics parameters OOD. Further, the augmented neural network suffers from the same extrapolation issues discussed in Section 3.1 leading to poor OOD performance as seen in Figure 3b. Concurrent work by Park et al. (2023) aligns
4

---

Published as a conference paper at ICLR 2024
with our meta-learning architecture but lacks IRM-type penalty that empirically demonstrate as crucial for identifying the true causal structure.
With these key reasons identified for the fragiliity of existing methods to OOD scenarios, propose an approach (MetaPhysiCa) that includes basis functions in the architecture for algorithmic alignment and can inductively transfer knowledge from varied training tasks to OOD test task.
PROPOSED APPROACH: METAPHYSICA
We first describe a family of structural causal models describing dynamical systems (Section 4.1), then discuss the proposed architecture along with causal structure discovery approach (Section 4.2) to learn the correct causal model for improved OOD performance.
4.1 STRUCTURAL CAUSAL MODEL
We describe the dynamical system using deterministic structural causal model (Peters al., 2022) with measurement noise over the observed states and explicitly define the assumptions over the unknown function v of Definition 1. The causal diagram is depicted in Figure 4 in the plated notation iterating over time = lo, (T( for each task i. As before, the hidden state of the dynamical system is zi
(t) IR for task We define the causal process each time step & for task follows.
Let f=(: Rd + R.I S h I m, be m linearly independent basis functions each with separate set of parameters &s acting on an input state æt Examples of such basis functions include trigonometrie functions like
X;
(a( E1) = sin(C1,1=f_) (6) + E1,2), polynomial functions like f2 (4) .2) = (.) 212. and so on. The corresponding
f
fx
W
W2
Skt
zy
outputs from these basis are shown as =k,t = fic(zz Ex) in Figure 4. The derivative daL for particular di- mension j E {1.....d) d) is only affected by a few basis function outputs =k,t (green arrows in Figure 4). These selected basis functions along with their parameters & are shared across all the tasks i = 1, 1,..., M,, i.e., the causal structure remains the same. The derivative is linear com- bination of these selected basis functions with task-specific coefficients ( and dictates the next dynamical system state. We observe the dynamical system with independent
(i)
W
Not causally reliated future states
dx,
dt
Figure 4: SCM for dynamical system defined as an unknown linear combina- tion of basis functions. We allow arbi- trary interventions over initial conditions Ity and ODE parameters W
measurement noise X == x; et where e ~ o(
(6) (1)
We assume that with unknown & and prior knowledge of which {fr ) causally influence dæ) fat. The need for basis functions stems from our analysis in Section 3.1, where we show that appropriate basis functions must be incorporated within the architecture in order to extrapolate OOD scenarios.
given the collection of m possible basis functions fa(:: Ex), 1,.... ... m,
4.2 META LEARNING & MODEL ARCHITECTURE
Given the training data (x, (4) M generated from the unknown SCM described above, our goal is two-fold: (a) discover the true underlying causal structure, i.c., learn which basis functions fa causally affect dzs.//tt for truth task-specific parameters (i) that act coefficients in linear combination of the selected basis functions. In the following, we propose a meta-learning framework that introduces structure parameters & that are shared across tasks and task-specific coefficients Wi(i) that vary across the tasks:
..., d, along with their parameters Ex. and (b) learn the ground
dx! (i)
(W
where O is the Hadamard product and
E) = (x(")) ... fm (x :m) the vector of outputs from the basis func-
tions with parameters& (E::..., Emh
5

---

Published as a conference paper at ICLR 2024
Task-specific parameters
Basis functions :0) 8
Causal structure
Trajectories from unknown dynamical system
FXX
3
X sin LE
Wilh
+
Training
Shared across tasks
R
R
Pradictions fo...., Fr
Training Loss Equation (4)
Ground truth 40.--1
Timed
fy
{0,
3
+
W(M
R
1M0 :0 e RM
Timer
Small initial trajectory at test time
Test-time adapt
Obseerve
Predict
Test
1
Ground truth f fz obtain W
2 2
Equation (5)
110
E {0, 13
Wilte) E Bdaio
(Fixed at test)
FOX!" :))
@ and w/ forecast
Figure 5: Schematic diagram of MetaPhysiCa and its training/test methodologies. We observe M trajectories in training from the same dynamical system with different initial conditions and ODE parameters. In training, Φ, denoting the causal structure, is shared among all tasks while ((() are the task-specific parameters, both learnt using Equation (4). During test, given small initial observations of the test trajectory from time 4o ... 4 we use Equation (5) to adapt W keeping the learnt causal structure & fixed. Final forecast for the test trajectory for t > t, is given by solving the ODE in Equation (3) with optimal W and $
1,..., M,
(0, 1) are the learnable parameters governing the global causal structure across tasks such that 63 = iff edge = - dæts/át exists in Figure 4,
( e Rdd are task-specific parameters that act as coefficients in a linear combination of the selected basis functions.
Figure 5 shows schematic diagram of MetaPhysiCa along with the training/test methodologies (described next). Additional implementation details are presented in Appendix C.
Next describe procedure learn the structure parameters Q. Finding whether edge exists not in the causal graph is known as the causal structure discovery problem (e.g., Heinze-Deml et al. (2018)). We use score-based causal discovery approach (e.g., Huang al. (2018)) where we assign a score to each possible causal graph. We wish to find the minimal causal structure, i.c., with the least number of edges, that also fits the training data. This balances the complexity of the causal structure with training likelihood, and avoids overfitting the training data.
Prediction loss & sparse structure. A sparse structure for & implies fewer terms in the RHS of the learnt equation for the derivatives in Equation (3). Several causal discovery approaches have been proposed that learn such minimal causal structure via continuous optimization (Zheng et al., 2018; Ng et al., 2022). We use the log-like logllikelihood of the training data with e-reguulaiizttion term to induce sparsity that is known to perform well for general causal structure discovery tasks (Zheng et al., 2018). Note that since the direction of all the edges are known (i.e., =k,t - deL./att,, we do not need the acyelicity constraints, and the causal graph is uniquely identified by its Markov equivalence class (Pearl, 2009, Chapter 2).
x where x(") x(" + f (w O F(X,) dr are the predictions obtained using an ODE solver to integrate Equation (3). In practice however, we found the squared loss directly between the predicted and estimated ground truth derivatives, i.e., RO)( Φ. r941 Eisto ||ax|0/ / 'x(") /ae l
The prediction error is given by (W((), +FE) = .:
=
leads to stable learning procedure with better accuracy in-distribution and OOD. As discussed before, we use an E lrregullarization term |/ to learn a causal structure with the fewest possible edges Zh,t - deia/dt, j 1,...,d, d, while minimizing the prediction error in training.

---

Published as a conference paper at ICLR 2024
Learning & that induces equal risk across all tasks. Our structure discovery task comes with an additional challenge as each training task is obtained with potentially different initial condition (8) Pta and ODE parameter W Existing causal discovery approaches via continuous optimization (e.g., (Zheng et al., 2018)) may not learn the correct structure. For example, if the (unknown) ground-truth weight for basis function f' is nonzero only for a small percentage of training tasks, standard approaches may learn structure without this basis function (i.e., sacrificing accuracy for a few training tasks to learn simpler model). This will impact OOD performance if many of the test tasks now have (unknown) nonzero ground-truth weight for the basis function fk-
Our goal then is learn structure that minimizes the prediction error across all training tasks simultaneously, similar to learning robust representations via invariant risk minimization-type meth- ods (Arjovsky et al., 2019; Krueger al., 2021). We use V-REx regularization (Krueger et al., 2021) that minimizes the variance of prediction errors across tasks. We empirically show in an ablation study (Table 7 in Appendix D.2) that V-REx penalty necessary to learn the true causal structure.
Meta-learning objective. Now we are ready to describe our final optimization objective. Similar to standard meta-learning objectives (Finn al., 2017; Franceschi et al., 2018; Hospedales et al., 2021), we propose a bi-level objective that optimizes the structure parameters and the global parameters in the outer-level, and the specific parameters W in the inner-level as follows
$ arg min 2,6
R (W
P. Ao 1111 + ARES Variance (W $ (.E))
S.L.Vi W
arg min R Wio
2.E).
(4)
where Åς and ÀREx are hyperparameters. We approximate the discrete structure parameters using deterministic binarization techniques (Courbariaux et al., 2015; 2016). We reparameterize $ ,,
1 (a(5).)) > 0.5) where O' E TRdxx,a({ is the sigmoid function, and the gradients are estimated via seaat-ttrough-etimmaae alternate optimization steps for (Φ, {) and (W("), J==1 in outer and inner loops respectively (Borkar, 1997; Chen al., 2021). In our experiments, jointly optimizing all parameters instead resulted in comparable performance with considerable computational benefits (Appendix D.2.3). We choose the
The bi-level optimization in Equation (4) can be approximated by
hyperparameters Åą, and AREx that result in sparsest model (i.e., with the least || ||0) while achieving validation loss within 5% of the best in-distribution validation loss. Use of in-distribution data for validation is a key requirement for OOD tasks as we do not have access to the test distribution.
In what follows, we show that MetaPhysiCa learns the true causal structure in Figure 4 given a set of assumptions described in Assumption 1 (Appendix A).
Theorem (MetaPhysiCa identifies the true causal structure). Under Assumption I (Appendix A) there exists a Àạ > 0 in the optimization objective of Equation (4) such that the optimal 4 learns the true causal structure in the SCM of Figure 4, i.e. 6 = 1 iff there exists an edge 2k,t - des../ in the true causal graph.
Proof in Appendix A shows that, given small enough Ao > O, the true causal structure Φx is the unique minimizer of the objective in Equation (4). While Theorem holds only for SCM in Figure 4, MetaPhysiCa can be extended to more expressive SCMs by composing basis functions (c.g., obtain sin (in ). MctaPhysiCa with such an expressive SCM shows OOD performance gains on a complex ODE task (Appendix D.44, but can suffer from learning stiff ODEs due the complexity of such 2-layer composition procedure. Better optimization techniques may help alleviate this problem.
4.3 TEST-TIME ADAPTATION
After learning the optimal causal structure Φ we adapt the model's 's task-specific parameters for each test task using the initial test observations available up to time f,- The test task may involve (unknown) out-of-support ground-truth parameters W((e)* and out-of-support initial conditions æt, Keeping क. & fixed, optimize the following adapt the model's task-specific parameters W(()
(tee) arg min
1
((e)
x
(5)
Will
tr
t=to
7

---

Published as a conference paper at ICLR 2024
Table I: (Epidemic model results) Test NRMSE + for different methods. NaN indicates that the model returned errors during test. MetaPhysiCa outputs 28x and 9 more robust 0OD predictions for the two OOD scenarios respectively.
Test Normalized RMSE (NRMSE) +
Methods
ID
OOD X
OOD X and W
Standard Deep Learning
NeuralODE (Chen et al., 2018)
0.005 (0.000)
1.139 (0.031)
1.073 (0.102)
Meta Learning
DyAd (Wang al., 2021b)
0.006 (0.001)
1.147 (0.044)
1.207 (0.202)
CoDA (Kirchmeyer et al., 2022)
0.004 (0.001)
1.341 (0.389)
1.090 (0.274)
Physics-informed Machine Learning
APHYNITY (Yin et al., 2021)
0.151 (0.150)
0.544 (0.249)
0.898 (0.211)
SINDy (Brunton et al., 2016)
1.999 (0.046)
2.746 (0,476)
NaN
EQL (Martius & Lampert, 2016)
NaN
NaN*
NaN
MetaPhysiCa (Ours)
0.009 (0.004)
0.019 (0.002)
0.100 (0.080)
MetaPhysiCa
NeuraIDDE
Aphynity
SINDY
DyAd
EQL
True
10.0
30.0
7.5
75
15
5.0
5.0
25
25
75
(a) in-distribution predictions
(b) Predictions under OOD X
Figure 6: (Epidemic model results) Ground truth curves (blue stars) in- and out-of-distribute along with corresponding predictions. Only MetaPhysiCa (orange) closely follows the true curve OOD,
where x = x + S (Wile) @ ()F((( dd are the predictions obtained using the optimal structure $. The final predictions for >> i, are obtained with the adapted parameters W7 and the fixed global parameters ④ f (see Figure bottom).
Two key aspects of the test-time adaptation: (a) Only the task-specific model parameters Wite are adapted whereas the meta-model & learnt during training is kept fixed, and (b) only the observations from time f, of the given test trajectory is used to adapt the parameters Wiit)) This allows the model be robust out-of-distribution ground-truth ODE parameters Wi(te).
EMPIRICAL EVALUATION
We evaluate MetaPhysiCa' in synthetic forecasting tasks based on 3 different dynamical systems (ODEs) from the literature (Yin et al., 2021: Wang et al., 2021a) adapted to our OOD scenario, namely, (i) Damped pendulum system, (ii) Predator-prey system and (iii) Epidemic model. We compare against: (a) NeuralODE (Chen et al., 2018), a deep learning method for learning ODEs, (b) DyAd (Wang et al., 2021b) (modified for ODEs), that adapts to different training tasks with a supervie encoder, (c) CoDA (Kirchmeyer et al., 2022), that learns modify its parameters to each task with a low-rank adaptation, (d) APHYNITY (Yin et al., 2021), that augments a known incomplete physics model with a neural network, (e) SINDY (Brunton et al., 2016), a transductive PIML method that uses sparse regression to learn linear coefficients over a given set of basis functions, (f) EQL (Martius & Lampert, 2016), transductive PIML method that uses basis functions within a neural network and learns sparse model (details on all models in Appendix C).
Dataset generation. For each dynamical system, we simulate the respective ODE as per Definition to generate M 1000 training tasks each observed over regularly-spaced discrete time steps t= 0 ...... by with At 0.1. For each training task i, we sample an initial condition (9) par (* and W prr W+ ). At OOD test, we generate M' = 200 test tasks by simulating the dynamical O, ...,, by with Al- 0.1 with r T/3 or T/10. For every test task
system over timesteps (J)
set initial conditions Vi and test ODE parameters W(i) outside their respective training supports.
Code available http:///iiiiiiccoo///ttttttnnnes
8

---

Published as a conference paper at ICLR 2024
PetaPh
NeuralODE
Aphynity
SINDy
DyAd
EQL
True
10.0
Predater g
Pro Prodstor
Pray
7.5
5.
23
8
10
0.0
2.5
5.0
7.5
10.0
3.0
0.0
2.5
5.0
7.
16.0
(a) In-distribution predictions
(b) Predictions under OOD X
Figure 7: (Predator-prey results) Ground truth curves (blue stars) in- and out-of-distributiin along with corresponding predictions. Only MetaPhysiCa (orange) closely follows the true curve OOD,
We consider two OOD scenarios: (a) OOD Xto when only the initial conditions are out-of-support, and (b) OOD X and W when both initial conditions and ODE parameters are out-of-support.
We consider three dynamical systems in our experiments, with 3 to 6 RHS terms in their respective differential equations: a damped pendulum system (Yin et al., 2021), predator-prey system (Wang al., 2021a), and an epidemic (SIR) model (Wang et al., 2021a), with following OOD shifts in their initial conditions respectively: acute initial angles in training to nearly vertical initial angles in OOD test, initial prey population 10x less in OOD test than in training, and initial population susceptible to a disease 10x more in OOD test than in training. For all three dynamical systems, all ODE parameters are R 1.5x higher OOD than in training (with non-overlapping support). We generate the damped pendulum dataset with 1% zero-mean Gaussian noise and the rest with no noise to show that OOD failure of baselines is unrelated to noise: existing methods fail OOD even with clean observations. Detailed description of the datasets is presented in Appendix B.
Results. We report normalized root mean squared errors (i.c., RMSE normalized with ground truth standard deviation) for the in-distribution (ID), OOD w.r.t. X and OOD w.r.. X,, and W in Tables 1, 2 and 3 (latter two in appendix) for the 3 datasets. Figures 3, 6 and 7 show example predictions in- and duttt--distribution from all models. NeuralODE, DyAd, CoDA and APHYNITY use neural network components and are able to learn the in-distribution task well with low errors. However, the corresponding errors OOD are high as they are unable to adapt to OOD initial conditions and OOD parameters. Example OOD predictions (Figures 3b, 6b and 7b) from these methods show that they have not learnt the true dynamics of the system. For example, for epidemic modeling (Figure 6b), most models predict trajectories very similar to training trajectories even though the number of susceptible individuals is 10x higher in OOD test. SINDy and EQL. being transductive, cannot use the training data and are fit on the initial test observations alone. Thus, they are unable to identify an accurate analytical equation from these few test observations, resulting in prediction issues due to stiff ODEs. MetaPhysiCa performs the best OOD across all datasets achieving 2x to 28x lower OOD NRMSE than the best baseline, and closely follows the true curve OOD.
Qualitative analysis. Appendix D.1 (Table 5) shows that MetaPhysiCa's 's performance gains stem from two factors: (i) The optimal meta-model & learns the ground truth ODE for all 3 dynamical systems, and (ii) the model adapts its task- speciffe parameters separately to each OOD test task. The former is key for robustness over OOD initial states and the latter helps to be robust over OOD parameters W Appendix 1.2.1 shows f--regularization and test-time adaptation are necessary components; OOD performance degrades significantly without either. Appendix D.2.2 shows that V-REx penalty is necessary to learn the true causal structure, without which OOD NRMSE is 24x worse. Appendix D.7 shows that MetaPhysiCa is able to learn the true dynamics for varying number of basis functions, and Appendix D.6 evaluates MetaPhysiCa without full algorithmic alignment.
6 CONCLUSIONS
We considered the out-of autoo-diistrbbutin task of forecasting a dynamical system (ODE) under new initial conditions and new ODE parameters. We showed that existing PIML methods do not perform well in these tasks and proposed MetaPhysiCa that uses a meta-learning framework to learn the causal structure for the shared dynamics across all tasks, while adapting the task-specific parameters. Results on three OOD forecasting tasks show that MetaPhysiCa is more robust with 2x to 28x reduction in OOD error compared to the best baseline. Limitations & future work: We believe that forecasting models should be robust to OOD shifts, and that our work takes step in the right direction with several potential avenues for future research: (i) Extending MetaPhysiCa to forecasting PDEs under OOD scenarios requires an expanded set of basis functions that includes differential operators, and considering OOD boundary conditions. (ii) Better optimization techniques to avoid learning stiff ODEs when extending MetaPhysiCa to more expressive SCMs.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was funded in part by the National Science Foundation (NSF) awards, CCF-1918483, CAREER IIS-1943364 and CNS-2212160, Amazon Research Award, AnalytiXIN, and the Wabash Heartland Innovation Network (WHIN), Ford, Nvidia, CISCO, and Amazon. Computing infrastruc- ture was supported in part by recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the sponsors.
(CloudBank). Any opinions, findings and conclusions or
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint urXiv:1907.0289) 2019.
Vivek S Borkar. Stochastic approximation with two time scales. Systems & Control Letters, 29(5): 291-294, 1997.
Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural pde solvers. arXiv preprint arXiv:2202.076433 2022a.
Johannes Brandstetter, Daniel E Worrall, and Max Welling. Message passing neural pde solvers. In International Conference on Learning Representations, 2022b.
Steven I.. Brunton, Joshua L.. Proctor, J. Nathan Kutz, and William Bialek. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings the National Academy of Sciences of the United States of America, 113(15):3932-33377 2016. ISSN 10916490. doi: 10.1077/pna.. 1517344113
Rick Chartrand. Numerical differentiation of noisy, nonsmooth data. International Scholarly Research Notices, 2011, 2011.
Gang Chen. Learning symbolic expressions via gumbel-max equation learner networks. arXiv preprint arXiv:2012.00921, 2020.
Gang Chen. Learning Symbolic Expressions via Gumbel-Max Equation Learner Networks. arXiv:2012.06921 [cs], May 2021.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Proceedings of the 32nd International Conference Neural Information Processing Systems, NIPS' 18, pp. 6572-6583, Red Hook, NY, USA, December 2018. Curran Associates Inc.
Tianyi Chen, Yuejiao Sun, and Wotao Yin. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. Advances in Neural Information Processing Systems, 34: 25294-25307, 2021.
Matthieu Courbariaux, Yoshua Bengio, and networks with binary weights during propagations. Advances in neural information processing systems, 28, 2015.
David. Binaryconnect: Training deep neural
Matthieu Courbariaux, Itay Hubara, Daniel Soudry, Ran Gl-yaniv, and Yoshua Bengio. Binarized neural networks: Training deep neural networks with weights and activations constrained to+ 1 or-1, arXiv preprint 1602.02600.02888 2016.
M. Cranmer, Alvaro Sanchez-Gonzalez, P. Battaglia, Rui Xu, K. Cranmer, D. Spergel, and S. Ho. Discovering Symbolic Models from Deep Learning with Inductive Biases. NeurIPS, 2020a.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2020b.
Yuchen Dang, Zheyuan Hu, Miles Cranmer, Michael Eickenberg, and Shirley Ho. Tnt: Vision transformer for turbulence simulations. arXiv preprint arXiv:2207.046166 2022.
10

---

Published as a conference paper at ICLR 2024
Arka Daw, Anuj Karpatne, William Watkins, Jordan Read, and Vipin Kumar. Physics-guided neural networks (pgnn): An application in lake temperature modeling. arXiv preprint arXiv:1710.1144 2017.
James H Faghmous and Vipin Kumar. A big data guide understanding climate change: The case for theory-guided data science. Big data, 2(3):155-163, 2014.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks. In Proceedings of the 34th International Conference on Machine Learning, pp. 1126-1135. PMLR, July 2017.
Marc Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivari- ant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine PMLR, 2021.
Learning, pp.
Luca Franceschi, Paolo Frasconi, Saverio Salzo, Ríccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-leaming. In International Conference on Machine Learning, pp. 1568-1577. PMLR, 2018.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 (11)):65-6111 673, 2020.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in neural information processing systems, 32, 2019.
Christina Heinze-Deml, Marloes H Maathuis, and Nicolai Meinshausen. Causal structure learning. Annual Review of Statistics and Its Application, 5:371-391, 2018.
Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 44(9): 5149-51699 2021.
Biwei Huang, Kun Zhang, Yizhu Lin, Bernhard Schölkopf, and Clark Glymour. Generalized Score Functions for Causal Discovery. KDD proceedings. International Conference on Knowledge Discovery & Data Mining, 2018:1551 1560, August 2018. ISSN 2154-817X. doi: 10.1145/
3219819.3220104.
Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normal- ization. In Proceedings of the IEEE international conference on computer vision, pp. 1501-1510, 2017.
Chiyu "Max" Jiang, Karthik Kashinath, Prabhat, and Philip Marcus. Enforcing Physical Constraints in Neural Neural Networks through Differentiable PDE Layer. September 2019.
Matthieu Kirchmeyer, Yuan Yin, Jeremie Dona, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing New Physical Systems via Context-Informed Dynamics Model. In Proceedings of the 39th International Conference on Machine Learning, pp. 1128--11301. PMLR, June 2022.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning-aceen accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21), 2021.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes in physics-informed neural networks. Advances in Neural Information Processing Systems, 34:26548-26560, 2021.
David Krueger, Ethan Caballero, Hbern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distriuution generalization via risk extrapola- tion (rex). In International Conference on Machine Learning, pp. 5815-5826. PMLR, 2021.
Raphael Leiteritz and Dirk Pflüger. How to Avoid Trivial Solutions in Physics-Inforaa Informed Neural Networks. arXiv: /2 2,0660 (cs, stat), December 2021.

---

Published as a conference paper at ICLR 2024
Julia Ling, Andrew Kurzawski, and Jeremy Templeton. Reynolds averaged turbulence modelling using deep neural networks with embedded invariance. Journal of Fluid Mechanics, 107:155-166, November 2016. ISSN 0022- 120, 1469-7645. doi: 10.1017/jfm.2066.611
Bethany Lusch, Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature communications, 91)::-11,, 2018.
Georg Martius and Christoph H. Lampert. Extrapolation and learning equations. arXiv:1610.02995 [cs], October 2016.
Viraj Mehta, lan Char, Willie Neiswanger, Youngseog Chung, Andrew Nelson, Mark Boyer, Egemen Kolemen, and Jeff Schneider. Neural dynamical systems: Balancing structure and flexibility in physical prediction. In 2021 60th IEEE Conference on Decision and Control (CDC), pp.
3735-3742. IEEE, 2021.
Ignavier Ng. Shengyu Zhu, Zhuangyan Fang, Haoyang Li, Zhitang Chen, and Jun Wang. Masked gradient-based causal structure learning. In Proceedings of the 2022 SIAM International Conference on Data Mining (SDM), pp. 424-432. SIAM, 2022.
MoonJeong Park, Youngbin Choi, Namhoon Lee, and Dongwoo Kim. Spreme: Sparse regression for multi-environment dynamic systems. In When Machine Learning meets Dynamical Systems: Theorr and Applications workshop, AAAI, 2023.
Judea Pearl. Causality. Cambridge university press, 2009.
Jonas Peters, Stefan Bauer, and Niklas Pfister. Causal models for dynamical systems. In Probabilistic and Causal Inference: The Works of Judea Pearl, pp. 671-690. 2022.
Maziar Raissi. Deep hidden physics models: Deep learning of nonlinear partial differential equations. The Journal of Machine Learning Research, 19(1):932-955, 2018
Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Physics informed deep learning (part i): Data-driven solutions of nonlinear partial differential equations. arXiv preprint 2017a.
Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Physics informed deep learning (part ii): Data-driven discovery of nonlinear partial differential equations. arXiv preprint arXiv: 2017b.
Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Multistep neural networks for data- driven discovery of nonlinear dynamical systems. arXiv preprint arXiv:1801.01233 2018.
Paul K Rubenstein, Stephan Bongers, Bernhard Schölkopf, and Joris Mooij. From deterministic odes to dynamic structural causal models. arXiv preprint 1608.:1608.0000 08028, 2016.
Leonid Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: nonlinear phenomena, 60(1-4):259-26.) 1992.
Michael Schmidt and Hod Lipson. Distilling free-form natural laws from experimental data. Science, 2009.
Sebastian Thrun and Lorien Pratt. Learning to learn: Introduction and overview. In Learning to learn, pp. 3-17. Springer, 1998.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1457-1466, 2020a.
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arXiv preprint arXiv:2002.0300 2020b.
Rui Wang, Danielle Maddix, Christos Faloutsos, Yuyang Wang, and Rose Yu. Bridging physics-based based and data-driven modeling for learning dynamical systems. In Learning for Dynamics and Control, pp. 385-398. PMLR, 2021a.
12

---

Published as a conference paper at ICLR 2024
Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. arXiv preprint arXiv:2102.1027,, 2021b.
Sifan Wang. Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks. arXiv preprint S2KV::2203.074044, 2022.
Jared Willard, Xiaowei Jia, Shaoming Xu, Michael Steinbach, and Vipin Kumar. Integrating physies- based modeling with machine learning: A survey. arXiv preprint arXiv:2003.04919 2020.
SHI Xingjian, Zhourong Chen, Hao Wang, Dit- Yan Yeung, Wii-Kin Wong, and Wang-chun Woo. Convolutional Istm network: A machine learning approach for precipitation nowcasting. In Advances in neural information processing systems, pp. 802-810, 2015.
Keyulu Xu, Mozhi Zhang. Jingling Li, Simon Shaolei Du, Ken-Ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference on Learning Representations, 2021.
Alireza Yazdani, Lu Lu, Maziar Raissi, and George Em Karniadakis. Systems biology informed deep learning for inferring parameters and hidden dynamics. PLoS computational biology, 16(11): e1007575, 2020.
Kyongmin Yeo and Igor Melnyk. Deep learning algorithm for data-driven simulation of noisy dynamical system. Journal of Computational Physics, 376:1212-1231, January 2019. ISSN 00219991. doi: 10.1016/J.ep..011 10.024.
Yuan Yin, Le Vincent, DONA Jérémie, Emmanuel de Bezenac, Ibrahim Ayed, Nicolas Thome, et Augmenting physical models with deep networks for complex dynamics forecasting. In International Conference on Learning Representations, 2021.
Jeremy Yu, Lu Lu, Xuhui Meng, and George Em Karniadakis. Gradient-enhanced physics-infone -nnormee neural networks for forward and inverse pde problems. Computer Methods in Applied Mechanics and Engineering. 393:1 14823, 2022.
Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing. Dags with tears: Continuous optimization for structure learning. Advances in Neural Information Processing Systems, 31, 2018.
13

---

Published as a conference paper at ICLR 2024
Supplementary Material of "MetaPhysiCa: Improving OOD Robustness in Physics-informed Machine Learning"
A PROOF OF THEOREM
Assumption 1. We are given M tasks generated from the structural causal model in Figure 4 with infinite observations per task, i.., T - oc, Vi (L.....MJ. Then, for all p E {l..... ......d) d} and every basis function fE active in the true (unknown) casual model with $ p.k = 1, we assume there exists v > O and at least one task j (I,...,,M such that,
-Ww p.k ;x(=))))))) for all and
"Ik is linearly independent from all other basis functions in the domain of task j.
where W is the (unknown) ground truth parameter corresponding to the basis function fx to p.k. predict ! tp//aa at.
Theorem (MetaPhysiCa identifies the true causal structure). Under Assumption (Appendix A) there exists a XØ > the optimization objective of Equation (4) such that the optimal & learns the true causal structure in the SCM of Figure 4, i.e,, D 1 iff there exists an edge Zh,t - diee./ at in the true causal graph
Proof. In what follows, we prove the statement for 1- I-dimensional trajectories, i.e., X, E Rd =1 The proof can be trivially extended for d 1. Our optimization objective
M
क = arg min L(Φ) = arg min
Σ
R (W)
10110111,
(6)
M
where
arg min R (WO, Φ, denotes the optimal task-specific parameter given a structure 4. We use the notation W (Φ) explicitly denote the dependence of the optimal task-specific parameter given structure Φ.
7 (W E = E | / ax(" /azl/ /a. with the predictions from the
proposed model given by
ax!1) de
(W E FX{
We assume that the parameters & the basis functions are known.
We set ARES = 0.
Let the true unknown causal structure be $. Note that from Assumption 40 has the property that 2 1 if and only if there exists task j E {I...., M) such that the ground truth parameters () # O, otherwise i 0.
The inner optimization (i = arg minw R (i),,, 22, is a linear regression problem. (1) 'T()
can define an ( design matrix F()) such that F Given the data for task i, (x)) tto
[a(2(()) Ex) for 11, b,, (T()- Also, define -enghhh vector (8)) of noisy ground truth derivatives where (1) = dxl /dl + Et, for (= .,..., (10) and e, is zero-mean Gaussian noise related to the noise in the ground truth observations. Then, with (
(+)) (()) for all i E M....,,). i.e., if we knew the true causal structure, the corresponding task-specific parameters converge to the ground truth parameters (i).
Assume & * $², i.e. the optimal structure found Equation (6) not the true causal structure. In the following, will show that this leads contradiction. For any b define the function b( (ʃx 11, 1) denoting the set of basis functions that active in T. Consider b( b(Φ) and
14

---

Published as a conference paper at ICLR 2024
b( IC b()) n (()²) denotes the basis functions that shared by the true causal structure and क. Then essentially, the basis functions in b((**) \ K are replaced by b(Φ) \ A. (at least of these sets is non-empty).
Case 1: 1b( K| = 0. i.e., all basis functions in the true causal structure are present in & but $ contains at least one additional basis function.
Then, with ((() + co, the optimal error R (W ( +..) {) 2 (W -- .) with equality achieved, for example, when for all k K. we have (l) ( and for all k' b(() /C, we have (i)()) ( O. Furthermore, || | = 11)) [b(() \ K0| >
Thus, for any d4 0,
M
((
1
Σ
R (W
(4), σ) 1
1 M
Σ
R (W))
E") =
[((*).
M
This is contradiction as is the optimal solution of Equation (6). Case 2: [b(2)) |K] > 0. There exists a task E (I,..., .M such that ( o T) FO) + (Wi() o Φ*) FÚ)
where F() the design matrix corresponding to task j (defined above). W (i(
1, fx (xx;i) >
for some k E b((*)) K. and fa cannot be written as a linear combination of other basis functions in the domain of task j. There always exists least one such task j, otherwise Assumption is violated. Let m)
(3) be the difference between LHS and RHS of Equation (7), then |n]")] > v/ for some fixed
d >> 0 and all l.
For this task j, the optimal error as T()) + 00,
R ($)) .) = Eg
(w
F y/- - -()
(
F()) - Yr - Ee + (111) 27
Ex
(a (a--" m 27
Eg
E. + Ee Ee(m 1)2
(Belea) 0)
=o² +
=
for some e > Then,
M
M
1
Σ
R
(W
(W 10 (0),
Σ R (w") 2- > M
1
e M
(8)
M
and ||0111
1 (in the extreme case, when Ø' contains all basis functions, but +b uses single basis function). Then choosing any 0 < da Im- we can rewrite Equation (8) as
|
1 M
M
(w
(0)
Σ
e M
M
M
111
Xo (m - 1)
|(011))
Rearranging the terms, we get L(Φ) > L(("), leading to contradiction.
In both cases above, we have shown there exists small enough As 0 such that the true causal structure 4x is the unique minimizer of L()).
15

---

Published as a conference paper at ICLR 2024
B DESCRIPTION OF TASKS
For each dynamical system, we simulate the respective ODE to generate M 1000 training tasks each observed over regularly-spaced discrete time steps (lo...., (T) where Vl, 4 = 0.11. Our data generation process is succinctly depicted in Table 4. For each dataset, the second column shows the state variables X, and the unknown parameters W For each training task T 1...., M, sample an initial condition x(() ~ (X.) (shown under ID columns of the table). We sample a different W ~ ~U(W param: 2W param) for each task i with W param shown in Table 4.
At test, generate M' 200 test tasks by simulating the respective dynamical system over timesteps (lu...,,<7 (,,,, where again Vl, = 0.1l. For each test task j = 1, ,,,.., M', we sample initial conditions X pit)(((,))) with completely different support for the initial conditions X!6 than in training. The distribution of the dynamical system parameters W is kept the same for "OOD Xto scenario but shifted for "OOD Xt, and W scenario. In the latter, we sample a different W U(W param: 3 param) for each test task j with param shown in Table 4.
Damped pendulum system (Yin et al., 2021). The state X, = [0,, we] E I² describes the angle made by the pendulum with the vertical and the corresponding angular velocity at time t. The true (unknown) function v describing this dynamical system is given by de, de - w/: dt
d -a 22 sin(0,)
p" or where W (a", p²) are the dynamical system parameters. We simulate the ODE over time steps (to:..... br] } with VI, tr = 0.11, T = 100 in training and over time steps {to..... t, } in test with ,T. In training, the pendulum is dropped from initial angles e/ U(0, 7/2) with no angular velocity, whereas in OOD test, the pendulum is dropped from initial angles 0/0 ~ U(T 0.1, rT) and angular velocity w/6) E ((-1,,)).
Predator-pre prey system (Wang et al., 2021a). We wish to model the dynamics between two species acting as prey and predator respectively. We adapt the experiment by Wang et al. (2021a) to our out-of- distribution forecasting scenario according to Definition 1. Let p and q denote the prey and predator populations respectively. The ordinary differential equations describing the dynamical system is given dg by de = a p- 3 pq. सt system parameters. We simulate the ODE over time steps {oo............... ..,,..,,,,]...... tT} with Vl, t = 0.11, T 100 in training and over time steps {fo..... } in test with r T.. We generate M = 1000 training tasks
= 8* pq >"q where W = (ax 3"* ,"",,,""""""""""""","" -*)) are the (unknown) dynamical
...,....
with different initial prey and predator populations with prey Pio ~ U(1000, 2000) and predator
Tto (10, 20) for each 1...., M.. At OOD test, we generate M' 200 dis---iiittioottion
(2)
()) (OOD) test tasks with different initial prey populations pia U(100 200) but the same distribution for predator population
Uta ()) ((10,20))
Epidemic modeling (Wang al., 2021a). We adapt the experiment by Wang et al. (2021a) to our out-of-distribution forecasting scenario according to Definition 1. The state of the dynamical system described by three variables: number of susceptible (5), infected (T) and recovered (R) individuals. di 3 71, dR 71, where W = (8,7) are the (unknown) dynamical system parameters and N S R is the total population. We simulate the ODE over time steps (to, ... &r with VI, = 0.11, 100 in training and over time steps = 10" T. We generate M = 1000 training tasks with different initial populations for susceptible (S) and infected (I) individuals, while the number of initial recovered
The dynamics is described using the following ODEs: dt - -3$% N di
ds
(R) individuals are always zero. In training, we sample s(0) ~ (9,9 10) and ( ~ U(11 5) for each = 1, M. At OOD test, we generate M' = 200 duitr--dist test tasks with different initial susceptible population, s(3) (90, 100), while keeping the same distribution for infected population.
C IMPLEMENTATION DETAILS
In what follows, we describe implementation details of MetaPhysiCa and the baselines.
16

---

Published as a conference paper at ICLR 2024
Table 2: (Damped pendulum results) Normalized RMSE L of test predictions from different methods in-distribution and two OOD scenarios. NaN" indicates that the model returned errors during test- time predictions, for example, because the learnt ODE was stiff (numerically unstable) solve.
Test NRMSE 1
Methods
ID
OOD X
OOD X and W
Standard Deep Learning
NeuralODE, (Chen al., 2018)
0.083 (0.033)
0.591 (0.119)
0.717 (0.210)
Meta Learning
DyAd (Wang et al., 2021b)
0.078 (0.051)
0.834 (0.263)
0.804 (0.267)
CoDA (Kirchmeyer et al., 2022)
0.052 (0.032)
0.764 (0.201)
1,011 (0.226)
Physics-informed Machine Learning
APHYNITY (Yin al., 2021)
0.097 (0.020) NaN*
0.970 (0.384)
1.159 (0.334)
SINDY (Brunton et al., 2016)
NaN*
NaN*
EQL (Martius & Lampert, 2016)
NaN
NaN
NaN
MetaPhysiCa (ours)
0.049 (0.002)
0.070 (0.011)
0.181 (0.012)
Table 3: (Predator-pree prey results) Test NRMSE | for different methods. NaN indicates that the model returned eтrors during test. Standard deep learning methods and physics-informed learning methods fail to forecast accurately out-o-diistributin.. MetaPhysiCa outputs 8 and 2x more robust OOD predictions in the two OOD scenarios respectively.
informed deep
Test Normalized RMSE (NRMSE) 4
Methods
ID
OOD X
OOD X and W
Standard Deep Learning
NeuralODE (Chen et al., 2018)
0.193 (0.024)
1.056 (0.141)
0.969 (0.172)
Meta Learning
Dy (Wang al., 2021b)
0.244 (0.025) NaN"
1.088 (0.373) NaN*
1.025 (0.403) NaN"
CoDA (Kirchmeyer et al., 2022)
Physics-ifformed Machine Learning
APHYNITY (Yin et al., 2021)
0.421 (0.332) N
3.937 (1.686) NaN*
1.281 (0.457)
SINDy (Brunton et al., 2016)
N
EQL (Martius & Lampert, 2016)
-
NaN'
N
MetaPhysiCa (Ours)
0.049 (0.008)
0.129 (0.030)
0.434 (0.128)
17

---

Published as a conference paper at ICLR 2024
Dalasets
State variables
ID
OOD X
OOD X,, and W
0o ~ ((0, x/2) = 0
Ou ~ 1 0.1, T)
06~ (T 0.1,
X, (4)
Damped pendulum
we ~ 1,
~ 21 (1,0)
W (a,
puram 1, Pparan
0.2
Po 4(1000, 2000)
No 24(100, 200)
Po ((00 200)
X, = (pc, 26)
Predator prey system
9o (10, 20)
To ((10, 20)
lo (10, 20)
W 3,, 7:6)
puuar
5
pamam
0.06, "/puram
0.5, 6pram 0.0005
So ~ (1,9, 10)
So (900 100)
So ~ (900 100)
X, = (St, I,, R.)
In ~ ((155)
1o ~ (1(15) R.
Io ~ (1155)
Epidemic modeling
Ro
Ro 0
W (3.7)
Jaram 4, Tparam 0.4
Table 4: Description of the dataset generation process. For each dataset, X₂ denotes the state variable of the dynamical system and W denotes its parameters. Column "ID" represents in-distribution initial states while the last two columns represent the two out-of aut-of--ditribution scenarios. In-distribution ODE parameters W sampled from uniform distribution W U(W param: 2W param)
and the out-of-distiibution ODE parameters are sampled as ( U2W param; 3W (aram)). For example, in the damped pendulum dataset, in-distribution parameters are sampled as (1) U(Uppram: parmm)) = 21, and p(0)+ (Ppprrmm 2ppurmm)
(0.2, 0.4) for each task i. Sim- ilarly, the out-of-distribution ODE parameters (in the last column) are sampled as ali)- (20ppamm 3pmmmm)) = (2, 3) and p(a> (2ppprrmmm 3p = (0.4, 0.6).
C.I METAPHYSICA
Recall from Equation (3) that the proposed model is defined as
dx
(w -))
(9)
where 2 is the Hadamard product and
F F(X 5) = [x(
Sm (x(1 Em) {
the vector of outputs from the basis func-
tions with parameters E
(0, 1) 1 are the learnable parameters governing the global causal structure across tasks such that 6,,x 1 iff edge = - dæt./dt exists,
( Rd are task-specific parameters that act as coefficients in linear combination of the selected basis functions.
In our experiments, we use polynomial and trigonometric basis functions. such that
F(X) EE
(02 1.1
(02 t.d
10 + E1,2) sin (() + Ed,2) X10 sin(Ei,1
polynomial order
polynomial order
trigonometric
Equation (4) describes bi-level objective that optimizes the structure parameters & and the global parameters & in the outer-level, and the task-specific parameters W in the inner-level as follows
$ غ arg min 3.5
R
W
No |1 + ARES Variance (( 2.
(W M)
s.t. W (4)
arg min R (W (0) (0,4, Vi = 1...,,M, M W6
18

---

Published as a conference paper at ICLR 2024
where Ay. and ARES hyperparameters. As discussed in the main text, the jointly optimizing $ E and W i ...... M, instead of alternating SGI) resulted in comparable performance with considerable computational benefits. We use the following joint optimization objective to approximate Equation (4),
M
$ W
((M)
arg min
P + Ao 1111111 W
(10)
e W. W(M M
1=1
+ AREx Variance (( (w 2, ())
We perform a grid search over the following hyperparameters: regularization strengths Ao (10-1 10-3, 10-3, 10-2), AREs E (0,10 10-1001111 10), and learning rates 1 {10-2, 1003, 1o--s10044 10-4). We use a linearly increasing scheduler for AREs to avoid the V-REx penalty dominating the loss during initial epochs. We choose the hyperparameters that result in sparsest model (i.e., with the least [|0) while achieving validation loss within 5% of the best validation loss in held-out in-distribution validation data.
Time complexity. For a single task i and time t, computing the predicted derivatives dx;" /tt in Equation (3) has ((md) complexity evaluate the Hadamard product and the matrix-vector product, where m is the number of basis functions and is the state dimension of the dynamical system. Since there are M tusks/trajectories of maximum length T, the overall complexity per epoch is O(ndMT). However MetaPhysiCa could require a higher number of epochs to converge for higher values of m (Figure 9).
C.2 NEURALODE (CHEN AL., 2018)
The prediction dynamics corresponding to the latent NeuralODE model is given by 4x. Fm. (xr,
dt zg; Wi) where zc, Fene (X,,..... X,,; W2) encodes the initial observations using a recurrent neural network Fenc (e.g., GRU), and Fm is a feedforward neural network. The model trained with ODE solver (dopri5) and the gradients computed using the adjoint method (Chen et al., 2018). We perform a grid search over the following hyperparameters: number of layers for Fm L {1, 2, 3), size of each hidden layer of Pon. dn {32, 64, 128), size of the encoder representation Zgr. d2 E {32, 64, 128]. J. batch sizes B E (32, 64), and learning rates 17 E {10-3 2, 1000 10 10-4).
C.3 DYAD (MODIFIED FOR ODES) (WANG ET AL., 2021B)
DyAd, originally proposed for forecasting PDEs, uses a meta-learning framework to adapt to different training tasks by learning per-task weak label. We modify their approach for our ODE-based experiments. Since we do not assume the presence of weak labels for supervision for adaptation, mean of each variable in the training task as the task's weak label. We use NeuralODE as the base sequence model for the forecaster network. The forecaster network takes the initial observations input and forecasts the future observations while being adapted with the encoder network. The encoder network is a recurrent network (GRU in our experiments) that takes as input the initial observations and predicts the weak label. The last layer representation from the encoder network is used to adapt NeuralODE via AdaIN (Huang & Belongie, 2017). We perform a grid search over the following hyperparameters: size of hidden layers for the forecaster and encoder networks dj E (32, 64, 128), number of layers for the forecaster network, L E {1, 3}, batch sizes B E (32, 64), and learning rates r) E {10-3 10 10-3, 10-4).
C.4 APHYNITY (YIN ET AL., 2021)
APHYNITY assumes that are given a (possibly incomplete) physics model d(,, Ophy with parameters Ophy- When the training data may consist of tasks with different ((),, APHYNITY predicts the physics parameters with respect to the task i inductively using a recurrent neural network Gm from the initial observations of the system as (() G.n (Xi,...., X,, Wz). Then, APHYNITY augments the given physics model with feedforward neural network component F... (t)
and defines the final dynamics
dx!"
Fin (x(00 Wi). APHYNITY solves phy
19

---

Published as a conference paper at ICLR 2024
constrained optimization problem minimize the norm of the neural network component while still predicting the training trajectories accurately. The model is trained with an ODE solver (dopri5) and the gradients computed using the adjoint method (Chen et al., 2018). In our experiments, provide APHYNITY with simpler physics models:
For damped pendulum system, use physics model that assumes no friction:
dē, dt
Wes
duis
-a sin(Ot) where pphy a-yyy is the physies model parameter.
For predator-prey system, we use a physics model that assumes no interaction between the two dp species: at = Ophy?, a
da 7phy? where Ppy (Aphy, (phy) are the physics model parameters.
For epidemic model, we use physics model that assumes the disease not infectious: as
de
df 0,
, II, where Phyy "phy is the physics model parameter.
dR
In each dataset, A APHYNITY needs augment the physics model with neural network component for accurate predictions.
We perform grid search over the following hyperparameters: number of layers for Fmn- L E {1, 2, 3). size of each hidden layer Fons dn. {32, 64, 128), batch sizes B E {32, 64), and learning rates E {10-2, 10-3, 10
C.5 SINDY (BRUNTON ET AL., 2016)
SINDy uses given dictionary of basis functions to model the dynamics as 4x O(,) W where + is feature map with the basis functions (such as polynomial and trigonometric functions) and W is simply weight matrix. SINDy trained using sequential threshold least squares (STLS) for sparse weights W. We perform grid search over the following byperparameters: threshold parameter used STLS optimization, To E (0.005, 0.01, 0.05, 0.1, 0.2, 0.5), and the regularization strength (0.05, 0.1,
C.6 EQUATION LEARNER (MARTIUS & LAMPERT, 2016)
Equation learner (EQL.) is neural network architecture where each layer defined as follows with input æ and output
Wx +
o ((()))) (zk, Ze4 14, 2843),....).
where fi are unary basis functions (such as sin, cos, etc.) and are binary basis functions (such as multiplication), We use id, sin and multiplication functions in our implementation. EQL is trained using sparsity inducing & érregularizztton with hard thresholding for the final few epochs. We perform a grid search over the following hyperparameters: number EQL lay- ers, L {1,2), number of nodes for each type of basis function, h 31,3, 5}. regulariza- tion strength a E {10-1, 10 10-3 10 10010-pp- $3. batch sizes B {32, 64), and learning rates n 110-2,10-3 10-1).
D ADDITIONAL RESULTS
D. QUALITATIVE ANALYSIS
Recall from Equation (3) that the proposed model is defined as
xx! (w o F)F(X (E),
(II)
dt
where F(X(() (1) {) is the vector of outputs from the basis functions, E {0, 1]dxd are the learnable parameters governing the global causal structure across all tasks, and (i e ddxm are tpsk-specific parameters that act coefficients in linear combination of the selected basis functions.
20

---

Published as a conference paper at ICLR 2024
Dalasets
State variables
Ground truth ODE
Learnt ODE (Irom 4)
d0, di Wiw
Damped pendulum
X = (8,, ws)
del
-a (in((nn())) p'w
del
W2 sin(O:) + W3as
dt
dt
Pt P.
d2.
Wipe Wapeo
Predator prey system
X, = -pp,,2))
d 8 P19: -44
di WaPs9: + Watt
d5. -3"
S.L.
45. W1ST
dt
Epidemic modeling
X, = (Sr, I1, Rt)
8 + 7
W2S,7; + W3 W.J.R.
dR: It
dR.
W&S: + Wal? + W,I.R
Table 5:
(Qualitative analysis.) Ground truth dynamical system vs learnt ODE the meta-model + Recall that & E {0, |dxm dictates which of the basis functions affect the output <X./at. The weights W; the learnt ODE column are learnable parameters that are optimized via test-time adaptation in Equation (5). MetaPhysiCa learns the exact ground truth ODE for Damped pendulum and Predator-pres prrttorrrpryyyyee system, and reparameterized version of the true ODE for epidemic modeling task.
After training, the ODE learnt by the model can be easily inferred by checking all the terms in & that are greater than zero, i.e., $ > 0 implies S& (æ:;; Ex) - dtt exists in the causal graph. In other words, RHS of learnt ODE for dms.s/tt contains the basis function [a( Ex).
Table shows the ground truth ODE and the learnt ODE for the three experiments. For each learnt ODE, we also depict the learnable parameters Wi that can be adapted using Equation (5) during test-time. For damped pendulum and predator-prey system, the RHS terms in the learnt ODE exactly matches ground truth ODE, and from Figures and 7, is clear that the method is able to accurately adapt the learnable parameters Wi during test-time. For epidemic modeling task, MetaPhysiCa learns reparameterized version of the ground truth ODE. For example, MetaPhysiCa learns dR. W + W3 + W which can be written as dR, Wal, (the ground dt
of truth ODE) if W W = W because S, + 4 + R. = N is constant denoting the total population. While the learnt reparameterized ODE is more complex because it allows different values for W Wy W the test-time adaptation of these learnable parameters with the initial test observations results in them taking the same values.
D.2 ABLATION RESULTS
D.2.1 ALL COMPONENTS
We present ablation study comparing different components of MetaPhysiCa in Table 6. Table shows -uitof--isttibution test NRMSE for MetaPhysiCa without each individual component on the three dynamical systems (OOD w.r.t Xt,). We observe that sparsity regularization (i.e., ||@]]) and test-time adaptation are the most important components. For two out of three tasks, the method returns prediction errors without sparsity regularization.
When testing MetaPhysiCa without test-time adaptation, we simply use the mean of the task-specific weights learnt for training tasks as the task-specifer weight for the given test trajectory, i.e., M+1+ M Σ, W This results in high OOD errors showing the importance of test-time adaptation. V-REx penalty (Krueger al., 2021) helps in experiments and performs comparably in others.
D.2.2 V-REX PENALTY
We repeat the damped pendulum experiment described in Appendix B with modified training distributions. Recall that the true function vb describing this dynamical system is given by
de, = Wi, de dt
dals a2 sin (0.) pp"w, where p² is the damping coefficient. In training, the pendulum
dropped from initial angles o U0, π a/2) with no angular velocity, whereas in OOD test, the pendulum is dropped from initial angles 0(6) ((T 0.1, :T)
21

---

Published as a conference paper at ICLR 2024
Test Normalized RMSE (OOD X
Method
Damped Pendulum
Predator- Prey
Epidemic Modeling
MetaPhysiCa
0.070 (0.011)
0.129 (0.030)
0.019 (0.002)
without ||
NaN*
1.806 (0.736)
NaN
without test-time adaptation
1.223 (0.741)
1.404 (3.794)
0.358 (0.554)
without V-REx penalty
0.070 (0.014)
0.129 (0.030)
0.042 (0.065)
Table 6: (Ablation.)
test NRMSE for MetaPhysiCa without each individual
component the three dynamical systems (OOD w.r.t X alone). Sparsity regularization (i.e., ||0] ()) and test-time adaptation are the most important components.
Table 7: (Damped pendulum.) True casual model for the damped pendulum is given by d0,//t wi; / = aa "2 sin (0,) - p' where p" is the damping coefficient. In training, we have (pt) > 0) 10-2, i.e., most training tasks have damping. In OOI) test, with arbitrary interventions allowed on p", we consider (p > 0) = 0.5. MetaPhysiCa without V-REx penalty sacrifices performance on 1% of training tasks and learns the incorrect ODE with no damping. With V-Rex penalty, the risks of all tasks are forced to be close to equal, and the model learns the true causal structure.
Test Normalized RMSE
Method
Learnt +
ID OOD X and W
MetaPhysiCa
without V-REx penalty
de,
Wiw
0.053 (0.009)
0.829 (0.106)
gt
W2 (in
dt
with V-REx penalty
d9
W
0.049 (0.008)
0.034 (0.007)
de
W2 sin(
Wyws
In training, we sample the damping coefficient p² from a mixture distribution p(t) such that pt = 0 with probability 0.99 and p² ~ (0(..,, 0.4) with probability 0.01. This means that very few training tasks have damping, but the model should still learn the true causal structure with the damping term.
In OOD test, with arbitrary interventions allowed on p², we consider (( > 0) 0.5.
Table 7 shows that MetaPhysiCa without V-REx penalty sacrifices performance on 1% of training tasks and learns the incorrect ODE with no damping. On the other hand, MetaPhysiCa with V-REx penalty forces the risks of all tasks to be close to equal (including the 1% tasks with damping), and thus learns the true causal structure.
D.2.3 BI-LEVEL OPTIMIZATION JOINT OPTIMIZATION
Recall that Equation (4) describes a bi-level objective that optimizes the structure parameters & and the global parameters & in the outer-level, and the task-speeific parameters ( in the inner-level as follows
M
$ غ arg min
Σ
2) + X0 + ARES Variance(( ( R (W
M
2
s.t. (
min R (W +)E) Vi = I,.,,
W
where da and AREs are hyperparameters.
We evaluate two ways to optimize the above: (a) bi-level optimization using inner and outer loops, and (b) joint optimization of all variables together.
Bi-level optimization. We use two loops where the outer loop optimizes the global parameters $ and f, and the inner loop optimizes the task-specific parameters (i+ for all while keeping
22

---

Published as a conference paper at ICLR 2024
1.75
NevsraDt
ReuraGDE
1.50
Ryed
APHINTY
2.0
RyAc APHNVIT
1.25
Mauutya
1.5
NRMSE
1.00
NRMSE
0.75
1.0
0.50
0.25
0.5
0.00
0.0
10
Percentage (%) of Gaussian noise
Percentage (%) of Gaussian noise
(a) Damped pendulum task
(b) Predator-prey system
Figure 8: (Performance with increasing noise.) Out-of-distribution NRMSE values for Damped Pendulum and Predator-prey pproy experiments with different percentages of Gaussian noise added (0%, 1%, 5%, 10%). MetaPhysiCa is relatively robust to < 5% Gaussian noise and outper- forms the baselines. With a larger amount noise, MetaPhysiCa is unable to identify the dynamical system accurately but performs comparable to the baselines.
the global parameters fixed. Rather than running the inner loop till convergence, only perform a few gradient steps S of the inner optimization (S in our experiments). This is standard meta-learning technique to improve computation complexity.
Joint optimization.
We use the following joint optimization objective to approximate Equation (4),
M
Φ E (1)
(M))
arg min
R()
(W + + X 1111
(12)
&. (())) W(M) M
1=1
+ ARES Variance (f ((.()) (W D 1)
Both bi-level optimization with S = 10 and joint optimization in Equation (12) result 6 learning the true dynamics for all tested dynamical systems. However, joint optimization is 8.2x faster per epoch than the bi-level optimization (taking 90ms vs 744ms on average on one Intel(R) Xeon(R) CPU core).
D.3 ROBUSTNESS TO NOISE
We repeat the Damped pendulum and Predator- prey experiments with increasing amounts of noises. Specifically, add 1%, 5% and 10% Gaussian noise all the trajectories, both in training and in test. We use Total Variation Regularization (TVR) (Rudin et al., 1992; Chartrand, 2011) for estimating derivatives from noisy data as done by Brunton et al. (2016). We report the normalized RMSE for different models trained the noisy versions of data in Figure 8. SINDy and EQL are not shown as they returned errors during test-time predictions similar to the case with no noise because the learnt ODE too stiff (numerically unstable) to solve. In both tasks, the proposed method is relatively robust to small amounts of noise and outperforms the baselines. With 10% noise, MetaPhysiCa is unable to identify the dynamical system accurately, but performs comparable the baselines.
D.4 COMPLEX ODE TASK
In this section, we extend MetaPhysiCa to consider significantly more expressive structural causal models (compared to Figure 4) that allow for composition of the basis functions. This is achieved with a 2-layer learnable basis function composition procedure. For example, given basis func- 2 1,,,, one can construct more ex- pressive basis functions with compositions: fa( E3) = sin(Es.3 sin(fa,1X1,1 + Ea,2) + Ea,4). fa {1) = tt, sin(E1,1 R + E1,2), etc., where are global parameters that remain con- stant for all training/test tasks. The rest of the SCM remains the same and the derivative dx ds for particular dimension j E {1...., ...... d} is a sparse linear combination of the original basis functions and the more expressive second layer ones.
tions fi (xx;&)
sin(f1,101,1 + {,22), and fe(
23

---

Published as a conference paper at ICLR 2024
Test Normalized RMSE (NRMSE) +
Methods
ID
OOD X
NeuralODE (Chen et al., 2018) APHYNITY (Yin et al., 2021)
0.034 (0.008)
0.296 (0.064)
0,027 (0.010)
0.684 (0.117)
MetaPhysiCa (Ours)
0.188 (0.035)
0.203 (0.046)*
Table 8: Test NRMSE for different methods.
indicates that the method returned errors during
predictions due to learning stiff ODE.
Initial conditions X
ODE parameters W
Learnt +
ID-1
0g ~ ((0 x/))
a' a',, ~1(0.2.0.4)
ID-2
0g ~
a' "'= ~l(0.2,0000
dd, Wiw: di (2ssin(n)) Wawe
ID-3
~ ~00, 7/2), ~ 1))
a² 21"((,,)))'~~0(((22,,11
ID-4
8g ~ ((0, (T(0, =
a² ~²,1(223),, ~2((0,))
ID-5
8g ~ 1/2), =0
a" p000,,, 0.6)
Table 9: (Damped pendulum) Under various training ranges for initial conditions and parameters, MetaPhysiCa learns the true structure for the damped pendulum system.
We evaluated MetaPhysiCa on a more complex ODE task from Chen (2020) adapted to our setting. We consider two-dimensional ODE with state X, = [pt, q: E B22 dp a sin(p:) +-* sin("?); des dt
de e" sin(p,) cQ:( where W (a", D" c²) are the dynamical system parameters. We simulate the ODE over time steps {to..... tr with Vl, t 0.11, T = 100 in training and over time steps {fo, ....... ...,.... in test with out- -ut-of-disttibution test, we sample Pt, q: ~ (1, 1.5). The dynamical system parameters are sampled as a (1) ,0()) c(i'+ ~ (1.00,11,1)) 1.5).
T. In training. we sample initial states pr. q: ~ (0(5,, 1), whereas in
Table 8 shows the results for this task. First, we note that due to the complexity of 2-layer learnable basis function procedure, we sometimes need to use validation data (held out from training) to cross- validate the learned model (and reject meta-models that do not do well in validation). MetaPhysiCa learnt a stiff ODE for 2 out of 5 folds of cross-validation, resulting in no predictions for in-distribution validation data, which were rejected (marked as superscript *). In these experiments MetaPhysiCa performs 1.5x to 1.7x better than the competing baselines. We believe there is room for improvement in the optimization procedure of these more complex models.
D.5 METAPHYSICA WITH DIFFERENT IN- AND OUT-OF-DISTRIBUTION DATA
In this subsection, we evaluate MetaPhysiCa for different in- and -itto---iitribttion ranges for the initial conditions X,, and ODE parameters W Table 9 shows that under variety of training ranges for initial conditions and parameters, MetaPhysiCa is able to learn the true structure of the damped pendulum dynamical system. Furthermore, given the learnt structure, Table 10 shows that MetaPhysiCa outputs robust predictions for different OOD ranges of initial conditions and parameters.
Initial conditions X
ODE parameters W
OOD test NRMSE
00D-1
0o ~ U(x m., 2
a' U(2.3) =((04,00.6
0.181 (0.012)
00D-2
0n ~ ())
a' ~ a' ~ ~ u(11100) a* ~ U(1,2) (0066,,, 0406000
0.103 (0.016)
00D-3
6o =1
0.029 (0.005)
00D-4
0.039 (0.002)
Table 10: (Damped pendulum) OOD test performance of MetaPhysiCa for various 00D test ranges for initial conditions and parameters. Using test-time adaptation, MetaPhysiCa robust to different OOD ranges for the ODE parameters W
24

---

Published as a conference paper at ICLR 2024
Test NRMSE +
Methods
OOD Xto
OOD X1, and W
NeuralODE
0.591 (0.119)
0.717 (0.210)
MetaPhysiCa without sine/cosine basis
0,290 (0. 137)
0,468 (0.059)
MetaPhysiCa with all basis terms
0.070 (0.011)
0.181 (0.012)
Table 11: (Damped pendulum results) Normalized RMSE 4 of test predictions from MetaPhysiCa with and without appropriate basis functions. While MetaPhysiCa without sine/cosine basis terms performs 1.5x to 2x better OOD than the best baseline, it is unable to extrapolate properly and worse than MetaPhysiCa with these basis terms included.
D.6 METAPHYSICA WITHOUT APPROPRIATE BASIS TERMS
In this subsection, we evaluate MetaPhysiCa without algorithmic alignment, i.e., appropriate basis functions required to learn the ground truth dynamics are not present. We repeat the damped pendulum experiment without sine/cosine basis functions.
In the absence of sin(0,) term to learn the true dynamics of the damped pendulum system, MetaPhysiCa learns truncated Taylor series approximation of this term via Og and 0f terms: de W20, + W30 + W4wr. Table 11 shows that after test-time adaptation, the learnt model achieves 1.5x to 2 better OOD test NRMSE than the best baseline, but worse than MetaPhysiCa with sinusoidal basis functions included.
d8, Wiw dwy d
D.7 METAPHYSICA WITH DIFFERENT NUMBER BASIS FUNCTIONS
We repeated our damped pendulum experiments with increasing numbers of basis functions. We begin from basis terms (sinusoidal terms and polynomial terms up power 1) 32 basis terms (sinusoidal terms and polynomial terms up to power 6) for the two output dimensions in the damped pendulum system (0, and w.).
Figure 9 shows the loss and ||6 over epoch. Note that when number of basis functions is m, | begins from 2m as all basis functions are initialized to be active for both the input dimensions. MetaPhysiCa converges to the true dynamics of the damped pendulum system with 3 basis terms for all the different values of m tested. However, as m increases, model requires a higher number of epochs reach convergence.
25

---

Published as a conference paper at ICLR 2024
Number of terms in learnt structure
m=32
m=25
m=19
m=14
m=10 m=7
60
50
40
30
20
10
Epochs
2k
4k
6k
8k
10k
12k
14k
(a) ||+1]
Training Loss vs Epochs
m=32
m=25
m=19
m=14
m=10
m=7
L05
Epochs
0
2k
4k
6k
8k
10k
12k
14k
(b) Training loss
Figure 9: (Damped pendulum.) || and training loss vs epochs for different number of basis functions. MetaPhysiCa learns the true dynamics of the damped pendulum with 3 terms in all the cases; however takes longer to converge for higher values of m.
26


Reviews:
[{"review_id": "YJV9vMYptk", "reviewer": "Reviewer_HxZg", "summary": "The paper focuses on the significant challenge of forecasting tasks within dynamical systems where the underlying ordinary differential equation (ODE) parameters may vary. The key contribution is the application of a meta-learning procedure for causal structure discovery, which aims to improve model performance even when faced with initial conditions and ODE parameters that lie outside the training distribution. The proposed method is tested across three different OOD tasks, and the results indicate a substantial performance improvement over existing state-of-the-art physics-informed machine learning and deep learning methods.", "questions": "- Is it possible for you to run extra experiments for multiple in- and out-distribution \"pairs\" and show that the results are consistent? I'd be happy to increase my score after seeing a more robust evaluation of the method.\n\n- A real epidemic does not follow the SIR model but there are some characteristics of the SIR model that are useful. How can we adapt your work to handle such real-world dynamical system?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper's approach to modeling dynamical systems by interpreting them as structural causal models within a meta-learning framework is both novel and significant. This perspective opens new avenues for understanding and forecasting complex systems.\n\n- The framework's construction is innovative, and the incorporation of the V-REx penalty to uncover causal structures appears effective.", "weaknesses": "- My main concern with this paper is the robustness and completeness of the empirical results. Authors developed two OOD scenarios and picked one for each dynamical system. The selection of OOD scenarios for each dynamical system requires further clarification, as the current rationale is not provided. In addition, the authors use limited in- and out-distribution \"pairs\": only one single distribution for training, and another single distribution for testing. This limits the assessment of method generalizability. Expanding the experimental design to include both OOD scenarios across all dynamical systems and a broader range of distribution pairs would likely enhance the validity of the findings.\n\n- The methodological scope is limited by not accounting for interactions between basis functions, potentially restricting the model's capability to learn more intricate dynamics. Additionally, by confining the approach to ODE-style dynamics, it may not accurately reflect the complexity of real-world systems, such as those found in real world epidemics.\n\n- While the method can integrate prior knowledge of dynamical systems into the creation of basis functions, it appears that this is the extent to which such insights can be utilized. Expanding the method to incorporate prior knowledge more deeply could further improve model performance and applicability.", "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": "", "extra_scores": {}}, {"review_id": "mX2bXocKFp", "reviewer": "Reviewer_jJvs", "summary": "The authors propose a meta-learning based method for physics-informed out-of-distribution (OOD) generalization. Specifically, the proposed method comprises a set of basis functions that are assumed given where each basis function is governed by its set of parameters (unknown). Additionally, it is assumed that the proposed system is trained on a set of related training instances (from some dynamical system e.g., pendulum) to learn the parameters of the basis functions (specifically a sub-set of the basis functions that govern the task) as well as to learn parameters that are involved in the linear combination of these basis functions. These basis functions (with appropriate learned parameters) and their linear combination comprises the causal structure discovery mechamism (CSM) proposed in the paper. Additionally, the authors adopt a meta-learning approach where the CSM model (trained on multiple trajectories with different initial conditions and PDE parameters) is additionally also trained with a invariant-risk minimization (IRM)  type of loss (specifically minimal variance of loss across all training tasks).  Once the model is trained with the CSM + IRM based losses, the test-time comprises a few-shot adaptation (only of a subset of parameters of the model) to the  related but new (initial, PDE parameters) condition of the dynamical system.  \n\nThe authors demonstrate that the proposed model in multiple contexts (i) epidemic modeling (ii) predator-prey systems (iii) pendulum system  and showcase better OOD adaptation compared to traditional physics-informed approaches as well as other approaches like Neural ODE.", "questions": "1. Have authors tested on more challenging (stiff PDE, ODE settings and on settings like 1D convection, convection-diffusion where traditional physics-informed models fail)? \n\n2. Could you please expound on the assumption of the m possible basis functions? Are there contexts where all sub-parts of the full applicable basis might not be present in the pre-trained network? How can the proposed method adapt to this OOD scenario? Has this been tested?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The problem of OOD generalization is important and is a significant challenge (as highlighted in the paper) for physics-informed neural networks (transductive or inductive) as well as traditional neural network models to accomplish. However, any useful neural network model applied to scientific domains needs to have good OOD generalization capabilities. Hence, the authors develop an effective solution to an important problem. \n \n\n- Overall, the paper is well written, and the related work and methodology as well as results are well organized and clear.", "weaknesses": "- Testing is required on more challenging settings (e.g., 1D convection, convection diffusion other “stiff” PDE / ODE settings where physics informed approaches are known to fail). \n \n\n- The assumption that the collection of m possible (appropriate) basis functions are always available seems too strong to the reviewer. It would be helpful if more clarity about this can be provided by the authors. Basically, this strong assumption might significantly reduce the impact of the proposed method as the full set of basis functions might not always be available to select from.", "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": "", "extra_scores": {}}, {"review_id": "qvCjfLycVK", "reviewer": "Reviewer_qpSw", "summary": "In the paper, the authors propose an approach to prove more robust dynamical system forecasting, for physics-informed algorithms. The paper first describes the out-of-distribution (OOD) setting with transduction and inductive setting, where the algorithm alignment is needed for the accurate forecasting. For the proposed model, MetaPhysiCa, the authors first describe the deterministic underlying structural model, where at each time step, the hidden state $x$ can go through different bases functions. The derivative is a combination of the task-specific coefficients and selected bases functions. Given the training data, the authors extract the underlying causal structure from causal structure discovery problem (minimal causal structure with the last number of edges). The structure parameters, global parameters and task-specific parameters are minimized using the binarization tricks to approximate. The authors have provided the proof to show that MetaPhysiCa can correct identify the causal structure. During the test time, the task-specific parameters can be adjusted. In the experiments, the authors show that the MetaPhysiCa can perform better than existing PIML method in the tasks.", "questions": "Q1.  The authors have described the MetaPhysiCa primarily for ODE tasks. However, could the MetaPhysiCa be adapted to PDE setting, for example, for finding the parameters for Burger’s Equation? How could MetaPhysiCa be adapted to methods such as PINN?\n\nQ2. Why are only the task-specific parameters updated during the test-time? Why would causal structure not needed to be updated? For fraudulent detection system, it is possible that a ODE function shifts during the test time.\n\nQ3. Theorem 1 is only guaranteeing that MetaPhysiCa identified the true causal structure. What is the theoretical guarantee that MetaPhysiCa discovers the good task specific parameters, especially for few-shot updates during the test time?\n\nQ4. Could the authors provide more insights into the computational complexity introduced by MetaPhysiCa, especially for more larger pools of $f$ basis functions and increasing number of observations?\n\nQ5. Would MetaPhysicCa be adapted or be integrated into PIML methods that contain different loss functions?  Is that as simple as adding the additional loss terms into Equation 4?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "(S1) The paper introduces a novel concept, MetaPhysciCA, which focuses on out-of-distribution learning in the PIML setting. This approach is distinct from the state-of-the-art methods and addresses an important robustness problem that is yet to be solved in the research community.\n\n(S2) The paper is clearly written. The authors give a clear problem definition (inductive setting and transductive setting) with compelling examples (Figure 1-2). There are several components in the methodology, yet each component is clearly addressed. The author has also provided a theoretical proof to show that the method can correctly identify the underlying causal structure.\n\n(S3) The empirical results presented in the paper show that the proposed method outperforms standard PIML methods in the OOD setting, showcasing the quality and robustness of MetaPhysiCa. Given the ubiquity and increasing reliance on PIMLs in real-world applications, the capability to adapt to real-time data shifting and providing robust estimation under underspecified physics prior is important. MetaPhysiCa's methodology has potential impact for ensuring PIML’s reliability.", "weaknesses": "W1. While MetaPhysiCa showcases success in the datasets mentioned, it is usually the case that the the correctly specified basis function is included in the search space of $f$. MetaPhsysiCa’s performances in the set of incorrectly specified basis function is not known.\n\nW2. There should be an ablation study that shows MetaPhysiCa’s performances for a relatively small pool of $f$ basis functions and a relatively large pool of $f$ basis functions. For example, for the large pool of $f$, does learning the causal structure more difficult?\n\nW3. The paper mentions that joining optimizing all parameters result in comparable experiments than bi-level optimization. But the experiments in the main body of the paper forego showing this. The authors should consider adding experiment results that shows bi-level optimization results as well as the jointing optimization.", "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": "", "extra_scores": {}}, {"review_id": "MmjO3HF1zy", "reviewer": "Reviewer_AhF2", "summary": "This paper adopted a meta-learning technique to improve the robustness of OOD. The experimental results demonstrate the proposed method exhibits a better generalization ability.", "questions": "Please see the comments above.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Develop a meta-learning technique for identifying the underlying governing equations of dynamical systems.\n\n2. Formulate it as a bi-level optimization problem\n\n3. The evaluation results have demonstrated the good performance of the proposed method over baselines.", "weaknesses": "1. How many basis functions are used in the experiments? If there is no prior knowledge, the number of basis functions will be very large.\n\n2. The proposed method still adopts the SINDy-like approach in which all key terms should be included in a set of basis functions. It may not work on complex equations like $dx/dt = 1/(x_1^2 + sin (x_2^3+x_3))$\n\n3. Conduct experiments on more complex systems in the CoDA paper, such as reaction-diffusion system and Navier-Stokes system\n\n4. I am curious about the causal structure in Fig. 4. Which two terms in each of the three dynamical systems have causal relationships? Please give an example to explain it. It seems that the authors did not discuss this in the Evaluation section.\n\n5. Please conduct experiments on noisy data.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)