First-order Context-based Adaptation for Generalizing to New Dynamical Systems
Authors: Junyoung Park
Federico Berto
Arec Jamgochian
Mykel Kochenderfer
Jinkyoo Park
Venue: Submitted to ICLR 2023 (2023)
Keywords: physical system modeling
differential equation
generalization
context
adaptation

Abstract: In this paper, we propose FOCA (First-Order Context-based Adaptation), a learning framework to model sets of systems governed by common but unknown laws that differentiate themselves by their context. Inspired by classical modeling-and-identification approaches, FOCA  learns to represent the common law through shared parameters and relies on online optimization to compute system-specific context. Due to the online optimization-based context inference, the training of FOCA involves a bi-level optimization problem. To train FOCA  efficiently, we utilize an exponential moving average (EMA)-based method that allows for fast training using only first-order derivatives. We test FOCA  on polynomial regression and time-series prediction tasks composed of three ODEs and one PDE, empirically finding it outperforms baselines.

Under review as a conference paper at ICLR 2023
GENERALIZING TO NEW DYNAMICAL SYSTEMS THROUGH FIRST-ORDER CONTEXT-BASED ADAPTATION
Anonymous authors Paper under double-blind review
ABSTRACT
In this paper, we propose FOCA (First-Order Context-based Adaptation), learn- ing framework to model sets of systems governed by common but unknown laws that differentiate themselves by their context. Inspired by classical modeling-and- identification approaches, FOCA learns to represent the common law through shared parameters and relies on online optimization context. Due to the online optimization-base FOCA involves bi-level optimization problem. To train FOCA efficiently, we utilize an exponential moving average (EMA))-ased method that allows for fast training using only first-order derivatives. We test FOCA on polynomial regres- sion and time-series prediction tasks composed of three ODEs and one PDE, em- pirically finding outperforms baselines.
compute system-specific context inference, the training of
INTRODUCTION
Scientists and engineers have made tremendous progress modeling the behavior of natural and engineering systems and optimizing model parameters to best describe the target system (Ljjung 2010). This modeling and system identification paradigm has made remarkable advances in modern science and engincering (Schrödinger, 1926; Black & Scholes, 1973; Hawking, 1975). However, applying this paradigm complex systems is difficult because the mathematical modeling of sys- tems requires a considerable degree of domain expertise, and finding the best system parameters requires massive experimentation.
The availability of large datasets and advances in deep learning tools have made it possible to model a target system without specific mathematical models, relying instead on flexible model classes (Brunton et al., 2016; Gupta et al., 2020; Menda et al., 2020; Jumper et al., 2021; Kochkov et al., 2021; Degrave et al., 2022). However, when the characteristics of target systems change (c.g., system parameters, boundary conditions), the flexibility of data-driven models makes them difficult to adapt.
Deep learning approaches typically handle the contextual change by collecting data from the new behavioral mode and re-training the model on the new dataset. However, this approach can be impractical, especially when the system is complex and context change is frequent. We are interested in developing a framework that learns a common shared model of the systems and inferring the context that best describes the target system to predict response. Our study considers a target system whose input x and response can be described by f(z c). where ff denotes the function class shared by the target systems and denotes the system-specific context.
One possible approach for modeling such target systems is meta-learning (Hospedales et al., 2021) I), which learns how to adapt new systems. Meta-learning is typically combination of an adap- tive mechanism and training for the adaptation. One typical meta-learning approach is to use an encoder that takes the adaptation data and returns the learned context (Santoro et al., 2016; Mishra et al., 2018; Garnelo et al., 2018; Kim et al., 2019). Although sacoder-based adaptation schemes require constant memory usage, their parameterized encoders limit the adaptation capability. Other approaches (pre) train the parameters on the dataset collected from the various modes and update all parameters using gradient-descent (Finn et al., 2017; Nagabandi et al., 2018; Rajeswaran et al., 2019). Despite their effective adaptability, those approaches (Antoniou et al., 2019), especially when the adaptation target is complex and adaptation data is scarce. Instead of updating all parameters, Raghu al. (2019); Zintgraf et al. (2019) update subset
often prone (meta) over-fitting

---

