Neural Context Flows for Meta-Learning of Dynamical Systems
Authors: Roussel Desmond Nzoyem
David A.W. Barton
Tom Deakin
Venue: ICLR 2025 Poster (2025)
Keywords: meta-learning
OOD generalisation
physical sciences
neural ODEs
Primary area: learning on time series and dynamical systems
TL;DR: We propose a robust meta-learning framework for dynamical systems to quickly generalise to OOD parameter values.

Abstract: Neural Ordinary Differential Equations (NODEs) often struggle to adapt to new dynamic behaviors caused by parameter changes in the underlying physical system, even when these dynamics are similar to previously observed behaviors. This problem becomes more challenging when the changing parameters are unobserved, meaning their value or influence cannot be directly measured when collecting data. To address this issue, we introduce Neural Context Flow (NCF), a robust and interpretable Meta-Learning framework that includes uncertainty estimation. NCF uses Taylor expansion to enable contextual self-modulation, allowing context vectors to influence dynamics from other domains while also modulating themselves. After establishing theoretical guarantees, we empirically test NCF and compare it to related adaptation methods. Our results show that NCF achieves state-of-the-art Out-of-Distribution performance on 5 out of 6 linear and non-linear benchmark problems. Through extensive experiments, we explore the flexible model architecture of NCF and the encoded representations within the learned context vectors. Our findings highlight the potential implications of NCF for foundational models in the physical sciences, offering a promising approach to improving the adaptability and generalization of NODEs in various scientific applications.

Published as a conference paper at ICLR 2025
NEURAL CONTEXT FLOWS FOR META-LEARNING OF DYNAMICAL SYSTEMS
Roussel Desmond Nzoyem
David A.W. Barton
School of Computer Science
School of Engineering Mathematics and Technology
University of Bristol
University of Bristol
rd. wet.nncyeenguuguuinnee.. www.:brrrtool.... .uc.uk
david. Sevid.bbrttom. www.bbarttttriieeeoee istol. ac. uc.uk
Tom Deakin
School Computer Science University of Bristol
tom. deak n@rr stol. ac.uk
ABSTRACT
Neural Ordinary Differential Equations (NODEs) often struggle to adapt to new dynamic behaviors caused by parameter changes in the underlying physical system, even when these dynamics are similar to previously observed behaviors. This problem becomes meaning their value or influence cannot be directly measured when collecting data. To address this issue, introduce Neural Context Flow (NCF), a robust and inter- pretable Meta-Learning framework that includes uncertainty estimation. NCF uses Taylor expansion enable contextual self-modulation, allowing context vectors influence dynamics from other domains while also modulating themselves. After establishing theoretical guarantees, we empirically test NCF and compare related adaptation methods. Our results show that NCF achieves state-of-the-arr Diist--Dist performance on 5 out of linear and non-linear benchmark problems. Through extensive experiments, we explore the flexible model architec- ture of NCF and the encoded representations within the learned context vectors. Our findings highlight the potential implications of NCF for foundational models in the physical sciences, offering a promising approach to improving the adaptability and generalization of NODEs in various scientific applications.
challenging when the changing parameters unobserved,
INTRODUCTION
prototypical autonomous dynamical system describes the con- tinuous change. through time E IR+, of quantity x(t) E Rd. Its dynamics are influenced by its parameters c E R according to the (ordinary) differential equation
Neural ODE
Neural Contect Flow
dr
dt (t) = f(x()).c).
(1)
Figure 1: Predicted angle for the simple pendulum & g sin a = O. The OFA Neural ODE that disregards context fails to generalize its oscillation fre- quency vT/2- to unseen environ- ments, due in part to merged in- homogeneous training data. Our work investigates Neural Context Flows and related Deeta-Le methods to overcome this issue.
where f Rd X Rd + Rdd is the vector field. Learning dynamical system from data is synonymous with approximating f. task neural networks have been remarkably good at in recent years (Chen et al. 2018 Kidger 2022).
da
Consider the challenge of reconstructing the mechanical motion of an undamped pendulum given limited data from two distinct environments: Earth and Mars. Disregarding the physical vari- ations between these environments (e.g., gravitational constants, ambient temperatures, etc.), one might employ One-For-All (OFA) approach to lear a single pevironment-aggnoste vector field from all available data (Yin al. 2021a). This model would struggle adequately fit such

---

Published as a conference paper at ICLR 2025
heterogeneous data and would face difficulties generalizing to data from novel environments, e.g., the Moon. Fig. iillustrates this issue, using gravity as the underlying physical parameter. Alternatively, learning individual vector fields for each environment with a One-Per the-Per-Env (OPE) approach would miss inter-domain commonalitiess, proving both time-intensive and inhibiting rapid adaptation to new environments (Yin et al. 2021a). Given these constraints, it becomes imperative to develop methodology that can effectively learn what to learn from the aggregate data while simultancously accounting for the unique properties of each environment.
In Scientific Machine Learning (Cuomo et al. 2022) (SciML), the problem of generalization has largely been tackled by injecting domain knowledge. in Eg. (1) that captures as much of the dynamics as possible leads to lower evaluation losses (Yn et al. 202Tb). For such terms to be added, however, is essential to have knowledge of the physical parameters that change, which may then either be directly estimated, or predicted by a neural network within the vector field (Rackauckas et al. 2020). We are naturally left to wonder how to efficiently learn generalizable dynamical system when such physics are absent.
commonly understood that adding a term
Under constantly changing experimental conditions, two major obstacles to learning the parameter dependence of vector field can be isolated: + (P1) Limited data SciML models can be data- intensive (Hey et al., 2020 Yin et al.) 2021b), and limited data in each environment may not be enough to learn vector field suitable for all environments; (P2) Unobserved parameters during the data collection and modeling processes, one might be unfamiliar with the basic physics of the system. Solving these two problems would contribute to the efficient generalization of the learned models, particularly to related but massively reduce cost and complement recent efforts towards foundational models in the physical
(OoD) datasets. Fast OoD adaptation would
sciences (Subramanian et al., 2024; McCabe et al. 2023 Herde et al. 2024).
Neural Ordinary Differential Equations (Neural ODEs) Chen al. 2018 Weinan 2017 Haber
& Ruthotto 2017) have emerged as powerful backbone for learning ordinary, stochastic, and controlled differential equations (Kidger al. 2020). Trained using Differentiable Programming techniques (Nzoyem et al., 20233 Ma et al. 2021 Rackauckas et al. 2020). they have demonstrated broad utility with outstanding results in areas like chemical engineering (Owoyele & Pal (0022). geosciences (Shen et al. 22023), and climate modeling (Kochkov et al. 2024). In the increasingly
popular SciML subfield of solving parametric PDEs (Li et al. 2020 Takamoto et al., 2023 Li
et al. 2023 Subramanian et al. 2024 Koupai et al. 2024). Neural ODEs occupy place of choice due to their flexibility and performance (Yin et al. 2021a Lee & Parish| [202 Kirchmeyer et al.)
2022). That said, existing methods seeking generalize Neural ODEs to various parameter-defined environments fail to leverage information from environments other than the ones of interest. Not to mention the pervasive lack of interpretability and ways of accounting for the model's 's uncertainty.
This work presents Neural Context Flows (NCFs), a novel approach for multi-environment gener- alization of dynamical systems based on Neural ODEs. By leveraging the regularity of the vector field with respect to unobserved parameters, NCFs parameterize an environment-agnossi field and environment-speede is Taylor- expandee about these context vectors, effectively allowing information to flow between environments. Our contribution is threefold:
vector latent context vectors to modulate the vector field. The vector field
(1) We introduce a Meta-Leamnng methodology for enhancing the generalizability of dynamical systems. Our approach effectively addresses problems (P1) and (P2), challenging the prevailing notion that standard Deep Learning techniques are inherently inefficient for cuttof-Diisributtion (OoD) adaptation (Mouli et al. 2024).
(2) We present an interpretable framework for Multi-Task representation learning, incorporating a straightforward method for uncertainty estimation. This work extends the emerging trend of explainable linearly parameterized physical systems (Blanke & Lelarge 2024) to non- linear settings, thus broadening its applicability. For affine systems, we provide a lean proof for the identifiability of their underlying parameters.
(3) We provide a curated set of benchmark problems specifically designed for Meta- Metannng of dynamical systems. This collection encompasses a diverse range of problems frequently encountered in the physical science literature, all accessible through a unified interface.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
Learning data-driven Neural ODEs that generalize across parameters is only recent endeavor: To the best of the authors' knowledge, all attempts to solve (P1) and (P2) have relied on Multi-Tast and Meta-Learning to efficiently adapt to new parameter values, thus producing methods with varying levels of interpolation and extrapolation capabilities.
Multi Maski Learning. Multi-Tass Learning (MTL) describes a family of techniques where model is trained to jointly perform multiple tasks. In Scientific Machine Learning, one of the earliest methods attack this generalization problem LEADS (Yin al., 2021a). In LEADS, the vector field is decomposed into shared dynamics fø and environment-speene specciccctttttcccccc Iú components
dr
di = So(x) + 96
(2)
where the superscript e identifies the environment in which the dynamical system evolves, and {Ø, learnable neural network weights.
While LEADS excels at interpolation tasks, it performs poorly during extrapolation (Kirchmeyer
al. 2022p. Furthermore, it requires retraining new network ge each time a new environment be constraining in scenarios where adaptation is frequently required. Before LEADS, other MTL approaches had been proposed outside the context of dynamical systems (Caruana 1997 Rebuffi et al. 2017 2018. Lee & Parish 2021). Still, they do not address the crucial adaptation to new tasks which is our focus.
encountered, which
Leannng. Another influential body of work looked at Meta-Learning: a framework in which, in addition to the ATT,joint training scheme, shared representation is learned for rapid adaptation to unseen tasks with only minimum data (Wang et al. 2021). The seminal MAML (Finn et al. 2017)) popularized Gradient-B Bradee Meta-Learning (GBML) by nesting inner gradient loop in the typical training process. Since then, several variants aimed at avoiding over-fitting and reducing cost have been proposed, e.g. ANIL (Raghu et al. (209)) and CAVIA (Zintgraf al.) 2019). The latter is contextual learning approach (Garnelo et al.) 2018)) that partitions learnable parameters into some that are optimized on each environment, and others shared across environments, i.e., mela-trained.
Ad Wang et al., 2022) is one of the earliest Meta-LLearning approaches to target generalizable dynamical systems. It learns to represent time-invariant features of trajectory by an encoder network, followed by forecaster network to learn shared dynamics across the different environments, DyAd only performs well under weak supervision, that when the underlying (observed) parameters are made known to the loss function via penalization.
Arguably the most successful Meta- Detanne method for physical systems is CoDA (Kirchmeyer et al. 2022). which assumes that the underlying system is described by parametrized differential equations whose form is shared by all environments. However, these equations differ by the values of the vector field's weights, which are produced by (linear) hypernetwork. For the environment indexed by e, these weights are computed by
0- = + WEC,
where 0f and W are shared across environments, and 5 E Rd is an environment--pecifee latent context vector (or simply context). While achieves offte-of-the-art performance on many physical systems, the main limitation of CoDA is its hypernetwork approach, which might hinder parallelism and memory scaling to large root or target networks. In practice, methods based on hypernetworks require more computational resources to backpropagate and train, and exhibit more complex optimization landscape (Chauhan et al. 2023).
Taylor Expansion. This local approximation strategy finds numerous applications in SciMI., notably helping establish single-trajectrry geometrical requirements for lincar and affine system identification (Duan et al. 2020). Recently, Blanke & Lelarge (2024) modeled the variability of
'The general CoDA formulation encompasses the GMBI. adaptation nule (Kirchmeyer et al. 2022 Further-
more, MTL models he identified Eq. (3) with 0º 0.
3

---