Under review as a conference paper at ICLR 2023
Table I: Comparison of context-based generalization approaches. The memory column specifies the additional memory requirements for adaptation during the training phase of the algorithms. | denotes the number of elements of
Training models (parameters)
Adaptation mochanism
Momory
Encoder
(Mishra ol.,, 2018; Lee et al., 2020a)
fo.9s
go(D)
Oilel
6412 m 0g a Eiraeen (fola
(o (a, de), ()).))
CAVIA (Zinigral al., 2019)
Io
O100
èu 0,8 egs where is the adapeation steps.
CoDA (Kirchmeyyr al., 2022)
Jo.W w
== arg min, Σ (awi: [(fo (x))
((W)
A any min, ) xteD
slen C (f((s,,)) 0),9)
FOCA
fo
where fa an EMA copy of fo-
0(1)
of parameters, which call context é, while fixing the remainder. Although this modeling approach effective, the training requires computing higher-order derivatives.
Many training method for meta-learning have been proposed (Finn et al., 2017; Nichol al., 2018 8; Rajeswaran et al., 2019; Deleu et al., 2021). Typically mncoerr-based meta-learning trains an en- coder and prediction model jointly and no online optimization-basee adaptation occurs. The train- ing of gradient-based meta-learning is typically cast as bi-level optimization. Typical gradient-basee meta-learning often carried out by propagating through the update steps. which requires higher- order derivative calculations. To avoid such issues, Finn et al. (2017); Nichol et al. (2018) propose first-order approximation of derivatives to update the (meta) parameters. Alternatively, the implicit gradient method can be used to lower the computational burden (Rajeswaran et al., 2019), but gra- dient computation errors can significant and result performance degradation (Liao et al., 2018; Zhou et al., 2019; Blondel et al., 2021; Chen et al., 2022).
Adaptation), a context-based meta- learning method that is specialized for complex dynamical systems whose behavior can be charac- terized by a common mathematical model f and a context c. Specifically, FOCA considers target systems of the form y fo(x, é). where fo denotes the learned function class shared by the target systems and è denotes the inferred system-specific context. FOCA learns the function class fo dur- ing training and solves the system identification problem through numerical optimization find the proper é. The online context optimization bypasses the limitation of encoder-based approaches, but entails higher computational burden. To train FOCA efficiently, we thus propose using an ex- ponential moving average (EMA) based training method, which operates with first- order derivatives and no additional memory usage. From our experiments, we also confirm that EMA-based train- ing decreases the computational burden of training and improves the generalization performance of FOCA. The contributions of this work are summarized as follows:
In this paper, we propose FOCA (First-Qrder Context-
We propose FOCA, learning framework specialized for modeling complex dynamical system families that are described by an unknown common law and system-specific context.
We propose an EMA-based training method for training FOCA that overcomes the burden of second-order derivative calculations while showing better generalization results compared to other training methods.
We empirically demonstrate that FOCA outperforms is competitive to various meta-learning baselines in static function regression and time-series prediction tasks evaluated both in- distribution and out-of-distribution..
2 RELATED WORK
Learning to generalize to the new systems. Transfer learning (Zhuang et al., 2020) attempts to generalize learned models to new systems (or tasks) by fine-tuning a small number of parameters to the new tasks. As more direct way of generalization new systems, meta-learning learns how to adapt quickly to new systems. In particular, gradient-based meta-learning (GBML) approaches perform few-step gradient descent updates of the model parameter 0 for adaptation (Finn et al., 2017: Nichol et al., 2018). Focusing on the empirical evidence that the plain GBML is prone to overfit the training tasks, CAVIA (Zintgraf et al., 2019) performs gradient-based updates for small
2

---

Under review as a conference paper at ICLR 2023
dimensional context vector c. On the other hand, ancoder-based meta-learing methods jointly train an encoder network 9e alongside the prediction model fa (Santoro et al., 2016; Mishra et al., 2018; Gamelo et al., 2018; Kim et al., 2019; Gordon et al., 2020).
A few works investigate how to solve our target problem (i.e., learning a common unknown rule f of static dynamic systems with the learned adaptation mechanism). Lee al. (2020b) applies GBML to generalize over Hamiltonian systems. LEADS (Yin et al., 2021) jointly learns the system specific networks 9a and shared network fa. and predicts the response of the dynamic system with fu + 9s- CoDA (Kirchmeyer et al.. 2022) learns the common model parameters O, a linear basis W, and context vectors è while defining fo(-. é) as fo+we(-). To adapt to the new system with the observation dataset D, it solves an optimization, mine ((,w)eD L (fo+we(z)) y). In contrast, FOCA adapts by solving Eq. (6), which is an effective method to generalize to systems that do not require any structural assumptions. This approach is different from LEADS, where the adaptation is sum-decomposable, and from CoDA, where the adaptive parameters are span of linear basis. We summarize the differences between FOCA and baselines Table
Solving bi-level optimization to learn. Training FOCA involves solving bi-level optimization. Differentiable optimization layers (Amos & Kolter, 2017; Agrawal et al., 2019; Meng et al., 2021) encapsulate an optimization problem as a layer of neural networks. As they employ optimization as a layer of the network, the training becomes a bi-level optimization problem. The backward path of differentiable optimization layers often implemented by backpropagating the optimization steps or implicit gradient methods. As a special case, Bertinetto et al. (2019) employs inner optimization problems that can be solved analytically.
In the context of meta- meta-learning, Bertinetto et al. (2019) considers solving the analytically solvable inner optimization problem. However, such approaches limit the flexibility of network parameteri- zation or adaptation mechanisms. As more flexible approach, Implicit MAML (Rajeswaran et al., 2019) extends MAML to overcome the extensive memory usage of adaptation and allows for more inner adaptation steps without memory issues. However, implicit gradient methods can calculate inexact vreta-gradiens (Deleu et al., 2021), and, as we show in the experiments, such errors can de- grade performance. Instead, we propose to solve the bi-level optimization training problem with the EMA method, which also has constant memory usage. We empirically show that the EMA-based training method outperforms implicit gradient-based methods on test metrics.
PRELIMINARIES
This section reviews classical parameter identification and meta-learning as preliminaries for ex- plaining the proposed method.
System identification. Consider the response y from input of system f with an unknown (or unobservable) context
y f(x,
Classical parameter identification aims inferring the (optimal) context c* with data from task D = {E yn)) ) Under the assumption that (In, Un) pairs in D are generated from the same system, the identification process can be formulated an optimization problem as follows:
c² = arg min
Σ V).).
where L is a discrepancy metric. By using the optimized context c", we can query responses of the (identified) system f(),e)) et to inputs.
Meta-learning. Meta-learning is collection of algorithms that learn to adapt quickly to new task by learning from experiences of related tasks. Assuming that tasks (i.e., the systems of dif- ferent contexts) drawn from distribution p(T) and each task includes inference data DJ and
3

---

Under review as a conference paper at ICLR 2023
Training
f
Trii data Prodictionn
Pidt
0
-10
10
(a) Initial fo(:))
(b) Task adaptation
(c) Prediction
(d) Trained fō (:)
Figure 1: Training of FOCA: (a) the function class captured by the initial (o(-.) and task distribution p(7).
(b) inference by solving Eq. (6) with context inference data and identification of fo(;,c),, c.), (c) prediction of fo(:.') and the outer loss [(fo [1fo c"). v.)), (d) the optimized (o())) FOCA is trained by repeating (b), (c), and minimizing C(fo(xi,, c²), y.).
evaluation data D the meta-learning objective can be defined follows:
min TT-p(T) [ (D). Ser)
(3)
0,0
s.t. 0' = gv (Dirr.
(4)
where Ie (.,0) 0) is a (eatred adaptation algorithm that parameterized by and a, for is the task-specific model that is parameterized by the adapted parameter 6', and L is loss metric. In summary, meta-learning learns 0 and vo to infer a model from D that minimizes the error on DJ Optimized meta-parameters 0* and vp can then be used "quickly" adapt tasks from p(T).
4 METHODOLOGY
In this section, we present First-Order Context- based Adaptation (FOCA) to learn to jointly identify and predict systems in the same family. We first introduce the training problem of FOCA and then discuss the EMA-based training method for efficiently solving the training problem. We also provide our code as supplementary material.
4.1 PROBLEM FORMULATION
We formulate the learning problem of FOCA as bi-level optimization where the inner optimization Eq. (6) is for inferring the context, and the outer optimization Eq. (5) is for learning fe conditioned on the identified context. The proposed optimization problem is as follows:
min Σ
Σ
C (fo(ahh e),,
(5)
DED (ar B...)
subject
é arg min
Σ
(folah,, v),
(x
where Dr and Di, are the context inference and evaluation datasets of the ith system, which satisfy Di nDi = and Di UDi C D', D = (D' is a collection of meta-training datasets from multiple systems, fo prediction model parameterized by 0, and is a loss metric.
The prediction model fa is trained by solving Eqs. (5) and (6). The trained fü describes the common structure of the systems D. Identification of specific function (i.e., adaptation) done by solving Eq. (6) with the data points sampled from the specific system. After finding c², we can query for meta-test predictions with fo(:, e²) as visualized Fig. (b) and (c).
FOCA can be viewed special case of the meta-learning recipe explained in Section 3 by defin- ing 0' [0, øt] and setting Eq. (4) as Eq. (6). Even though FOCA formulated similarly to optimization-bar based meta-learning methods, the optimization variables for the adaptation is e rather than 8. In terms of implementation, this algorithmic selection allows us to batch-solve inner-level optimization with standard automatic differentiation tools. FOCA can therefore be trained effi- ciently. We can also interpret this clear separation of 0 and c as capturing function classes and their coefficients separately. We further inspect this view of FOCA in Section 5.1.
4

---

Under review as a conference paper at ICLR 2023
Algorithm 1: Training FOCA with exponential moving average (EMA)
Input: Prediction model fa. training tasks DD inner optimization steps K, inner optimization
step size a, weighting factor т, context regularization weight // initialize the target model
< fo < for 1, ... do
Sample batch of tasks D ~ D for D D' do construct Dj and Di, from D E² +-0 for k: = 1,... do
[(c') = ((Jo(nn,c'), A||cs1|
è' + ē² - -VVal(l)
// Solve Eq. (6)
10
end end
11
Evaluate (()) = EDi E(ohwieeD (ff(=",,)).
12
+ 0 - (V)(())
+- 70 + (1 (1-
// Update the target model
end
4.2 TRAINING FOCA
Training FOCA with gradient descent involves computation of the partial derivatives of Lupper EDIED= E(r L (Jo(r,, ēt), vi) with respect to 0. We express ac /80 using the chain rule taking čt as the intermediate variable,
C de
00 ae
00
Here, OC"F"/06" /ast can be computed via an automatic differentiation package. However, the computa- tion of 8e fee is less straightforward because of the implicit relations between č and 0. That is, their relations are defined through an optimization problem rather than a computational chain.
One possible approach to calculate de' /ao is to backpropagate through the inner optimization steps, which calculates the derivatives by unrolling the inner optimization steps through the computational graph. However, this method is not memory efficient, especially with a large number of model parameters |6] or with many inner optimization steps. Another method is to approximate 80%an with first-order methods as similarly done Finn et al. (2017); Nichol et al. (2018). However, as reported by others (Zhou et al., 2019; Jayathilaka, 2019; Chen et al., 2022), this approach often requires techniques to mitigate the approximation errors. Another method calculates 0e/0e using implicit gradients and bypass the memory usage of the first approach. However, calculating the gradient without error requires solving the inner optimization problem optimally (Liao et al., 2018; Blondel et al., 2021). In practice, this restricts the expressivity of So to architectures that can be solved optimally Eq. (6). We propose an efficient approach to train FOCA that uses no additional memory and allows for flexible architecture of fe.
EMA-based training method. Inspired by recent developments in self-supervised learning (Grill al.. 2020; Caron al., 2021) and the sarget-network of reinforcement learning (Lillicrap al., 2016), we use a simple method based on the exponential moving average (EMA) to train FOCA. The core idea is to consider as an independent input given from outside of the gradient chain. First, we copy to create delayed target model fø parameterized by 0. We use fa to infer e with gradient descent on Eq. (6). We then update 0 by using V VgC òc /00 evaluated on č As the parameters of fo and fa separate, this avoids differentiating the argmin operator to calculate VoL(O). After updating 8, we update as 0 +- 70 + (1 - 7) with 0 < < 1. The training procedure is summarized in Algorithm
The proposed training algorithm utilizes fo (i.c., the delayed copy of fo) to infer a by numerically solving Eq. (6). When train fo, it takes ?t as input that contains meaningful information that
5

---

Under review as a conference paper at ICLR 2023
ctx point
True
Prediction
FOCA
Encoder
CAVIA
CoDA
MSE: 0.006
MSE: 0.011
MSE: 0.261
MSE: 0.192
1
1
Figure 2: Polynomial regression results. All models infer the polynomial from the given context points The prediction results and ground truth values are visualized with the blue-solid (-) and gray-dashed (--) lines, respectively.
differentiates the system i from the other systems. At the early training phase, èt might have no meaning (i.e., the different č's do not change the prediction results). However, this phenomenon is resolved during training because So is trained as if è is (learned) gystem-differentiating context. Furthermore, disentangling the model for prediction and inference results in a practical advantage: employ any optimization algorithm (even non-differentiable ones) to solve Eq. (6). The proposed method shows its advantages in memory usage and implementation.
We also observed that EMA-based training can increase the generalization performance the other training methods. We discuss this aspect with the experimental results in Section 5.3.
EXPERIMENTS
In this section, we discuss experiments that show the properties of FOCA with static regression problems (Section 5.1), demonstrate its effectiveness in predicting the response of dynamics systems (Section 5.2), and investigate the proposed EMA-based training method (Section 5.3).
Baselines. We compare FOCA to three baselines that adapt to the tasks by inferring é. The first two baselines are representative methods that find ět via the inference schemes discussed in Section 1. The last baseline, CoDA (Kirchmeyer et al., 2022), is our closest baseline in terms application: meta-learning method tailored for generalization of physical systems. To make fair comparisons, all models use the same network architecture fa(-) for each task. The baselines are as follows:
Encoder employs an encoder network extract the context è' from Di That is, E² go (Dinr). where 9s is the encoder network. We employ SNAIL (Mishra et al., 2018) and CaDM (Lee et al., 2020a)-based Je for static and time-series tasks, respectively.
CAVIA (Zintgraf et al., 2019) performs a few gradient steps on Eq. (6) to extract the context č from Di and train fo via the backpropggation-througeg shrouuh-oppimmizatton-sttp scheme.
CoDA (Kirchmeyer et al., 2022) learns the common model parameters 6, linear basis W, and context vectors č while defining fo(, 2¹ as fo+was(-). For adaptation to a new system, solves the optimization problem mine EettaheED L (Jo+we
5.1 STATIC REGRESSION
10a
We first examine polynomial function regression to under- stand the properties of FOCA. In this task, the objective is to identify a polynomial and make predictions from the identified functions as shown Fig. 2. We prepare 4th or- der polynomial datasets by sampling the coefficients from (0.11, 2.5), and set [Di and [Di, as 5 and 15, respec- tively. Fig. 3 shows the average of the test root mean squared error (RMSE) values of models. As shown in Fig. 3, Figure 3: Static regression test RMSE. FOCA adapts the different polynomials better than the baseline models. We provide the details of the model architectures and training scheme in Appendix A. We then investigate the properties of FOCA by using the polynomial dataset.
poSSE
POCA
Encodeer
CAVIA

---

Under review as a conference paper at ICLR 2023
FOCA
Encoder
State
2
2
CAVIA
CoDA
State
2
time [s]
time [s]
Figure 5: MS prediction results. All models infer the physical constants from historical (shaded in yellow) and current (shaded in gray) state observations, and then predict rollouts of future state trajectories (colored in bluc). The gray dashed lines visualize prediction targets. FOCA successfully identifies the target systems and predicts long-term futures.
Do 0 and contain function class and parameter information, respectively? The roles of 0 and are to capture function class and to identify spe- cific function from that class, respectively. For ample, 0 should express the fact that functions are polynomials and should capture the coefficients of a specific polynomial. To analyze whether and c fulfill these roles, we train two FOCA models: one with parameters 01 meta-trained on linear func- tions, and one with parameters 02 meta-trained on quadratic functions. We then define models that in- terpolate between these two sets parameters with 0(A) (1 A)0, + AO2 with O I X < 1. Fig. shows predictions from (e())(,,e) with dif- ferent samples ~ u(-0.025, 0.025) and dif- ferent values of A. We conjecture that fe(xy(-,)) represents linear functions when X = O, quadratic functions when À = 1, and functions that interpolate between linear and quadratic functions when 0 < À 1. We verify the conjectures by observing that within single function class 0(A), each parameterizes a different function from that class (e.g., ., the red curves are all linear but each of them a different linear function).
9(0)
8/0.5)
9(1)
8(0.25)
8(0.75)
0.5
0.0
0.5
-1.0
-05
0.0
0.5
1.0
Figure 4: Role of è within model classes (())
that interpolate between quadratic linear.
5.2 TIME TIME-SERIES PREDICTION
We evaluate FOCA on time-series prediction in three ODE datasets, including mass-spring systems Lotka- equations (LV), glycolytic cillators (GO), and one PDE dataset derived from the Navier-Stokes equations (NS). In this setting. the objective is first to identify target systems from historical and current state observations and then to predict rollouts of future state trajectories as shown in Fig. 5. We provide details for dataset genera- tion. model architectures, and training schemes in Appendices E.
0.08
0.6
0.7
0.3
0.9
1.0
1.2
1.3
14
Viscosiity
X10
ROCA
Encoder
CoDA
Figure 6: RMSE for different models on the Navier-Stokes equations at different viscosities v. The orange area denotes the training region.
Task generalization. We study the generalization performances of the models systems whose context parameters are generated from inside as well as outside the training distribution (i.c., in- training and out-of-training distribution adaptation tasks). We evaluate the models by measuring the average and standard deviation of the rollout prediction root mean squared errors (RMSE) over the test systems. As shown in Table 2, FOCA generally outperforms the baseline algorithms, indicating that FOCA adapts well to new systems regardless of the parameters of the system. To further analyze the performance variation of models to the change of system parameters, we visualize the prediction

---

Under review as a conference paper at ICLR 2023
Table In/out-of training distribution generalization results. Smaller is better (U). We measure the RMSE values per trajectory. Average RMSES and their standard deviations are reported. Best in bold; second underlined.
In-training
Out-of training
In-training
Out-of training
In-training
Oat-of training
In-iraining:
Out-ol training
Uncodex:
2:272 +: 1.422
2,379 =: 1.506
0.137 0.093
0.326 =: 0.362
0.062 1: 0.092
0.156 0.174
0.057 1: 0.016
0.114 1: 0.026
CAVIA
1.743 :: 0.595
1,793 0.628
0.312 0,316
0.706 =: 0.551
0.046 0.061
0.157 0.161
0.057 :: 0.011
0.114 0.028
0.051
P0CA
05858
0.156
0.478
0.295
0.079 0.061
0.303 2: 0.291
0.043 0.066
0.147 0.168
0.042 0.026
0,070 0.038
Target
2.5
Vorticity
0.0
Encoder
-2.5
Context
CAVIA
CoDA
MSE
0.2
0.1
FOCA
Figure 7:
Target solutions and RMSE propagation on the Navier Stokes equations for different models. FOCA manages to better contain crrors even after long autoregressive rollouts and out-of of-distribution context samples.
errors of the models on the Navier-Stokes equations for different viscosity values in Fig. 6, while in Fig. 7 show sample of target predictions and error propagation for out-of-distribution viscosity = 10 FOCA shows better prediction results than baselines, not only for new systems with in- distribution contexts but also for ouist--distribution ones. We provide the additional studies for MS, LV. and GO in Appendices B to D, including out-of-distribution error comparison visualizations.
5.3 TRAINING ABLATION STUDY
this section, we perform ablation studies to understand the effect of the proposed training scheme on the performance of FOCA.
Is the EMA target network scheme effective? To train FOCA, we employ the delayed target network fū to infer čt, which we consider an independent input when we train fo (EMA). However, we can consider employing other training methods as mentioned in Section 4.2. We compare against training FOCA by calculating gradients by backpropagating through the unrolled inner optimization steps (BPTO), approximating the inner loop gradient with a first-order approximation (FO), or em- ploying the implicit method (e.g.. Amos & Kolter 2017) for calculating the gradient of the inner optimization (Implicit). We compare the test set performance of each method on the LV dataset while applying the same hyperparameters. Fig. 8a visualizes the test performances of each method over the training steps.
As shown in Fig. 8a, EMA has the lowest test set errors. We conjecture that EMA outperforms Implicit because the implicit gradient estimator can have error in practice, the inner optimiza- tion problem cannot be solved optimally, especially when fo is parameterized with general neural network (Blondel et al., 2021; Liao et al., 2018). This issue might be resolved by employing spe- cial structure to fo such as linearity or convexity. However, such treatment can prevent the use of arbitrary network architecture as fo.
8

---

Under review as a conference paper at ICLR 2023
2,75 10
RPTO
2,75
Test RMSE
2.20
2.5
KO
EMMY=
Implict
2
EMMT 0.5)
x 10-1
1.75 10
1.5 16-1
Tost
1.5
1.25 x 10-1
1.25 10
10
10
500
1000
2000
2500
3000
500
1000
1500
2000
25000
3000
Epoch
(a) Training method ablation results
(b) ablation results
Figure 8: Training method ablation results. We train each method five times with different random seeds. The bold lines and shadow region illustrate the averages and standard deviations of training runs. We smooth the curves using moving average.
We also note that EMA outperforms BPTO. Because BPTO calculates the gradient by chained com- putation the computational graph (at the cost of memory usage), it is difficult to claim that the performance degradation is related to the gradient computation. We conjecture this phenomenon is related the consistency of è estimation. It is noteworthy that BPTO is equivalent CAVIA with larger inner gradient step. Hence, FOCA is also superior to CAVIA with the same number of inner optimization steps. The results indicate that EMA trains FOCA effectively while allowing us to employ arbitrary network architecture and lower memory usage.
How does affect the performance? We investigate the effect of the delay parameter in our training scheme. Our training scheme assumes contains enough information to differentiate specific system. As it is considered an input, the consistency of over the training step becomes crucial to the stable training of fe. In this sense, т controls the degree of consistency of è over the training steps. Higher (~ 1.0) introduces relatively drastic changes of è during training, while smaller r changes conservatively. Fig. 8b illustrates the test performance of EMA with 7 = {1.0, 0.5, 0.1} against BPTO. It is noteworthy that the forward processes of EMA with BPTO are identical, while the backward processes (i.e., the gradient calculation) are different. By comparing EMA 1.0) and BPTO, can observe that using consistent mechanism is crucial for successful training FOCA. Furthermore, as we decrease r to 0.1, FOCA exhibits better predictive performance.
1.0 and
6 CONCLUSION
Motivated by classical modeling-and-inn and-idettiifiiittionn approaches, we introduced in this work new contextual meta-learning method: FOCA. We separate the role of different parameters into a model, that captures context-invariant behavioral laws, and context, that specializes in finding instance- specific parameters. FOCA jointly learns model and context parameters by solving online bi-level optimization efficiently. Instead of relying higher-order derivatives, it employs EMA for faster and more stable training. We empirically demonstrated in both static and dynamic regression tasks complex dynamical systems that approach outperforms is competitive baselines when evaluated on both in- and out-of out-of-distribttion contexts. We also performed ablation studies to exam- ine the efficacy of proposed training method.
FOCA is limited by modeling assumptions for the target system. Namely, FOCA seeks to model dynamical system families that share underlying mathematical model Se- If the target systems do not share consistent behavior, it would be hard to characterize specific model family, and the performance of FOCA might degrade. Future work could develop a more general algorithm that allows for the adaptation of the function class at test time by changing @ when the common is not identifiable.
REFERENCES
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.

---

Under review as a conference paper at ICLR 2023
Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization a layer in neural networks. In International Conference on Machine Learning (ICML). pp. 136-145, 2017.
Antreas Antoniou, Harrison Edwards, and Amos Storkey. How to train your MAML. Interna- tional Conference on Learning Representations (ICLR), 2019.
Luca Bertinetto, Joao F Henriques, Philip HS Torr, and Andrea Vedaldi. Meta-learning with differ- entiable closed-form solvers. In International Conference on Learning Representations (ICLR), 2019.
Fischer Black and Myron Scholes. The pricing of options and corporate liabilities. Journal of Political Economy, 81(3):637-653 1973.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares- López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint 1wt.:i::205..51111 2021.
Johannes Brandstetter, Daniel E Worrall, and Max Welling, Message passing neural PDE solvers. In International Conference on Learning Representations (ICIR), 2021.
Steven L Brunton, Joshua L Proctor, and Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 113(15):392-33377, 2016.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In International Conference Computer Vision (ICCV), pp. 9650-9660 2021.
Dong Chen, Lingfei Wu, Siliang Tang, Xiao Yun, Bo Long, and Yueting Zhuang. Robust meta- learning with sampling noise and label noise via eigen-reptile. In International Conference on Machine Learning (ICML), pp. 3662-3676 2022.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties -ecoderr of neural machine translation: Encoder decoder approaches. In Workshop on Syntax, Semanties and Structure in Statistical Translation, pp. 103-111, 2014.
Bryan C Daniels and Ilya Nemenman. Efficient inference of parsimonious phenomenological mod- els of cellular dynamics using s-systems and alternating regression. PloS One, 10(3):e0119821, 2015.
Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al. Mag- netic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):414- 419, 2022.
Tristan Deleu, David Kanaa, Leo Feng. Giancarlo Kerg, Yoshua Bengio, Guillaume Lajoie, and Pierre-Luc Bacon. Continuous-time meta-learnig with forward mode differentiation. In Inter- national Conference Learning Representations (ICLR), 2021.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International Conference on Machine Learning (ICML), pp. 1126-1135, 2017.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes. In International Conference Machine Learning (ICML). pp. 1704-1713, 2018
Jonathan Gordon, Wessel P Bruinsma, Andrew YK Foong, James Requeima, Yann Dubois, and Richard E Turner. Convolutional conditional neural processes. In International Conference on Learning Representations (ICLR), 2020.
dean-Bastion Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Beraard Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to superviee learning. Advances Neural Information Processing Systems (NeurIPS), 33:21271 1177712-1244 2020.
10

---

Under review as a conference paper at ICLR 2023
Jayesh K Gupta, Kunal Menda, Zachary Manchester, and Mykel Kochenderfer. Structured mechan- ical models for robot learning and control. In Learning for Dynamics and Control, pp. 328-337, 2020.
Stephen W Hawking. Particle creation by black holes, In Euclidean Quantum Gravity, pp. 167-188. World Scientific, 1975.
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-5169, 2021.
Mirantha Jayathilaka. Enhancing generalization of first-order meta-learning. 2019.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-5588
2021.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnel,, Ali Eslami, Dan Rosenbaum, Oriol
Vinyals, and Yee Whye Teh. Attentive neural processes. In International Conference Learning Representations (ICLR), 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Matthieu Kirchmeyer, Yuan Yin, Jeremie Dona, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning (ICML), pp. 11283-1 11301, 2022.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine (eamiing-cceelerateedee computational fluid dynamics. Proceedings of the National Academy of Sciences, 118021))e2101184111 18, 2021.
Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynam- ics model for generalization in model based reinforcement learning. In International Conference on Machine Learning (ICML), pp. 5757-5766, 2020a.
Seungjun Lee, Haesang Yang. and Woojae Seong. Identifying physical law of hamiltonian systems via letanne. International Conference on Learning Representations (ICLR). 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. 2021.
Renjie Liao, Yuwen Xiong, Ethan Fetaya, Lisa Zhang, KiJung Yoon, Xaq Pitkow, Raquel Urta- sun, and Richard Zemel. Reviving and improving recurrent pack-propagation. In International Conference on Machine Learning (ICML), pp. 3082-3091, 2018
Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforeement learning. In Inter- national Conference on Learning Representations (ICLR), 2016.
Lennart Ljung. Perspectives on system identification. Annual Reviews in Control, 34(1):1-12,
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In Inter- national Conference on Learning Representations (ICLR), 2017.
Alfred J Lotka. Contribution to the theory of periodic reactions. The Journal of Physical Chemistry, 14:271-277 1910.
Kunal Menda, Jean De Becdelievre, Jayesh Gupta, Ilan Kroo, Mykel Kochenderfer, and Zachary Manchester. Scalable identification partially observed systems with certainty-equvvalee em. In International Conference on Machine Learning (ICML), pp. 6830-6840, 2020.

---

Under review as a conference paper at ICLR 2023
Zihang Meng, Sathya N Ravi, and Vikas Singh. Physarum powered differentiable linear program- ming layers and applications. In AAAI Conference on Artificial Intelligence (AAAI), pp. 8939- 8949, 2021.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta- learner. In International Conference on Learning Representations (ICLR), 2018.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt in dynamic, real-world environments through meta- reinforcement learning. In International Conference Learning Representations (ICLR), 2018.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-leaming algorithms. arXiv preprint 1033:v:88030029900 2018.
Jaideep Pathak, Shasbank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Four- castnet: A global data-driven high-resolution weather model using adaptive fourier neural opera- tors. arXiv preprint arXiv:2202. 12114, 2022.
Michael Poli, Stefano Massaroli, Federico Berto, Jinkyoo Park, Tri Dao, Christopher Re, and Ste- fano Ermon. Transform once: Efficient operator learning in frequency domain. In International Conference on Machine Learning (ICML), AI for Science Workshop. 2022.
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse? towards understanding the effectiveness of maml. In International Conference Learming Rep- resentations (ICLR), 2019.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with im- plicit gradients. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta- learning with augmeete neural networks. In International Conference Machine Learning (ICML), pp. 1842-1850 1850, 2016.
Erwin Schrödinger. An undulatory theory of the mechanics of atoms and molecules. Physical Review, 18(6),1049, 1926.
George Gabriel Stokes. On the effect of the internal friction of fluids on the motion of pendulums. Transaction of the Cambridge Philosophical Society, 9:8-106, 1851.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- fion Processing Systems (NeurIPS), 30, 2017.
Yuan Yin, Ibrahim Ayed, Emmanuel de Bézenac, Nicolas Baskiotis, and Patrick Gallinari. Leads: Learning dynamical systems that generalize across environments. Advances in Neural Informa- tion Processing Systems (NeurIPS), 34:7561-7573, 2021.
Pan Zhou, Xiaotong Yuan, Huan Xu, Shuicheng Yan, and Jiashi Feng. Efficient meta learning via minibatch proximal update. Advances in Neural Information Processing Systems (NeurIPS), (2.iiP)),2. 32, 2019.
Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE, 109(1): 43-76, 2020.
Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. In International Conference on Machine Learning (ICML), pp- 7693-7702,
12

---

Under review as a conference paper at ICLR 2023
A DETAILS OF STATIC REGRESSION EXPERIMENTS
In this section, we provide the model architectures, training process, and additional experimental results for the polynomial regression task.
Model architecture employ Encoder, CAVIA, and CoDA as baselines. All models use the same multi-layer perceptron (MLP) for fo and are trained with the same data batches for fair comparisons.
For brevity, refer to multi-layer perceptron (MLP) with hidden dimensions R1, 72, ... ny for each layer and hidden activation act, as MLP(11, ?2, ny; act). We refer to a cross multi-head attention block (Vaswani al., 2017) with h heads and x hidden dimensions as R.HHA(h 2).
Table A.1: Polynomial regression model architectures
Context encoder
Jo
Inner step
Inner step
parameter generator
Eucoder
MLP(2.32/MLP(1332)) (BMA)/)))
NLP(I+32, 64, 32, [[LaakReLL))
CAVIA
NLP(L+32, 64, L:LeakyReLU)
CoDA
W R R!
HLP(L, 64, l:LeskyßeLU)
0.001
FOCA
MLP(1+32, 64, 1:LeakyReLU)
0.001
0.1
Table A.1 summarizes the network architectures. The green colored values indicate the dimension of context è. For CAVIA, we also perform hyperparameter search to optimize K. We found that CAVIA with K 5 underperforms, as compared K For CoDA, context as 2 (i.e., |c| 2), following the default setting of Kirchmeyer et al. (2022).
the dimension of
Training details We train all models with mini-batches of 256 polynomials for 4,048 epochs using Adam (Kingma & Ba, 2015) with an initial learning rate of 0.001. The learning rate is scheduled by the cosine annealing method (Loshchilov Hutter, 2017).
B DETAILS OF MASS SPRING (MS) EXPERIMENTS
In this section, we provide the model architectures, training process, and additional experimental results for the mass-spring systems.
A3s
MM
M
MM
m2
Figure A.l: Target mass-spring system. (ki, k2, ks) masses (mi, m2g). The models require adapt the change of spring constants
We consider a frictionless three-spring two-mass system shown in Fig. A.1 with mass positions 21 and 22 governed by the following second-order ODE:
d2r
kn +
hiz
h2
+
x2
d/2
(A.1)
d2r
h +
his
1z
D2
+
dt2
m2
where is a coefficient matrix with spring constants (k1, k2, ki) and masses (m1, mz).
Data generation For training data generation, generate 128 mass-spring systems whose pa- rameters (k), 22 K3, my m2) are sampled from ((077,, 1.25) and numerically solve the mass- spring systems with Runge-Kutta 45 for 7 10 seconds with At 0.15 second time intervals.
13

---

Under review as a conference paper at ICLR 2023
Model architecture For fe. employ the ID-CNN model from Brandstetter et al. (2021). which stacks an MLP, ID-CNN, and consistency decoder (Brandstetter et al., 2021) with the bundling refer to ID-CNN layer with input channels, y output
parameter N 25. For brevity, channels, and filter size w as Conv(x; y, w), a consistency decoder as C.Dec, and the bi-directional GRU (Cho al., 2014) with hidden dimension GRU(x)).
Table A.2: Mass-spring prediction model architectures
Context encoder
fg
purameter generator
Encoder
(RDT644)-MPP(6 64)
MEP(4x 64, 32, 4x 4x255LaakyRettt)CCoomm Convv(,, 4, 8) 1ontot0
CAVIA
MLP(4X25+64, 64, 62, 1x225LeeskyteelU)Coo//// 4, 8) 1y-eealyRelU--oorrr........
CoDA
W Rld 12
MLP(4x25, 64, 32, 4x2551eeayRe11 Convi, 1yLeaayyeelU--CCov4444..C...
FOCA
MLP(4X25464, 64, 32, 1x225LeeakyReLU)Coovlll has.://wwwwaaoooo....aaaaaaaaaaa 4, 3yLeaayyaeuu-Conr/4444.........
Table A.2 summarizes the network architectures. The green colored values indicate the dimension of context é. We the inner step K and step size a 100/1 and 0.001/1 for FOCA and CAVIA, respectively, and T as 0.1. For CoDA, we set the dimension of context as 2 (i.e., c) = 2), following the default setting of Kirchmeyer et (2022).
Training details We train all models with mini-batches of size 512 for 1,000 epochs using Adam (Kingma Ba, 2015) with the initial learning rate of 0.001 and the pushforward regularization (Brandstetter et al., 2021). We use the past observations from the previous 2N to N steps as the input the adaptation process. For CoDA, we performed a hyperparameter search and found that L1 regularization on the norm of as le-3 and W le-7 performs well.
Evaluation setting For in-training evaluations, we generate 2,048 mass-spring systems whose parameters (ka, R2, lss my, m2) sampled from (00.75, 1.25) and numerically solve the mass- spring system with Runge-Kutta 45 with T 5.0 and Al = 0.01. For out-of-training evaluations. generate 2,048 mass-spring systems whose parameters (ki, İ2, ka, m), sampled from 4(0.60, 0.75) U ((1.25, 1.35) for T = 5.0 with At = 0.01. All models take the first and second 0.25 seconds of observation for the task adaptation and predict 4.5 seconds of future states via the model rollout.
Additional results We provide the visualization of the model generalization errors to the change of mi, M2 on the mass-spring systems, as shown in Fig. A.2.
FOCA
Encoder
CAVIA
CoDA
1.25
4.8
1.00
3.6
3.0 mg
m
0.50
1.2
0.25
0.25
0.50
0.75
1.00
1.25
0.25
0.50
0.75
1.00
1.25 0.25
0.50
0.75
1.00
1.25
0.25
0.50
0.75
1.00
1.25
m
+m
-my
Figure A.2: In/out- distriibbt losses (in RMSE) on MS. Two parameters of MS (m²1, m2) are generated from |0.25, 1.25 while the other parameters fixed to 0.75. The orange box indicates the boundaries of the training parameter distribution.
C DETAILS OF LOTKA- VOLTERRA (LV) EXPERIMENTS
In this section, we provide the model architectures, training process, and additional experimental results for the Lokta-Volterra systems (Lotka, 1910).
14

---

Under review as a conference paper at ICLR 2023
FOCA
Encoder
CAVIA
CoDA
1.25
2.8
1.00
1.6
0.75
1.2
0.50
0.8
0.4
0.25
0.0
1:25
+
8
-
Figure A.3: In/out-of-distriuttion losses (in RMSE) on LV. Two parameters of LV (8, 7) are generated from [0.25, 1 1.25] while the other two parameters are fixed to 0.5. The orange box indicates the boundaries of the training parameter distribution.
The Lokta- Volterra system describes the interaction between prey-predator pair in ecosystem, formalized into the following ODE:
dr
ax - Bry,
dt
(A.2)
dy
ōxy - TA..
dt
where z, y respectively the quantity of the prey and the predator and a, 3,8,y, 8 define how two species interact.
Data generation For training data generation, generate 128 Lokta-Volterra systems whose parameters (a, 6) (0.5, 0.5) and (3, 7) ~ 4(0.5, 1.0) following Kirchmeyer et al. (2022). We then numerically solve the systems with Runge Kutta 45 for T 50.0 seconds with At 0.5 second time intervals.
Model architecture We employ the same network architectures of the MS experiments except with different input dimensions of 2 and bundling parameter N 20.
Training details We train all models with Mani-batches of size 128 for 5,000 epochs. Other details kept the same as mass mass-spring experiments.
Evaluation setting For in-training evaluations, we generate 2,048 Lokta- Volte-Volterr systems sim- ilarly to the training data generation. For oft-of-training evaluations, generate 2,048 Lokta- Volterra systems whose parameters (a, 5) (0.5, 0.5) and (B.r) ~ 140.25, 0.5) U 1.25)² for T 50.0 with Al 0.5. All models take the first and second 10.0 seconds of observation for the task adaptation and predict 30.0 seconds future states via the model rollout.
D DETAILS OF GLYCOLYTIC OSCILLATOR (GO) EXPERIMENTS
In this section, we provide the model architectures, training process, and additional experimental results for the glycolytic oscillators (Daniels & Nemenman, 2015).
15

---

Under review as a conference paper at ICLR 2023
The glycolytic oscillators describe yeast glycolysis dynamics with the following ODE:
dSi
K1S158
Ju
df
(1/K))S@
dS2 = 2 1 (1/K?) K2 S5) - k25255
S
df
dSg
1 N - 5s) k3 - S6)
dt
4Sp
SA (6) - K4SSS - r(S1 ST)
(A.3)
df
SSp
(N - S6) K444, K2
dp
dSp
S + 2k3 56) Kg
dt
1 (1/K4)
dS,
vu((S. S7) k:S7,
dt
where S1, S2, S3 S4, Ss, So, S7 (states) represent the concentrations of 7 biochemical species and Jo. k1, 22, 33, k1, Ks, ka, K1, q, N, A, and k are the parameters of the glycolytic oscillators.
Data generation For training data generation, generate 128 glycolytic oscillators with the fixed
parameters Jo = 2.5, 12 6, kg = 16, h₂ 100, kı 1.28, kş = 12. q = 4, N A, 4, 13, d 0.1 and k = 1.8 by sampling integer k 4/(00, 100) and K₁ 2/(0.5, 1.0). We adopt the values or ranges of the parameters from Kirchmeyer et al. (2022). We then numerically solve the systems with Runge-Kutte Kutta 45 for 5.0 seconds with At 0.05 second time intervals.
Model architecture We employ the same network architectures of the MS experiments except with a different input dimensions of 7 and bundling parameter N = 10. For CoDA, we set the context dimension as
Training details
We train all models with mini-batches of size 512 for 5,000 epochs. Other details are kept the same as mass-spring experiments.
Evaluation setting For in-training evaluations, generate 2,048 glycolytic oscillators similarly to the training data generation. For cut-of-training evaluations, we generate 2,048 glycolytic oscil- lators whose parameters hn U775, 80) U 105) and K2 (0.445, 0.5) U (1100, 1.05) and (Jo, k2, kg, ka, ks, kis ka, q, NN A, K, 1h, k) = (2.5, 6, 16, 100, 1.28, 12, 4, 1, 13, 0.1, 1.11,,1,, 1.8) for T 5.0 with At 0.05. All models take the first and second 0.5 seconds of observation for the task adaptation and predict 4.5 seconds future states via the model rollout.
Additional results We provide the visualization of the model generalization errors to the change of k2 K2 on the glycolytic oscillators, as shown Fig. A.4.
FOCA
Encoder
CAVIA
CoDA
0.80
0.5
0.75
Figure A.4: Infout-of-disrribtiio losses (in RMSE) on GO. The orange box indicates the houndaries of the training parameter distribution.
16