Published as a conference paper at ICLR 2025
linearly parameterized dynamical systems with affine function of low-dimensional environment- specific context vectors. They empirically showed that this improved interpretability, generalization abilities, and computation speed. In our work, we similarly explain non-lineat systems, thus generalizing existing work with higher-order Taylor expansion. Additionally, we extract benefits such massive parallelizability and uncertainty estimation.
NEURAL CONTEXT FLOW
A training dataset Dr = {=f()) is defined as a set of trajectories collected from m related environments, with S trajectories per environment, each of length N N" over time horizon T >>> 0. Given Dr. we aim to find the neural network weights 0 that parameterize a vector field fa, along with several context vectors { Jee[ 1,m] that modulate its behavior such that
dx dt (8) = (t)
e [0, Vi [1,S], Ve [1, m.
We learn a single vector field for all environments our training set Dr.. The same vector field will be reused, unchanged, for future testing and adaptation to environments in similarly- Das-
The vector field fe is assumed to not only be continuous, but also smooth in its second argument ç. Exploiting this constraint, we "collect" information from other environments by Taylor-expanding fo around any other {E' her, where P S [1, m ] is a context poof containing == [P] 2 environment indices. This gives rise, for fixed and i, to p Neural ODES
dx ed
(1) = (1,55))
di
VjEP,
(5)
1 ( (0) (0),
where (-) E Rd, and {",&) E Rd,, T (", denotes the k-th order Taylor expansion of fo at ę" around ę In particular, T), can be written
(x (5,5))))' = folo;" (2) Vef(a ee)((" -
where o(-) captures negligible residuals. Eq. T6) directly consists of a Jacobian- Vector Product (IVP), making its implementation memory-efficient Since higher-order Taylor expansions of vector-valued functions do not readily display the same property, we provide the following proposition to facilitate the second-order approximation.
Proposition 1 (Second-order Taylor expansion with JVPs), Assume Rd x Rd + Rd is C² wrt its second argument. Let x E Rd E E Rd and define g ; g + vef(z, - E). The second-order
Taylor expansion of around any E TR is then expressed as
gg(t)) + Vo())E - E) +oll =
((x,))=
Proof. We refer the reader to Appendix A
By setting <= : E. and é : E, Proposition I yields an expression for T7, (-, ( in terms of JVPs, an implementation of which we detail in Appendix E. During training (as described below) trajectories from all Neural ODEs are used within the loss function.
This new framework is called Neural Context Flow (NCF) is the resulting model. The "flow" term refers to the capability the context from one environment, i.e., j. to influence predictions in another environment, i.c., e. by means of Taylor expansion. It allows the various contexts to not only modulate the behavior of the vector field, but to equally modulate theirs, since they are forced to remain close for the Taylor approximations be accurate. This happens while the same contexts are pushed apart by the diversity in the data. This self-modulation process and the beneficial friction
We note that P might include itself, and is reconstituted at each evaluation of Fq. 9 Its size is constant, each element indexing distinct environment for computational efficiency (see Appendix A.6)
4

---

Published as a conference paper at ICLR 2025
a
(b)
Candidate trajectionis
content vistoos
condition
trajectory
#f(1)
(1)
fo(o(())
(4,5)
41
Figure 2:
(a) Illustration of the Neural Context Flow (NCF). Given an initial condition (0) for training trajectory the environment e, NCF predicts in parallel, several candidate trajectories { her that are all compared to the ground truth x; upon which {0, (+, ) is updated. (b) Depiction of the 3-networks architecture for fo. where the state vector and arbitrary context vector projected into the same representational space before they can interact inside the main network.
creates is notably absent from other contextual Meta-Learning approaches like CAVIA ZZintgraf et al. 2019) and CoDA (Kirchmeyer et al. 2022).
We depict the NCF framework in Fig. along with a 3-networks architecture that lifts the contexts and the state vectors into the same representational space before they can interact. Concretely, 2 (t) and the context & are first processed independently into ãद (t) (by a state network in blue) and (by a context network in green) respectively, before they are concatenated and fed to a main network (in purple) to produce fe (zf (t), E). This explicitely allows the model to account for the potentially nonlinear relationship between the context and the state vector at each evaluation of the vector field.
3.1 META-TRAINING
Starting from the same initial state 2 (0) x x() the Neural ODEs in Eq. ( are integrated using differentiable numerical solver (Kidger 2022 Nzoyem al. 2023 Poli al. Chen 2018):
(1) 2, + [ T% (a) (7),,",
Ej
(8)
The resulting candidate trajectories are evaluated at specific time steps (ta (talneli,n) such that = 0 and IN = T. We feed these to a supervised (inner) loss function
N
X2 dg
E00.5". 1 22)
(t.)) (tt)ll
11012.
(9)
N
de
where d, dg, and de are the dimensions the state space, context vectors, and flattened network weights, respectively. || ll1 and || ||z denote the regularizing L' and I norms, with \ι and X their penalty coefficients, respectively. To ensure its independence from enviromment count, context pool size, and trajectory count per environment, the overall MSE loss function is expressed as in Eq. (10b; after which it is minimized wrt both weights 0 and contexts cl.m == {E" Je via gradient descent, alternating between updates:
S
L(0.E c:m Dw) mXSxp
H0.E0 E 2
(10)
The most effective way to train NCFs is via proximal alternating minimization as described Algo- rithm I Although more computationally demanding compared to ordinary alternating minimization (see Algorithm 3) it is adept at dealing with non-smooth loss terms like the L' norm in Eq. (Parikh et al. 2014). Not mention that ordinary alternating minimization (Algorithm 3 can easily lead to sub-optimal convergence (Attouch et al. 2010 Li et al. 2019),, while its proximal counterpart converges, with random initialization, almost surely second-order stationary points provided mild
5

---

Published as a conference paper at ICLR 2025
assumptions satisfied (see Theorem Appendix A.3]
l). We provide short discussion those assumptions in
The above comparison demands the definition of two variants of Neural Context Flows. NCF-tz uses second-ored order Taylor expansion in equation [S] and is trained via preximal alternating minimization (Algorithm [ NCF-11 the other hand, is implemented using first-order Taylor expansion, and trained using ordinary alternating minimization (Algorithm 3 Although less expressive than NCF-12, NCF-1 is faster and serves powerful baseline in our experiments.
Theorem 1 (Convergence to second order stationary points), Assume that C(., Dr)) satisfies the Kurdyka-Loye Lojasiewiez (KL) property, is bi-smooth, and v. Dir) Lipschitz continuous on any bounded subset of domain Re x I& Under those assumptions, let (fo, (p-)) be a random initialization and (0g. 5q Elmm) be the sequence generated by Algorithm 1 If the sequence (0,. el:m)
bounded, then converges to a second-order stationary point of C( Drr) almost surely.
Algorithm 1 Proximal Alternating Minimization
Algorithm Sequential Adaptation of NCF
1: Input: Du == (e") 1.m]
1: Input: Dau (D (D"
2: 0o E Rd randomly initialized
A e Rd learned
3: <0-Tm U 5. where 5" e R
3: g 0 Rdde, Ve' E [a,bb
f >
4: Amnx : 2 R ; 0
5: for e' t a, ao
5: for < 1, fma do
6: L(E D
)=
6:
clam ((0) == L(A, 5q Du) + 2/11 11-00p-110 0
81 Σ 1(0,E" , ,24))
7:
0, 0
8:
repeat
7: repeat
9:
Og 0g n/0 ()0gg)
8:
5d + - )"')
10:
until Oy converges
9: until &d converges
11:
H((5 tm == [(, : 5 mm.m Du)
10: end for
+ $ 11c
1:m
5q
cl:m
12:
cl:m 2
Sg
Proof. The proof of Theorem | is straightfor- ward by adapting Assumptions 1 and 4, then Thcorem 2 of (Li ct al. 2019).
13: repeat
14: 5T
+- E
178 VH(E).m)
until el:m converges
16:
end for
3.2 ADAPTATION (OR META-TEST HESA-TES
Few-shot adaptation to new environment e' [[,b]] in or out of the meta-training distribution requires relatively less data, i.c., its trajectory count < Here, the network weights are frozen, and the goal is to find a context 5d such that
dx dt (t) = fo.
Vi [1,
(11)
Our adaptation rule, as outlined in Algorithm extremely fast, converging in seconds for trainings that took hours. In scenarios where we want to adapt to more than one environment, we outline bulk version in Appendix A..5 (see Algorithm 4). Although the bulk adaptation algorithm is parallelizable and framed in the same way as during meta-training, does not allow for flow of contextual information, since this causes significant accuracy degradation. Most importantly, disabling the Taylor expansion this stage limits the size of the context pool to p = and significantly improves memory efficiency, an important resource when adapting to large number of environments.
3.3 IND AND OoD TESTING
We distinguish two forms of testing. For In-Domain (InD) testing, the environments are the same the ones used during training. In InD testing data, the underlying parameters of the dynamical system we aim to reconstruct are unchanged, and so the meta-learned context vectors are reused. fut-of-Distribution (OoD) testing considers environments encountered during adaptation. The data

---

Published as a conference paper at ICLR 2025
is from the same dynamical system, but defined by different parameter values (either interpolated or extrapolated). In all forms of testing, only the main predicted trajectory corresponding to = j = e' is used in the MSE and MAPE metries computation (see Appendix B.4 for definitions), thus returning to a standard Neural ODE. Other candidate trajectories are aggregated to ascertain the model's uncertainty. The initial conditions that define the trajectories always unseen, although their distribution never changes from meta-training to meta-testing.
MAIN RESULTS
In this section, we evaluate the effectiveness of our proposed framework by investigating two main questions. (i) How good are NCFs at resolving (P1) and (P2) for interpolation tasks (Section 4 (ii) How does our framework compare to SoTA Meta-Learning baselines (Section (.3)) Further questions regarding interpretability, scalability, and uncertainty estimation are formulated and addressed in Appendices A.2 A.7 and B.2 respectively.
4.1 EXPERIMENTAL SETTING
The NCF framework is evaluated on five seminal benchmarks. The Simple Pendulum (SP) models the periodic motion of a point mass suspended from a fixed point. The Lotka-Volterra Volterra (LV) system
models the dynamics of an ecosystem in which two species interact. Additional ODEs include a simple model for yeast glycolysis: Glycolytic-Oscillan
(GO) (Daniels & Nemenman 2015). and
the more advanced Sel'kov Model (SM) (Strogatz) 2018 Kov 1968). Like GO. SM is non- linearly parameterized, but additionally exhibits starkly different behaviors when its key parameter consider three PDEs, all with periodic boundary conditions, and cast as ODEs via the method of lines: the non-linear oscillatory Brusselator (BT) model for autocatalytic chemical reactions (Prigogine & Lefever 1968), the Gray-Scott (GS) system also for reaction-diffusia in chemical settings (Pearson 1993). and Navier-Stokes (NS) for incompressible fluid flow (Stokes et al. 1851).
is varied (see Fig. 4(a). Finally,
For all problems, the parameters and initial states are sampled from distributions representative of real-world problems observed in the scientific community, and the trajectories are generated using a time-adaptive 4th-order Runge-Kutt Kutta solver (Virtanen et al. 2020). For LV, GO, GS, and NS, we reproduce the original guidelines set in (Kirchmeyer et al. 2022). while exposing the data for ODE and PDE problems alike via common interface. Such use of synthetic data is a common practice in this emerging field of generalizable dynamical systems, where the search for unifying benchmarks remains an open problem (Massaroli et al. (020). This need for shared datasets and APIs has motivated the Gen-GennDnnenn -aynami cs open- open-source initiative, our third and final contribution with this paper. Further details, along with the data generation process, are given in Appendix B
We now highlight a few key practical considerations shared across experiments. We use the networks architecture depicted in Fig. 2 to suitably process the state and context variables. The dimension of the context vector, the context pool's size and filling strategy. and the numerical integration scheme vary across problems. For instance, we set de = 1024 for LV, de 202 for NS, and de = 256 for all other problems; while = for LV and SM, p = 4 for LV and GO, and 3 for all PDE problems. Other hyperparameters are carefully discussed Appendix D
4.2 INTERPOLATION RESULTS
This experiment explores the SP problem dis- cussed in Fig. [ During meta-training. we use 25 environments with the gravity g regularly spaced [2, 24). Each of these environments contains only 4 trajectories with the initial con- ditions x; (0) ~ (u(-5 3. ((-,,)))) Dur- ing adaptation, we interpolate 2 new environ- ments with g E (10.25, 14.75}. each with 1 trajectory. For both training and adaptation testing
Table I: Training and adaptation testing MSEs (L) with OFA, OPE, and NCF-t, on the SP problem.
TRAIn (x10-1)
ADAPT (x10-
OFA OPE NCF
9.49 1 0.04
115000 + 3200
0.18 1 0.02
459.0 + 345.0
0.10 1 0.03
0.0356 0.001
scenarios, we generate 32 separate trajectories.
7

---

Published as a conference paper at ICLR 2025
Table In-Domain (InD) and adaptation (OoD) MSEs (1) for the LV, GO, SM, BT, GS and NS problems. The best is reported in bold. The best of the two NCF variants is shaded in grey
LV (x10-3)
GO (x10-1)
#PARAMS
IND
OoD
#PARAMS
IND
OoD
CAVIA CODA
305246
91.01163.6
120.1
130711
64.01114.1
463.4184.9
305793
1.40=0.13
2.19
135390
5.06+0.81
4.22=4.21
NCF-h h
308240
6.73+0087
1.92+1.04
131149
40.3+9.1
19,4=1.24
NCF- to
308240
1.68+0.32
1.99
131149
3.33:0.14
2.83+0.23
SM (x10-)
BT (x10
#PARAMS
IND
OoD
#PARAMS
IND
OoD
CAVIA
50486
979. 10.1+111..
859. 77007
116665
21.93+1.8
22.6+7.22
CODA
50547
156.01140.52
8.28
119679
25.40
19.47
NCF-4
S0000
680.6.1320.1
677.2.18.7
117502
21.5318.99
20.8
NCF-
50000
6.42.10.4 0.41
2.03.10.12
117502
3.46.10.09
3.77
GS (x10
NS (x10
#PARAMS
IND
OoD
#PARAMS
IND
OoD
CAVIA
618245
69.9+21.2
68.014.2
310959
128. 199
126.4=20.7
CoDA
619169
1.23++0.14
0.75++0.65
309241
1.69
7.08+0.07
NCF-1 610942
7.64+0.70
5.57+0021
310955
2.98+00.09
2.83=0.06
NCF-12
610942
0.124
3.40+0.51
310955
2.92-0.08
2.79+0.00
Table emphasize the adaptation-tine time merits of Meta-Learning approaches like NCFs in lieu of baselines where one context-agnostic vector field trained for all environments indiscriminately (One-For-All or OFA), or for each environment independently (One-Per-Env or OPE), Additional results and further analysis of these differences is done in Appendix B.3
4.3 EXTRAPOLATION RESULTS
Meta-Lest Leta-Learning algorithm is only as good as its ability to extrapolated to unseen environments. In this section tackle several one-shot generalization problems, i.e. S' 1. We consider the LV, GO, SM, GS, BT, and NS problems, which all involve adaptation environments outside their meta- teeta-training distributions. We compare both variants of NCF to baselines: CAVIA (Zintgraf et al. 2019 which is conceptually the closest GBMI. method to ours, and CoDA-E, (Kirchmeyer et al. 2022). Their hyperparameters. laid out in Appendix D are tuned for a balance of computational efficacy and performance, all the while respecting each baseline's i key assumptions.
At similar parameter counts Table 2 shows that CAVIA is the least effective for learning all six physical systems, with tendency to overfit on the few-shots it receives during both meta- veta--taining and meta-testing (Mishra et al.) (017). NCF-1z achieves SoTA OoD results on out problems. While CoDA retains its superiority on GS, we find that it struggles on non-linear problems, most notably on the SM problem whose trajectories go through distinct attractors in a Hopf bifurcation.
This shows that CoDA's ability to automatically select low-rank adapta- tion subspaces via linear hypernetwork decoder is limited, particularly for highly nonlinear systems. We note that NCF-t, equally fails to ac- curately resolve SM and other non-linear problems, thus illustrating the value of second-order Taylor expansion. This suggests that higher-order Taylor-base based regularization prevents NCF-tz from learning spurious as- sociations, which is problem commonly associated with poor Ool) generalization (Mouli et al. 2024).
402
Figure 3: Grid-wise adap- tation on the LV problem showing low MAPES ()).
Grid-wide adaptation on L.V. We report in Fig. 3 how well NCF per- forms 625 adaptation environments obtained by varying its parameters
Thhis accounts for the total number of learaalee parameters in the framework excluding the context vectors. whose size might vary with the method explained in Appendix
8

---

Published as a conference paper at ICLR 2025
and ó a 25x uniform grid. It shows consistently low MAPE below 2%, except in the bottom right corner where the MAPE rises to roughly 15%, still remarkably low value for this problem. We highlight the 9 metaa-rrinnng environments in yellow, and the 4 environments used for OoD adaptation for Table in indigo. Remarkably, our training environment's convex-hull where adaptation is particularly low is much larger compared CoDA's (Kirchmeyer et al. 2022 Figure 3).
Nonlinear adaptation on SM. We investigate how the models perform across the SM attractors
depicted in Fig. [a] training environments ej and 62 fall into limit cycle (L1), 6g and Ca collapse
to a stable equilibrium (E), while eş and és fall into another limit cycle (L.2.. Fig. [4cc presents the adaptation MSE for each environment in the OoD testing dataset. While CoDA equally fails
to capture all 3 attractors, we observe that CAVIA and NCF-tz tend to favor E. But unlike CAVIA,
NCF-t2 succeeds in capturing limit cycles as well, as evidenced by the low MSE on all environments.
NCF
CODA
CAVIA
107
16a
10a
10-1
10
0
250
500
750
1000
52
1500
10
Iterations
02
(a) Trajectories
(b) Training losses
(c) OoD MSEs per environment
Figure 4: (a) Sample trajectories from the SM problem illustrating the Hopf bifurcation; (b) Losses during meta-training; (c) Subsequent adaptation MSE per method per environment.
DISCUSSION
5.1 BENEFITS OF NCFS
Neural Context Flows provide a powerful and flexible framework for learning differential equations. They can handle irregulallyssappled sequences like time series, and can easily be extended to general regression tasks (Finn et al. 2017). Some of their other desirable properties are highlighted below.
Massively parallelizable. Eq. (10) indicates that NCFs are massively parallelizable along di- rections. Indeed, evaluations L can be vectorized across m x p environments, and across all S trajectories. This leads to better use of computational resources for meta-training, We provide details on such vectorized NCF implementation in Appendix along with a thorough discussion on its scalability in Appendix(A.7
Interpretable. Understanding how a model adapts to new physical settings is invaluable in many sci- entific scenarios. Moving information from one en- vironment to another via context-affine transforma- tions provides a powerful framework for explaining Multi-Task Learning (Blanke & Lelarge 2024).. With contextual self-modulation via Taylor expansion, we generalize this framework while maintaining compu- tational efficiency. In Fig. 5) for instance, we show- case system identification with NCF-1, where the underlying physical parameters & and 5 of the LV problem are recovered up to linear transform. The corresponding experiment is detailed in Appendix A.2 which further illustrates NCF's robustness to noise during system identification.
Kladddie 3.
CoDA
NCF
Figure 5: Interpretability with NCF.
Proposition 2 (Identifiability of affine systems), Assume de 2 de that P is full-rank, and that fave is differentiable in its second argument. In the limit of zero training loss in Eq. (10), fo trained with first-order Taylor expansion and the Random-All pool pool-filling strategy is affine on an open region of de Furthermore, there exists Q E Rddxx xdx and q E TR such that for any meta-trained E { and its corresponding underlying parameter e E {ee m we have c = QE +
"A discussion which neighboring contexts should go into the context pool is provided in Appendix A.6

---

Published as a conference paper at ICLR 2025
To theoretically demonstrate the capacity identify physical systems, consider learnable vector field fuve x + R is that is affine i.c. 3P E Rddd, and p E R such that (: = Pctp. Under these assumptions, the predictor fo R x R IR satisfies Proposition 2 above, inspired by Proposition of (Blanke & Lelarge 2024). A detailed proof is provided in Appendix A.2
Extendable. The relative simplicity of our formulation makes NCFs easily adaptable to other works involving Neural ODE models. For instance, it is straightforward to augment the vector field as in APHYNITY (Yin et al. 2021b), augment the state space as in ANODE (Dupont et al. 2019 or control the vector field as in (Massaroli et al. 2020).
Uncertainty requirements are increasingly critical in Meta-Leearned dynamical systems (Liu et al.). NCFs provide, by their very definition, a measure of the uncertainty the model. During testing, all candidate trajectories stemming from available contexts (meta-trained, adapted, or both) can be used to ascertain when or where the model is most uncertain in its predictions. In Fig. Fig.[ffor forr instance,
Provides uncertainty.
Figure 6: Uncertainty estimation with NCF.
visualize the means and standard deviations (scaled 10 folds for visual exposition) across 9 and 21 candidate forecasts for the Lotka- Volterra and Sel'kov Model problems, respectively. Additional results with quantitative metrics are provided in Appendix B.2
Lotka-Volterra
5.2 LIMITATIONS
Further theory. While massively parallelizable, extendable, interpretable, and applicable to count- less areas well outside the learning of dynamics models, Neural Context Flows still need additional theoretical analysis to explain their effectiveness and uncover failure modes. While the provided Propositions I and partly compensates for this shortage, we believe that this creates an interesting avenue, along with other limitations outlined below, for future work.
Regularity assumptions. Another limitation faced by Neural Context Flows lies within the assump- tions they make. While differentiability of the vector field f wrt is encountered with the majority of dynamical systems in science and engineering, some are bound to be fundamentally discontinuous. NCFs break down in such scenarios and would benefit from the vast body of research numerical continuation (Allgower & Georg 2012).
Additional hyperparameters. Finally, NCF introduces several new hyperparameters such as the context size dç, the context pool size p, the pool-filling strategy, the proximity coefficient / in Algorithm [ and many more. Although we offer insights into their roles the Appendix C we acknowledge that tuning them complicates the training process.
5.3 CONCLUSION AND FUTURE WORK
This paper introduces Neural Context Flows (NCFs), an innovative framework that enhances model generalization across diverse environments by exploiting the differentiability of the predictor in its latent context vector. The novel application of Taylor expansion NCFs facilitates vector field modu- lation for improved adaptation, enhances interpretability, and provides valuable uncertainty estimates for deeper model understanding. Our comprehensive experiments demonstrate the robustness and scalability of the NCF approach, particularly with respect to its most demanding hyperparameters. Future research will explore the limits of NCFs and their adaptation to even more complex scenarios. This work represents promising step toward developing foundational models that generalize across scientific domains, offering a fresh and versatile approach to conditioning machine learning models.
ETHICS STATEMENT
While the benefits NCFs are evidenced in Section 5.1
its negative impacts should not be neglected. For instance, malicious deployment of such adaptable models in scenarios they were not designed for could lead to serious adverse outcomes. With that in mind, our code, data, and models are openly available at httpt::pp.... / ttt..b.ttuu....gg.g com www/ddoooo/onnttteeeooeo
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by UK Research and Innovation grant EP/S022937/1: Interactive Artificial Intelligence, EPSRC program grant EP/R006768/1: Digital twins for improved dynamic design, EPSRC grant EP/X039137/1: The GW4 Isambard Tier-2 service for advanced computer architectures, and EPSRC grant EP/TO22205/1: JADE: Joint Academic Data science Endeavour-2. We extend our sincere gratitude to the anonymous reviewers whose insightful feedback and rigorous commentary substantially enhanced the quality and clarity of this manuscript. We also acknowledge Amarpal Sahota for valuable discussions on uncertainty quantification methodologies.
REFERENCES
Eugene L Allgower and Kurt Georg. Numerical continuation methods: an introduction, volume 13. Springer Science Business Media, 2012.
Hédy Attouch, Jérôme Bolte, Patrick Redont, and Antoine Soubeyran. Proximal alternating minimization and projection methods for nonconvex problems: An approach hased on the kujdyka-tnjasiewicr inequality. Mathematics of operations research, 35(2):438 457, 2010.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Jesse Bettencourt, Mathew Johnson, and David Duvenaud. Taylor-mode automatic differentiation for higher- order derivatives in jax. In Program Transformations for ML Workshop at NeurIPS 2019, 2019.
Matthieu Blanke and Mare Lelarge. Interpretable meta-lcarning of physical systems. In ICLR 2024-The Twelfth International Conference on Learning Representations, 2024.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula. Adam Paszke. Jake VanderPlas, Skye Wanderman Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. htt.hhtp::/// htttt::://tttt tht.btt.t....... comm gowp://oolle////g
Rich Caruana. Multitask learning. Machine learning, 28:41-75, 1997.
Vinod Kumar Chauhan, Jiandong Zhou, Ping Lu, Soheila Molaci, and David A Clifton. brief review of hypereetwors deep learning. arxiv preprint arXiv:2306. 06955, 2023.
Ricky T. Chen. torchdiffeq. 2018 URL
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations, Advances neural information processing systems, 31, 2018.
Salvatore Cuomo, Vincenzo Schiano Di Cola, Fabio Giampaolo, Gianluigi Rozza, Maziar Raissi, and Francesco Piccialli. Scientific machine learning through thysics-informen neural networks: Where we are and what's 's next. Journal of Scientific Computing, 92(3):88, 2022.
Bryan C Daniels and Ilya Nemenman. Efficient inference of parsimonious phenomenological models of cellular dynamics using s-systems and alternating regression. PloS one, 10(3):e01 19821, 2015.
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keek, Turii Kemaev, Michael King. Markus Kunesch, Lena Martens, Hamza Merzie, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan. Roman Ring. Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojciech Stokowiec, Luyu Wang. Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http://githh com hamm:/gooleeeeeee deepmiid
P Kingma Diederik. Adam: A method for stochastic optimization. (No Tirle), 2014.
Xiaoyu Duan, JE Rubin, and David Swigon. Identification of affine dynamical systems from single trajectory. Inverse Problems, 3688)445 2020.
Emilien Dupont, Arnad Doucet, and Yee Whye Teh. Augmented neural odes. Advances in nearal information processing systems, 32, 2019,
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1135. PMLR, 2017.
11

---

Published as a conference paper at ICLR 2025
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan. Yoe Whye Tch, Danilo Rezende, and SM Ali Eslami. Conditional neural processes. International conference on machine learning, pp. 1704-1713, PMLR, 2018.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34(1):014004 2017.
Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes. arXiv preprint 22t::::005.9900 2024.
Hansika Howamalage. Klaus Ackermann, and Christoph Bergmeir. Forecast evaluation for data scientists: common pitfalls and best practices. Data Mining and Knowledge Discovery, 37(2):788 832, 2023,
Tony Hey, Keith Butler, Sam Jackson, and Jeyarajan Thiyagalingam. Machine learning and big scientific data. Philosephical Transactions of the Royal Society A, 378(2166))2010003)) 2020.
Patrick Kidger. On neural differential equations. arXiv preprint arXiv:2202.02433. 2022.
Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time series. Advances in Neural Information Processing Systems, 33:6696- 6707,, 2020.
Matthieu Kirchmeyer, Yuan Yin. Jérémie Donå, Nicolas Baskiotis, Alain Rakotomamoniy. and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning, pp. 1283-11301. PMLR, 2022.
Dmitrii Knchkov, Janni Yuval, Ian Langmore, Peter Norgaard, Jamie Smith, Griffin Mooers, Milan Klöwer, James Lottes, Stephan Rasp, Peter Düben, et al. Neural general circulation models for weather and climate. Nature, 632(8027): (11/80027)1160-10000 1066, 2024.
Armand Kassal Koupai, Jorge Misfut Benet, Yuan Yin, Jean-Noël Vittaut, and Patrick Gallinari. Geps: Boosting generalization in parametric pde neural solvers through adaptive conditioning. arXiv preprint arXiv:2410.23889. 2024.
Kookjin Lee and Eric Parish. Parameterized neural ordinary differential equations: Applications to computa- tional physics problems. Proceedings of the Royal Society A, 477((223):202101622 2021.
Qing Li, Jiahui Geng, and Steinar Evje. Identification of the flux function of nonlinear conservation laws with variable parameters. Physica D: Nonlinear Phenomena, 451:133773, 2023.
Qiuwei Li, Zhihui Zhu, and Gongguo Tang. Alternating minimizations converge to second-order optimal solutions. In International Conference on Machine Learning, pp. 3935-3943 3943. PMLR, 2019.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arX0v:2010.08895, 2020.
Jiaqi Liu, Jiaxu Cui, Jiayi Yang, and Bo Yang. Stochastic neural simulator for generalizing dynamical systems across environments.
Yingbo Ma. Vaibhav Dixit, Michael Innes. Xingjian Guo, and Chris Rackauckas. A comparison of automatic differentiation and continuous sensitivity analysis for derivatives of differential equation solutions. 2021 IEEE High Performance Extreme Computing Conference (HPEC). pp. 1-9. IEEE, 2021.
Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissecting neural odes. Advonces Neural Information Processing Systems, 33:3952-3963, 2020.
Michael McCabe, Bruno Régaldo-Sain Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg. Siavash Golkar, Geraud Krawezik, Francois Lanusse, al. Multiple physics pretraining for physical surrogate models. arXiv preprint arXiv:2310.0299 2023.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-learner. arXiv preprint arXii:1707.03144 2017.
S Chandra Mouli, Muhammad Alam, and Bruno Ribeiro. Metaphysica: Improving ood robustness in physics- informed machine leamning In The Twelfth International Conference Learning Representations, 2024.
12

---

Published as a conference paper at ICLR 2025
Roussel Desmond Nzoyem, David AW Barton, and Tom Deakin. A comparison of mesh-free differentiable programming and data-driven strategies for optimal control under pde constraints, In Proceedings of the SC'23 Workshops of The Interaational Conference on High Performance Compuring, Network, Storage, and Analysis, pp. 21 28, 2023.
Opeoluwa Owoyele and Pinaki Pal. Chemnode: A neural ordinary differential equations framework for efficient chemical kinetic solvers. Energy and AI, 7:100118, 2022.
Neal Parikh, Stephen Boyd, al. Proximal algorithms. Foundations and trends@ Optimization, 1(3):127-239) 2014.
Junyoung Park, Federico Berto, Arec Jamgochian, Mykel Kochenderfer, and Jinkyoo Park. First-order context- based adaptation generalizing to new dynamical systems. 2023.
John Pearson. Complex patterns in simple system. Science, 26 261(5118):189-11111111 1922 1993.
Michael Poli, Stefano Massaroli, Atsushi Yamashita, Hajime Asama, Jinkyoo Park, and Stefano Ermon. Torch- dyn: Implicit models and neural numerical methods in pytorch.
Ilya Prigogine and René Lefever. Symmetry breaking instabilities in dissipative systems. ii The Journal of Chemical Physics, 48(4):1695-1770, 1968.
Tiexin Qin. Hong Yan, and Haoliang Li. Generalizing to new dynamical systems via frequency domain adaptation.
Christopher Rackauckas, Yingho Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Dominic Skinner, Ali Ramadhan, and Alan Edelman. Universal differential equations for scientific machine learning. arXiv preprint orXtv:2001.04580 2020.
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse? towards understanding the effectiveness of maml. arXiv preprint 1wXi:::909.09157, 2019.
Prajit Ramachandran, Barret Zoph, and Quoe Le. Swish: a self-gated activation function. arXiv: Neu- ral and Evolationry Computing, 2017. http::/: heaat/pp...emmttt.cccho......r ott.aaaaaiaiaac...oooooooooooaaa org/CozpuslD
196158220
Sylvestre-Alvise Rebufli, Hakan Bilen, and Andrea Vedaldi. Learning multiple visual domains with residual adapters. Advances in neural information processing systems, 30, 2017.
Sylvestre-Alvise Rebuffi, Hakan Bilen, and Andrea Vedaldi. Efficient parametrization of multi-domain deep neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp- 8119-8127, 2018.
BE Sel' Kov. Self-oscillations in glycolysis I, simple kinetic model. Luropean Journal of Biochemistry, 4(1): 79 86, 1968.
Louis Serrano, Armand Kassaï Koupai, Thomas Wang, Pierre Erbacher, and Patrick Gallinari. Zebra: In-context and generative pretraining for solving parametric pdcs. arXiv preprint 44Xiv2410..03477. 2024.
Chaopeng Shen, Alison P Appling, Pierre Gentine, Toshiyuki Bandai, Hoshin Gupta, Alexandre Tartakovsky, Marco Baity-Jest Fabrizio Fenicia, Daniel Kifer, Li Li, et al. Differentiable modelling to unify machine learning and physical models for geosciences. Nature Reviews Earth & Environment, pp. 1-16. 2023.
Michael Spivak. Calculus on manifolds: a modern approach to classical theorems of advanced calculus. CRC press, 1965.
George Gabriel Stokes al. On the effect of the internal friction fluids on the motion pendulums. 1851.
Steven H Strogatz. Nonlinear dynamics and chaos: with applications physics. biology. chemistry, and engineering. CRC press, press 2018
Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael W Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. Advances in Neural Information Processing Systems, 36, 2024.
Makoto Takamoto, Francesco Alesiani, and Mathias Niepert. Learning neural pde solvers with parameter-guided channel attention. International Conference on Machine Learning, pp. 33448-33467, PMLR, 2023.
Ezio Venturino. The influence of diseases on lotka-volterra systems. The Rocky Mountain Journal of Mathematics, pp. 381-402, 1994.
13

---

Published as a conference paper at ICLR 2025
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy. David Cournapeau. Evgeni Burovski, Pearu Peterson, Warren Weekesser, Jonathan Bright, Stéfan J. van der Walt, Marthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, CJ Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, António H. Ribeiro, Fabian Pedregosa, Paul Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing Python. Nature Methods, 17:261-272, 2020. doi: 199--441922001-006666
Andrey Voynov and Artem Babenko. Unsupervised discovery of interpretable directions in the gan latent space. International conference on machine learning, pp. 9786-9796. PMLR, 2020.
Haoxiang Wang, Han Zhao, and Bo Li. Bridging multi-task learning and meta-learning Towards efficient. training and effective adaptation. International conference on machine learning, pp. 10991-11002. PMIR, 2021.
Wang, Robin Walters, and Rose Meta-learning dynamics forecasting using task inference, Advances in Neural Information Processing Systems, 35:21640-21653, 2022.
Gerhard Wanner and Ernst Hairer. Solving ordinary differential equations II, volume 375. Springer Berlin Heidelberg New York, 1996.
Be Weinan. proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 1(5):1-11, 2017.
Lifeng Wu, Sifeng Liu, and Yinao Wang. Grey vootero model and its application. Technological Forecasting and Social Change, 1720-1770-, 2012.
Yuan Yin, Ibrahim Ayed, Emmanuel de Bézenac, Nicolas Baskiotis, and Patrick Gallinari. Leads: Learning dynamical systems that generalize across enviromments. Advances in Neural Information Processing Systems, 34:7561-7573, 2021a.
Yuan Yin, Vincent Le Guen, Jérémie Dona, Emmanuel de Bézenac, Ihrahim Ayed, Nicolas Thome, and Patrick Gallinari. Augmenting physical models with deep networks for complex dynamics forecasting. Journal of Statistical Mechanics: Theory and Experiment, (021(12)122012,, 2021b.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvorekk, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in 2020.
neural information processing systems, 33:18795--888
I Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation meta-lear,ing. In International Conference on Machine Learning, pp. 7693-7702, PMLR. 2019.
14

---

Published as a conference paper at ICLR 2025
SUPPLEMENTARY MATERIAL FOR NEURAL CONTEXT FLOWS FOR META-LEARNING OF DYNAMICAL SYSTEMS
ALGORITHMS & PROOFS
16
A.1 SECOND-ORDER TAYLOR EXPANSION WITH JVPS
16
A.2 IDENTIFIABILITY OF PHYSICAL PARAMETERS
16
A.3
CONVERGENCE OF PROXIMAL ALTERNATING MINIMIZATION
18
ORDINARY ALTERNATING MINIMIZATION
19
A.5
BULK ADAPTATION ALGORITHM
19
A.6 WHLAT'S IN CONTEXT POOL
20
A.7 SCALABILITY OF THE NCF ALGORITHMSS
20
DATASETS & ADDITIONAL RESULTS
22
B. GEN-DYNAMICS
22
B. UNCERTAINTY ESTIMATION
22
B. 3 SIMPLE PENDULUM (SP)
23
LOTKA VOLTERRA (LV)
27
B. 5
GLYCOLYTIC- OSCIILATTR (GO)
28
B.6 6 SEL' KOv MODEL (SM)
29
B. 3.7 BRUSSELATOR (BT)
29
B.8 GRAY-SCOTT (GS)
31
B.9 NAVIER -TOOES (NS)
31
ABLATION STUDIES
32
C.I LIMITING THE CONTEXT POOL SIZE
32
RESTRICTION OF THE CONTEXT SIZE
................................
C.3 VARIATION OF THE POOLING STRATEGY
33
ABLATION OF THE 3-NETWORKS ARCHITECTURE
D EXPERIMENTAL DETAILS
36
D.I NCFS
36
D.2 CODA
38
D.3 CAVIA
38
EXAMPLE IMPLEMENTATION NCFS
39
TRAJECTORIES VISUALIZATION
40
15

---

Published as a conference paper at ICLR 2025
A ALGORITHMS & PROOFS
A. SECOND-ORDER TAYLOR EXPANSION WITH JVPs
For ease of demonstration, we propose an equivalent formulation of Proposition that disregards the first input of f not necessary for its proof. Although we will consistently write terms like fz),, we emphasize that is meant to stand for the context, not the state variable.
Proposition (Second-order Taylor expansion with JVPs). Assume f Rd Rd is C², Let x E Rd and define + (z y). The second-order Taylor expansion of f around any Fo Rd is then expressed
f(x) = f(zo) + 2 g00) + Vg(zo)(= 2o) + o(1 x xo 2-.l
Proof. I, E Rdc, The second-order Taylor-expansion of f includes its Hessian that we view as 3-dimensional tensor, and we contract along its last axis such that
f(x) f(xo) + (z(o)( 2o) + 2 [22 mo)] (x ro) + o(|2 =ool)). (12)
Next define gi : Vz(v)(r v), and consider small perturbation h Rd write 5ly + h) = V + (z y h)) [V)(y) + (h) o ((lA)))
(x y--v-h) var Taylor expansion of V V v) Vf(y)h [())( (a v) olI o(Hl
The Hessian is by definition symmetric along its last two axes, which allows us to rewrite the third term [ f(y)(x y)] h. We then have
g + h) g(v) + [V2 s())(= - v) - v5(r) ((Il)),,
which indicates that Vg(y) = v²f(o) v) Vf(y), from which we derive
E Rd,, f2f(y)(z y) = Vg(y) Vf()).
In particular,
V' mp) Vg(zo) Vf(o)).
(13)
Plugging this into Eq. 12 we have
(xx) ((00 + (i(o)) mo) + 1 (g((zo + V/(oa)) (x 1o) + o(l1o zoll?)
= f(0o) + 3 (x0) x0)
Vg(zo)(= xo) + o[|ı xo.l²).
g(zo)
This concludes the proof.
While this expression of the second-order Taylor expansion of a vector-valued function makes its implementation memory-efficient via Automatic Differentiation (AD), it still relies on nested derivatives, which scale exponentially with the order of the Taylor expansion due avoidable recomputations. For even higher-order Taylor expansions, this scaling is frightfully inefficient. Toylor-Mode AD a promising avenue address this issue (Bettencourt et al. 2019).
A.2 IDENTIFIABILITY OF PHYSICAL PARAMETERS
Identifying the underlying physical parameters of system is of enduring interest to scientific machine learning practitioners. Our work critically builds on CAMEL (Blanke & Lelarge 2024) which extensively studies model similar NCF-1 linearly parametrized system. Its Proposition
16

---

Published as a conference paper at ICLR 2025
I states, informally, that in the limit of vanishing training loss C(...-) = 0. the relationship between the learned context vectors and the system parameters is linear and can be estimated using ordinary least square. This forces the model to learn meaningful representation of the system instead of overfitting the examples from the training tasks.
We now reformulate Blanke & Lelarge (2024)'s identifiability result for linear systems trained with first-order Taylor expansion and the Random-A pool-filling strategy (see Appendix (A.6), then we provide an alternative proof suited our setting. Building notations from Eqs. (1) and (4), restate that our goal is to approximate the true vector field based on
dæ dt = fine fime
dæ
and
dt (t) - (t)=
V [0,2].
We drop the dependence on Rd to ease notations in the sequel (the V will henceforth indicate gradients wrt Q). As such, the predictor So IRdc - Ied is parametrized as a neural network and learned, while fine 12de + Fd is known and affine, i.c. 3P Rdxde and p Ra such that
five (c) Pc + p.
Proposition 2 (Identifiability of affine systems). Assume de 2 de. that P is full-rank, and that Suue differentiable. In the limit of zero training loss in Eq. (10). fo is affine on an open region of RK., Furthermore, there exists Rdexds and Rde such that for any mela-trained E {{t))mm its corresponding underlying parameter E {
Je=1 and
Je= we have = QE + 4
Proof. Let e E [. In the limit of zero training loss, fo coincides with its first-order Taylor expansion neighborhood UE") which contains all other {5* 1- We can write
fo(E) = fa( + A(E - )
VEEU(E")
(14)
where A = (5( is constant.
Similarly, for [1, ml I. there exists open set U(E)) which contains all other {5)) such that
fo( = fo((') B(g - {),, V(E
(15)
where B = fe()) is constant.
To show that necessarily is non-empty since both şe and & are included in both sets) and proceed by contradiction, assuming + B. Let R& sufficiently small, we can set & = 5 + to Eq. 14 and write the directional derivative
B. let's consider without loss of generality UEE") n ((') (which
fo + tv) - fo(5")
fo(5") A(E + tv Ed) fo(5")
lim 10
lim (+0
t
tAv
lim +0
Av.
(16)
We also write, using Eq. (15). the same directional derivative
folE + +tv fo(
fa(E)) B + tu - (3) (c(-)
lim 00
lim
t-
fo(5") - fo(5") ) + B(f-- -T
tBu
lim
lim 6-0
(-+0
Bo.
(17)
For v & ker (B - A), we have Av # Bo which contradicts with the uniqueness of directional derivatives for differentiable functions (Spivak 1965). This shows that fo affine on the open set
U= 0 U({)).
Furthermore, using Eq. 14 we have
fo(E) Ag = fo(E ACe,
AyE
(18)
17

---

Published as a conference paper at ICLR 2025
which is valid for all [11 m. This indicates that the right hand side of Eq. (8 is constant, i.e. B E IR's such that V E [1, m ]. fo(5") - A q. We then use the fact that in the limit of zero training loss, the predicted and true vector fields coincide for context Ç U and its corresponding underlying parameters c::
firue (c) folE) = Pe re+p p= A(S + folE)) for [1.m] Ag + q-
Since de 2 de and full-rank, its rows are linearly independent, guaranteeing the existence of a pseudo-inverse. We can thus write Q + 9 with
Q= (P" P) PT A,
and
q 8=(PT PPTP)-1pp 1p'lpp p].
(19)
The closed form expression Eq. 19 can be challenging to derive, especially since P and p might not be fully known when collecting data. So similar to (Blanke & Lelarge (024), one can perform post-training, ordinary least squares regression on observed {e Je=1 (with m' S m) to estimate the optimal Q and q":
2²,4" argmin Q.4
Σ
10 + 12---11 cl.2.
(20)
Experimental validation. We validate Proposition by modifying the Lotka-Volterss (LV) experi- ment. For CoDA (Kirchmeyer et al. (222) and NCF-1, we use the exact same network architecture: a 4-layer MLP with 224 hidden units. This means no context nor state network is used in NCF-1 the context vector of size de 2 is directly concatenated to the state vector done by Zintgraf et al. (2019). We note that the results obtained using this configuration further indicate the superiority of NCF, even when model comparison centers their main/root networks (see also (Park al. 2023) for a similar model comparison based on parameter count] 5
After meta-training and meta-testing, set Volterra systems via a linear transformation of the learned context vectors. We fit a linear regression model to the 9 meta-training context vectors, using the true physical parameters as supervision signal. We test on the adaptation contexts. The results, displayed in Fig. 5 adequately illustrate interpretability as stated in (Blanke Lelarge 2024, Proposition 1) and our Proposition 2 They show that our trained meta meta-parameterr recover the underlying system parameters up to a linear transform, and thus enable zero-shot (physical parameter -nduuc)d adaptation via inverse regression.
recover the underlying parameters of the Lotka-
Robustness to noise. Additionally, system identification with NCFs is robust to noise in the trajectory. We show this empirically by corrupting the single trajectory each adaptation environment with Gaussian noise scaled by factor of 1). Upon addition of this noise, sequential adaptation is performed to recover new & which are then transformed into e and plotted. The weight Q and bias of the affine transform are fitted on the training environments and their corresponding underlying parameters, which are unchanged across all noise levels. Fig. 7 shows that the reconstruction MSF remains low despite the noise (compared to CoDA's MSE 1.57 10 2 when 7) shown in Fig. and the physical system remains visually identifiable, especially in the convex hull of training environments. Outside the convex hull, the identifiability is notably worse with r 2 0. indicating that this noise level is excessively high given the range of the LV state values.
A.3 CONVERGENCE OF PROXIMAL ALTERNATING MINIMIZATION
For clarity of exposition, Theorem stationary points is repeated below.
I expressing the convergence of Algorithm T to second-order
$ As reported in Appendix the primary model comparison approach in this work counts all learnalee parameters, including hypernetworss if involved. Only the context vectors are exempt from this count. "Sample context vectors pre-reconsttuction can be observed in Fig. 22 the structure is preserved with changing seeds.
18

---

Published as a conference paper at ICLR 2025
0=0.01
MSE=5.2e-04 1.4e-04
0=0.05
0=0.10 MMSE=6.e-03
0.35
1.5
0.5
0.50
0.55
3.52
6.75
1.30
3255
232
3.25 0.23 0.75
125
Figure 7: Robustness of NCF-4 when identifying physical parameters with varying levels of noise y injected into the adaptation trajectory. The MSE is reported with the standard deviation across 10 runs with different seeds for the Gaussian noise. (Left) = 0.01, (Middle) 77 = 0.05, (Right) 7 0.1.
Theorem (Convergence to second-order stationary points). Assume that C(., -, Drr) satisfies the Kurdyka- Lojasiewicz (KL) property, is bi-smooth, and V Drr) is Lipschitz continuous on any bounded subset of domain IRdv x xm Under those assumptions, ler (0o, ()-m)) be random initialization and (8g. 5q elum) be the sequence generated by Algorithm / If the sequence (0g. el:m) is bounded, then converges to a second-order stationary point of C(., -..-. Dir) almost surely.
We use this section to briefly emphasize that the assumptions of KL property (Attouch et al. 2010) and Lipschitz continuity are mild and easily achievable with neural networks. Interestingly, boundedness of the sequence (0p. </:m)) is guaranteed priory if L is coercive (Li et al. 2019), property we encouraged by regularizing wrt weights and contexts as in Eq.
A.4 ORDINARY ALTERNATING MINIMIZATION
Here provide an additional procedure (Algorithm 3 for training Neural Context Flows which reserved for the NCF-11 variant. While stronger assumptions (Li et al. 2019) are needed to establish convergence guarantees like its proximal extension, is relatively easier implement, and exposes fewer hyperparameters to tune.
Algorithm 3 Ordinary Alternating Minimization
Algorithm 4 Bulk Adaptation of NCF
Input: Du == (Dr)ee( 1.m]
1: Input: Daa (DD be'E[a,o]
2: 0 E IR de randomly initialized
2: 0 E Re learned
3: elim Ü 8, where ge =0 E Rd e=1
3: ga:b
U ge where ge' = 0- E Rdce c'=a
4: Nj, 7
4: ")
5: repeat
5: repeat
6:
+- 1/0 ofL( (0,g010000)) gl:m Drr)
6:
goib t- eaib n55 Daa)
7:
el:m t- El:m ng C(, elim Du)
until garb converges
8:
until (1mm) converges
A.5 BULK ADAPTATION ALGORITHM
Neural Context Flows designed to be fast adaptation time. However, might want adapt to hundreds or thousands of environments, perhaps to identify where the performance degrades on downstream tasks. In such cases, Algorithm 2 could be slow due to its sequential nature. We provide Algorithm that leverages the same parallelism exploited during training (while restricting information flow from one adaptation environment to the next).
Although highly parallelizable, we realize in practice that Algorithm 4 is susceptible to two pitfalls:
1. Memory scarcity: the operating system needs to allocate enough resources store the data
Dat, the combined context vectors go:b IR ((b=++1) and backpropagate their gradients
19

---

Published as a conference paper at ICLR 2025
in order to adapt all environments once, which might be impossible if the context pool size p is set too high. To avoid this issue, we recommend using p < m, or ideally p = 1 contextual self mofullooulttio is disabled.
2. Slower convergence: the bulk algorithm could take longer converge poorer context vectors when the jointly adapted environments are all very far apart. This is because the contextual self-modulation process would be rendered useless by such task discrepancy, and sampling j = e in the context pool P to return to standard a Neural ODE in Section [3 would be harder. We find in practice that manually setting j = e by adjusting the vector field to
forego the Taylor expansion works well (see Fig. 3) A more direct way of achieving the
same result disregard and retain fo (rather than its Taylor expansion) in Eq. (
All adaptation results in this paper use the sequential adaptation procedure in Algorithm except for the grid-wise adaptation in Section 4
A.6 WHAT'S IN A CONTEXT POOL? ?
The content of the context pool P not only defines the candidate trajectories we get from Eq. (8. but also the speed and memory cost of the meta-training process. Based on intuitive understanding the role of P. we outline 3 tunable pool-filling strategies for selecting neighboring contexts:
Random-All (RA): all distinct contexts can he selected by randomly drawing their indices from [1, m. By repeatedly doing so, we maximize long-range interactions to provide the broadest form of self-modulation since information can always (in the stochastic limit) flow from any environment into e.
Nearest-First (NF): only the p closest contexts 8" are selected, thus encouraging en- vironments to form clusters. (Note, however, that if p 1, then P = {ee itself, and no sef-moouulttion occurs.) In one ablation study, we observe that this strategy is the most balanced with regard to training time and performance (see Appendix C.3
Smallest-Firs First (SF): the smallest contexts in L,- norm are selected first. Since environ- ment with context close to can be interpreted as an envirommennt-gnootio feature like in (Kirchmeyer et al. 2022 Blanke & Lelarge 2024), this strategy prioritizes the flow of information from that base or canonical environment to the one of interest e.
A.7 SCALABILITY OF THE NCF ALGORITHMS
The Neural Context Flow (NCF) framework demonstrates excellent scalability with respect to various hyperparameters, owing to its innovative design and implementation. This scalability is evident in three key aspects: distributed training capabilities, efficient handling of large context vectors, and utilization of first-order optimization techniques.
Primarily, NCF's training process can be distributed and parallelized across environments and trajectories (see Eq. (10), Fig. 12 and Algorithm 4 feature that distinguishes it from baseline methods which are limited to parallelization across trajectories. Furthermore. the framework employs an efficient approach, as outlined in Proposition 1. to avoid materializing Jacobians or Hessians wrt potentially large context vectors, thereby significantly enhancing scalability. Additionally, NCF utilizes only first-order gradients wrt model weights 0, in contrast to methods like CAVIA that require second-order information in their bi-level optimization loop.
Scalability can also be evaluated terms of component size, particularly in meta-learning adaptation rules that incorporate additional contextual parameters beyond shared model weights. These compo- nents may include encoders, hypernetworks, and other mechanisms for generating or refining context vectors necessary for task adaptation. The complexity and memory requirements of contextual meta-learning rules, which are directly related to the size of these components, can be quantified by their parameter count among other metrics. In this regard, NCF maintains constant memory cost (1), while baseline methods such as CAVIA and CoDA require additional memory to produce better contexts (see (Park et al. 2023
Table 1)).
Despite these advantages, the NCF framework may face challenges related to memory and compu- tational efficiency due to the requirement of solving p Neural ODEs in equation 3 opposed to
20

---

Published as a conference paper at ICLR 2025
single In scenarios where all training environments utilized (i.e., p m). this results in quadratic cost O(m²) for I and However, our ablation studies in Appendix C.1 be achieved most problems using few p
demonstrate that competitive performance neighboring environments. Additional studies in Appendices C.2 and C.4 [C..esstblist the necessity of expressive context vectors (Voynov & Babenko 2020) and validate the efficacy of the 3-networks architecture, respectively.
It is worth noting that limiting these quantities could directly contribute to improved parameter counts and more interpretable models. Moreover, restricting the total number of environments contributing the loss equation 10 each iteration may further enhance efficiency.
21

---

Published as a conference paper at ICLR 2025
B DATASETS & ADDITIONAL RESULTS
B.1 GEN-DYNAMICS
Given the lack of benchmark consistency in Scientific Machine Learning Massaroli et al. 2020). we launched Gen Gen-Dynamii cs:
com homm/ddross/geer dynamics cs This is a call for fellow authors to upload their metrics and datasets, synthetic otherwise, while following a consistent interface. In the context of OoD generalization for instance, we suggest the dataset be
split in 4 parts: (1) train: For In-Domain meta-training: (2) test: For In-Domain evaluation;
(3) ooditrain: For Out-of-Distribution adaptation to new environments (meta-testing); (4)
oodtest: For OoD evaluation.
Each split should contain trajectories X and the time points at which the states were recorded. The time I-dimensional array, and we recommend 4-dimensional tensor with dimensions described as follows: (I) nb envs: Number of distinct environments; (2) nb trajs per env: Number of trajectories per environment; (3) nb www.sttepppppp.. per-traj Number of time steps per trajectory (matching the size of t); (4) state_size: Size of the state space.
While the suggestions apply mostly to dynamical systems' Meta-Learning, we believe they are generalizable to other problems, and are open to suggestions from the community. All problems described below are now represented in Gen-Dynam is..
cs.
Finally, we define the MSE and the mean absolute percentage error (MAPE) criteria as they apply to all trajectory data found in Gen Gen-Dynama DynaDmmmmmc cs. Unless stated otherwise, the following are the metrics used throughout this work, including in Table
MSE(2, xz, 2)
| - (tnn)
(21)
N x d
n=1
N
(ttt.)) (tn))
MAPE(2, 2)
x 100.
(22)
N x d
x(tn)
B.2 UNCERTAINTY ESTIMATION
Neural Context Flows can provide uncertainty about their predictions. To show this, we calculate (i) the relative mean squared error (MSE), (ii) the mean absolute percentage error (MAPE), and (iii) the 3-d Coverage Confidence Level (CL) (Serrano et al. 2024), with the following formulae applied
In- In-Domain:
100
m
S
(tn(tt) D (t.))l1
Rel. MSE
(23)
m
X S x
N X
d
(t,
100
1 (tn) al x (t.)
MAPE
(24)
m x S N
100
ΣΣΣΣ
CL
(t))ec((c..nndd))
(25)
m mx5xNxd
i=1 n=1 d==1
with the mean and standard deviation across candidate trajectory predictions defined as
(4)))) = Σ 5 (In), and (1))))) =
(L)) (1m) - (L,)))
(26)
N
1=1
1=1
and the pointwise empirical confidence interval defined as
CI(e,,,,,)) [A] 36 jif + 3+
(27)
where S indicates the number of trajectories used per environment, the length of each trajectory. and d the dimensionality of the problem. The number of InD environments is denoted by m, to be
replaced with b a + for OoD cases We set p = m for InD uncertainty metrics calculation,
"onsequently. the first summation symbol's 's bounds and its corresponding factor in the denominators of Eqs. 23 to 25) should he adjusted accordingly for OoD formulae.
22

---

Published as a conference paper at ICLR 2025
and we guaranteed the existence of those same environments for OoD cases. However, if our model performs well across all adaptation environments it encounters -as the case with NCF-tz fz observed in Table 2 then all training and adaptation environments m + b a + instead of only p = m environments at our disposal for Eq. 26 This produces more sample predictions, allowing for more reliable population statistics.
be used, resulting in
The results show low relative MSE on ODE problems, and remarkably low MAPE scores on all problems. These indicate that the empirical mean of the predictions is indeed close to the ground truth. We also notice that the confidence levels vary from very low values on NS to very high on BT. Based on Eq. 25). we hypothesize that this is primarily due standard deviations across predictions
To test our hypothesis, we plot the standard deviations as they evolve with time in Fig. When compared to per-problem InD and OoD CLs from Table [3 Fig. 8 reveals that higher confidence levels align with higher standard deviations, which particularly noticeable on forecasts beyond training time horizons. Additionally, in Fig. / relate to the corresponding absolute errors. Naturally, we observe a non-negligible correlation of the two, especially on the GO problem. Focusing on OoD behavior, our model successfully avoids undesirable regions of low uncertainty but higher-than InD errors (along the lop left corners.. Instead, some OoD predictions for ODE problems fall in a region of higher- higher-than- than-InD uncertainty, but still low error (along the bottom right corners), which stresses the well-suitedness of our approach for OoD generalization. Despite not knowing the underlying source of uncertainty wish to model, these results suggest that our framework is capable of providing meaningful uncertainty estimates, This said, we emphasize that the calculation and interpretation of aforementioned uncertainty metries (e.g. the width of CI) should be grounded on knowledge and goals of the problem at hand.
we plot the pointwise standard deviations they
Table 3: Uncertainty estimation metrics with NCF-t2, all expressed in percentage points (%). The star indicates cases where the close to zero denominators had to be filtered out to retain state values greater than 10 10-3 The Relative MSE is the most sensitive to these instabilities due to squaring the denominator.
LV
GO
REL. MSE
МАPE
CL
REL. MSE
MАPE
CL
IND OoD
0.032 0.183
0.80 1.86
56.22 78.24
5.119
10.36
82.06
1.653
7.17
70.47
SM
BT
REL. MSE
MАPE
CL
REL. MSE
MAPE
CL
IND OOD
2.005* 0.158
4.07
65.64 94.70
42.199 46.028
18.95
92.25
1.99*
22.28
90.63
GS
NS
REL. MSE
MAPE
CL
REL. MSE
MAPE
CL
IND
2118.51 281.986*
58.33*
14.58
166.696* 1
17.30*
11.19
OoD
33.43*
11.89
16.72*
11.09
B.3 SIMPLE PENDULUM (SP)
The autonomous dynamical system at play here corresponds to a frictionless pendulum suspended from stationary point by string of fixed length I.. The state space (a, w). comprises the angle the pendulum makes with the vertical, and its angular velocity, respectively
da
W.
dt
(28)
dw
sin(a).
'e remark that the empirical confidence level CL. is metric that favors models that uncertain in giving the right mean prediction.
23

---

Published as a conference paper at ICLR 2025
1rpuawn ham-ard stat
Std.
12-
36-
36-
34-
22
02
3.
3.0 02
04
0.0
32
02
0.5
0,0
3.0
32
Normalized
Figure 8: Average standard deviations indicating how uncertainty grows with time, including when the model forecasts in time domains not seen during training. Higher standard deviations correlate with higher confidence level metrics observed in Table 3
13
275-
027-
135-
302
333
161-
185
353
6.-
383
Erooe le- 10l
1.0 15 Stel Dev.
Std.
Dav.
Std. Dev.
Figure 9; Pointwise absolute errors against standard deviations for all problems.
For this problem, each environment corresponds to different gravity With L = 1 set, the goal is to learn dynamical system that easily generalizes across the unobserved g. Trajectories are generated with Runge-Kutts Kutta 4th order solver, and fixed step size of At 0.25. Only the NCF-1 variant is used all experiments involving this problem.
OFA OPE NCF. The training and adaptation MSE metrics were reported in Table I during our favorable comparison to two baselines: One-For-AN (OFA) - one context-free model trained for all environments: and One-Per-Env (OPE) - one model trained from scratch for each environment encounter 10 Here, we expand on said comparison with our loss values during training reported in Fig. 10 Additionally, we probe the training time of the different methods. Given that OFA and OPE do not require contextual information, we increase the capacity of their main networks to match the NCF total parameter count. Each method then trained until the loss stagnates, and we report amortized training times in Table
While exhibiting much larger adaptation times, OPE overfits its few-shot trajectories as evidenced in Fig. 10 The same figure shows that the OFA training loss quickly stagnates to a relatively high
TT give an intuition behind the "environment", might consider the surface of celestial body in the solar system (see Section
19The OFA paradigm offers mechanism adapt to new environments, unless fine-tune the vector field's weights, thus returning OPE-like setting.
24

---

Published as a conference paper at ICLR 2025
value during training. Unsurprisingly, learning context-agnostic vector field for all enviromments (OFA) is suboptimal given the vast differences in gravity from one environment to the next. These observations align with the more complete study OFA and OPE Yin al. (2021a). Leveraging Meta-l lerning our method, trained on environments at once like OFA, effectively learns to discriminate between them and produces low validation MSE metrics and accurate trajectories, one of which is presented in Fig. [p. Taken together, Tables a and Fig. [ that Meta-Learning delivers on training time, adaptation time, and most importantly, testing accuracy.
Table 4: Meta-training and adaptation times (4) with OFA, OPE, and NCF-4 the SP problem. We report the amortized times (in minutes) corresponding to fitting one single environment (of which we count 25 when meta-training, and for adaptation).
#PARAMS
TRAIN
ADAPT
OFA
49776
0.34
OPE
49776
4.63
5.72
NCF
50000
2.96
0.51
10-1
OFA
10"9
OPE
NCF
100
Tranning Progress (96
Figure 10: MSE loss values when training NCF-1 on the SP problem, compared with the baseline OFA and OPE formulations. The crosses indicate validation eurves across 3 runs, color-coded to match the training curves. OFA fails catastrophically since the diversity of environments in the training dataset prevents the approximation of any meaningful vector field, while OPE overfits its 4 training trajectories.
Sample efficiency with the number of trajectories. We compare model to the two OFA and OPE baselines as the number trajectories in each training environment is increased from to 12. The results reported in Fig. [T indicate that NCF is indeed the best option when data is limited, while the improvements in OFA MSES are barely noticeable. As S increses, we observe that OPE is able overcome overfitting to ultimately achieve the best results. These results demonstrate that NCF effiently uses its few-shots trajectories. However, if neither data nor training time are not concern (cf. Table 4). then the traditional One-Per-Env should be prioritized.
Scaling with the training environments. The computational speed of any method based on Neural ODEs depends the numerical integrator uses. To provide consistent number of function evaluations (NFEs), we switch the adaptive time-stepper Dopri5 for the fixed time-stepper RK4, then we measure the duration of epochs the training progresses. These times used to produce Fig. 12 which indicates training times per epoch (in seconds) as the number of training environments increased while keeping the range of gravities unchanged between and 24. We observe excellent scaling, with the training time only increasing by roughly 23% (from 0.38 to 0.47 seconds) when the number of environments is scaled by 10 (from 50).
Probing the context vectors. encode useful representations. In Fig. 13 visualize the first two dimensions of the various
Beyond serving as a control signal for the vector fields, the contexts
25

---

Published as a conference paper at ICLR 2025
10-1
MSE
10
OFA
OPE
10 3
NCF
Trajectories per Environment
Figure 11: In-Domain MSEs on the SP problem comparing the sample eficiency NCF-1 against OPE and OFA. NCF is effective in low-data regimes, and OPE overcomes the gap data increases.
Time per Epech (SecS)
0.55
0.50
0.45
0.40
0.35
25
Training Environments (m)
Figure 12: Training time the number of training environments m is increased from 5 to 65. The vertical bars are proportional to the standard deviation across epochs. OOM indicates that our workstation ran out of memory.
{5") 12ses2 after training. We observe that environments close in indices are equally close in the two context dimensions. Similarly, distant environments are noticeably Tar apart in this view of the context space. The same observation is made during adaptation, where, for instance, e' 02 (corresponding to 14.75) gets context close to = 15 (corresponding to 14.83). This observation indicates that the latent context vector encoding features related to gravity, which may be used for further downstream representation learning tasks.
330
-00
622 2.2
Figure 13: Representation of the first and second dimensions of the learned contexts for the SP problem. The labels 1 to 25 identify the training environments (in shades of blue), while 0, and 02 (in red), indicate the adaptation environments. We observe that environments close indices (and thus in gravity) share similar contexts along these dimensions.
The indices of the training environments correspond to their ordering increasing values of gravity.
26

---

Published as a conference paper at ICLR 2025
B.4 LOTKA- VOLTERRA (LV)
With dynamics that are of continued interest to many fields (epidemiology Venturino 1994). economy (Wu et al. 2012), etc.), the Lotka-Volteryy (LV) ODE models the evolution of the concentration of preys x and predators in a closed ecosystem. The behavior of the system controlled by the prey's natural growth rate a, the predation rate B, the predator's increase rate upon consuming prey 6, and the predator's natural death rate Y
dæ
dt ax fry,
(29)
dy ōxy ry.
dt
We repeat the experiment as designed in (Kirchmeyer et al.] 2022). All synthetic ground truth data is generated with the two initial states both following the (11 (1) distribution. Once sampled, we note that the same initial condition used to generate trajectories for environments. The parameters that vary across training environments are B E {0.5, 0.75, 1} and ò E (0.5, 0.75, 1). In each training environment, we generate trajectories with a Runge-Kuut Kutta time-adaptive 4th-order scheme, while generate 32 for In In-Domain evaluation. For one-shot adaptation, extrapolate to {0.625, 1.125] and E {0.625, 1.125), with only trajectory per environment, and 32 for OoD testing. The observed parameters and y are always fixed at 0.5.
100
3.00
1e03
MAPE (%)
MSE
0775
1e04
1e05
0e06
622
630
1.30
1.35
648
$30
4.15
1.0
1.5
(a) MAPE
(b) MSE
Figure 14: Results for the Large adaptation of the LV problem to grid, showcasing the MAPE and training MSE. Compared to Fig. [] the colorbar range in (a) is shrunk to focus on the MAPEs between () and 8%.
The large grid-wise adaptation experiment conducted on the LV problem in Section 4.3 showed that NCFs were powerful at extrapolating. Here, through the MSE training losses in Fig. 14(b) highlight the fact that the bottom and left edges the grid display quite high errors. This illusrates the importance of observing several loss metrics when learning time series (Hewamalage
2023
27

---

Published as a conference paper at ICLR 2025
B.5 FLYCOLYTIC-OSCILLAAT (GO)
This equation defines the evolution of the concentration of biochemical species {s:).=12....77 according to the ODE:
ds
81 $6
Jo k1
dt
1 1+ (s6/ K1)
ds2
$186
2h1
K232(N 85)) -k68283
dt
+ (56/ K1)9
ds3
2 -36) -333
-s6)
-66)
ds4 A383(A A 86) k48485 K(84 87)
dt
dss dt 1282(N 8s) KAs4S5 k68285
dsc 2k1 1 (s0/ K1)
S1
2k s6) kg9e
dl
ds? vik(84 87) - k87
dl
where the parameters either vary or are set fixed as per Table 5
Table 5: Physical parameters and their values for the GO problem
Parameter
Jo
ki
lz
ha
ks
ka
Ki
N
A
+
h:
Value
2.5
(varies)
6
16
100
1.28
12
(varies)
13
0.1
1.8
Again, we follow the same procedure in (Kirchmeyer al.] 2022) generate trajectories for each environment. Namely, we sample initial conditions from a specific distribution (Daniels & Nemenman 2015] Table 2). We vary k1 E {100, 90, 80} and K₃ {1, 0.75, 0.5] to create 9 training environments with 32 trajectories each, both for InD training and testing. We use trajectory for OoD adaptation environments defined by E (85, 95) and E (0.625, 0.875). We use 32 trajectories for Ool evaluation.
Intuitive post-pros pratepoccesing of the contexts offers many insights into the meta-training process, particu- larly because it clusters the 9 environments into groups of 3, as illustrated in Fig. [5
360
200
-200
200
1-56
Figure 15: Illustration of 256-dimensional context vectors consistently clustered with t-SNE embed- dings with perplexity of on the Glycolytic Oscillator (GO). The training environments are in green and labeled as their IDs, while the adaptation environments in purple. This clustering mirrors the distribution the system parameters.
28

---

Published as a conference paper at ICLR 2025
B.6 SEL' KOv MODEL (SM)
We introduce the Sel' kov model in dimensionless form Strogattz 2018 a highly non-linear ODE mainly studied for its application to yeast glycolysis:
dx x. + ay + x²y
dt
dy
dt ay x²y
where a 0.1 and b the parameter that changes. The trajectories are generated with a time horizon of 40, with 11 regularly spaced time steps. We sample each initial condition state from the distribution [4003), 3), and we observe the appearance of a limit cycle (LI), then an equilibrium point (E), and another limit cycle as changes (see Fig. 4(a).
Specifically, our 21 training environments are a union 7 environments evenly distributed in b e [-1, 0.255 then 7 evenly distributed in b E [-0.1, -0.1]. and finally 7 others with b E [0.25, 1]. We generate 4 trajectories for training and 4 for InD testing. As for adaptation, we choose environments, with 6 {-1.25, -0.65, 0000, 0.02, 0.6, 1.2] (see Fig. 16). We set aside 1 trajectory for adaptation and for OoD testing.
In Section 4 we showed that NCF-12 outperformed other adaptation rules. We note, however, that training such systems is very difficult, even with NCFs. We observed in practice that convergence is extremely sensitive to weight initialization, as the shaded regions on the loss curves of Fig. [4bb attest. The curves are complemented with Fig. 17 emphasizing that the model performs best in environments the equilibrium.
Di -2
px13)
D-03
03
Q%
D-13
p-12
Figure 16: Visualization a trajectory with the same initial condition in various environments of the Sel' kov Model's dataset. In blue are the meta-training environments, in red the meta-testing ones. We observe three attractors: the limit cycle (LI) along the top row, the fixed equilibrium (E) along the second row, and another limit cycle (L2) along the third row.
B.7 BRUSSELATOR (BT)
The Brusselator model (Prigogine & Lefever 1968) describes the cnectiio-diffusion dynamics of two chemical species, and V, and given by the following system of partial differential equations (PDEs) defined on 8 x 8 grid:
aU D,AU A - (B + I)U + UVV,
Øt
0V
0t = D,AV BU - U
where:
U and V are the concentrations of the chemical reactants.
D, and D. 0.1 are the diffusion coefficients for U and V, respectively.
29

---

Published as a conference paper at ICLR 2025
10
In-Domain
OOD
0179
< b
Figure 17: Per-environment In- on-Domain and adaptatation MSEs for the SM problem. In blue are the meta-training environments, in red the meta-testing ones. Central environments in the attractor (E) are better resolved than the others.
A and B are constants representing the rate parameters of the chemical reactions, both varying across environments.
Using an RK4 adaptive time-step solver, the system is simulated up to T = 10 (excluded), with the step reported every time step At 0.5. The initial condition for each trajectory for all environments is sampled and broadcasted as follows:
Uo A
B
+ 0.17k Vo Á
where Ā ~ ((0.5 2.0), B ((1.225,550) and Mjj ~ ~N(0,) 1) for each grid position (i.j).
We aim to keep all environments involved in this problem outside the Brusselator's oscillatory regime Ba > 1 + A². For Meta-training, A and B are selected from {0.75, 1.25] x {3.25, 3.5, 3.75) yielding 12 environments, each with trajectories for training, and 32 for InD testing. For adapta- tion, we select 12 environments from (0.875, 1.125, 1.375] x (3.125, 3.375, 3.625, 3.875J. with trajectory for training, and 32 for OoD testing.
As observed with the metrics in Table the BT dataset is one of the most challenging learn Indeed, the Non-Meta-Learning baselines OFA and OPE equally struggle on To show this, we vary the number of trajectories in each trainnig environment between 1 and 8. Then, we plot in Fig. [8] the sample efficiency of our NCF approach against the non-Meta-Learning baselines. Like with the SP problem (cf. Fig. [T], we observe good NCF performance when data is scarce, underlining the suitability of our approach for few-shot learning.
MSE
OFA
OPE
NCF
10-1
12
Trajectories per Environment
Figure 18: In-Domain MSEs on the BT problem comparing the sample eficiency of NCF-/2 against
OPE and OFA. NCF is effective in low-data regimes, and OPE closes the gap as data increases.
30

---

Published as a conference paper at ICLR 2025
B.8 GRAY-SCOTT (GS)
We aim to validate the empirical potential of our method on spatiotemporal systems beyond ODEs. We thus identify the Gray-Scott (GS) model for reaction-diffusion, partial differential equation (PDE) defined over a spatial 32 x 32 grid:
aU
0L D.AU UV2 F -U)
2V
D.AV + - k)V
öt
Like in the Brusselator case above, and V represent the concentrations of the two chemical components in the spatial domain with periodic boundary conditions. Du, Dy denote their diffusion coefficients respectively, while F and are the reaction parameters. We generate trajectories on temporal grid with At 40 and temporal horizon T 400.
The parameters we use to generate our environments are the reaction parameters and k. We consider 4 environments for meta-training: F E (0.30, 0.39). k E (0.058, 0.062); and 4 others for adaption: F (0.33, 0.36), k {0.59, 0.613. Other simulation parameters, as well the initial condition-generatinn distribution, are inherited from (Kirchmeyer et al. 2022). where we direct readers additional details.
B.9 NAVIER-STOKES (NS)
Like the Gray-Scott case, the 2D incompressible Navier-Stokes case inherited from (Kirchmeyer et al. 2022). The PDE is defined on 32 x 32 spatial grid as:
dw Ne + DAw f
V 0
where:
w V x u is the vorticity,
v is the velocity field,
v/ is the kinematic viscosity.
The trajectory data from to T' 10 with At is obtained through custom Euler integration
10 10-3) we gather meta-training environments, each with trajectories for training and 32 for testing. Similarly, we collect 4 adaptation environments with {8.5 10-1 1,9.5 10-4, 10-1,
scheme. By varying the viscosity from E {8- 10-4,9 9- 10-4, 1.0 10 -.1-10 10-3,11111 1.2
1.05 10-3, 1.15 10-3 each with 1 trajectory for training, and 32 for testing. Other parameters of the simulation, as well as
the initial condition generating distribution, inherited from (Kirchmeyer et al. 2022 where encourage the readers to find more details.
31

---

Published as a conference paper at ICLR 2025
C ABLATION STUDIES
The ablation studies described in this section are designed to investigate how the context pool size, the context size, the pool-filling strategy, and the 3-networks architecture affect the performance of NCFs.
C.I LIMITING THE CONTEXT POOL SIZE
The context pool from which environments are randomly sampled contributes significantly to the computational and memory complexity of the algorithm. As evidenced in Algorithms [nd 3 the computation of the loss and hence its gradient can be parallelized across the contexts vectors & With the goal of assessing the associated computational burden, we vary the size of the pool size for the LV problem from to 9, reporting the In-Domain and OoD metrics, and computational time in Fig. 19
In-Domain
In-Domain
150
OOD
COD
125
6.5
Wall Time [min)
100
MSE
75
5.5
50
5.0
25
4.5
0
1
2
4
6
8
2
4
6
7
8
Context Pool Size
Context Pool Size
(a) MSE Losses
(b) Training Times
Figure 19: (a) In-domain and OoD MSEs, and (b) training and adaptation times when varying the context pool size on the LV problem with NCF-12.
10-1-
MSE
MSE
6x 10-5
(a) In-Domain
(b) OoD
Figure 20: Average MSEs for the ablation of the context pool size for NCF-1,, with evaluation carried over many seeds, hence the vertical bars indicating the standard deviations.
We dig further into the influence of the context pool size. While ideal value might difficult to conclude based on Fig. [9] Fig. 20 investiigats NCF-1 1 to highlight how important any value bigger than is (with p = 6 proving to be exceptionally adequate). Indeed, 0** indicates no context pool was used, and the vector field is evaluated without Taylor expansion. This ablation results in significant drop in MSE of about one order of magnitude. While computational metaa-taining time scales linearly with the pool size (see Fig. 19(b). the various losses do not. Overall, these results suggest that the context pool size p largely remains a hyperparameter that should be tuned for maximum balance of accuracy and training time.
32

---

Published as a conference paper at ICLR 2025
Our motivations for the above conclusion lies in that Taylor expansion is the key information flowing from one environment to another, along with the concept of "task relatedness" or "context proximity" Indeed, depending from one another may be required to interact in context space. As a result, the Taylor approximation incurs non-reducible residual error term. To counter this effect, could restrict the pool size to limit the impact of those far-apart environments since they will be less often sampled from the pool the training progresses.
the pool-filling strategy used, some environments that are prohibitively far
C.2 RESTRICTION OF THE CONTEXT SIZE
Rather than limiting the pool size, what if the context vectors themselves were limited? The latent context vectors are the building blocks of NCF, and having shown that they encode useful represen- tations vital for downstream tasks Appendix [.3 we now inquire as to how their size influences the overall learning performance. This further provides the opportunity to test the contextual self- modulation process. Indeed, over-parametrization should not degrade the performance NCFs since context vectors must be automatically kept small and close each other (in L norm) for the Taylor approximations to be accurate.
Like in CoDA (Kirchmeyer al. 2022). the context size de is directly related to the parameter count of the model; and limiting the parameter count bears practical importance for computational efficiency and interpretability. Thus, perform the LV experiment described in Appendix [.44 with dg E {2, 4, 8, 16, 32, 64, 128, 256, 512, 1024). The results observed in Fig. (1 align with our intuitive understanding of increased expressiveness with bigger latent vectors. Together, Figs. (1) and 21(b) shed light on the relationship between & and the underlying physical parameters pair (B, suggesting that NCFs would benefit from the vast body of research representation learning.
We appreciate that while selecting de small is more interpretable (as can be observed by the clustering of the training environment Fig. 2I(b)), it important choose de sufficiently big. Indeed, using the JVP-based implementation we provide in Appendix large context vectors come at reduced extra cost in both speed and memory.
In-Domain
5.322
0.10
5572
10-4
2.403
60
0.07
1232
0.955
16
32
120
256
512
1024
-0.0 -0.05
0.00
0:15
2.5
Context
size dr
dim
(a) MSE Loss
(b) Loss landscape
Figure 21: (a) InD and OoD evaluation MSES for the LV problem as the context size is increased. (b) Loss landscape plotted in context vectors space for de = 2. The clustering of the training contexts mirrors the parameters pairs (3,8) from Fig.
C.3 VARIATION OF THE POOLING STRATEGY
The context pool plays important role in the NCF training process. In order to complement our comments in Appendix A.6 we test the effect choosing various pooling strategies, namely focus
Nearest-First (NF), Random-All (RA), and Smallest-Firs (SF). In this experiment, the LV case with de = 2 due to its linearity and our knowledge of the interpretable behavior of its contexts (see Appendix A.2. We train all three strategies with all hyperparameters identical
33

---

Published as a conference paper at ICLR 2025
(including the neural network vector field initialization). We monitor the validation loss to select the best model across all 10000 epochs. The experiment is repeated 3 times with different seeds.
Run
Run
Run
basen in
02
Figure 22: Both dimensions of the context vector post-training (indicated by dots), and post-adaptation (crosses) across several runs. Given the same neural network initialization, all pooling strategies are able to find the structure mimicking that of the underlying physical parameters observed in Fig. 5
Although all strategies capture the underlying structure of the LV physical system (see Fig. 22 Table 6 indicates that RA and NF have the best performance, contrasting with high metries and uncertainties displayed by SF. The training and validation dynamics in Fig. 23 provide clearer insight into this discrepancy, highlighting
important spread of validation values for SF. Interestingly, Fig. 23 shows that RAA takes much longer to converge during training, despite ultimately producing the best OoD performance. This slow convergence is because very early on, the model is still forced to find commonalities between all pairs of environment, despite some pairs being less related than others. All in all, NF is the most balanced strategy on the LV problem since it converges fast, shows signs of increased stability, and produces excellent results InD and OoD.
Table 6: Comparison of pooling strategies on the LV problem. The reported mean and standard deviation training and sequential adaptation times are expressed in minutes. All strategies are trained for the same number of epochs across 3 separate runs with identical hyperparameters.
RANDOM-ALL (RA)
NEAREST-FIRST (NF)
SMALLEST-FIRST (NF)
TRAIN TIME (mins)
22.90 + 0.04
22.82 + 0.03
22.94 + 0.04
ADAPT TIME (mins)
2.80 0.01
2.80 0.01
2,80 : 0.01
IND MSE (x (x10-1)
2.65 + 0.16
2.52 + 0.66
5.03 + 3,01
OoD MSE (x10-)
2.38 1 0.43
2.69
3.34 0.96
Random All
Nearest First
Smallest First
-17
Epocte
Spochs
Figure 23: Loss curves with varying pooling strategies the LV problem. (Left) Training: (Right) Mean validation losses across 3 runs, color-coded to match training curve labels.
C.4 ABLATION OF THE 3-NETWORKS ARCHITECTURE
Another key element of the NCF framework the 3-networks architecture described in Fig. 2p. In it, the vector field consists of state and context-specific networks that help bring the inputs into the same latent representational space before passing to another network to approximate the local derivative. Like the context size, its ablation directly contributes reduction in learnable parameter count.
34

---

Published as a conference paper at ICLR 2025
Table MSE upon ablation of the 3-networks architecture (NCF#) both SP and LV problems. highlighting difficulties adapting to new environments.
(x10-))
IV (x10-B)
IN-DOMAIN
ADAPTATION
IN-DOMAIN
ADAPTATION
NCF
11.0
0.0003 + 0.00001
6.73 0.87
7.92 + 1.04
NCF
79.8 0.8
1.6402 + 0.7539
11.03 + 1.1
69.98 1 84.36
For this experiment, we consider NCF-4 We remove the data- and context-specific neural networks in the vector field, and directly concatenate (1), and E; the result of which passed (single) neural network. In this regard, NCF* (which designates NCF when its 3-networks architecture is removed) analogous CAVIA (Zintgraf et al. 2019). That said, NCF still benefits from the flow of contextual information, which allows for self- modulatinn...
109
NCF Completa
10-1
MSE
10
100000
200000
100000
a00000
500000
600000
Epoch
Figure 24: MSE loss curves when training the LV problem on a complete NCF, and on NCF*, the NCF variant deprived the 3-networks architecture.
This study is of particular significance since, if in addition to foregoing the 3-networks architecture, we eliminated the Taylor expansion step, and NCF would turn into the data-controlled neural ODE (Massaroli al. (2220). So we both dynamics forecasting problems without the 3-networks architecture, and we report the training MSE in Fig. 24 In-domain and adaptation MSEs for both
and SP problems are reported in Table 7
While Fig. 24 highlights performance discrepancy of nearly order of magnitude during training, the key insight is hidden in the adaptation columns of Table Indeed, the removal of the data- and context-specific networks in the vector field considerably restricts the model's ability to generalize to unseen environments, both for the SP and LV problems. We observed in experiments, similar performance drops when only one of the two state or context networks was removed. This consolidates the 3-networks architecture essential piece of the NCF framework.
35

---

Published as a conference paper at ICLR 2025
D EXPERIMENTAL DETAILS
This section shares crucial details that went into the experiments conducted in Section 4 and expanded upon in Appendices B and C Unless otherwise specified, the hyperparameters in the next 3 paragraphs were applied to produce all figures and tables in this work.
Each method was tuned to provide an excellent balance of accuracy and efficiency. The main motivation behind our choices was to restrain model sizes both to allow efficient training of all baselines and fair comparison at roughly equal parameter counts. Irreconcilable differences with the baseline adaptation rules made it difficult to perform a systematic comparison, which is why we focused on parameter count as done (Qin et al.); all the while noting that all models involved had sufficient capacity to approximate the problems hand. For NCF and CAVIA, we counted the number of learnable parameters in the main network being modulated. For CoDA, we counted the root network plus the hyperneworr's 's learnable parameters outlined in Eq. (3)
Our main workstation for training and adaptation was fitted with an Nvidia GeForce RTX 4080 graphics card which was used for the SP. LV, SM, and NS problems. Additionally, we used an RTX 3090 GPU of the same generation for the GO problem, and an NVIDIA A100 Tensor Core GPU for the GS problem as its CNN-based architecture was more wemorry-inenssine training hardware and deep learning frameworks (JAX for NCFs, and PyTorch for CAVIA and CoDA) that made apple-to-apple comparison challenging, our network architectures varied slightly across methods, as the subsections below highlight. This is part of the reason we decided to focus on total parameter count as the great equalizer.
In addition to the
To ensure fair comparison, we made sure each model was sufficiently large to represent the task at hand (by comparing to commonly used hyperparameters in the literature, c.g., in (Kirchmeyer et al. 2. Since an "epoch" meant something different for the frameworks under comparison, we simply made sure the models were trained for sufficiently long, i.e., the lowest pewest-pefforming method for each task was trained for at least as long as the second-best for that task, and its checkpoint with the best validation error was restored at the end.
D.1 NCFs
Model architecture The MLPs used as our context, state and main networks (as depicted in Fig. 2p) typically had a uniform width. When that was not the case, we use the notation [him - h1 +
h + hmme] to summarize the number of units in each layer of an MLP of depth for brevity. We used a similar notation to summarize the number of channels in a Convolutional Neural Network
(CNN). For instance, on the SP problem, the context network was the MIP 255 + 64 - 64 -> 64] (which indicates context size of de 256), the state network was - 64 64 + 64], and the main network was [128 + 64 - 64 + 64 + 2]. all with Swish activations (Ramachandran et al. 2017). For the OFA and OPE experiments that do not require contextual information (cf. Tables and we delete the state and context networks then increase the hidden units of the main network to 156 to match the NCF parameter count. Below, we follow the same convention to detail the networks used for other problems in this paper's 's main comparison (see Table
LL (context) [1024- 256 + 256 - 64]; (state) [2 - 64 64 64 > 64 - 2] GO (context) [256- - 64 - 64 + 122]: (state) [7 -> 122 + 122 - 122]: (main) [244 - 64 + 64 + 64 + 7 SM (context, state, main) Exact same architectures as with SP above BT (context) linear layer with 256 input and 128 output features followed by an appropriate reshaping then 2D convolution with channels and circular padding to maintain image size, a kernel of size 3 x 3, and Swish activation (state) appropriate reshaping before 2D convolution with channels; (main) CNN [16 + 64 - 64 - 64 + 2] GS Same network architectures as for the previous BT problem, with the exception that the linear layer in the context network outputted 2048 features NS (context) A single linear layer with 202 input and 1024 output features reshaped and concatenated to the state and the grid coordinates before processing; (state) No state network was used for this problem; (main) 2-dimensional Fourier Neural Operator (FNO) (Li et al., 2020) with spectral convolution layers, each with 8 frequency modes and hidden layers of width 10. Its lifting operator was a x 1 convolution with 4 input and 10 output channels respectively - functionally
64]: (main) [128 + 64
Except for the number of channels, these were default convolution settings used throughout. The circular padding enforced the periodic boundary condition observed in the trajectory data.
36

---

Published as a conference paper at ICLR 2025
similar to fully connected layer with weight sharing. The projection operator had two such layers, with 10 inputs, output, and 16 hidden channels in-between.
Training hyperparameters As for the propagation of gradients from the loss function to the learnable parameters in the right-hand tihht-hand-sde of the neural ODEs (5), we opted to differentiate through our solvers, rather than numerically solving for adjoint states (Chen et al., 2018). Since our main solvers were the less demanding RK4 and Dopriš solvers, we didn t incur the heavy memory cost Differentiable Programming methods are typically known for (Nzoyem et al. 2023 Kidger et 2020).
Specifically, we used the following integration schemes: DopriS (Wanner & Hairer 1996) for GO, SM, GS, and BT, with relative and absolute tolerances of 10 and 10 respectively; RK4 for LV with At = 0.1: Explicit Euler for NS with At 1. Across all NCF variants, we kept the context pool size under 4 reduce computational workload
2 for LV and SM, and p = 4 for LV and GO, and = 3 for all PDE problems). We used the Adam optimizers (Diederik) 2014) for both model weights and contexts ODE problems, and Adahelief (Zhuang al. 2020) for PDE problems. Their initial learning rates were as follows: x 10-4 for LV and NS, 10- for GO, BT, and GS, 10-4 for SM. That learning rate was kept constant throughout the various trainings, except for BT, GS, and NS where it was multiplied by a factor (0.1, 0.5, and 0.1 respectively) after third of the total number of training steps, and again by the same factor two-thirds completion. The same initial learning rates were used during adaptation, with the number of iterations typically set to 1500.
As per Eq. (IO), the NCF implementation was batched across both environments and trajectories for distributed and faster training. This was in contrast to the two baselines, in which predictions were only batched across trajectories. Using full batches accelerate training meant that with LV for instance, all 4, 32, or trajectories were used at once depending on the data split (meta-training vs. meta-testing, support vs. query). For regularization of the loss function Eq. (10), we set A = 10-3 for all problems; but x 0 for ODE problems and X 10-3 for PDE problems. For NCF-12 we always used proximal coefficient 3 = 10 except for the LV where B = 100.
As for the simple pendulum SP problem used in this work, it leveraged NCF-t, with p 4 and dg = 256. The integrator was Dopri5 with default tolerances, a proximal coefficient 3 = 10?. a constant learning rate 10-4, and 12000 epochs. The OFA and OPE comparisons used the same hyperparameters, but with 6000 epochs for OPE and 2000 for OFA. For the other problems highlighted Table the remaining crucial hyperparameters (not stated above) are defined in
Tables g and
Table 8: Hyperparameters for NCF-1
Hyperparameters
LV
GO
SM
BT
GS
NS
Context size de
1024
256
256
256
256
202
Pool-filling strategy
NF
RA
RA
NF
RA
NF
# Epochs
10000
2400
24000
10000
10000
5000
Table 9: Hyperparameters for NCF-
Hyperparameters
LV
GO
SM
BT
GS
NS
Context size de
1024
256
256
256
256
202
Pool-filling strategy
NF
NF
RA
NF
RA
NF
# Inner iterations
25
10
10
20
20
25
# Outer iterations
250
1000
1500
1000
700
250
The inner iterations used in NCF-12 not be confused with the inner gradient update steps in CAVTA, since NCF performs alternating rather than bi-level optimization.
37

---

benchmarks introduced in Published as a conference paper at ICLR 2025
D.2 CoDA
The reference CoDA implementation from (Kirchmeyer et al. 2022) was readily usable for most problems. We adapted its data generation process to incorporate this paper (the SM and BT problems).
Model architecture Neural networks without physics priors were employed throughout, all with the Swish activation function. - LV: 4-layers MLP with 224 neurons per hidden layers (width) GO: 4-layers MLP with width 146 SM: 4-layers MLP with width 90 BT: 4-layers ConvNet with 46 hidden convolutional filters and x 3 kernels. Its output was rescaled by 10 to stabilize rollouts GS: Same ConvNet as BT, but with 106 filters. NS: 2-dimensional Fourier Neural Operator (FNO) (Li ct al. 2020) with 4 spectral convolution layers, each with 8 frequency modes and hidden layers of width 10. Its lifting operator was single fully connected layer, while its projection operator had two such layers, with 16 hidden neurons.
Training byperparameters Using TorchDiffEq Chen 2018). CoDA backpropagates gradients through the numerical integrator. We used the same integration schemes as NCF above. We stabilized its training by applying exponential Scheduled Sampling (Bengio et al. 2015) with constant & = 0.99 and initial gain 60 0.99 updated every 10 epochs (except for the NS problem updated every 15 epochs). The Adam optimizer with constant learning rate of 10-4 was used. The batch size and the number epochs are given in the Table 10
Table 10: Hyperparameters for CoDA
Hyperparameters
LV
GO
SM
BT
GS
NS
Context size de
2
Minibatch size
4
32
4
1
1
16
# Epochs
40000
40000
12000
10000
120000
30000
We note that in line with CoDA's fundamental low-rank assumption, maintained context size de = throughout this work since no more than 2 parameters varied for any given problem.
D.3 CAVIA
The reference implementation of CAVIA- ZZnngraf et al. 2019) required substantial modi- fications to fit dynamical systems. Importantly, we incorporated the TorchDis ffEq (Chen 2018) open-source package and we adjusted other hyperparameters accordingly to match NCF and CoDA on parameter count and other relevant aspects for fair comparison. The performance of CAVIA-Con was heavily dependent on the context size, which we adjusted depending on the problem.
Model architecture Like CoDA, the Swish activation was used for all neural networks LV: 4-layers MLP with width 278 GO: 4-layers MLP with width 168 SM: 4-layers MLP with width 84 . BT: 7-layers ConvNet with 46 hidden convolutional filters and 3 x 3 kernels. Its outputs was rescaled by 0.1 to stabilize rollouts GS: 4-layers ConvNet with 184 hidden convolutional filters and 3 x 3 kernels. Its outputs was rescaled by 0.1 stabilize rollouts. NS: Same FNO as with NCF above, including the single-layer context network with 1024 output neurons.
Training hyperparameters We used full batches to accelerate training. We use the Adam optimizer with learning rate of 0.001 for the meta-update step. The single inner update step had a learning rate 0.1. The number of iterations is given in Table
Table 11: Hyperparameters for CAVIA
Hyperparameters
LV
GO
SM
BT
GS
NS
Context size de
1024
256
256
64
1024
202
# Iterations
200
100
1200
500
50
1200
38

---

Published as a conference paper at ICLR 2025
E EXAMPLE IMPLEMENTATION OF NCFS
highly performant JAX implementation (Bradbury et al. 2018 of our algorithms available at https: //ubbbb.uutuuuuu thub. com coon/ddrous /nffoooos///oo We provide below few central pieces of our codebase using the ever-growing JAX ecosystem, in particular Optax (DeepMind et al. 2020 for optimization, and Equinox (Kidger Garcia 2021 for neural network definition.
Vector field The vector field takes center-stage when modeling using ODEs. It is critical to use JVPs, since we never want to materialize the Jacobian (the context size can be prohibitively large). The lacobian-Vector Product primitive from Equinox (Kidger & Garcia 2021) filtes filter-jvp as illustrated below, coupled with jit-compilation, enabled fast runtimes for all problems.
ContextPlowlectoztPe (et.e.iiiiidde........d))) :edule):
physics: eqx. Module sugmentation: eqx Module
def _nnit_ (self, sugmentation, physica=Nons)> aon.eatgnett augmentation self.physics physics
def coall (self, ctx8): 10
etx, etx_ ctxs
Axie,, while (xinj
12
self.physica
14
vf lambda xi: ausmentgmnttsiinnt x xi)
16
vI lanada xi: self. physiihy xi) sel..auugeettaiinnt xi)
17
gradvf lembda xi_: egx. filter_ (xirrl (stx-xi_.)) sod_order_ten tere www.fiiier_jjvpeeraaae
19
totx))) (etx-otx_,11111
21
return v 1.5rgradvf (ctx_) o.dweed_
Listing 1: Second-order Taylor expansion in the NCF vector field
Loss function Eq. 10 provides a structured summation layout particularly suited for a function transformation like JAX's vmap. During implementation. that loss functions be defined in two stages, the innermost summation term along with Eq. (9) making up the first stage; while the two outermost summations of equation make up the second stage. Once the first stage complete, the second is relatively easy to implement. We provide vectorized implementation of the first stage below.
daf loss_ ft___________________________e__ (model, traje, t_eval, all_atx_a, key):
Tnner loas function Eg.
ind = jax pew.unatioemmttiiiiiiiio.. (key, stL_ct_ssshage (01) Betpnttex__ppo________ ctx_s all_stx_s (ind, !]
construstion of the concext pool
batcheed_nodel vmag mmoell, in_axes- (None,, None, None, 0))
vmag
traje_hat, nb_steps batched (trajell;, 6, :h, t_eval, ctx, ctx_a) new_traje hap.bbreadcast_ttttttssss. traje_hat..
11
12
terml jop.meen (mow__roop--traaa__attt
resonatruction error
13
14
Lermž jap.neea jap.abaaa (ctx))
context regularisation
16
Lermö Barame_norm_squee (model)
veights regulacisation
18
loss_val terml 10-34term2 1o-2+term3
20
return loss_val
Listing 2: Inner NCF loss function with vectorization support
A4A PyTorch codebase equally made available httt:pp://grrt.t....e//////////t
39

---

Published as a conference paper at ICLR 2025
F TRAJECTORIES VISUALIZATION
067
Figure 25: (Right) Visualizing the first ground truth and predicted testing trajectories and phase spaces in all 4 adaptation environments of the LV problem. The initial condition is the same across all 4 adaptation environments. (Left) Visualizing the first ground truth testing trajectory in 4 meta- training environments found in various attractors of the SM problem: the first in (L1), the second and third (E), and the fourth (I.2).
CAMA
CoDA
NCF (Oun)
Figure 26: Sample absolute reconstruction error of trajectory during adaptation of the Gray-Scott
system with NCF-12. Trajectories begin = 0 (left) and end 400. This describes the sole trajectory in the fourth adaptation environment.
Ground Truth
CAMA
CoDA
NCF (Ous)
Figure 27: Sample absolute reconstruction error of trajectory during adaptation of the Navier-Stokes
system with NCF-12. Trajectories begin at (left) and end at = 10. The viscosity reported environment is 1.15 10 3
40


Reviews:
[{"review_id": "V86ViDi30b", "reviewer": "Reviewer_oqLH", "summary": "The paper introduces Neural Context Flows, a method for meta learning. The main contributions of the work are focusing on how to combine context vectors in a way that allows OoD generalization and interpretability.", "questions": "Question:\n\n- With respect to the 3-network model, can you remove the state-network but keep the context network? Does it make any difference compared to the one-network model which performs similar to CAVIA?\n\n\nSuggestions:\n\n- Include sample efficiency experiments for some example ODEs and PDEs. For example MSE for a model trained from scratch vs one finetunes via a new context vector.\n- Include uncertainty as a function of forecast. One would assume that uncertainty increases as the forecast becomes longer. Could you provide such an estimate?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Writing is very clear and the methodology is well explained. This allows readers to understand the differences between this method and previous ones.\n- Interesting use of context vectors through the 3-network model. Ablation studies in supplementary material show the need for such an architecture. \n- The combination via context through the Taylor expansion seems to be an interesting and novel application, which I can see being used in other fields outside of ODE and PDE simulations.\n- The estimation of uncertainty via different context vectors is very simple yet very clear and useful.", "weaknesses": "Manuscript makes reference to sample efficiency of using such adaptive models for new context. However the manuscript does not include any experiments to support such a statement.", "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": "C9fZ6AQSNj", "reviewer": "Reviewer_wQzk", "summary": "This work proposes a new meta-learning strategy for learning dynamical systems governed by PDEs. It introduces a new multi-environment framework, where environments are defined by specific PDE coefficients, each describing a specific behavior. To do so, the paper proposes a Taylor expansion of a forecaster network at a context vector $\\xi^e$ around other context vectors $\\xi^j$. It thus collects information from other environments via a context pool of P indices, modulating the forecaster network and allowing also contexts themselves to be self-modulated. The method has other properties: parallelization, interpretability, uncertainty estimation, ... The method is evaluated both for in-domain and out-domain environments agianst two baselines (CAVIA & CODA) on multiple datasets and show competitive or sota performance.", "questions": "You introduce different context pool strategy (RA, NF, SF). What are the differences in terms of performance for your method? Some ablations should have been done to compare the different methods, e.g.:\n- Comparing the performance (e.g. MSE, adaptation time) of RA, NF, and SF for different datasets.\n- Analyzing how the choice of strategy impacts the learned context representations.\nThen, propose guidelines for choosing the appropriate strategy.\n\nYou mentioned LEADS, a multi-task learning problem. It would have been nice to add this baseline to the different generalization experiments, especially for the new datasets that are not present in [3].\n\nCan you provide training time details for the different meta-learning frameworks?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The targeted problem is important. Building neural-ode like solvers able to generalize to changes in the PDE coefficients is important, often referred to solving parametric PDEs.\n\nThe method seems novel for learning dynamical systems with changes in pde coefficients. The use of taylor expansion is intuitive and natural. I particularly liked the intuition given line 202-205. Existing context-based methods do not try to leverage information from each context vectors, each describing the environment information. NCF fills this gap.\n\nThe method is evaluated on a wide range of PDE problems and is SOTA or competitive when considering a second order talyor expansion.", "weaknesses": "Regarding the writing style of the paper:\n- I think there is room for improvement. In the introduction, I think the problem of solving parametric PDEs / learning dynamical systems with varying PDE coefficients should be stated more clearly and explains what is an environment in your specific setting. The introduction should 1) clearly defines the problem of building generalizable neural PDE solvers, 2) what are the different directions taken to do so [1, 2, 3] and 3) explain how your work fits into these different directions and advances the field.\n- There are especially 2 paragraphs, where neural ODEs and physics-based (hybrid) approaches are introduced, that are not necessary or too much detailed in my opinion.\n\nThe authors state that the method can be interpretable, provides uncertainty quantification and is parallelizable. These are important properties that lack for instance for CoDA, as you mentioned. I think that the paper should have exploit these properties in more depth, provide more ablation studies to show that NCF can exploit these properties such as:\n- a detailed analysis showing how the learned context vectors relate to physical parameters, demonstrating interpretability\n- benchmarks showing how parallelization impacts training time as the number of environments increases\n\n[1] Subramanian et al.,  Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior, NeurIPS, 2023.\n\n[2] Takamoto et al., Learning Neural PDE Solvers with Parameter-Guided Channel Attention, ICML, 2023.\n\n[3] Kirchmeyer et al., Generalizing to New Physical Systems via Context-Informed Dynamics Model, ICML, 2022.", "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": "5IJmYj1gH2", "reviewer": "Reviewer_6smX", "summary": "The paper introduces a new method for meta-learning of dynamical systems by enforcing context vector to be smooth and close to each other using a training method based on Taylor expansion of the vector field.", "questions": "1) Did you used **different** trajectories to compute the metrics in the tables than the one used to perform the adaptation in OoD setup?\n2) Could you please explain why your model performs better in OoD than InD ? This is an unexpected resuls : we would expect the model to perform always better on seen data.\n3) Could you clarify what *robust* means in this context ? (robustness to OoD, to noise, to unseen initial condition ? ) and point to the corresponding experiment that support the claim ?\n4) Could you elaborate on the motivation of using Taylor expansion for training the model ? If the main reason behind this is to force context to remain close, then a crucial ablation is missing where the Taylor expansion is replaced by a simple distance loss between context vector. \n5) The context vector is a learned embedding, so what is the point of learning a supplementary dedicated network to convert it into $\\tilde \\xi$, since you could directly learn $\\tilde \\xi$ ? In the paper, you explains that \"this allows the framework to automatically balance the potentially non linear influence of the context with that of the state vector\". What does this means ? Could you please elaborate on this ?\n6) Could you please justify the claim line 248 that $L^1$ regularization is crucial ? Did I miss the ablation of this regularization in the paper ?\n7) I am not sure to understand fig. 6. It seems that you collected the prediction of NCF for a same initial condition with different context vectors (including unrelated ones) and consider the max/min of these prediction as uncertainty measure. Hence my questions:\n\t- What is the probability of the true solution to lie within these bounds ?\n\t- It seems that all model gives fairly similar outputs, which is surprising since they are adapted to different dynamics. How do you explain that ?\n\n**Minor**\n - I do appreciate code release. May I ask you to document and clean the code before publication ?\n - Could you please fix the url link to the associated github repo in the paper ?\n - Could clarify if Theorem 3.2 should be considered as a contribution ? It seems to be a straightforward application of Theorem 2 in Li et al, 2019. If not, it might be interesting to provide more details about what changed, and for which reason.\n\n# Rebuttal update\nThe authors have provided a thorough reply to most of my concerns and have clearly demonstrated that my suspicion of overfitting was unfounded. This significantly impacts my evaluation, and I have raised my score from 1 to 5.\n\nHowever, I still recommend rejecting this paper. While the proposed method is interesting and well-supported, it would benefit from further development, particularly in the analysis of its behavior. Although the authors made commendable efforts to improve their manuscript, the short timeframe allocated for rebuttals limited the depth of the supplementary results included. Specifically:\n\n- The emergence of clusters of \"similar\" environments due to the context pool-filling method, coupled with the Taylor expansion constraint, is a compelling idea. However, it warrants deeper exploration, perhaps with more environments and a detailed analysis of the content of these clusters. These behavior is at the heart of the success of NCF and deserve a better, deeper exploration.\n- Some design choices remain unclear, particularly the use of a projection network on the context embeddings.\n- The authors have strengthened the \"uncertainty\"-related experiments, but this discussion should have been made clearer in the main paper. However, addressing this would require significant restructuring of the paper, which is impractical within the rebuttal timeline.\n\nFor these reasons, I have updated my evaluation to a borderline rejection. While the paper presents interesting results, I believe it is not yet ready for acceptance.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper introduces a method that seems reasonable, with several interesting analysis of the behavior of the learned vector field. The related works section does present the most popular baselines and methods for this problem. The datasets used to evaluate the method are in par with what is currently used in the litterature.", "weaknesses": "I have several major concerns on this paper that I will try to rank from higher to lower importance.\n1) **Overfit**: I am very confused by the OoD adaption protocol used in the paper. From what I understood, Algorithm 2 is used on the validation set to tune the value of $\\xi$ using gradient descent. Then, the prediction error is computed **on the same trajectories as the one used to pick $\\xi$**. Given the size of $\\xi$, it is most certainly overfitting to the (small) set of trajectories used to tune the context vector. I suspect that the value of $\\xi$ is specifically tuned to match a small set of trajectories during validation, hence the results are not representative of the true performance of the model. A quick look at the code seems to confirm that, although it is hardly readable (many commented functions, no comments, empty README, residual files from project development). Moreover, the model performs almost systematically better in OoD setup than in In-domain, (table 1 and 2), which is unexpected, and very unusual.\n2) **Unsupported claims**: \n\t- \"we introduce a *robust*  ...\" ? The word \"robust\" is never used in the experimental section.\n\t- line 248 : \"*The $L^1$ regularization term is particularly vital in promoting sparsity in possibly high dimensional context vector*\". Sparsity is indeed used a lot for equation retrieval, such as SinDy-like methods, but I don't see why this is crucial here. Yet it probably is, but I don't see which experiment in the paper justify this.\n\t- \"*straightforward method for uncertainty quantification*\" : seems like an overstatement according to my understanding of the experiment. Quantifying the uncertainty means that the model can provides an interval to which the true trajectory must belong (up to some probability score). From what I understand from your experiment, you consider that the predictions from your model obtained with different contexts (including unrelated to the current environment) gives an uncertainty on the prediction.\n1) **Motivation of the method**: the use of Taylor development is motivated (section 3, l.200 to 206) by the fact that it forces context vector to remain close. It is not clear to me why closeness is important for the task, and more importantly why simpler regularization to force $\\xi_i$ to remain close from each other would not perform well.\n2) **Complexity of the method**: training NCF requires to compute the second order derivative of a neural network, and then back-propagate through the entire computation graph, including the use of a Neural ODE. The method seems utterly complex and computationally demanding. Moreover, some design choices are unclear to me (see questions). Many hyper-parameters have to be set, and the paper provides no clear indication on how to set them (mostly, size of the context vector and context pool mode (RA, NF or SF)). However, I did appreciated the experiments on the size of the context pool. \n\n\tFinally, Figure 7 and 16 shows training curves of the model exhibiting high spikes, including one (figure 7) from which the model never recover. This seems to indicate a highly unstable and difficult training, which is not a good sign for extrapolation to more complex dynamics.\n\t\n1) **Experimental section**: It seems that the order 1 NCF model is trained with a different algorithm than the order 2 (l.264-268). I am not sure to see why, and more importantly, if the gap between the two models in table 1 is due to the supplementary order of the Taylor expansion, or the different training algorithm.\n\n**Minor** (no effect on rating)\n- in the introduction: \"*Its dynamics are heavily depending on its parameter*\" looks like an overstatement. The dependance of a dynamical system on its parameter can vary significantly from a system to another.\n-  In section 2, you mentioned that CoDA involves two networks instead of one, hence requiring more computational resources to train. Your approach uses three networks. You might clarify this statement to explicitly mentioned hyper-networks as the bottleneck, and not the use of two networks ?\n- in the beginning of section 3 (l.168), $\\mathcal{D}_{ad}$ is mentioned before being introduced.\n- L. 171, you refer to the smoothness assumption as an inductive bias. It's closer to a constraint than a real inductive bias (or at least, it is an inductive bias on the type of dataset you will test your model on, but not an inductive bias of general physical systems).\n\n**Motivation for my grade**\nMy grade is mostly driven by my suspicion of overfitting. However, I am also concerned by the (in my opinion) poorly motivated design choices (use of Taylor expansion to promote closeness of context vector, supplementary encoder for context vector, back-propagating through the Hessian) and my difficulty to understand several claims and experiments.", "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": "HQRibIzZiA", "reviewer": "Reviewer_fPSV", "summary": "This paper proposes NCF to solve parametric ordinary differential equations. It is designed to improve the adaptability and generalization of learning dynamical systems across various environments. NCF introduces a meta-learning approach that employs a context-modulation mechanism, incorporating uncertainty estimation. Specificaly,  NCF uses k-th order Taylor expansion to enable contextual self-modulation. The authors demonstrate the performance of NCF on a variety of ODEs and PDEs problems and illustrate the effectiveness of the proposed methods.", "questions": "Refer to weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "(1) Introducing a meta-learning framework to study complex dynamical systems driven by ODEs and PDEs is a novel and interesting direction.\n(2) This article includes both experimental validation and some theoretical analysis.", "weaknesses": "(1) Currently, there are several methods for parameterizing equations, such as those for ordinary differential equations [1] and partial differential equations [2]. In the paper, the authors mainly compared their approach with meta-learning methods like CAVIA and CODA. So, does it have a competitive advantage over other non-meta-learning methods?\n[1] Parameterized Neural Ordinary Differential Equations: Applications to Computational Physics Problems\n[2] Identification of the flux function of nonlinear conservation laws with variable parameters\n\n(2) The authors mentioned that their method is robust. One important aspect of proving robustness is how the model performs when the observed data contains noise. However, this was not demonstrated in the experiments.\n\n(3) The authors converted three PDE problems into ODE problems for their study. They should provide a detailed description of how this conversion was done and how the errors were controlled during the process. PDE problems are quite sensitive to the choice of numerical schemes, and different discretization methods can significantly affect the accuracy of the solution. Converting them into ODEs and solving them using an ODE solver will inevitably introduce errors.\n\n(4) The link to the code provided in the paper is inaccessible.", "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 (Poster)