---

Under review as a conference paper at ICLR 2023
==
[= 10
== 15s
<= 20s
J = 25s
== 30
Target
Contod
Predition
Figure A.5: Out--ffutistiito distribution generalization on the Navier-Stokee equations: FOCA is trained with viscosity E [8 x 102
1.2 x 10-3] and tested with v = 6 x 10-4. The model can reconstruct states even after long autoregressive rollouts in unseen conditions.
E DETAILS OF
(NS) EXPERIMENTS
In this section, we provide the model architectures, training process, and additional experimental results for modeling the 2D Navier-Stokes equations (Stokes, 1851).
The Navier-Stori Stokes equations describe the dynamics incompressible flows with two-dimensional PDE. In vorticity form they can be written as:
dw v vAw + Øt
(A.4)
Vr
0
V X
where is the velocity field and is the vorticity, is the viscosity, and forcing term. The domain is subject to periodic boundary conditions.
Data generation We generate trajectories with temporal resolution of At and time horizon of / = 10. The space is discretized on a 32 x 32 grid and we set f(x, y) = 0.1 (s1(ni(22 (x + w)) + (2 y))), where x, y coordinates on the discretized domain (Yin al., 2021). For training data, similarly to Kirchmeyer et al. (2022), we consider training environments with
E (8- 10 4, 10- 1.0 10 1.1 10 1.2 10 3] respectively. Each environment contains a total of 100 different initial conditions for a total of 500 training sequences.
Model architecture For this experiment, we we employ the Fourier Neural Operator (FNO) (Li al., 2021). Thanks to their layers in the spectral domain and frequency mode pruning, frequency domain models have been proven to be well suited to model complex dynamical systems characterized by multitude of natural frequencies (Pathak et al., 2022; Poli et al., 2022). We employ FNO model with spectral convolution layers, hidden layers with width 10, and frequency mode pruning set to the 12 highest frequencies. Temporal bundling is set to N = 1. We set the context dimension of CoDA to 3. For models whose input includes the context, i.e. Encoder, CAVIA and FOCA, we keep the same context dimension |o] 64 in the previous experiments: the inferred context passes through a linear layer and resized to the grid size of 32 x 32 as an additional input channel of the FNO.
Training details We train all models with mini-batcher of size 16 for 100 epochs. Other details are kept the same as in the mass-spring experiments.
Evaluation setting For in-training adaptation, we consider environments with viscosity v E {8.5 1 9.5 10-4, 1.05 10 1.15 10-33 while for out-of-distribution adaptation we generate data from 4 environments with unseen ranges of viscosity, iee. и E (7.0 10 7.5 10 1.25 10 10-3, 1.30 10 10-31 00-3. Each adaptation environment is generated out of 30 different initial conditions. All the models take as input the first two steps (i.e. the first seconds) and predict the rest of the sequence with autoregressive rollouts.
17

---

Under review as a conference paper at ICLR 2023
Additional results We provide in Fig. A.5 an additional qualitative visualization showing the outto--diitributtin generalization on the Navier-Stokes equations. FOCA is trained with viscosity [8 x 10 +, 1.2 x 10 'I and tested with = 6 x 10 demonstrating its adaptation capability in challenging and unseen settings.
18


Reviews:
[{"review_id": "uzRTORufZR", "reviewer": "Reviewer_fp24", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors  propose FOCA (First-Order Context-based Adaptation), a learning framework to model sets of systems governed by common but unknown laws that differentiate themselves by their context. FOCA  learns to represent the common law through shared parameters and relies on online optimization to compute system-specific context and its training involves a bi-level optimization problem. This paper use an exponential moving average (EMA)-based method that allows for fast training using only first-order derivatives. ", "summary_of_the_review": "I would suggest to accept this paper. It certainly is above the average. ", "strength_and_weaknesses": "Pros: \nThis paper is well-written and present a new contextual meta-learning method FOCA.  FOCA jointly learns model and context parameters by solving online bi-level optimization efficiently. Instead of relying on higher-order derivatives, it employs EMA for faster and more stable training. The authors empirically demonstrated in both static and dynamic regression tasks on complex dynamical systems that FOCA outperforms or is competitive to baselines when evaluated on both in- and out-of-distribution contexts.\nCons: \nAs the authors say in the last part,  FOCA only models systems that are identifiable from the given dataset. In the setting where the systems are sampled from nearly identical systems, FOCA may not be able to infer context properly. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gQZZb5tkqZ", "reviewer": "Reviewer_dXrp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper investigates the problem of generalizing to new dynamical systems through an efficient first-order context-based algorithm (FOCA). The algorithm is based on a bi-level learning objective obtained through the modeling-then-identification approach. An EMA step is taken to accelerate the training process.", "summary_of_the_review": "The manuscript is overall in good quality but I'm not aware of its relevance towards dynamical systems and its novelty thereafter.", "strength_and_weaknesses": "The most confusing part of this manuscript is the *dynamical system* that is claimed to be adapting. But throughout the derivation I'm not seeing why it is even a dynamical system and what makes the adaptation of a dynamical system special in the derivation. I suspect that this is the largest weakness of the manuscript. Some detailed questions follow.\n\nQuestions:\n\n1. As suggested by the title of this paper, FOCA should be a method that focuses on the generalization of new dynamical systems. The referred related works such as GBML and CoDA are also works that are specific to dynamical systems. However, from the problem setting described in the paper, it is not very clear to me how $f$ is a dynamic system (or perhaps the relationship between $x$ and $y$). This characterization will be very important as it will set apart this problem from the general setting of Meta-learning, where the datasets are often required to have i.i.d. data. \n2. Following the last question, when $f$ represents a general dynamic system where data are not i.i.d. Solving (6) or line 9 of the algorithm with mean squared loss is perhaps non-trivial as one can no longer directly uses ERM which assumes data are i.i.d. It is indeed discussed in the CoDA paper that in general the MSE objective of (6) is intractable and can only be approximated. \n3. In section 3, it is stated that the learning objective for meta-learning is to “minimize the generalization error on $D_{tr}$. This is confusing as usually the generalization error is defined with the testing dataset while $D_{tr}$ denotes the training dataset. The paper actually did not define the test dataset for the generalization problem studied, thus this left me confused. \n4. For the time-series prediction tasks, FOCA is compared against CoDA, CAVIA, and encoder. However, it seems that CAVIA and encoder were originally designed for the general meta-learning setting and thus it is reasonable to have suboptimal performance for generalization to new dynamical systems. Therefore FOCA can be evaluated more comprehensively by comparing it with other algorithms designed specifically for learning dynamical systems, such as [1] and [2]. Besides, in the CoDA paper, two versions of the algorithm are proposed (with $\\ell_1$ or $\\ell_2$ constraints), could you state which version of the CoDA is used in this work as a baseline?\n\n[1]: Yin, Y., Ayed, I., de Bezenac, E., Baskiotis, N., and Gallinari, P. LEADS: Learning dynamical systems that generalize across environments. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. Advances in Neural Information Processing Systems, 2021a\n[2]: Wang, R., Walters, R., and Yu, R. Meta-learning dynamics forecasting using task inference. CoRR, abs/2102.10271, 2021c.\n\nMiscellaneous questions:\n1. For Table 1, it is better to include references to each algorithm in the caption. Otherwise, it can be a bit confusing (e.g. encoder can be a general terminology used while here it is used to refer to a specific algorithm?) The use of absolute value ($| \\cdot |$) to donate memory usage is also confusing. To my understanding, $\\theta$, $\\psi$, and $W$ are parameters that can be a vector/matrix. Therefore it is unclear what the definition of the absolute value of a vector/matrix is. \n2. For Figure 5, could you report the standard deviation of the prediction outcome obtained with the algorithms?\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "OYMt_CVTfIB", "reviewer": "Reviewer_t6gz", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper investigates a first order optimization-based framework for context-based meta-learning of static or dynamical systems. In the considered problem setup, the authors look into settings where a large class of models can be expressed as a combination of a function parametrized by shared parameters $\\theta$ and task-specific context parameters $c$. The authors propose FOCA as the framework which uses typical first order optimization methods for online optimization for training with a key ingredient being the use of asymmetric configurations of $\\theta$ to learn the shared parameters and the setting $\\bar\\theta$ used to identify the best context parameter set $c^i$ for any task $i$ during training. To maintain stability over training, the authors propose to use a slowly exponentially moving average of $\\theta$ as a proxy for $\\bar\\theta$ instead of drastically moving best values of $\\theta$. Finally, the authors test their proposed framework on a toy static regression problem and dynamical system datasets based on well-known ODE/PDE systems.", "summary_of_the_review": "Overall, the paper investigates an interesting idea and shows promising empirical results. However, I feel that the paper needs improvement on the thoroughness of experiments, adequate comparison with potential baselines and a better ablation study of the key component providing the improvement (if stability is the main issue, it can be combined with other baselines as well). Further, the tests for whether the context captures task specific information is rather weak and needs further evidence.\n\nNOTE: I want to take this space to note that the review is being submitted outside the stipulated timeframe, and I apologize from my end for the delay. I hope that the authors do find my comments helpful.", "strength_and_weaknesses": "Strengths:\n1. The overall idea behind FOCA is well-motivated along with its connections of using a target network in RL and other works in supervised learning. Addressing the stability of bilevel optimization problems, even outside of meta-learning is a common issue, and this work provides empirical evidence of the efficacy of a simple method.\n2. The empirical results for FOCA look quite promising on the datasets considered in the paper. However, there are related issues for comparison (see below).\n\nWeaknesses:\n1. Overall, despite the empirical focus of this paper, the discussion about weaknesses of other approaches to handle the dependence of best fit $\\hat c^i$ is rather weak. Backpropagating through the optimization trajectory in MAML style methods is an issue but various alternatives have been studied in recent years, some of which are even referenced in this paper. Did the authors investigate any such variations empirically on these datasets? Further, implicit gradients 'ideally' would require solving the inner optimization accurately but can be used as an approximate way of solving the $\\hat c$ dependency issue. However, I do agree that adding an exponentially moving average with large values of $\\tau$ in this case would benefit as well. Also, can a simpler two timescale learning rate scheme for updating the inner variable $c$ at a slower rate work? It will be worth investigating that.\n2. In the experiments, it is not clear as to what extent the context variables capture the task specific parameters vs shared parameters capturing the underlying functional form. For instance, is there any result on how the dimensionality of the $c$ affect performance when it is known that the actual cardinality of the context set if small. Similarly, changes in the overall architecture and investigating the distribution of errors across the tasks in the test set will indicate any undesired behavior.\n3. Evaluating CAVIA: How is CAVIA run during test inference? Have the authors ran a few gradient steps $K$ during test time as well? If yes, can the alternative of finding the best fit $\\hat c^i$ be tested as well? That way, we can actually decouple the effect of meta-training procedure and not induce indirect effects of the meta-test step.\n4. Architectural differences with CODA: In the experiments, the authors use $\\theta^i = \\theta + W\\hat c^i$ as the task-adapted parameter for CODA. This is starkly different from how the counterpart for FOCA is structured: $\\hat f^i(x) = f_\\theta(x,\\hat c^i)$. The main difference between CODA and FOCA can be said to be joint training of $\\theta,c$ vs interleaved training using EMA respectively. As such, I do not see the linear approximation for the model as anything inherent to CODA. I recommend the authors to have the same setup for comparison.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject