In-Context Learning of Stochastic Differential Equations with Foundation Inference Models
Authors: Patrick Seifner
Kostadin Cvejoski
David Berghaus
Cesar Ojeda
Ramses J Sanchez
Venue: NeurIPS 2025 poster ()
Keywords: SDE inference
equation discovery
zero-shot inference
in-context learning
zero-shot parameter estimation
inference of diffusion process
stochastic differential equations
diffusion process
foundation models
Primary area: deep_learning
TL;DR: We introduce a framework for in-context (zero-shot) inference/estimation of drift and diffusion functions underlying SDEs from empirical data of different dimensionalities.

Abstract: Stochastic differential equations (SDEs) describe dynamical systems where deterministic flows, governed by a drift function, are superimposed with random fluctuations, dictated by a diffusion function. The accurate estimation (*or discovery*) of these functions from data is a central problem in machine learning, with wide application across the natural and social sciences. Yet current solutions either rely heavily on prior knowledge of the dynamics or involve intricate training procedures. We introduce FIM-SDE (Foundation Inference Model for SDEs), a pretrained recognition model that delivers accurate *in-context* (or zero-shot) estimation of the drift and diffusion functions of *low-dimensional* SDEs, from noisy time series data, and allows rapid *finetuning* to target datasets. Leveraging concepts from amortized inference and neural operators, we (pre)train FIM-SDE in a supervised fashion to map a large set of noisy, discretely observed SDE paths onto the space of drift and diffusion functions. We demonstrate that FIM-SDE achieves robust *in-context* function estimation across a wide range of synthetic and real-world processes --- from canonical SDE systems (*e.g*., double-well dynamics or weakly perturbed Lorenz attractors) to stock price recordings and oil-price and wind-speed fluctuations --- while matching the performance of symbolic, Gaussian process and Neural SDE baselines trained on the target datasets. When *finetuned* to the target processes, we show that FIM-SDE consistently outperforms all these baselines.

In-Context Learning of Stochastic Differential Equations with Foundation Inference Models
Patrick Seifner Kostadin Cvejoski"., David Berghaus
César Ojeda' & Ramsés J. Sánchez
Lamarr Institute's University of Bonn?, Fraunhofer IAIS* & University of Potsdam" seifnerccs. uni mani-bonn.com bonn.de, de, sanchez@es. uni-bonn.de
Abstract
Stochastic differrnniaal equations (SDEs) describe dynamical systems where de- terministic flows, governed by drift function, are superimposed with random fluctuations, dictated by a diffusion function. The accurate estimation (or dis- covery) of these functions from data is a central problem in machine learning, with wide application across the natural and social sciences. Yet current solutions either rely heavily on prior knowledge of the dynamics or involve intricate training procedures. We introduce FIM-SDE (Foundation Inference Model for SDEs), a pretrained recognition model that delivers accurate in-context (or zero-shot) estima- tion of the drift and diffusion functions of min-dmmesiionll SDEs, from noisy time series data, and allows rapid finetuning to target datasets. Leveraging concepts from amortized inference and neural operators, we (pre)train FIM-SDE in supervised fashion to map a large set of noisy, discretely observed SDE paths onto the space of drift and diffusion functions. We demonstrate that FIM-SDE achieves robust in-context function estimation across a wide range of synthetic and real-world processes from canonical SDE systems (e.g. double-well dynamics weakly perturbed Lorenz attractors) to stock price recordings and oil-price and wind-speed fluctuations while matching the performance of symbolic, Gaussian process and Neural SDE baselines trained on the target datasets. When finetuned to the target processes, we show that FIM-SDE consistently outperforms all these baselines. Our pretrained model, repository and tutorials are available online!
Introduction
Stochastic differential equations (SDEs) govern dynamical systems featuring deterministic flows superimposed with random fluctuations. The deterministic components of these systems dictated by a drift function, and can used to model the slow, tractable or accessible degrees of freedom of dynamic phenomena. The random components be used to represent fast or intractable degrees of freedom. Such description in terms of tractable versus intractable, slow versus fast, is abstract enough be approximately valid across disciplines and observation scales, for phenomena in and out of equilibrium. However, the central obstacle limiting its applicability to many real-world scenarios lies in the accurate estimation of the drift and diffusion functions that best characterize collection of empirical observations. In other words, the problem of SDE discovery from data, whose solution is a key step toward the automation of scientific discovery. Yet current machine learning approaches either rely heavily on prior knowledge of the underlying dynamics or involve intricate training schemes.
instead controlled by diffusion function, and can
Likely inspired by the proliferation and advancement of foundation models in natural language processing and time series forecasting, there has been recent shift toward (pre)training neural
"ntcpe://finnscieence github ho/DppenIM/intrr.....moooooo html
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

P prior (f. G)
p Pay
(1) Sampling
Payz
G
VFF
Co
Go FIM-SDE
&ar
x KiL 1)
Loss
[3) (Prajtraining
(2) Simulation
Comruption
Figure 1: Left: Three-step pretraining strategy for SDE discovery problem: (1) sample target drift and diffusion f, G from the prior distribution Pprior: (2) simulate and corrupt SDE paths, with the
circles denoting noisy observations; and (3) compute neural estimators fe. Gg. and match them to the target functions. Right: Foundation Inference Model for SDEs (schematic representation). The input (context) consists of [LL 1) tuples of the form (y, Ay, Ay², AT) that are projected by the linear layers. The result is processed by a Transformer encoder V that returns the Context Matrix (Eq. 7). This matrix is used as keys and values to M Functional Attention layers: Lr., the and dy. The input queries are the (embedded) location where we evaluate the estimated functions.
network models on large, synthetic datasets to perform in-context (or zero-shot) estimation of functions governing the evolution of unseen dynamical systems. Examples include models that estimate in-context the vector fields of ordinary differential equations (d' (d'Ascoli et al., 2024), the rate matrices of Markov jump processes (Berghaus et al., 2024), the intensity functions of point processes (Berghaus et al., 2025), and parametrie functions for time series imputation (Seifner et al., 2024) and forecasting (Dooley et al., 2024). This emerging paradigm is characterized by a general pretraining strategy, which typically unfolds in three steps. First, constructs broad prior probability distribution over the space of target functions and, consequently, over the space of dynamical systems they define. This distribution is to represent one's beliefs about the general class of systems one expects to encounter in practice. Second, one samples target functions from this distribution, simulates the corresponding dynamical systems, and corrupts the simulated paths to generate a dataset of noisy observations and target function pairs. This step effectively defines type of supervised, meta-learning task that amortizes the estimation process". Third, one trains a neural network model to match these observation- (trrgett))))uuuttio pairs in supervised way. In this work, we apply this general strategy to the SDE discovery problem (see Figure 1) and make the following contributions:
1. Introduce a data generation model for sampling and corrupting SDE paths, restricted to SDEs with polynomial drift and polynomial diffusion. We empirically show that data generated by this model encodes strong prior, enabling models trained on it to generalize to both in-distribution SDEs and real-world datasets.
2. Train the first neural recognition model capable of estimating in-context the drift and diffusion functions of low-dimensional SDEs from noisy data. We demonstrate that its zero-shot perfor- mance matches symbolic, Gaussian process, and neural SDE methods across eight canonical systems, and four real-world datasets, covering processes of one to three dimensions, dataset sizes of 5K-20K samples, and inter-observation times ranging from 10-1 10-3 (arbitrary units).
3. Show that the model can be rapidly finetuned to context data, both dense and sparse, converging up to 50x faster than neural SDE methods while ultimately outperforming all baselines.
In what follows, we adopt the definition of Bommasani et al. (2021) and name our pretrained model FIM-SDE: Foundation Inference Model for (lw-dimennsiona)) SDEs. We briefly revisit related work in Section 2, and formally define the data-driven SDE discovery problem in Section 3. We then present our methodology in Section 4, and demonstrate its capabilities in Section 5. Finally, we discuss limitations and future directions in Section
We invite the reader to refer to Appendix A, where we discuss how these ideas relate to other approaches, like the neural process family, or the recently introduced prior fitted networks.

---

2 Related Work
The machine learning community has primarily relied three approaches tackle the data-driven SDE discovery problem, namely symbolic regression, Gaussian process (GP) regression, and neural SDEs. Symbolic regression assumes that drift and diffusion can be expressed as linear combinations of basis functions, drawn from a carefully specified dictionary (Boninsegna et al., 2018; Frishman and Ronceray, 2020; Wang al., 2022). GP regression, by contrast, employs nonparametric priors, modeling the unknown functions as samples from judiciously chosen Gaussian processes (Garcia et al., 2017; Batz et al., 2018). Both approaches suffer from two major limitations. First, they are highly sensitive to the quality and correctness of prior knowledge about the functions they seek to estimate. Second, they assume access to densely sampled data, or else must resort to a variational approximation in order to regularize the inter-observation dynamics (Batz et al., 2018; Duncker et al., 2019; Course and Nair, 2023). While effective in principle, this latter strategy often undermined by convergence problems (Verma et al., 2024; Adam et al., 2021).
Neural SDE methods (Tzen and Raginsky, 2019; Li et al., 2020; Kidger et al., 202 la; Zeng et al., 2024) offer a prominent deep learning alternative by representing drift and diffusion functions with neural networks, thereby reducing the reliance on prior knowledge. Yet their training involves costly simulations and backpropagation through SDE solvers, typically via adjoint sensitivity methods, which are known suffer from accuracy issues and slow convergence (Kidger et al., 2021b). More recently, simulation free Neural SDEs have been proposed (Bartosh et al., 2025), alleviating this computational bottleneck. However, like most traditional machine learning methods, they still require separate optimization for every newly observed system. This constraint severely limits their suitability in real-world scientific workflows, where prior knowledge scarce and repeated retraining, whether via complex training pipelines models with slow convergence, quickly becomes impractical.
The goal of this paper is to provide zero-shot, deep learning solution the data-driven SDE discovery problem that circumvents (most of) these limitations, and thus can readily be deployed in scientific settings. Appendix A provides a more nuanced discussion of related work on the SDE discovery problem.
3 Preliminaries
In this section, concretely define the data-driven SDE discovery problem. First, let us properly introduce the SDE class we will focus on, namely first-order Markov SDEs in the Ito form.
Ito Stochastic Differential Equations. A d-dimensional SDE in the Ito form is defined as
dx f(x)dt + G())dW(I).
(1)
The vector-valued function f: Rdd Rd denotes the state-dependent drift function of the process and characterizes the deterministic components of the dynamics. The matrix- function G Rd -> Rdxm denotes the state-depe deeendee diffusion matrix and controls the stochastic components which, in turn, are generated through an m-dimensional Wiener process W R+ + I Given some initial condition x0) in Rd, the SDE solution corresponds d-dimensional stochastic process x(I), which we call the state of the system.
Formally, the drift and diffusion functions are defined as (Gardiner, 2009)
f((x) = lim
(ac) 10:===xxxxxt t At|x, t)dx',
(2)
St- At
GG(x) A0 At
lim 1
(a' 2z)(
xx)p(x'' t -,)((,,)))))p(''' At|x, t)dx',
(3)
where p(x', t- - At t) denotes the transition probability for the state of the system to evolve from x into x' under Eq. over the infinitesimal time At. In what follows, we will restrict our attention to diagonal diffusion matrices of the form'
Gxx = diag( / g1 (x), (2(x),
(x))
'We opt for diagonal diffusion matrices to limit the number of functions to be inferred. Nothing in the architecture we present below prevents non- -on-diigonal diffusion, and extending pretraining to that setting is straightforward. We leave this extension for future work.
3

---

Appendix B provides additional details.
Data-driven SDE Discovery Problem. Suppose we given ordered sequence of I. observations }i)} on some empirical process y (1) : R+ + R2d recorded at non- equidistant times 0 S o
DD = [(y, (i); -i),... (yL,
r I T with T denoting the observation horizon. Now assume generated conditioned on hidden process x' (t) R + R governed by an
these observations SDE of the form defined in Eq. 1, with unknown drift f* (x) and diffusion G (x). Our objective is to construct non-pann parametric estimators fg (x("") and Gx(x(D")
parametrized by 0 and conditioned on the context data D that best approximate the putative drift f* (x) and diffusion G' (x( over some "reasonable" domain' X' E Rd. Conventional machine learning approaches invariably address this problem by fitting 0 directly to "D" As a result, traditional models cannot generally be reused to estimate the SDE of a second empirical process y" ((), even when the latter is "similar" (i.e., with respect to some metric) to y (t). The methodology we introduce in the next section strongly departs from this tradition.
Notation. We use x(t) to denote simulated SDE paths, and x" (t) to denote hidden SDE processes assumed to underlie observed data. Similarly, (t) refers artificially corrupted SDE paths, while y" (t) denotes target data. Furthermore, we distinguish between the simulation discretization step At and the empirical onter-observation times Ar. Finally, use the same symbols to denote probability distributions and their densities, as well as random variables and their realizations.
4 Foundation Inference Models For In-Context Learning
In this section, we propose methodology for in-context estimation of drift and diffusion functions from data, which involves training neural network model to map large set of corrupted SDE paths into their a priori-knows drift and diffusion functions. The latter being sampled from a heuristically constructed synthetic distribution. Thus, the methodology consists of two components: a data generation model and a neural recognition model.
The success of this methodology- namely, that model (tre)traiined solely on synthetic data can help understand and predict unseen empirical processes D - will hinge not only on the inductive biases encoded in the architecture of the recognition model, but also on the family of SDEs represented in the synthetic dataset. Our primary assumption is that, even though we can only consider a very restricted family of drift and diffusion functions, the solution of the SDEs they define can exhibit sufficient complexity to mimic real-worid empirical datasets. This assumption resonates with a classical idea popularized by Stephen Wolfram (Wolfram and Gad-el Hak, 2003), but that can be traced back to Kadanoff himself (Kadanoff, 1986, 1987), namely that simple rules can create complex patterns. In what follows, we first introduce our data generation model and then present FIM-SDE, our transformer-basee neural-operatee recognition model.
4.1 Data Generation Model
In order to easily explain the different steps involved in our synthetic data generation process, let us define the probability of observing the noisy sequence y1,.-.:y1 Yy Rd at the observation times 0S < ... < TL I T for given time horizon T. We write this probability as follows
II 1=1
Proise
Pm G, x(77=))))(((()))paa
TL) TL)Pair( Gld))ana (f(d), (5)
where each term represents one step in the data generation process from the drift and diffusion function generation, to the corruption of the simulations. Let us specify each of them.
Drift Function Generation. Drift functions are sampled from the distribution Pain (f()))d), which is conditioned on the system dimensionality d, and is defined to factorize as Pant (f1 (x)) Pana (x)). Following Berghaus et al. (2024) and d' Ascoli et al. (2024), we generate drift functions for processes of different dimensionalities. ranging from one until some maximum dimension den. This ensures the applicability of methodology across systems whose dimensions lie within this range. The set
By "reasonable" we loosely mean the subdomain X E _ that contains the typical set (Thomas and Joy, 2006) of the empirical process y" (t).

---

Gesend-Truth
SparsedP
BISIDE
FIM-SDE (Zeno-Shol)
Drift
Difkusion
Drift
Diffasion
20
1
2
Figure 2: Drift and diffusion estimation for two canonical SDEs with state-dependent diffusion in the uncorrupted setting (p - 0.0, ÅT 0.002). Left: Double-well system (Eq. 37). Right: Synthetic 2D system (Eqs. 38 and 39). FIM-SDE infers the target functions in-context (i.e., zero-shot mode), showing excellent agreement with the ground truth.
of distributions Pant (x))). for E 1.... d, should reflect our beliefs about the class drift functions likely to arise in nature. In this work, we define them over the space of polynomials of degree at most three, with randomly sampled coefficients. Albeit simple, this choice covers a broad range of commonly used physical models, including chaotic systems. See, for example, the large collection of drift functions in ODEBench (d' Ascoli et al., 2024). We invite the reader to check our discussion about the limitations of this choice in Section 6.
Diffusion Function Generation, Diffusion functions are sampled from the distribution Par (Gx)) which also factorizes Par (x)) Par (ar()a(x)), where the (x),, for E 1... d, correspond to the arguments of the square roots in Eq. 4. State-dependent diffusion common in financial applications and also plays role in transport problems (see e.g. the classical work by Büttiker (1987)). We define Pant (a (x)) over the space of positive polynomials of degree at most two, with random coefficients. This construction covers both geometric Brownian motion and constant diffusion, among others.
f, G, x(T:-1)) denotes the transition probability between the states x(T;-1) and x(r;), and corresponds to the solution of the Fokker- Planck equation defined by f and G over the interval [T--1. This equation evolves the stochastic process at the "macroscopic" (or path-ensemble) level (Gardiner, 2009). In practice, we define p(x(To 0)) the standard normal distribution and simulate individual sample paths by numerically solving the corresponding SDE via the Euler- scheme, using discretization step At until the time horizon T'
SDE Simulation. The term
Empirical data noisy, often recorded at irregular time intervals, and can feature inter-observation gaps much larger than the microscopic timescale of the dynamics (i.e. AT >> At). The distribution (T, 7)) represents the uncertainty in the sequence of recording times and we implement it by subsampling our SDE solutions through different schemes. Similarly, the distribution Proie (y: (t),, f, G) is defined to represent additive Gaussian noise the maximum entropy attractor of empirical noise distributions (Jaynes, 2003) with variance that depends on the characteristic scale of the system's 's dynamics.
SDE Corruption.
Please refer to Appendix C for the implementation details of all these distributions.
We use the data generation model defined in Eq. 5 to sample synthetic dataset consisting of tuples of the form (D, (f. G)). Here, D denotes a set of K corrupted time series {ykı...., .KL ..L > k==1 where each series contains L observations sampled from the SDE defined by the pair (f, G). Each tuple represents one distinct SDE, and we generate SDES of varying dimensionalities.
4.2 FIM-SDE: a Transformer-based, Neural Operator Recognition Model
We now introduce FIM-SDE, neural recognition model designed process instances of D, and produce estimates of the target pair (f, G) over our "reasonable" domain X. To implement this map, we draw on concepts from neural operators, specially DeepONets (Lu et al., 2021). However, different SDEs and thus, each tuple (D, (f, G)) in our dataset are characterized by distinct spatial and temporal scales. ensure that FIM-SDE is scale-agnostic, we first normalize every D and renormalize the associated (f, G) pair accordingly. Appendix D. provides the details.
5

---

Table 1: MMD evaluation canonical SDE systems under varying noise levels p and inter- observation times ÅT. FIM-SDE is compared against Symbolic (BISDE) and GP (SparseGP) regres- sion. Results show the mean and standard deviation over five runs (lower is better). Superscript Roman numerals indicate the number of failed estimations, while the character "_" denotes that all runs failed. Bold entries mark the best results. All values scaled by factor of 100.
Double Well
2D-Synt
Damped Linear
Damped Cubic
A
Model
Duffing
Glycolysis
Hopf
(Wang)
SparseGP
3(4)
0.3(3)
0.9(2)'
0.5(4)
0.68(0)
0.4(3)
0.002
BISDE
0.2(4)
0.4(5)
0.0(0)
0.03(6)" "
0.0(0)
0.3(3)
0.4(4)
FIM-SDE
0.5(5)
(.5
0.0(0)
0.9(5)
0.6(5)
1(2)
1(5)
0.0
SparseGP
12(2)'
12(1)
6.67(0)
13.6(5)
0.02
BISDE
0.0(0)
0.1
0.0(0)
0.7(5)
0.0(0)
0.02(3)
0.0(0))
FIM-SDE
1(1)
0.2(4)
0.0(0)
0.1
0.0(0)
0.2(2)
0.23(9)
SparaeGP
5(2)
9(2)
9(2)"
18(2).
11(2) tMI
10(1)
0.002
BISDE
3.1(8)
9(2)
11.8(2)
18(2)"
11.9(2)
14(4)"
13(2)"
FIM-SDE
2.6(6)
2.2(8)
2(2)
6(2)
1.8(9)
3(2)
1.8(9)
0.05
SparseGP
14(2)"
17(1)
17(8)
19.42(0)
18(1)
0.02
BISDE
0.3(5)
0.6(4)
3(1)
1.7(5)
1(1)
2.1(3)
0.4(3)
FIM-SDE
1(1)
(.4
1.1(5)
0.5(3)
1.4(8)
0.8(2)
0.2(2)
direct inspection of Eqs. 2 and reveals that the values of the drift and diffusion functions given location, say x, only depend on transitions that take place in the "neighborhood" of x. In other words, the sequential nature of the time series in D does not provide additional information for estimating the local expectations. We therefore reorganize each set DD into a collection of K(L -) transition tuples D {y,, Ay, Ay?, (r.)k((-1) where each tuple only keeps information about one-step transitions and where explicitly include both Ay; Ay? enable the network define below to learn separate feature channels for drift and diffusion estimation. Intuitively, a model that estimates or G at location x from D, should dynamically query every transition in D, and weight them according their spatial location (i.e. the y:) with respect to x, to compute the relevant local expectation. Thus, attention mechanisms seem like the natural choice.
We embed all transitions in D using a self attentitee encoder, and introduce a functional astention mech- anism that performs docation-dppender querying over the set of embedded transitions. Let 6 (-) and Φ6 (.) denote linear projections and feed-forward neural networks, respectively. Let (., denote attention layers, and let (, denote Transformer encoders with linear attention (Katharopoulos al., 2020: Shen et al., 2021), both of which take three arguments as inputs (i.e. queries, keys and values). Finally, let 0 denote model parameters.
Context Matrix (DeepONet's Branch-net equivalent). Let us define the n-dimensional embeddings
de = concat o6 (y (Ay)), 2 (Tr)).
where runs from to K(I. 1) for every element in D, and the linear projections d map their inputs onto R=/4, Let now D [df....,d'', do (L)))))) denote the n x K(L - 1) matrix of linear embeddings. We define the (self-attentive) context matrix
D v (D' DD.D"))........ so that 1 E BXXK(L--)
(7)
which encodes the entire context data D.
Functional Attention Mechanism (DeepONet's Trunk-net equivalent). Let us assume we want to estimate the drift function at location x. Let x be the input query to sequence of M attention networks wa and let De be both keys and values to every attention network in that sequence. We define a set of M + location-dependen embeddings for the estimator of f as
h (D) = vi,, (x(D). 5 DD),
(8)
'That said, note that the Ar in D are typically much larger than the At characteristic of our simulations.

---

Table 2: MMD evaluation data sampled from the Lorenz attractor under three initial conditions: (O 1). (1, 1) and (0, 2). LatentSDE was trained for 5000 iterations on data from A ((, 1), while FIM-SDE, by contrast, required only 100 finetuning iterations the same dataset.
Model
[O 1)
(1, 1,)
.V 2)
LatentSDE
0.051708
0.135003
0.054335
FIM-SDE (Zero-Shot)
0.403933
1.151080
0.378515
FIM-SDE (Finctuned)
0.003838
0.037159
0.003335
where runs from 1 until M, and the first embedding is defined (x([)) (x) = (x).. We
then compute (x Φ (x,,,( The calculation of (x( is analogous, and Figure illustrates the FIM-SDE architecture.
Target Objective. We adopt the mean-squared error as target objective, and define
4 (x, ((,,)) Σ {
fi(x)) (5())) -Vorotti))-Vor (x(xx)
2
which can be computed by marginalizing D. and G under the generative model (Eq. 5), and x under e-g. a uniform distribution over X. Now, can readily show that this objective reaches global minimum when the estimators fg and Gg respectively match the mean of the true posterior distributions Paen (f]D) and Pai (G)D). However, the reader will immediately notice that marginal- izing Ci under uniform distribution
result in numerical instabilities, caused by some local contributions dominating the integral, not due to the objective itself, but because of the norm of the involved functions in those regions. We therefore need weighting mechanism to balance these contributions over X. Consider the loss
(o G) +
Ca E x-l((x) (f.
E
(10)
where U denotes uniform distribution, and Ppeior labels our generative model (Eq. 5). The learnable function (x, D) provides adaptive weights, and can be interpreted as epistemic uncertainty estimator (see e.g. Appendix B in Karras al. (2024)). Indeed, assuming that C is location- dependent Gaussian variable, Up acts as its log-variance. From Eq. 10, we see that large Up values weighweight regions where the model is uncertain, while the additive +Ug term prevents it from collapsing to +.. Empirically, we have observed that Uo peaks in sparsely observed or highly noisy areas, acting proxy for epistemic uncertainty. We implement Ug with third set of M attention functions ".. .t,
but only back-propagate up D (i.e. we detach Uo).
Finetuning. FIM-SDE supports rapid finetuning on both densely and sparsely sampled data². For densely sampled data, the model is finetuned by maximizing the likelihood of short-time transitions induced by the estimators fo. Go with respect to the target data. For sparsely sampled data, fine- tuning proceeds instead by simulating the inferred SDE (defined by fe, Ge) for l steps between observation pairs, and minimizing the discrepancy between the last simulation step and its target value. Appendix D.4 gives further details and the explicit loss functions used for finetuning.
Let us close this section by noting that, unlike traditional approaches. both neural and not, that rely on variational methods to regularize the dynamics between sparse observations, FIM-SDE requires no such explicit regularization. Its pretraining implicitly enforces it. We demonstrate this empirically below.
5 Experiments
What follows outlines the experimental setup, including pretraining details, datasets, evaluation metrics, and baselines.
Pretraining. We (pre)trained single 20M-parameter instance of FIM-SDE on a synthetic dataset of 600K low-dimensional SDEs. The dataset contains three-, two-, and one-dimensional systems in
TThe target data may consist of context data D. additional samples from the same underlying process y
7

---

Ground-Truth
LatentSDE
IZen-Shot)
FIM-SDE
FIM-SDE
0.
MMD
0-
04-
0--
001
730
1500
teration
Figure 3: Sample paths (first three panels) and MMD convergence (right) for the Lorenz experiment. FIM-SDE in zero-shot mode already captures the global dynamics, though errors accumulate over time. With finctuning. FIM-SDE quickly refines its estimates of the dynamics, requiring only a few iterations and converging much faster than LatentSDE.
3:2:1 ratio, with context sizes (i.e., DD) ranging from 128 to 12800 tuples, and inter-observation times (i.e,, År) spanning 10-1 to 10-3 (arbitrary units). Further details on the (pre)training data, architecture and hyperparameters, as well as ablation studies, are provided in Appendix C, D, and H, respectively.
Datasets. We evaluate FIM-SDE on two classes of systems: eight canonical SDEs of varying dimensionality and complexity, and four real-world systems derived from experimental data. The canonical systems are defined in Appendices G.1 and G.2, and belong to the SDE class used to construct the prior distribution (Eq. 5). That is, SDEs with polynomial drift of dimension at most three, and polynomial diffusion of dimension at most two. We denote this class by C in what follows. The experimental data consist stock prices of Facebook and Tesla, price fluctuations, and wind speed recordings (Wang et al., 2022).
We compare against three baselines: the GP regression model of Batz et (2018), the symbolic regression model of Wang et al. (2022), and the neural SDE model of Li et al. (2020), Let us referred to them SparseGP, BISDE, and LatentSDE, respectively. Appendix E contains details about their implementation, hyperparameters and training.
Baselines.
Metrics. The quality of the inferred SDEs is assessed using three complementary metrics. First, compute the maximum mean discrepancy (MMD) between simulated and held-out path ensembles, using the signature kernel of Király and Oberhauser (2019). Second, when ground-truth functions are available, we measure the mean-squared error (MSE) between estimated and true driftddiffusion functions predefined grid. And third, qualitatively assess the results by inspecting the inferred functions on uniform grids and comparing sampled paths. Implementation details for both MMD and MSE are provided in Appendix F.
5.1 Canonical SDE systems
Comparison against Symbolic and GP regression. We begin by studying seven canonical SDEs, defined in Appendix G.I, spanning diverse non-linear behavior such Mouble-welll dynamics and Hopf bifurcations. Each system is simulated with an Euler-Maruyama step of A/ 0.002, which fixes the "microscopic" time scale of the dynamics. To assess robustness of FIM-SDE and the baselines under different forms of corruption, we design four experimental settings per target SDE. Specifically, we introduce additive noise with variance p and vary the inter-observation time. The first setting preserves the microscopic scale, AT At = 0.002, while the second uses a coarser resolution, Ar 0.02, but keeps the number of observations per path fixed. Both baselines assume (by construction) access to the clean system states, and are therefore expected to perform well on dense data (i.e., AT At) without noise corruption. In contrast, FIM-SDE is (pre)trained on corrupted paths characterized by broad inter-observation time distribution p( AT) (see Appendix C.5). We use FIM-SDE and the baselines to estimate the drift and diffusion functions that best characterize sampled paths, repeating each experiment five times with independently generated datasets. Both baselines must be retrained from scratch for every system and configuration, whereas FIM-SDE applied off-the- the--hhlf,, without modification, across all experiments.

---

Table 3: MMD in four real-world systems. Per system, we report the mean and standard deviation across five cross-validttion splits. LatentSDE has been trained for 5000 iterations, while FIM-SDE has been finetuned for only 100 iterations. FIM-SDE (Finetuned) achieves the overall best performance across the different systems.
Model BISDE
Wind
Oil
Facebook
Tesla
0.000 + 0.000
0.137 = 0.097
0.053 + 0.060
0.013 0.015
LatentSDE
0.350 + 0.212
0.235 1 0.167
0.014 +: 0.018
0.016 + 0.023
FIM-SDE (Zero-Shot)
0.330 + 0.114
0.199 + 0.137
0.012 + 0.026
0.000 1 0.000
FIM-SDE (Finetuned)
0.029 + 0.034
0.068 = 0.042
0.006 1006=0000 0.007
0.004 0.008
Table reports the MMD evaluations across all four experimental setting, and all seven SDE systems, averaged over five runs. The emergent picture is clear: FIM-SDE consistently outperforms both baselines on systems corrupted by noise. On clean systems, BISDE tends to perform best, although FIM-SDE attains comparable MMD values in roughly half of the cases. Surprisingly, both BISDE and SparseGP frequently produce estimates that correspond to invalid SDEs (even the uncorrupted settings!). FIM-SDE is stable in this regard. Table 6 (Appendix) reports MSEs between estimated and target functions, providing a consistent and complementary picture. Moving on to a qualitative view, Figure 2 shows the inferred drift and diffusion for the double-well system (Eq. 37), and synthetic 2D system (Eqs. 38, 39), in the uncorrupted setting. FIM-SDE estimations are in excellent agreement with the ground-truth. Additional inferred functions and sampled paths for the remaining systems shown in Figure 5 (Appendix).
The results presented here show that FIM-SDE successfully generalizes to unseen SDEs within C (the SDE class in our prior distribution). Appendix H provides ablation studies on the effects of varying the parameter count, the definition of C, and the context size of the model on these canonical systems.
Comparison against Neural SDE methods. We now turn to the celebrated Lorenz attractor, weakly perturbed by constant diffusion. The system is simulated with initial conditions sampled from ((001), 1). observed sparsely in time with gaps of A = 0.025, and perturbed by additive Gaussian noise with variance p = 0.01. This setup reproduces the experimental conditions studied by Li et al. (2020) (see Appendix G.2 for details). We use this data to train LatentSDE, which handles data sparsity by modeling the dynamics in latent space, using an inhomogeneous SDE posterior together with a homogeneous SDE prior. The latter is what enables LatentSDE to tackle the SDE discovery problem. FIM-SDE, in contrast, treats the same data as context, and is evaluated both in zero-shot mode and after finetuning on its context (using the finetuning cost in Eq. 35 for sparse data).
For evaluation, we first generate dense ground-truth trajectories under three initial conditions, namely (0, 1). (1, 1,, and (0, 2). Then, we compute the MMD between these and samples drawn from the trained LatentSDE prior, as well as from FIM-SDE and its finetuned counterpart. Table 2 summarizes the results: LatentSDE outperforms FIM-SDE in the zero-shot setting, but after less than 100 finetuning iterations, FIM-SDE surpasses LatentSDE on all three tests. We attribute the weak zero-shot performance of FIM-SDE to the fact that, despite being in C, the perturbed Lorenz attractor and its samples are drift-dominated (Figure 3 makes this feature apparent). Because our pre-training set contains mostly non-zero diffusion functions, FIM-SDE is biased toward inferring finite diffusion, which increases the variance of its generated paths leading to higher MMD scores. Finetuning on the context data quickly removes the bias, yielding accurate estimates. The right panel of Figure further shows that fine-tuning FIM-SDE converges more than 5x faster than LatentSDE, and over 15x faster than training PIM-SDE from scratch, underscoring the strong initialization provided by the pretrained parameters.
5.2 Real-Worke Systems
In this section, we study how well FIM-SDE generalizes to four real-world stochastic systems obtained from experimental data. Specifically, we consider stock prices of Facebook and Tesla, as well as wind speed and oil price fluctuations, originally collected and analyzed by Wang al. (2022). As shown in Figures 4 and 6 (Appendix), these datasets exhibit inherent stochasticity and are densely recorded.
We partition each dataset into five folds, train BISDE and LatentSDE on four of these, and compute the MMD on the held-out fold, repeating this process across all folds and averaging the results (see Appendix G.3 for details). Within the same cross- noss-validation protocol, FIM-SDE uses the data as

---

Gruund
FIM-SDE (Zero-Shuti
FIM-SDE (Finetuned)
Tesla
sEe-
3.000
0:002
3.034
0.002
6004
Figure 4: Paths sampled from SDEs inferred by PIM-SDE (dark blue) and by FIM-SDE finetuned on oil (left) and stock price data (right). Finetuning noticeably improves sample path quality in the oil dataset, while the strong zero-shot performance on stock prices is retained.
context, and is evaluated both in zero-shot mode and after finetuning with the dense-data objective (Eq. 34). Table reports the resulting MMD scores.
The zero-shot performance of FIM-SDE is competitive overall, and on the stock datasets it even outperforms both BISDE and LatentSDE. On the wind and oil datasets, zero-shot FIM-SDE lags behind BISDE, but finetuning substantially improves its estimates, yielding the best performance across all systems. In particular, finetuning adjusts the inferred diffusion in the oil and wind datasets, as illustrated in Figures and 6 (Appendix). This shows that pretraining on our synthetic, prior distribution over the SDE class C (Eq. 5) encodes useful inductive biases into the model, while finetuming adapts it to -ptaset-specific features.
Finally, Table 9 (Appendix) compares the training trajectories of LatentSDE and FIM-SDE. Whereas LatentSDE requires thousands of iterations to converge, FIM-SDE surpasses it after only 100 fine- tuning iterations, representing 50x reduction in compute. Taken together, these findings demonstrate that FIM-SDE functions as a true foundation model: it achieves competitive zero-shot performance and can be rapidly adapted to complex real-world data through finetuning.
Conclusions
In this work, we introduced a methodology for in-context learning of stochastic differential equations (SDEs), built two components: (i) a prior distribution (C') over the SDE class C, consisting of polynomial drift of degree < and polynomial diffusion of degree I 2: and (ii) FIM-SDE, a neural recognition model pretrained on (C). We empirically showed that the zero-shot performance of FIM-SDE is on par with symbolic, Gaussian process, and neural SDE methods across eight canonical systems and four real-world datasets. Moreover, showed that FIM-SDE be rapidly finetuned on both dense and sparse data, consistently outperforming all baselines after only few iterations.
The main limitation of our methodology lies in the prior distribution, which is currently restricted to (i) the SDE class C of low-degres polynomial drift and diffusion, and (ii) low-dimensional systems. The practical impact of the first restriction depends on whether C is rich enough to capture SDEs relevant in real-world applications. Our experiments indicate that C offers not only useful inductive biases but also strong initialization for finetuning. Future work will explore more expressive priors, such as random function compositions via unary-binary trees (Lample and Charton, 2020), to generate more complex drift functions.
The restriction to now-dimensional systems stems mainly from the rejection rates in data generation, which increase sharply with dimension: 50% in ID, 78% in 2D, and 92% in 3D (curse of dimension- ality). In future work, we aim to reduce rejection rates by generating SDEs known a priori to admit solutions, e.g. by randomly sampling globally Lipschitz drift and diffusion functions with linear growth (Øksendal, 2010). This would enable the simulation of higher-dimensional SDEs. Still, many complex and relevant dynamical systems are, in fact, low dimensional.
Finally, future work will also explore leveraging the pretrained FIM-SDE weights to regularize equa- tion discovery from high-dimensional data, with applications ranging from neural population and chemical reaction dynamics (Duncker et al., 2019) to the evolution of natural language content (Cve- joski et al., 2023, 2022).
10

---

Acknowledgments and Disclosure of Funding
This research has been funded by the Federal Ministry Education and Research of Germany and the state of North-Rhine Westphalia as part of the Lamarr Institute for Machine Learning and Artificial Intelligence.
References
Vincent Adam, Paul Chang, Mohammad Emtiyaz Khan, and Arno Solin. Dual parameterization of sparse variational gaussian processes. Advances in neural information processing systems, 34: 11474-11486, 2021.
Cedric Archambeau, Dan Cornford, Manfred Opper, and John Shawe-Taylor. Gaussian process approximations stochastic differential equations. Gaussian Processes in Practice, pages 1-16. PMLR, 2007a.
Cédric Archambeau, Manfred Opper, Yuan Shen, Dan Cornford, and John Shawe-Taylor. Variational inference for diffusion processes. Advances in neural information processing systems, 20, 2007b.
Grigory Bartosh, Dmitry Vetrov, and Christian A Naesseth. Sde matching: Scalable and simulation- free training of latent stochastic differential equations. arXiv preprint arXiv:2502.02472 2025.
Philipp Batz, Andreas Ruttor, and Manfred Opper. Approximate bayes learning of stochastic differential equations. Physical Review E, 98(2):02210,, 2018.
David Berghaus, Kostadin Cvejoski, Patrick Seifner, Cesar Ojeda, and Ramses J Sanchez. Founda- tion inference models for markov jump processes. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL htttss://pperreviee wwtp:/eee/ooorooeeeeeeeeeeeoeeee formmiii f4v7 cmm
David Berghaus, Patrick Seifner, Kostadin Cvejoski, César Ojeda, and Ramsés J. Sánchez. In- context learning of temporal point processes with foundation inference models. arXiv preprint arXiv:2509.24762, 2025.
Marin Biloš, Kashif Rasul, Anderson Schneider, Yuriy Nevmyvaka, and Stephan Günnemann. Modeling temporal data as continuous functions with stochastic process diffusion. In International Conference on Machine Learning, pages 2452-20 24700. PMLR, 2023.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint arXiv:2108.07255 2021.
Lorenzo Boninsegna, Feliks Nüske, and Cecilia Clementi. Sparse learning of stochastic dynamical equations. The Journal of chemical physics, 148(24), 2018
David B Brückner, Pierre Ronceray, and Chase P Broedersz. Inferring the dynamics of underdamped stochastic systems. Physical review letters, 125(5):058105 103, 2020.
Steven L. Brunton, Joshua L Proctor, and Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 113(15):3932-393) 2016.
M Büttiker. Transport as a consequence of state-dependent diffusion. Zeitschrift für Physik Condensed Matter, 68(2):161-167, 1987.
Kevin Course and Prasanth B Nair. State estimation of a physical system with unknown governing equations. Nature, 622(7982):261 2023.
Kostadin Cvejoski, Ramsés Sánchez, and César Ojeda. The future is different: Large pre-trained language models fail in prediction tasks. arXiv preprint arXiv:2211.00388 2022.
Kostadin Cvejoski, Ramsés J Sánchez, and César Ojeda. Neural dynamic focused topic model. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 12719-12727, 2023.

---

Stéphane d' Ascoli, Sören Becker, Philippe Schwaller, Alexander Mathis, and Niki Kilbertus. ODE- Former: Symbolic regression dynamical systems with transformers. In The Twelfth International Conference on Learning Representations, 2024. TzoHLiGVMo.
wwt/ id=
Samuel Dooley, Gurnoor Singh Khurana, Chirag Mohapatra, Siddartha V Naidu, and Colin White. Forecastpfn: Synthetically-trie trained zero-shot forecasting. Advances in Neural Information Process- ing Systems, 36, 2024.
Lea Duncker, Gergo Bohner, Julien Boussard, and Maneesh Sahani. Learning interpretable continuous-time models of latent stochastic dynamical systems. In International conference on machine learning, pages 1726-1734. PMLR, 2019.
Harrison Edwards and Amos Storkey. Towards a neural statistician. arXiv preprint 120.2256 2016.
Anna Frishman and Pierre Ronceray. Learning force fields from stochastic trajectories. Physical Review X, 10(2):021009, 2020.
Ting-Ting Gao, Baruch Barzel, and Gang Yan. Learning interpretable dynamics of stochastic complex systems from experimental data, Nature communications, 15(1))6029, 2024.
Constantino A Garcia, Abraham Otero, Paulo Felix, Jesus Presedo, and David G Marquez. Nonpara- metric estimation of stochastic differential equations with sparse gaussian processes. Physical Review F., 96(2):022104, 2017.
Crispin Gardiner. Stochastic methods, volume 4. Springer Berlin, 2009.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes, In International conference on muchine learning, pages 1704-1713. PMLR, 2018a.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint 1rXiv:1807.01622 2018b.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel tampla test. Journal of Machine Learning Research, 13(25):723-775 2012. URL btttp:/jjmli orggoooggg papprs/v13//eeppppppppppppaassaa watpar///v///reetoononee
http//geetttttinnn html.
Nicolas Heess, Daniel Tarlow, and John Winn. Learning to pass expectation propagation messages. Advances in Neural Information Processing Systems, 26, 2013.
Luke B Hewitt, Maxwell I Nye, Andreea Gane, Tommi Jaakkola, and Joshua B Tenenbaum. The variational homoencoder: Learning to learn high capacity generative models from few examples. arXiv preprint arXiv:1807. 180.00899, 2018
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfn: transformer that solves small tabular classification problems in second. arXiv preprint 2207:0000000144 2022.
Yunfei Huang, Youssef Mabrouk, Gerhard Gompper, and Benedikt Sabass. Sparse inference and active learning of stochastic differential equations from data. Scientific Reports, 12(1):21691, 2022.
Edwin Thompson Jaynes. Probability Theory: the Logic of Science. Cambridge University Press, 2003.
Leo P. Kadanoff. On two levels. Physics Today, 39(9):7-99 09 1986. ISSN 0031-9228.
Leo P. Kadanoff. On complexity. Physics Today, 40(3):7-9, 03 1987. ISSN 0031-9228.
Tero Karras, Miika Aittala, Jaakko I Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEF/CVE Conference Computer Vision and Pattern Recognition, pages 24174-24184 2024.
12

---

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnss Fast autoregressive transformers with linear attention. International conference on machine learning, pages 5156-5165. PMLR, 2020.
Patrick Kidger, James Foster, Xuechen Li, and Terry J Lyons. Neural sdes infinite-dimensional gans. In International conference on machine learning, pages 5453-5463. PMLR, 2021a.
Patrick Kidger, James Foster, Xuechen Chen Li, and Tenry Lyons. Efficient and accurate gradients for neural sdes. Advances in Neural Information Processing Systems, 34:18747-18761, 2021b.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv:1901.05761 2019.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980. 2014.
Franz Király and Harald Oberhauser. Kemnlss for sequentially ordered data. Journal of Machine Learning Research, 20(31):1 45, 2019.
Guillaume Lample and François Charton. Deep learning for symbolic mathematics. International Conference on Learning Representations, 2020. URL http:p/s//gpmmeeeee..eeee/eeeeee S1eZYeHFDS.
Xuechen Li, Ting- Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pages 3870-3882 PMLR, 2020.
Xuanqing Liu, Tesi Xiao, Si Si, Qin Cao, Sanjiv Kumar, and Cho-Jui Hsich. Neural sde: Stabilizing neural ode networks with stochastic noise. arXiv preprint 10X.v3190. 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17Xiv:1711.05101, 2017.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218 229, 2021.
Sarthak Mittal, Yoshua Bengio, Nikolay Malkin, and Guillaume Lajoie. In-context parametric inference: Point distribution estimators? arXiv preprint arXiv:2502.11617, 2025a.
Sarthak Mittal, Niels Leif Bracher, Guillaume Lajoic, Priyank Jaini, and Marcus Brubaker. Amortized in-context bayesian posterior estimation. arXiv preprint arXiv:2502.06mm 266000 2025b.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers do bayesian inference. In International Conference Learning Representations, 2022. URL https:://// / wwwp/eweeeeeeeeee//e/ee/e.
Samuel Müller, Arik Reuter, Noah Hollmann, David Rügamer, and Frank Hutter. Position: The future bayesian prediction prior-fitted. arXiv preprint arXiv:2505.23947, 2025.
B. Øksendal. Stochastic Differential Equations: An Introduction with Applications. Universitext. Springer Berlin Heidelberg, 2010. ISBN 9783642143946. URL http:::/bookk...oggggg gowgge.. de/ books?id-EQQZEAMAQQE.
Brooks Paige and Frank Wood. Inference networks for sequential monte carlo in graphical models.
In International Conference on Machine Learning, pages 3040-3049. PMLR, 2016.
Patrick Seifner, Kostadin Cvejoski, Antonia Körner, and Ramsés J Sánchez. Foundational inference models for dynamical systems. arXiv preprint arXiv:2402.07594, 2024.
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Atten- tion with linear complexities. Proceedings of the IEEF/CVE winter conference on applications of computer vision, pages 3531-3539, 2021.
13

---

Andreas Stuhlmüller, Jacob Taylor, and Noah Goodman. Learning stochastic inverses. Advances in neural information processing systems, 26, 2013.
MTCAJ Thomas and Thomas Joy. Elements of information theory. Wiley-Intersciennee, 2006.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference Learning Theory, pages 3084-31 31 14. PMLR, 2019.
Csaba Tóth, Danilo Jr Dela Cruz, and Harald Oberhauser. A user's guide KSig: Gpu-accelerated computation of the signature kernel, 2025. URL. https:/ /tappss:////xvvv/////eesse/smmv/ wogg/ba///aa 07145.
Prakhar Verma, Vincent Adam, and Amo Solin. Variational gaussian process diffusion processes. In International Conference on Artificial Intelligence and Statisties, pages 1909-1917. PMLR, 2024.
Michail D. Vrettas, Dan Cornford, and Manfred Opper. Estimating parameters stochastic systems: A variational bayesian approach. Physica D: Nonlinear Phenomena, 240(23):1877-1000 2011.
ISSN 0167-2789.
Yasen Wang. Huazhen Fang, Junyang Jin, Guijun Ma, Xin He, Xing Dai, Zuogong Yue, Cheng Cheng, Hai-Tao Zhang, Donglin Pu, Dongrui Wu, Ye Yuan, Jorge Gonçalves, Jürgen Kurths, and Han Ding. Data-driven discovery of stochastic differential equations. Engineering, 17: 244-252, 2022. ISSN 2095-8099, doi: https://doi.org//00 wwwp////ee...22.......... www.:acceenediree.. com/ sttp//ccccccceee//e//e//eee/ecee /ttpe//ariaeeee/a/////000022222/
URL https:
Christian Wildner and Heinz Koeppl. Moment-based variational inference for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pages 1918-199 1926. PMLR, 2021.
Stephen Wolfram and M Gad-el Hak. A new kind of science. Appl. Mech. Rev., 56(2):118BB1 2003.
Sebastian Zeng, Florian Graf, and Roland Kwitt. Latent sdes on homogeneous spaces. Advances in Neural Information Processing Systems, 36, 2024.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction claim the development of FIM-SDE, a pretrained model for in-context (zero-shot) estimation of drift and diffusion functions for SDEs from noisy, sparsely observed time series. The experiments in section support these claims by showing FIM-SDE's performance diverse systems. The scope (low-dimensional SDEs, diagonal diffusion) is also implicitly and explicitly (in limitations) acknowledged.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We specifically discuss the limitations section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
15

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The paper is primarily empirical, introducing novel model architecture (FIM-SDE) and methodology for its training and application. While relies on established SDE theory (e.g., Ito form, definitions of drift/diffusion, etc.), it does not present new, core theoretical results or theorems that would require formal proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Our source code is open source and all experimental details are mentioned.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibilityy can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
16

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Our source code is open source.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips..
public/guides//ee
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. ccc/pbllc/guuiii
ssi ssionPolicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Extensive details are provided, particularly the appendices
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report the error bars that correspond to the standard deviation of 5-fold cross-validation
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We specified the hardware and amount of compute.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics pettps://neurjpe cc cc/public//atte htt.eiiichiiieeenneeeeeee..
Answer:
[Yes]
Justification: Our work is purely based on foundational science without any negative ethical implications.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is purely based on foundational science without any negative societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
18

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The FIM-SDE model, while applicable to potentially sensitive domains like finance, is a tool for understanding dynamical systems rather than model directly generating high-risk outputs. The training data used is purely synthetic.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All external datasets that have been used for evaluation have been properly cited and publicly available.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
19

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperewithcode. thcode. com/ con/datasete
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's ' creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The primary new assets are the FIM-SDE model and the synthetic data generation pipeline which are extensively explained.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing experiments research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The research does not involve new data collection from human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
20

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The research is focused on stochastic differential equations and does not use Large Language Models
component of its methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: //nurrupp..... scc/oonferences/20000 for what should or should not be described.
21

---

A Related Work
Variational Inference and Parametric Priors. In general, when the observed data is noisy, sparse in time, or both, one faces uncertainty not only in determining the drift and diffusion of the putative SDE, but also in the state of the system itself. Therefore, a Bayesian treatment requires the estimation of the posterior distribution over these states, conditioned on the noisy data (i.e. the so-called smoothing problem). Starting with the seminal works of Archambeau et al. (2007a.b), which approximated the posterior over the states with a variational and inhomogeneous Gaussian process, most proposals have mainly focused on devising different strategies to infer the smoothing distribution - while assuming prior parametric forms for the drift and diffusion functions. See e.g. Vrettas al. (2011), Wildner and Koeppl (2021), or the recent work by Verma et al. (2024).
non-parametric Priors. A notable exception is the proposal of Duncker et al. (2019), which extended the variational trick of Archambeau et al. (2007a) by imposing a non-parametric prior over the drift of the prior process. However, Archambeau et al. (2007a)'s trick has been shown to suffer from significant convergence issues (Verma et al., 2024), which are inherited by Duncker et al. (2019)'s model. In contrast, Batz et al. (2018) framed the drift-diffusion estimation problem from uncorrupted (i.e. clean and dense) data as a Gaussian process regression problem, and extended this (non-parametric) approach to observations that are sparse in time, by leveraging Drstein-Uhlenbeck bridges optimized with an expectation maximization algorithm. Nevertheless, this extension can only deal with non- -orrammmtriic drifts (i.e. the diffusion is restricted to parametric forms) and clean data, and is highly sensitive to the choice of prior hyperparameters.
Symbolic Regression. Symbolic regression methods for drift and diffusion estimation mainly extend the SINDY algorithm (Brunton et al., 2016) - which performs sparse linear regression on a predefined library of candidate nonlinear functions to SDEs. The first of these extensions corresponds to the work of Boninsegna et al. (2018), which sets the regression problem by approximating the local values of the drift and diffusion functions with the empirical expectations of Eqs and 3 of Section 3. However, the calculation of these local expectations generally requires significant amounts of data, even for one-dimensional systems. To (somewhat) alleviate this issue, Huang al. (2022) and Wang et al. (2022) recently resorted to sparse Bayesian learning, but their solutions still require sizable dataset sizes and, most problematically, assume access to clean and dense observations. Other key contributions came from the physics community. For example, Frishman and Ronceray (2020) leveraged symbolic regression too, but utilized the connection between the Ito and Stratonovich formulation of SDEs to devise an unbiased estimator for the drift and diffusion functions under noisy observations. They were also able to use their methodology to estimate non-equilibrium observables, like the stochastic entropy production. Nevertheless, their empirical results still showed significant sensitivity to noise. In response to these limitations, Course and Nair (2023) proposed a hybrid solution that leveraged the variational trick of Archambeau et al. (2007a), while allowing the drift of the prior process to be approximated by a sparse, linear combination of known basis functions. However, their model inherits the slow convergence problems of the variational approximation. What is more, all symbolic regression methods are limited by construction to linear combination of the functions in their library, which therefore makes their performance highly dependent on the preselected functions within the library. Besides these works, other recent efforts tackle underdamped stochastic systems (Brückner et al., 2020) and stochastic networked (that is, interacting) systems (Gao et al., 2024), and we refer the reader to them for details.
Neural SDEs. Most deep learning approaches for SDE learning rely mainly on neural network parameterizations of the drift and diffusion functions - the so-called Neural SDE models - and primarily focus on path generation (that is, smoothing or forecasting tasks). Prominent examples include the works by Tzen and Raginsky (2019), Liu et al. (2019), Li et al. (2020), Kidger et al. (2021a). Biloš et al. (2023) and Zeng al. (2024). The training of these models is however known to be a formidable hurdle, which many have attempted to overcome (Kidger et al., 2021b). We nevertheless note the very recent work by Bartosh et al. (2025).
Early Amortization Inference. Our methodology can be understood as an amortization of the probabilistic inference process (of the drift and diffusion functions) through a single neural network, and is therefore akin to the works of Stuhlmüller et al. (2013), Heess et al. (2013) and Paige and Wood (2016). Rather than treating, as these previous works do, our (rreetrained models as auxiliary to Monte Carlo or expectation propagation methods, we employ them to directly estimate the drift. and diffusion functions from various synthetic and experimentll datasets.
22

---

Conditional Neural Network Models and Meta-Learning. Different from "foundation models" trained on synthetic datasets (like FIM-SDE or Prior Fitted Networks discussed below), conditional neural network models such as the neural statistician Edwards and Storkey (2016); Hewitt et al. (2018) or members of the neural process family Garnelo et al. (2018b,a) Kim al. (2019) - are trained across sets of different, albeit related datasets, each assumed to share a common context latent variable. These meta- aeeta-learning models are trained exclusively on (sets of) datasets from their larget domains, rendering both their optimized weights and the representations they can infer (that is, their latent variables) problem/data specific. In contrast, our method maintains the same network parameters and representation semanties throughout all experiments. The representations consistently correspond to drift and diffusion functions of SDEs, regardless the target dataset.
Prior-Fitted Networks, A recent line of work introduces Prior-Fitted Networks (PFNs) (Müller et al., 2022; Hollmann et al., 2022; Müller et al., 2025) as powerful framework for amortized Bayesian inference. PFNs are typically implemented as transformers trained on large collections of synthetie, supervised-learning tasks sampled from a user-specified prior. Once (pre)trained, these models approximate posterior predictive distributions in single forward pass, effectively mimicking Gaussian processes and other Bayesian inference procedures at test time. Our foundation inference models (FIMs) (Berghaus et al., 2024; Seifner et al., 2024) share the philosophy of pretraining on synthetic data, but target a fundamentally different goal. Rather than learning posterior predictive distributions for supervised tasks, FIMs are designed to infer dynamical systems (i.e. mathematical models) from data. In other words. FIMs infer latent structure underlying target dataset. Let us close this section by noting that, very recently, the problem of amortized in-context Bayesian posterior estimation was framed very precisely and studied by Mittal et al. (2025b,a).
B Background: Stochastic Differential Equations
In this section we provide some background information for stochastic differential equations, that we already touched on in Section 3. For thorough discussion, we refer the reader to (Gardiner, 2009) and (Øksendal, 2010). We recall the definition of stochastic differential equations, detail path sampling scheme and specify the kind of equations we consider in this work.
B.1 Definition
d-dimensional stochastic process x(t follows an Ito stochastic differential equation (SDE), it satisfies
(1 = =I;(()= + [. ((x(1))))) + Σ G=(((())))aww(()
(11)
j=1
for all i < d. and some vector-valued function f: Rd x R+ + Rd and diffusion function G Rdd x 1 Rdxmm where W : R - R. is standard m-dimensional Wiener process. Such process is denoted as a differential equation
dx(l) f(x(1), D)dt + G(x().))d(()
(12)
Equation (11) contains deterministic Riemann integral and stochastic Ito integrals (which are defined as Riemann-Stieltjes integrals). As such, (11) is understood as the infinitesimal limit in discretization of the interval [t. ? into multiple short-time steps. Considering a short-time step [t, + At), the right-hand side of (11) is approximately
TR
(t) + f((((t) + Σ ((W At) Wt),),
(13)
by the definition of the integrals. The infinitesimal limit is well defined drift and diffusion are globally Lipschitz and feature (at most) quadratic growth (Øksendal, 2010).
stochastic process following an Ito SDE is Markovian, emitting a transition probability distribution p(., x, t). Its relationship to the drift and diffusion functions is characterized by two equations:
/ f(e) 11x))x)pxx.>> t + At x, t)dx'
f((xt))
lim A0 Al
(14)
23

---

1 (a4 A4-0 A/
/
G ()x,
Jim
( x;)(da) xx,)px))(xxtt + At t)dx'
(15)
Hence, drift and diffusion can be reinterpreted by infinitesimal limits of mean and covariance of the short-time transitions.
B.2 Simulating SDEs: Euler-Maruyama Scheme
The formal solution construction via discretization and short-time step approximation (13) is also scheme to simulate paths (i.e. realizations) of the stochastic process. Consider again a short-time step [1, + AtJ. Then given x(t), the solution at & + Al is approximately
x(t At) x() + f(x(t), t) + (xx(t), ()( + At)
W(()
(16)
because of (13). By definition of the Wiener process, its increment follows a the multivariate normal distribution
(W(t + At) W(t)) ~ NV Atmm)
(17)
which enables the sampling of x(t + At) x(1) in (16). In other words, the short-time transition probability is approximately Gaussian:
x A
f(x, t) G(x, G(x,))
(18)
To conclude, the called Maler-Marayama scheme for sampling SDE x(t) interval, say [0, 2'., from an initial state x(0) involves discretization of the interval into M small time-steps of T and iteratively sampling from (18). Finer discretizations, ie. smaller At, yield more
size At X7 accurate sample paths.
B.3 Our Setup, the Shot-time Transition Probability and its Logarithm
Our work focuses on the inference of SDEs with purely state dependen drift and diffusion functions f Fd + Rd and G Rd + Rdxm. We denote these differential equations by
dx = f(x)dt G(x)dW(tt.
(19)
Moreover, train data distribution and inference model assumes diagonal diffusion. functions of this form are uniquely defined by
Diffusion
G = diag( 91 92
9a(x)),
(20)
with component functions gi IRd R+ for i < d.
This diffusion function modeling assumption is not as restrictive might seem. All baseline models (see Appendix E) make the same assumption. In principle, our framework can handle non non-diagonal diffusion. We make this simplifying assumption for a fairer comparison to the basclines and lower computational load during data generation and model training.
Under these assumptions, the approximation of the short-time transition probability from (18) reduces to
p(: A/ x) R (x f(x)At, digg(n())).....(x(())))))) -9(()))
(21)
More explicitly, the transition density is
(xt 24 -x(x))1122
p(x', 1 At x,t) (2m21)a/22 IIL V (x(x)
exp
(22)
2At
(x(
whose logarithm reads
d
1 d
(x) 12 A(x)
Dxgp(x(+A/
(+At x, = log 1ogg(2 At)
log (x(g (x)
Σ
(23)
24
(x(x)
(x)
i=1
24

---

C Synthetic Data Generation Model
In this section we provide more details about the generation of synthetic training data for FIM-SDE. First, we describe a scheme to sample SDEs of up to dimension d using polynomial coefficients of restricted degrees. Then, we specify how we sample from these SDEs to generate paths, which are potentially corrupted by spatial or temporal noise. Finally, summarize the explicit hyperparameter choices made for the train set FIM-SDE.
C.1 Sampling Scheme for Polynomials
We begin with our sampling scheme for n-variate polynomials with degree bounded by Mass Below, this scheme is deployed to generate drift and diffusion functions that define the stochastic systems of our training data distribution.
polynomial is uniquely defined by a finite set of monomials with non-zero coefficients. Hence, our sampling scheme derives such of monomials (using hierarchical sampling approach) and then samples (a.s.) non-zero coefficients for them. More precisely, our algorithm for sampling n-variate polynomial f is
1. Sample the size of the set of degrees Nor ~ -U,,,.... Mmax] with terms of non-zero coeffi- cients the polynomial.
2. Sample {mi...., M } U(PNNeg [10,.... max Hl] the set of monomial degrees with terms of non-zero coefficients the polynomial.
)) 3. For each E .L....,,NNes!! sample Ni ~ LI[L...... (m;tt=')) of degree m, with non-zero coefficients in the polynomial.
the number of monomials
4. For each i E {l,..., Nong sample {af..... a'y Nx J U[P Ne.. Hf N" |a| mm]ll, the n-variate exponents monomials with non-zero coefficients in the polynomial.
5. For each i {1,..., Nas) and j E (1....,N*, Nim. sample a coefficient Caj
N(O, 1) of
the n-variate monomial x
6. The n-variate polynomial f is then defined as f(x) = -/ E)er C
Naeg Ninon
The uniform distributions cover broad range of polynomials with degree bounded by m-mxs while the hierarchical sampling scheme ensures some sparsity (regarding monomials with non-zero coefficients) and correlation between monomials of the same degree.
As an example for 3 and max consider the samples Nag 2, {m,, m2 = (0,3), with N1 = 1, { {(0,0, 0, 0)) and N2. = 2, {af, a²) {(3, 0,0), (1, 1, 1)]. Then, for given cat: Caf: Ca3 ~ NO, 1), the sampled polynomial f(11, 22, 23) = Cot + Caza] + 22233
C.2 Distribution over Drift Functions
To sample the drift function of a d-dimensional system, where d < dmax we apply the sam- pling scheme from Appendix C.1 for each component separately, using n d and Miax mail. In other words, our distribution Pann(f(x d) factorizes over the d components f as Pain( (f ... Pana (fa(x)). Each f; is sampled from the distribution over d-variate polynomials with degree bounded by m max specified by the algorithm in Appendix C.I.
dritt
C.3 Distribution over Diffusion Functions
To define the distribution Par (x d) over the diffusion functions of a d-dimensional system, where d' < d recall from (4) that we only consider diagonal Gxx). That is, for all < d there exists non-negative function gi such that G(x);: / (x), and G(x);; 0 for all j * i. Hence, our distribution Pan (x) d) also factorizes over the d diagonal components as Pan(g (x) ... Pan (x))))
For set S and le I denote by P*[S] the set of subsets of S with & elements.
25

---

We define Par (x)()) via an extension of the sampling scheme in Appendix C.1. First we sample ğ.. d-variate polynomial with degree bounded by may using the sampling scheme from Appendix C.I. Because 9i can attain negative values, we define the component function by gx(x) max(0, 9; (x)). This then defines the component-wist distribution Pan ( (x)) and therefore, by our assumed factor- ization, Part (x) d).
Note that some diffusion component functions gi (x) sampled by this procedure contain regions with low, or even zero diffusion (e.g. when 9: is constant of negative value). In practice, found that our trained model can therefore also approximate (almost) deterministic systems.
C.4 SDE Simulation
To simulate paths from an SDE, with drift and diffusion functions sampled as described in Ap- pendix C.2 and C.3, we first sample set of K initial states from (O, 1). We solve the equation (i.e. sample paths) beginning at these K initial states using the Euler-Maruyama method with a fine discretization step size Al until time horizon T.
To facilitate smooth training. we employ two criteria to reject the equation, should (at least) one of them be satisfied. If equation is rejected, we sample new equation and repeat the simulation process until a predefined number of valid equations has been reached.
The first rejection criterion targets stability. It rejects equation if at least one of the K solutions contain NaN or Inf values. The second rejection criterion targets diverging systems. It rejects an equation if at least one component of at least one of the K solutions exceeds threshold value + B.
C.5 SDE Corruption
In this subsection we describe how we implement the distribution Ppia(T1..... 7i) (over the sequence recording times) and Punise ( (y: f, G) (over the optional corruption of the process values).
Regular Observation Grids. By default, consider regular, and potentially coarse, observation grids. To realize such observations of our sampled equations, we subsample the fine grid simulations from Appendix C.4 regularly. We vary the subsampling factor, and therefore the Regular inter- observation gap Åt, throughout the dataset, to cover wide range of application. Each subsampling yields I. remaining observations per path, still with time horizon of T. See Appendix C.6 for more details about the specific choices we made for our dataset.
Irregular Observation Grids. To accommodate applications with irregular inter- oiter-bbservatin gaps, we subsample these regular grids with an additional sampling scheme using the Bernoulli distribution. Given the regular, coarse observations of process, we sample a single Bernoulli survival probability per equation. The irregular observation grid for a given equation are then realized by sampling the Bernoulli distribution with survival probability 77 at each observation of all paths.
Additive Gaussian Noise. To make the model more robust, e.g. for real world applications, we corrupt the (clean) observations {x; ]s] KxL of a process (f, G) by additive Gaussian noise. The standard deviation is determined relative to the observed values. Concretely, let us define the component-wee wise range of the process:
(xx G) = 1 [(
xi)
}
min K. x,
(24)
i- 1 KxI.
To realize the additive Gaussian noise, first sample noise scale a (x((), G) for the given (clean) observations of process. Then, each observation is corrupted by samples from ((0,o(x(t), G)r((((), G)). In other words:
Puoise (y: (x)),
((()), ())
(25)
Note that a Course and Nair (2023) employ similar relative additive noise scheme in some of their experiments on synthetic data. However, they define the range of a process based on the associated ODE process, i.c with zero diffusion.
26

---

Table 4: The inter-observation times Ar, the number of paths K, the number of observations per path and the time horizon 7 present in our dataset, and the percentage data they occupy. Note that we vary the sampling step size At as well, speed up generation of paths with longer time horizon.
% of Data
At
AT
K
L
K x L
T
1/3
0.004
0.1
100
128
12800
12.8
1/3
0.002
512
12800
5.12
1/3
0.001
0.001
1024
12288
1.024
C.6 Data Generation Hyperparameters
Let us now summarize the data generation hyperparameters of the train dataset of FIM-SDE.
We sample SDEs of up to dimension dmax 3. The drift function components are polynomials of up degree m drif 3. The diffusion function components are based on polynomials of up degree
maa 2.
We vary regular inter-observation times Ar, the number of paths K, the number of observations per path L and the time horizon I throughout the dataset. The information is summarized in Table 4. Note that cover three orders of magnitude for Ât and adapt At for slightly faster generation. Moreover, we set K, L and T such that the number of observations per equation K x L is (roughly) fixed.
To corrupt the observation grids, we sample the Bernoulli survival probability as r7 ~ 4(0.9, 1] per equation. The additive Gaussian noise scale is sampled o (x()), G) 2/00, 0.1] per equation. Each corruption scheme is applied to third of the total dataset. Note that these corruptions overlap for ninth of the total dataset, i.e. observations of a process be noisy and on irregular grid.
The rejection threshold value set to 78% for 2D and 92% for 3D equations. Higher dimensional equations therefore require more equation samplings, to populate the training distribution with valid equations, illustrating a curse-of- dimensionality effect. Moreover, higher dimensional equations are also (intuitively) more challenging to estimate.
100, which leads to a rejection rate of 50% for 1D,
We hence generate 300K 3D equations, 200K 2D equations and 100k 1D equations for our synthetic train dataset, for a total of 600K equations. An additional validation dataset of size 60K is generated using the same ratios.
D FIM-SDE Model Details
This section contains the implementation details of FIM-SDE. We first describe model input pre- processing and output post-processing approach that allows the model to be used wide range of applications. We also provide more insights into the model architecture, training procedure and hyperparameters.
D.1 Instance Normalization and Vector Field Renormalization
Observation sequences on different dynamical process are naturally characterized by different spa- tial and temporal length scales. To handle such sequences in consistent manner, introduce component-wise instance normalization transformations. We pre-process the inputs to our model before application, and post-process the outputs according relevant change of variables formulas.
Continuing the notation from Section 4.2, let the data D {y,, Ay, Ay? Ar, (((-)))) and location x be the inputs of FIM-SDE. We consider linear normalization transformations. so the normalizations of Ay. Ay² and AT are implied by the normalization transformations of y and
Spatial Instance Normalization. Before applying FIM-SDE, we standardize the observations {y Ly.kK(L- ((L-)) component-wise. Let yiy denote the j-th component of y., then the mean and standard
27

---

deviation of the j-th spatial dimension are
K(L-1)
K(L-1)
bJ K(L-1 Σ Miy and
Sj V K(L - 1)
(y 8,)²
(26)
These define y = (Br...., .a....,al Da) E Rd and
= diag( $$1...., sa) E Rddd which specify the spatial
transformation
S(y) = S S - y)
(27)
and inverse
S yy) Sy +
(28)
important to apply the same standardization transformation to the location x, because they are elements of the now transformed domain.
Temporal Instance Normalization. The absolute time 7, of observation is not input of our model. We therefore normalize ÅT; directly. However, during application (as well as some part of our training data), AT; might be constant, i.e. the input observations are on regular grid. Moreover, AT; might differ (vastly) between applications, even if all of them are on a regular grid. We therefore only centralize for interpretability.
{=r (
(L-1) around target value, while keeping the inter-observation gaps positive
Let ATHr = 0.01 be target inter-observation gap after normalization and
K(L-1)
TIA K(L 1)
Σ
In An;
(29)
the mean of (In AT,
K(L-1), Then we normalize the inter-observation gaps Ar, with
TH exp( (la
(30)
In other words, we center (In Ar) K(L-1) at In ATu..
Note that the unique corresponding temporal transformation T of absolute time t that retains 0, i.e. T(0) = 0, is
T(r) AT exp(-Inxrr)
(31)
with inverse
(T) ATur exp( -Inna)
(32)
Inverse Process Transformation. FIM-SDE processes instance normalized inputs and therefore estimates process in the normalized domain. To yield the corresponding process in the original unnormaiized domain, we apply change of variables formulas according to the inverse of the spatial and temporal transformations from (28) and (32).
Let (S(x) (D) and (S))))) denote the estimated drift and diagonal diffusion functions at the transformed location S(x). Recall that by our assumption in (4), the estimated diffusion matrix is diagonal.
According Ito's formula and Theorem 8.5.7 in Øksendal (2010), the corresponding vector fields in the unnormalized domain at location are
Tux cxp( ImA, (S(()))))) and V cxp(- Go(x
(33)
where c denotes the element-wise product.
D.2 Model Architecture
In this section, provide some more details about the architecture of the single FIM-SDE, all presented experiments were conducted with.
We fix hidden size 256 for all embeddings throughout the different parts of the model.
28

---

All features are embedded with linear layers of to dimension 4 s.t. their concatenated embedding di is of the desired size The embedded features are further encoded by the Transformer encoder with linear attention Vp It consists of 2 layers and attention dimension n.
Each location is single input, encoded by another linear layer & dimensions.
The trunk-net equivalent functional attention mechanism first applies attention blocks (includ- ing residual feed-forward layers). The final embeddings h () are then projected to output dimension by feed-forward networks with 2 hidden layers of dimension
Following standard conventions, the hidden size of the residual feed-forward neural networks in all transformer layers is set to 4n = 1024.
We dropout of 0.1 and the GeLU activation function.
In total, FIM-SDE has roughly 20M parameters.
D.3 Training Procedure
We train FIM-SDE with AdamW (Loshchilov and Hutter, 2017), using learning rate 1e-5 and weight decay 4.
In each batch, we sample the total number of observations passed to the model from [[188, 12800). so the model is exposed to widely varying context sizes. For each equation in batch, randomly sample 32 locations to compute the loss L from (10).
Memory requirements during training are quite large, as we use up to 12800 observations (inputs) per instance in a batch. utilize four A 100 40GB GPUs to train with a batch size 64 for 1.3M optimization steps over roughly 6 days.
D.4
Finetuning
In this section, we provide the explicit expressions we use finetune FIM-SDE to some target data. Our finetuning approach depends on the assumed sampling rate of the observations.
Dense Observations. For densely sampled data, transitions between observations approximately Gaussian (see Eq. (18)). We compute the log-likelihood of the short-time transitions induced by the
estimators (fg. Ge). with respect to a target sequence y1....,yI E IRd with small inter-observation times ATt. Specifically, we write
Σ (1i++ - B1,1 - f.
L-1
d
CPT
+ log 9r, prx)
(34)
2y),
1=1
Sparse Observations. For sparsely sampled data, transitions between observations are not well approximated by Gaussian. We alternatively simulate the one-step transition with the estimated equation and compare its result to the observed transition. Concretely, for an observed transition from yı to Yi+1 in time An, let 1, be the result of a simulation of (fo, Go) from initial state y: for An. The finetuning objective is the mean squared error" between the ground-truth and estimated transition:
1 L
1y41 $141.0112
(35)
In both cases, the objective can be computed in parallel for all observations in the train split of the target data. Alternatively. the objective can be computed for a randomly sampled set of observations, if the set is sampled anew each iteration.
'The MSE could be replaced by other distance metrics, such as Gaussian negative log- likelihood, under the assumption of additive noise. The variance of said noise term could also be learned.
29

---

E
Baseline Models
Throughout our experiments, we compare FIM-SDE to three baseline models: SparseGP, BISDE and LatentSDE. This section offers a brief introduction to each method, provides relevant references and summarizes our hyperparameter choices.
E.1 SparseGP
We use the SparseGP approach of Batz et al. (2018) as a nonparametric baseline for stochastic differential equation (SDE) estimation. In this method, both the drift and diffusion functions are modeled with Gaussian processes (GPs), and GP regression is performed directly on the observed one-step increments (X Xt)/ Át and their squared values (Xt+ x / Åt. This enables simultaneous, data-driven estimation of the deterministic and stochastic components of the dynamics without specifying any parametric form. We treat the observation interval Âr as equívalent to the discretization step At, effectively assuming dense data, and always perform GP regression for both drift and diffusion.
For larger datasets, where direct GP regression becomes computationally expensive, we employ the sparse approximation proposed in the original work. In this setting, fixed number of inducing poines-chosen in regions of high trajectory density-arn are used to approximate the full GP posterior, significantly reducing computational cost while retaining accuracy. In our implementation, we use polynomial kernels for the drift functions and RBF kernels for the diffusion functions, with roughly ten inducing points per experiment. Kernel hyperparameters and diffusion noise levels are selected through evidence optimization.
E.2 BISDE
BISDE (Wang et al., 2022) is a symbolic regression method for the low-dimensional SDE inference problem. Using sparse Bayesian learning, it avoids binning the data, common technique in symbolic regression approaches. Binning a costly operation, which suffers from a curse of dimensionality and introduces approximation error.
The effectiveness of BISDE, compared to a standard (SINDy-like) symbolic regression approach, is demonstrated in by Wang et al. (2022) in range of experiments.
We the open-source implementation" of BISDE, released by the authors. As symbolic regression based method, requires a family of basis function. Our choices, per dimensionality of the equation, are:
ID: 1, x², x3, sin(x), exp(x)
2D: 1, I. y, x² 3² zy. x3, xzy. xy², y², sin(x), sin(y),
exp(x), exp(v), exp(xy), sin(ry)
E.3 LatentSDE
LatentSDE (L. et al., 2020) is a neural model that learns (prior) SDE in latent space from data via variational inference. The model is trained end-to-end on train split of the target datasets by maximizing an ELBBO objective using gradient descent.
We use the open-source implementation of LatentSDE released by the authors in the torchsde 10 package.
In all experiments, we use the LatentSDE setup and training hyperparameters from the Lorenz experiment described in al. (2020, Section 9.10. We report them here only for completeness.
The GRU encoder has layer with hidden size 100. The latent dimension is of size 4. Drift and diffusion functions implemented by layer MIPs, also with hidden size 100. The prior drift time-inhomogenous. The diffusion consists separate ML.s for each latent dimension. Latent equations solved with fixed discretization step size of 0.01. For the Gaussian emission model,
30

---

Table 5: Sizes of the evaluation datasets. We report the sizes of the train splits, which used as context for FIM-SDE, and the reference splits, used for MMD computations. Here, År is the regular inter-observation time gap, the number of paths, 7 the length of each path and the last column contains the total number of time points.
(a) Train data sizes.
(b) MMD reference data sizes.
Dataset
Dim.
Ar
Total
Dataset
Dim.
ÅT
P
T
Total
Can. Sys.
1, 2
0.002,
0.02
5000
5000
Can. Sys.
1, 2
0.002
100
500
50000
Lorenz
0.025
1024
41
41984
Lorenz
3
0.025
5248
Wind
0.167
5240
20960
Wind
0.167
10
521
5240
Oil
1
4
1584
6336
Oil
1
10
158
1580
Facebook
188--
4
4992
19968
Facebook
1
1.8e
10
499
4990
Tesla
1.8e
4
4992
19968
Tesla
1
1.8e
10
499
4990
the mean is implemented as linear projection from latent space. The standard deviation assumed to be diagonal with fixed standard deviation of 0.01. All MLPs employ the sof sof tpluss activation function.
We train for 5000 iterations, where the KL term of the ELBO objective is annealed linearly over the first 50 iteration. Each iteration is performed on the whole train set, i.e. gradient descent without minibatches. The learning rate of the Adam optimizer (Kingma, 2014) is set to 0.01 with decay of 0.999 per iteration. Gradients are computed through the solver, without the adjoint method.
Observation values are standardized per dimension and the observation times are normalized to [0, 1] before training. For evaluation, model outputs are rescaled accordingly.
F
Metrics
F1
Maximum Mean Discrepancy
The Maximum Mcan Discrepancy (MMD) is non-parametric statistical test used to compare two probability distributions (Gretton et al., 2012). MMD is a kernel based method that measures the difference between two distributions, P and Q, purely based on samples drawn from these distributions. The samples are mapped to high-dimensiooal feature space, where linear methods are effectively applied.
Signature kernels for time series are implemented in the KSIG library (Tóth et al., 2025). Using the RBF-kernel and signature levels, we obtain a signature kernel k(.,)). The MMD between two sets of samples {x,]== and {y,}リ= ~ can then by computed as
MMD2(P.Q))
Σ x(x,, "k(y)))) 141
Σ {()x,,))) n²
2
(36)
ann 1)
In our experiments, detailed in Appendix G, we generate or hold-out a set of reference paths specifically for the MMD computation. With each model, we simulate a set of paths of the same size and with the same time horizon. For consistency, reference set as the initial states for the simulation. The time horizon (and number of observations) of the simulated paths is also determined from the set of reference paths.
the first observation of cach path in the
To simulate these paths, BISDE and SparseGP learn an equation from a separate train set, that does nor include the set of reference paths for MMD computation. FIM-SDE uses the same train set as context infer an equation. Finetuning of FIM-SDE is always performed on the same train set. Finally, LatentSDE is trained on the train set, but the posterior initial states have to be inferred from the reference paths. We use the learned prior equation from LatentSDE to simulate paths, starting at these posterior initial states.
31

---

F22 Mean-squared Error on Vector Fields
In some experiments, we have access the ground-truth vector fields used to generate the data paths. This is leveraged to compute the equared error between the estimated and the ground-truth vector fields on a pre-defined grid. We define regular grid in hypercube that surrounds the observations. The mean-squared error is reported separately for drift and diffusion.
G Experiments
This section covers our experiments in great detail, including a description of our employed metrics and dedicated subsection for each dataset. We provide thorough description of each dataset, our experimental setup and offer additional results. Table 5 contains information about the number of observations, number of paths and inter-observation time gaps for each dataset.
G.1 Canonical SDE Systems
G.1.1 Description
We collect set of well known SDEs with accompanying initial conditions, which were previously used by Batz al. (2018), Course and Nair (2023) and Wang et al. (2022) to benchmark inference performance of their models. We simulate paths from each of these equations to generate datasets with fine and coarse grid observations and with optional Gaussian noise.
Here, we summarize the selected processes dx (where W a Wiener processes), the initial conditions x(D) and the source of the equation.
Double Well:
dr 4(x - x²) dt + /max(4-- - 1.25x², .) (t((t)
(37)
with (0) = 0, extracted from Batz et al. (2018).
2D Synthetic System:
dæj = (I1 12 x11 - dt + 1 23 dWi(t)
(38)
dxa
(z) + 12 - 1²-02 x dt + 1 dWa(t)
(39)
with x(0) = (1.5, 1.5), extracted from Wang et al. (2022).
Damped Linear Oscillator:
dr - -00.1 - -200z)) dt + (t)
(40)
drg -22.0 0.1 x2) di + 2
(41)
with x(0) (2.5, 5), extracted from Course and Nair (2023).
Damped Cubic Oscillator:
dr - -0.1 - 2.0 di + d1)
(42)
dr -(2.0 0.1 =3) dt dW2(t)
(43)
with x0) = -1), extracted from Course and Nair (2023).
Duffing Oscillator:
dx
z2 dt dWi(t) (t)
(44)
dx -a 31 + 0.35 22) dt d(2
(45)
with x(0) = (3,2) 2), extracted from Course and Nair (2023).
Selkov Glycolysis:
dr -(11 0.08 r2 22 dt + dWi(t)
(46)
dzg
(0.6 0.08 Ø2 27 dt + dW2(t)
(47)
32

---

with x() = (0.7,1.255) extracted from Course and Nair (2023).
Hopf Bifurcation:
dr
(0.5 21 + 22 11-xxx1- (x x3) dt dWi(1)
(48)
dxg (-I1 + 0.5 22 22 x²)) dt + dW2((t)
(49)
with x(0) = (2, 2), extracted from Course and Nair (2023).
We simulate a single path from each systems with the Maller-aruuyaaa method and discretization step size At 0.002. This path is subsampled by a factor of
to achieve onter-observation
gaps of Á 0.002 or Åt = 0.02. The number of observations is kept fix at 5000, yielding time horizons of T 1 or T 10 respectively. Each resulting trajectory is corrupted by relative additive Gaussian noise (as in Appendix C.5) with standard deviation p 0.0 or p 0.05.
For each system dx, each inter-observation gap AT and each standard deviation p, we repeat this sampling and evaluation five times. These five repetitions of each setup therefore yield a standard deviation of the models' performance.
We generate 100 additional paths with inter-observation gap At = 0.002 and 500 observations from each system dx, without additive Gaussian noise. These paths used as reference compute the MMD, as described in Appendix F.I.
Finally, we record the ground-truth vector fields of each system at 1024 locations in axis- aligned regular hypereube, roughly surrounding the observations. These hypereubes are defined by the span [lower; upper;] they occupy in dimension
More precisely, we use {[-2, 2 for Double Well, ([-4, [-,, [-,44] for 2D Synthetic System, {[-2,2), 2, 2.} for Damped Linear Os- 2,
cillator, {[ 2|, 1-2, 2}} for Damped Cubic Oscillator, { {[-4. 4, 4), [4, 4]} for Duffing Oscillator,
{1-2, 4], 22 4]} for Selkov Glycolysis, and {1-2, 2), 2, 2 for Hopf Bifurcation.
G.1.2 Experimental Setup
For each system d, each At and and each of the five repetitions, the data generation yields single path with 5000 observations. These are used as context by FIM-SDE, BISDE and SparseGP to estimate drift and diffusion function, specifying SDE.
We sample paths from these estimated equations compute the MMD. More precisely, we take the first observation of each of the 100 reference paths as initial states and simulate the equations on the same observation grid as these reference paths, i.c. with Ar 0.002 and 500 observations. These two sets of 100 paths are used for the MMD computation (see Appendix F.1). The results are reported in Table and include the standard deviation over the five repetition.
Moreover, we extract the values of the estimated vector fields in the 1024 locations of the regular hypercube. For each vector field, we compute the mean-squared error between the estimated and ground-truth values. Table 6 contains the results for the drift, Table 7 for the diffusion.
Process estimation or simulation can fail for the basline models BISDE and SparseGP in some setups or experiment repetitions. We remove those repetitions from the performance metric computation, and only compute the mean and standard deviation over the successful attempts. The number of failed repetitions are marked by Latin numerals the respective tables.
G.1.3 Additional Results
Mean-squared Error on Vector Fields. Tables 6 and 7 contain the mean-squared error between the ground-truth and estimated vector field values regular hypercube. In general, these results complementary to the MMD results in Table 1. The baseline models perform decently on fine observation grids with no noise. FIM-SDE also handles coarser observation grids, potentially with additive Gaussian noise, very well.
Figure of Vector Fields and Sample Paths. Figure 5 contains ground-truft -tuth vector fields and sample paths of one experiment for each system. We display the results from the experimental setups without additive noise, / = 0.0, and the fine onter-observation grid, AT 0.002. For the vector fields, estimations from SparseGP, BISDE and FIM-SDE are displayed. We only show the sample paths from FIM-SDE for clarity. Vector field estimates and sample paths from FIM-SDE agree with the ground-truth.
33

---

Ground-Truth
SparseGP
BISDE
FIM-SDE iZero-Shot)
Drift
Diffusion
Paths
Double Well
Damped Synthetic
Damped
Glycosis
Figure 5: Drift and diffusion function estimation from FIM-SDE and baselines in all canonical SDE systems. The vector fields are estimated from observations with setup = 0.0 and Ar 0.002. The sample paths from FIM-SDE resemble the ground-truth paths closely.
34

---

Table 6: Mean-squared
drift estimation in canonical SDE systems. The error is computed between values of the estimated and the ground-truth drift on the regular grid described Ap- pendix G.I.I. Drift functions estimated by SparseGP and BISDE might diverge, we discard all
values > 10a and indicate them by
Double Well
2D-Synt
Damped Lincar
Damped Cubic
AT
Model
Duffing
Glycolysis
Hopf
(Wang)
SparseGP
35.7
4750
3.92
1.39
1010
452
44.8
0.002
BISDE
4.94
0.203
0.461
1.56
49.0
10.5
FIM-SDE
8.34
1150
1.16
1.83
6.25
81.6
9.66
0.0
SparseGP
21.2
621
3.47
478
3210
28.2
6.31
0.02
BISDE
1.92
0.0243
0.657
0.437
1.71
1.94
FIM-SDE
4.83
453
0.180
0.437
1.04
14.0
8.10
SparseCP
59.5
7.80
74.8
648
466
0.002
BISDE
819
1060
87.1
3380
4150
FIM-SDE
50.6
746
3.57
10.9
89.0
110
99.2
0.05
SparaeGP
33.0
546
3.69
8370
180
127
10.5
0.02
BISDE
0.409
1.23
2.67
21.0
45.1
6.50
FIM-SDE 2.31
489
0.297
0.591
5.87
30.1
6.77
Table 7: Mean-squared error on diffusion estimation in canonical SDE systems. The error is computed between values of the estimated and the ground-trufe diffusion the regular grid described in Appendix G.1.I.
Double Well
2D-Synt (Wang)
Damped Linear
Damped Cubic
P
AT
Model
Duffing
Glycolysis
Hopf
SparseCP
0.0227
1.70
0.0958
0.0205
0.152
0.0411
0.00373
0.002
BISDE
0.00620
0.548
0.000105
0.000115
0.000114
0.0264
0.00197
FIM-SDE
0.0209
0.349
0.0111
0.00619
0.0945
0.112
0.114
0.0
SparseGP
0.205
1.37
0.0302
1.42
0.848
0,129
0.00749
0.02
BISDE 0.0460
1.26
0,0000661
0,000331
0.000325
0.00223
0.00484
FIM-SDE
0.0334
0.268
0.0177
0.127
0.128
0.123
0.106
SparseCP
3.21
2.18
115
10.5
89.5
7.04
4.26
0.002
BISDE
2.43
2440
75.8
10.2
72.8
7.38
3.96
FIM-SDE
1.26
2.19
2.87
0.825
2.33
0.224
0.252
0.05
SparseCP
0.553
1.08
6.48
2.12
7.58
1.00
0.151
0.02
BISDE 0.0922
0.395
5.52
0.881
8.01
0.744
0.113
FIM-SDE
0.0371
0.733
2.07
0.334
2.50
0.171
0.129
G.2 Lorenz
G.2.1
Descríption
We generate synthetic data from Lorenz attractor system with constant diagonal diffusion, following the setup from Li et al. (2020, Section 9.10.2). For completeness, we recall the setup in this section.
Li al. (2020) consider the Lorenz system
dri 0 x1) di + a d(
(50)
dr
(mi(p ra) x2) dt + dW2(t)
(51)
dr3 = (I1 (aa)) di + dWa(t)
(52)
with parameters o = 10, p = 28, 3 = 8/3 and a 0.15. Initial states are sampled from the standard Gaussian distribution (0,
We simulate paths the time interval [0, 1] on a fine grid and record observations in a coarse
regular grid with 41 points, corresponding inter-observation gap Ar 0.025. These paths are
35

---

standardized in each dimension. Afterwards, they are corrupted with additive Gaussian noise using standard deviation of 0.01.
We generate a train set of size 1024 and a validation set of size 128, which is standardized using the statistics from the train set.
For the MMD computation, we generate three additional reference datasets with 128 paths each, using three distinct initial distributions: (0, I). (0, 21) and fixed value of 1,) Simulation time interval and inter-observation gap remain [0, 1] and ÅT = 0.025 respectively. Standardization is performed with the statistics from the train set. No additive Gaussian noise is applied.
G.2.2 Experimental Setup
We compare FIM-SDE to LatentSDE with MMD, using the additional reference sets.
Implementing LatentSDE with the hyperparameters described in Appendix E.3, we train the model on the train split for 5000 epochs. To sample paths for the MMD computation, the trained model encodes the reference set and estimates posterior initial condition. Using the learned prior equation, we sample set of 128 paths on the same observation grid as the MMD reference paths.
FIM-SDE processes the train data as context and estimates an equation. Using the first values each MMD reference path as initial states, we also sample a of 128 paths per reference dataset.
The MMD is computed between the set of reference paths and these sampled model paths. We report the results in Table 2.
G.3 Real-World Systems
G.3.1 Description
The set of four emirical datasets, representing complex real-world phenomena, were collected and studied by Wang et al. (2022). They were released alongside the open-source implementation of BISDE We provide a brief description of each of them and refer to Wang et al. (2022) for more details.
Stock Prices: The prices Facebook and Tesla stock were recorded every minute from July to September 2020. As such, there are 24960 one-dimensional observations per stock. Following Wang et al. (2022), the models process log-transformed observations. For evaluation, the estimated equations transformed back the original data space with Ito's formula.
Wind Speed Fluctuation:
For this dataset, wind speeds at meteorological station in Wellington, New Zealand, were recorded every ten minutes in the first half of 2020. The fluctuations, i.e. difference between two consecutive wind speed recordings, are modeled. There are 26202 one-dimensional observations in total.
Oil Price Fluctuations: Daily crude oil prices from 1986 to 2017, collected from data by the US Energy Information Administration. The daily price fluctuations are modeled, yielding 7921 one-dimensional observations.
To obtain meaningful results and standard deviation for the MMD metric, we perform five-fold validaaiint..... Each dataset is split into five splits of equal size. Data from four splits is used as context for the models, while the last split provides reference paths. We split this reference split into ten paths of equal size, because MMD computation requires of reference paths.
G.3.2 Experimental Setup
We compare SIE to BISDE and LatentSDE on each of the four datasets with MMD, using our five-fold vross-validttion setup.
For each instance of the vross-validation, data from four splits is used as context for FIM-SDE and as training data for BISDE and LatentSDE. From the estimated (prior) equations, sample ten paths on the same observation grid as the ten paths from the MMD reference split. FIM-SDE and BISDE
Il "nttps://githuu http://ttttb....o///AAAAAAAA////
36

---

use the first observations of the reference paths as initial states. LatentSDE infers a posterior initial condition latent space by encoding the ten reference paths.
The MMD is computed between the set of MMD reference paths and this set of ten sampled paths from the respective model. The five-fold coss-validation yields mean and standard deviation of the MMD, aggregating the results per dataset and model.
Table 3 contains the results of this experiment. FIM-SDE performs better than BISDE in the two stock price datasets, but worse in the two fluctuation datasets. LatentSDEs performance similar to FIM-SDE.
G.3.3 Additional Results
BISDE Equations. We report the equations estimated by BISDE per dataset and cross-validation split Table 8.
Table 8: Symbolic drift and diffusion function estimates from BISDE in real-world systems. We report the functions per cross-validation split for each real-world system.
System
Split
Drift
Diffusion
-7.2480 x 0.6628 sin(x)
2.9061
1.5182 x.5683 x 0.1565 x
1.9734 sin (x) 0.07/6
2
-7.1237 x 0.4295 sin(x)
3.0896
0.1223 r
1.7693
- e
Wind
3
-6.9100 + 0.2030 sin(x)
2.8466
2.4500 1.456.452222 x² 0.2513 +
2.8331
sin(x) 0.0877
4
6.8015 + 0. 0.1054 sin(2)
3.1170 + 1.3119 0.3860 sin(x) 0.0601 e*
5
6.88415 0.1443 sin(x)
2.8354
1.6006 1.5005 0.13 +
2.0828 sin(x) 0.0616
1
-1.0449
1.0427 + 0.3162
2
-1.0447 ir
1.1164 + 0.3060
Oil
3
-1 -1.0426.r x
1.0448 + 0.3108 x
4
-1.0405 x
0.5797 + 0.3132 0.0900 sin(x)
5
-1.0426
0.7484 + 0.3346
0.1720
2
0
Facebook
3
0
0 -0.2075 sin(2))
1
O
5
0.2555 sin(x)
1
-0.6620 sin
2
-10.5399 sin(x)
0
Tesla
3
-9.1783 sin(x)
1.3294
4
-7.2514 sin(x)
0.6220
5
-6.5110 sin(x)
0.4163 sin(2)
Convergenee Speed. We additionally compare the performance on the four real-world systems over the course of training of LatentSDE and finetuning of F FIM-SDE.
The experimental setup remains the same: a five-fold cross-validation evaluated by the MMD the reference split. The results are summarized in Table 9,
We observe that LatentSDE requires thousands of iterations to converge, while FIM-SDE reaches comparable or superior performance after only 100 iterations.
Figure of Sample Paths. Figure contains the reference paths and estimations from LatentSDE and FIM-SDE for one cross-validation split each system. zero-shot mode, FIM-SDE performs roughly as well as the custom trained LatentSDE, For the oil and wind systems, finetuning FIM-SDE improves the sample path quality drastically. In the stock price data, finetuning does affect the already convincing paths from the zero-shot model.
37

---

Ground-Truth
LatentSDE
FIM-SDE Shro-Shot)
FIM-SDE (Finetuned)
Wind
10
10
10
30
30
60
30
60
Facebook
255
255
0.000
0.002
0.004
0.000
0.002
0.004
0.000
0.002
0.004
480
Tesla
0.000
0.002
0.004
0.000
0.002
0.004
0.000
0.002
0.004
Figure 6: Paths from one cross-validatie split in each real-world system. For each system we plot the reference paths of this split and the estimations of LatentSDE and SDE, in zero-shot mode and finetuned. The displayed trajectories for oil and wind are truncated for clarity.
38

---

Table 9: Convergence speed of training LatentSDE and finetuning FIM-SDE in four real-world systems. Depicted is mean and standard deviation of the MMD over five cross- volsdaaaloo splits per system at different training (resp. finetuning) iterations.
Model
Iteration
Wind
Oil
Facchook
Tesla
500
1.019 0.078
0.382 + 0.236 0.091 + 0.043
0.066 = 0.059
LatentSDE
1000
0.874
0.127
0.224
0.071
0.037
0.052 0.056
2000
0.318
0.164
0.205=00.11
0.034
0.023 = 0.028
5000
0.350 = 0.212
0.235 + 0.167
0.014
0.016 = 0.023
FIM-SDE (Zero-Shot)
0.330 1 0.11
0.199
0.012
0.026
0.000 + 0.000
20
0.114 = 0.06
0.071 0.094
0.020 0.009
0.021 = 0.014
40
0.075 + 0.05
0.048 0.057
0.019 1 0.026
0.020 1 0.028
FIM-SDE (Finetuned)
60
0.042 1 0.04
0.049
0.003
0.007
0.013 1 0.018
80
0.043
+ 0.07
0.002
1 0.005
0.011 1 0.016
100
0.029
= 0.08
0.068 0.006 0.042
0.004 = 0.008
H
Ablation Studies
In this section, we present ablation studies on the training and application of FIM-SDE. First, we investigate the scaling of FIM-SDE, by scaling the train data size and the number of model parameters. Then we consider the train data distribution, ablating the degree of polynomials it consists of. Lastly, we study the impact of the context size on the performance of FIM-SDE, exemplified by data from two canonical systems.
H.1 Ablating Size of Train Data and Model Parameter Count
Table 10: Ablation study on scaling FIM-SDE in terms of parameter count and train data size. We report the MMD, scaled by a factor of 100, on canonical systems, which can be compared to the results from Table G.I.
Param. Data Double 2D-Synt Damped
Damped Cubic
P
AT
Count
Size
Well
(Wang)
Linear
Duffing
Glycolysis Hopf
5M
30k 30(1)
10(3)
10(2)
13(4)
9(4)
11(4)
13(4)
0.002
10M
100k 067(6)
8(3)
8(2)
7(3)
7(2)
10(4)
6.3(8)
0.0
5M
30k
0.4(5)
2.5(5)
0.5(5)
3.8(6)
0.6(2)
1.5(6)
3.2(3)
0.02
10M
100k
0.5(7)
0.2(4)
0.01(2)
0.6(4)
0.07(8)
0.3(2)
0.1(2)
5M
30k
30(1)
10(4)
10(2)
12(4)
16(8)
12(7)
14(4)
0.002
10M
100k
3.9(6)
7(3)
10(2)
14(2)
14(3)
13(4)
9(2)
0.05
5M
30k
0.8(6)
1.4(6)
1.1(7)
3.8(5)
3(1)
1.2(5)
1.2(5)
0.02
10M
100k 0.4(8)
0.4(8)
1.9(5)
0.5(2)
2.9(10)
0.8(6)
0.0(0)
Recall from Appendix D that FIM-SDE has roughly 20M parameters and was trained on 600k equations. We train two smaller models on fewer equations to investigate the scaling behavior of FIM-SDE, One model has 10M parameters and is trained on 100k equations. The other, even smaller, model has 5M parameters and is trained on just 30k equations. All models are trained for 500k optimization steps.
We compare both models in terms of MMD on the set of canonical systems from Appendix G.I. The results are reported in Table 10 and can be compared the results Table 1.
Without relative additive noise, i.e. P = 0.0, the larger model trained on more data performs consistently better. Only in a few setups, results of the two models less than standard deviation apart, otherwise the larger model (significantly) better. Although less decisive, the same pattern emerges data with relative additive noise, i.e. p 0.05.
Note that both (smaller) models are worse than the (larger) FIM-SDE from Table These results validate our choice of scaling the model and train data size for better overall performance.
39

---

Table 11: Ablation study the train data distribution of FIM-SDE, in particular the maximal degree of drift polynomials in the train data. We report the MMD, scaled by a factor of 100, on canonical systems, which can be compared to the results from Table G.I.
Drift Degree
Double
2D-Synt
Damped
Damped
P
Well
(Wang)
Linear
Cubic
Duffing
Glycolysis
Hopf
3
30(1)
10(3)
10(2)
13(4)
9(4)
11(4)
13(4)
0.002
4
3(1)
7(2)
6(1)
6(2)
6(1)
10(3)
8(2)
0.0
3
0.4(5)
2.5(5)
0.5(5)
3.8(6)
0.6(2)
1.5(6)
3.2(3)
0.02
4
0.4(6)
2.7(7)
0.1(1)
2.9(8)
0.33(10)
1.(5)
1.4(4)
3
30(1)
10(4)
10(2)
12(4)
16(8)
12(7)
14(4)
0.002
4
9(3)
11(4)
9(2)
15(4)
12(4)
16(5)
10(2)
0.05
3
0.8(6)
1.4(6)
1.1(7)
3.8(5)
3(1)
1.2(5)
1.2(5)
0.02
4
0.2(2)
2.1(4)
1.7(9)
4.0(4)
2.5(9)
2(5)
1.1(5)
H.2
Ablating the Train Data Distribution
In Appendix C we describe the synthetic train data distribution for FIM-SDE. Recall that this data is generated from SDEs with polynomial drift of up to degree 3. In this Section, we ablate the polynomial degree of the drift functions and study its impact.
We train two models with 5M parameters on 30k SDEs. For one model, the 30k SDEs are sampled with mi = 3. as described in Appendix C.6. The other model is trained on 30k SDE using the same setup, apart from allowing degree drift functions, i.c. m din
We compare the two models on their performance in terms of MMD on the set of canonical systems from Appendix G.1 and report the results Table 11.
The results often differ by less than one standard deviation. Without noise, i.e. 0.0, the model trained on up degree drift seem perform slightly better than the other model, when only considering the mean performance. This advantage vanishes if noise is added, i.e. 0.05.
Considering the significantly higher sampling rejection rate when generating the data with degree 4 drift, while only providing minimal performance increases, we decided to train FIM-SDE on data with driit
3.
H.3 Ablating the Context Size during Evaluation
During training, FIM-SDE has been exposed to context sizes from 128 to 12800 observations (see Appendix D). In this section, we ablate the context size of FIM-SDE during evaluation and study its impact on performance.
We conduct this ablation study on two canonical systems from Appendix 1: the Double Well and the Damped Linear Oscillator. Sets of context paths of sizes ranging from 500 50000 are generated using the fine inter-observation gap AT 0.002 and no additive Gaussian noise = O. Using the experimental setup from Appendix G.I, we evaluate the performance of FIM-SDE on these different context sizes using the MMD and repeating each experiment five times.
In Table 12 we report the results of ablating the context size of FIM-SDE over two orders of magnitude. The model requires at least 2000 context points to infer the Damped Linear Oscillator correctly. The Double Well system, inherently more difficult because of its state-dependent diffusion function, requires 4000 context points. Increasing the context size, even beyond the train maximum of 12800, does not deteriorate the already good performance FIM-SDE.
40

---

Table 12: Ablation study on the context size FIM-SDE. We consider data from two canonical
systems (see Appendix G.1), generate with p 0 and AT 0.002. Ablating the number of observations, we report the mean and standard deviation of the MMD across five repetitions each experiment.
Context Size
Double Well
Damped Linear Oscillator
500
0.155 1 0.090
0.128 0.068
1000
0.115 0.094
0.014 0.005
2000
0.053 0.073
0.006 0.005
3000
0.042 + 0.086
0.002 + 0.003
4000
0.004 + 0.003
0.000 + 0.000
5000
0.004 + 0.003
0.000 0.000
50000
0.006 0.006
0.000 0.000
41


Reviews:
[{"review_id": "8ammqGpIK2", "reviewer": "Reviewer_Kq3M", "summary": "This paper seeks to pre-train a large foundation model on the supervised mapping of SDE realizations to their associated drift and diffusion coefficients so that later on, as an alternative to GPs and neural SDEs.\n\nThe authors generate a dataset of 600k low-dimensional SDEs, make their observations a bit noisy, and learn to regress the associated drift and diffusion coefficients giving rise to them. \n\nThey then show that the tool can be used to predict the coefficients of unseen SDEs that fall within the class, and compare to existing methods.", "questions": "- The reviewer is trying to understand how fair the comparison is to GPs. GPs also give you a nice measure of uncertainty on your output. Can you get the same from your approach, rather than just a pointwise mapping?\n\n- Can the authors comment more on why they only chose this sparse set of low dimensional SDEs? Is it the difficulty trying to come up with higher dimensional dataset?\n\n- The \"generalization\" performance that the authors discuss for zero-shot prediction on real world datasets seems not strongly well founded. Surely there are more than four stochastic systems of real data to test on, and matching the performance of one other method is not the strongest argument for generalization. Could the authors investigate this further for the rebuttal? This is a pretty limited set of experiments to justify one of the main interesting claims of the paper.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "*Strengths*\n\n- While the paper relies on a very constrained parametric class of drift and diffusion coefficients on which the synthetic datasets are constructed, they are very upfront on this, and the reviewer sees this paper more as a proof of concept on building these foundation models\n- The phrasing of the approach is interesting and seems novel as I understand it, and, with the right generalization of function classes and datasets, could be an interesting approach going forward.\n\n*Weaknesses*\n\n- My main criticism is something the authors will likely agree with, namely that the generalization of the \"foundation model\" to SDEs that it has not been trained on makes sense when the tested models are still within, in some sense, the parametric class of polynomial drifts of degree 3. One ablation that should be done but that I don't see in the paper is if the test SDEs have drift and diffusions that fall within the class of polynomials of degree 3. \n- My other main criticism is the low dimensionality of the target SDEs. Why is this the case? A computational limitation? It does not seem absolutely necessary to the reader that they proceed this way, as there exist many paradigms for learning high dimensional SDEs already (e.g. score-based diffusion). Is the issue coming up with interesting candidate training datasets? \n- Equation 8 is not well explained. You do not address why you are introducing $U$ other than that it can be interpreted as a measure of epistemic uncertainty and a citation to another paper. This needs to be better motivated.\n\n\n\n\n*Minor comments*\n- typo with unnecessary comma line 57\n- MSE is not a divergence measure, you should consider changing this verbiage (line 215)", "was_revised": true, "final_justification": "The authors provided more experimental justification of their method on real-world datasets, showing that their proposed training approach was expansive enough to encompass meaningfully interesting SDEs and not just proof-of-principle ones. Nonetheless, they also helped make clear that there is a curse of dimensionality in terms of number of training points needed to generalize that is pretty steep, so I'm keeping my score at a 4 (weak accept).", "rebuttal": "Before we begin, we would like to thank the reviewer for their detailed feedback, for identifying potential weaknesses and raising important questions, and for recognizing the contributions of our work. We address each of the raised points below.\n\n**@W1**:  We agree with the reviewer that FIM‑SDE can be expected to generalize only to *(unseen) SDEs that are within the SDE class used to build the synthetic training set of Section 3.1*. However, the practical impact of this limitation depends on whether that class is already rich enough to cover SDEs used in real-world applications.\n\nFIM‑SDE was trained on data drawn from the class\n\n\n$\\\\mathcal{C} = \\\\{\\\\text{SDEs with polynomial drift of degree} \\\\leq 3, \\\\text{ and polynomial diffusion of degree} \\\\leq 2\\\\}$,\n\nwhere by diffusion we refer to the $\\\\{g_1, …, g_d\\\\}$ functions in Eq. 4.\n\n*Why $\\\\mathcal{C}$ is already practically useful*:\n\n- The eight benchmark SDEs used in Section 4.1 all belong to $\\\\mathcal{C}$. The reviewer can readily check this by looking at their definition in Appendix G.4.  These SDEs are widely used across natural and social sciences alike.\n- In Section 4.2, we studied four widely different real-word datasets, namely Facebook and Tesla stock prices, and oil price and wind speed fluctuations.  Stock prices and oil price fluctuations are traditionally modelled as geometric Brownian motion (linear drift and diffusion, see e.g. [1, 2] for the latter), while wind speed fluctuations have been recently described with linear drift and quadratic diffusion [3]. All these models belong to $\\\\mathcal{C}$.\n- Pretraining FIM‑SDE on data sampled from $\\\\mathcal{C}$ provides a good initialization for fine-tuning the foundation model to data (See our replies **@Q3** to you, and **@Q6** to reviewer WGed below).\n\nIn short, many SDEs of practical interest already satisfy the polynomial restrictions, and FIM‑SDE shows consistent generalization to them. \n\n**[PROPOSED MODIFICATION:]** We will include the discussion above on generalization within class $\\\\mathcal{C}$ into our experimental section (Section 4). \n\n**[QUESTION TO REVIEWER:]** Does this address the missing ablation the reviewer requested? If further evidence is required, we can include an ablation on degree‑4 drift SDEs (outside $\\\\mathcal{C}$) and report the expected drop in performance.\n\n**@W2**: The intuition of the reviewer is correct. One of the main reasons why we, in this first study, restrict our attention to low-dimensional SDEs are the **data‑generation rejection rates**.\n\nAs we described in lines 582-587 (Appendix C3), we simulated $100$ paths per candidate SDE and rejected any SDE whose paths contain NaNs, Infs, or values exceeding a threshold value of $\\\\pm 100$. We have empirically found that the rejection rate rises sharply with dimension: $\\\\mathbf{50\\\\%}$ **for 1‑D,** $\\\\mathbf{78\\\\%}$ **for 2‑D, and** $\\mathbf{92\\\\%}$ **for 3‑D**. Higher dimensionality therefore demands far more SDE candidates to populate the training distribution, illustrating a *curse‑of‑dimensionality effect*.  \n\nIn future work, we plan to reduce the rejection rates of our algorithm by generating SDEs that are *a priori known to have a solution* (i.e. by randomly generating drift and diffusion functions that are globally Lipschitz and feature linear growth). Note that we already commented on this in lines 322-326 of our Conclusions.\n\n**[PROPOSED MODIFICATION:]** We will state this bottleneck more explicitly and include the rejection rate statistics in the Limitation Section of the manuscript.\n\n**@W3**: In lines 221-224 of our manuscript we explained that when one marginalises Eq. 7 over a uniform prior on $\\\\mathbf{x}$, the integrand can be dominated by regions where the signal‑to‑noise ratio is extremely high (or low), producing large gradients and numerical instability. In line 226 we then wrote that the learnable and state-dependent $U_\\\\theta$  function works as a weighting mechanism that is introduced to address these instabilities. \n \n*Interpretation as epistemic uncertainty*: Assuming that Eq. 7 is a state-dependent Gaussian variable, $U_\\\\theta$ plays the role of the learnable, state-dependent log-variance. If we look at Eq. 8, large $U_\\\\theta$ values down‑weight regions where the model is uncertain, while the additive $+U_\\\\theta$ term prevents it from collapsing to $-\\\\infty$. Empirically, we have observed that $U_\\\\theta$​ indeed peaks in sparsely observed or highly noisy areas, acting as a proxy for epistemic uncertainty.\n\n**[PROPOSED MODIFICATION:]** We will add this explanation to the main text of our manuscript and provide a plot of $U_\\\\theta$ for a representative example in the Appendix.\n\n**@Q1**:  The primary goal of our proposal is to perform SDE inference from noisy data *in a single forward pass*. In other words, FIM-SDE was trained to perform zero-shot pointwise mappings only. **However, uncertainty estimation is a natural next step**. In fact, as we discussed in **@W3**, the function $U_{\\\\theta}$ can also be interpreted as an uncertainty estimator *wrt. the loss function defined in Eq. 7*. Future work will consider two independent $U$ functions, one that weights the drift contribution (first term  in Eq. 7), and the other that weights the diffusion contribution (second term in Eq. 7). Each of these $U$ functions can then be used as uncertainty estimators for the drift and diffusion functions, respectively.\n\n**@Q2**: We kindly refer the reviewer to our reply **@W2** above.\n\n**@Q3**: In the $16$ papers that we reviewed in our related work, Appendix A, we unfortunately did not find many real-world datasets on which to test FIM-SDE, but the ones that we considered. In fact, *most works typically only test their proposals on synthetic datasets*. \n\nHaving said this, we have now extended our experiments on the $4$ real-world experiments, Table 2a of the manuscript, by including the scores of \n\n(i) LatentSDE [5] (one of the most prominent neural network models for SDE inference); and\n(ii) FIM-SDE *fine-tuned* to the training sets.\n\n| Model           | Wind              | Oil               | Facebook          | Tesla               |\n|:----------------|:-----------------:|:-----------------:|:-----------------:|:-------------------:|\n| BISDE           | $0.000 \\pm 0.000$ | $0.137 \\pm 0.097$ | $0.053 \\pm 0.060$ | $0.013 \\pm 0.015$   |\n| LatentSDE       | $0.350 \\pm 0.212$ | $0.235 \\pm 0.167$ | $0.014 \\pm 0.018$ | $0.016 \\pm 0.023$   |\n| FIM (zero-shot) | $0.330 \\pm 0.114$ | $0.199 \\pm 0.137$ | $0.012 \\pm 0.026$ | $0.000 \\pm 0.000$   |\n| FIM (fine-tuned)   | $0.029 \\pm 0.034$ | $0.068 \\pm 0.042$ | $0.006 \\pm 0.007$ | $0.004 \\pm 0.008$   |\n\n*Key take‑aways*:\n\n- The **zero-shot** performance of FIM-SDE on real-world data is competitive with *both* symbolic regression methods (BISDE) and neural models (LatentSDE). These results highlight the robustness of FIM-SDE’s generalization across markedly different dynamical systems.\n- Leveraging Eq. 12 (Appendix D), FIM-SDE can be readily fine-tuned to real-world data. The fine-tuned model outperforms both baselines across all systems. \n\nThese findings demonstrate that FIM-SDE functions as a true foundation model, exhibiting competitive zero-shot performance and the ability to be fine-tuned on complex data.\n\n**[PROPOSED MODIFICATION:]** We will include these new results into our experimental section. \n\n**@minor comments**: We thank the reviewer for pointing these issues out. We have corrected them in our manuscript.\n\n*References*:\n\n[1] Noh, N. M. et al. Analysis of oil price fluctuations. In AIP Conference Proceedings, volume 1750, 2016.\n\n[2] X. Yang and Z. He. Predicting the price of crude oil based on the stochastic dynamics learning from prior data. Stochastic Environmental Research and Risk Assessment, 2024.\n\n[3] Wang et al. Data-driven discovery of stochastic differential equations. Engineering, 2022.\n\n[4] d’Ascoli et al.  ODE-Former: Symbolic regression of dynamical systems with transformer. ICLR, 2024.\n\n[5] X. Li, et al. Scalable Gradients for Stochastic Differential Equations. AISTATS 2020."}, {"review_id": "dqpiqHjIcV", "reviewer": "Reviewer_WGed", "summary": "The paper introduces FIM-SDE, a neural network model designed for in-context learning of stochastic differential equations (SDEs). Its main contributions are\n1. Developing a pretrained neural method capable of estimating drift and diffusion functions directly from noisy data across diverse tasks.\n2. Framing the SDE discovery task as an in-context inference problem, enabling zero-shot generalization to different SDE systems.\n3. Demonstrating that FIM-SDE can handle noisy, sparse, and irregularly sampled data effectively, outperforming traditional approaches.\n4. Providing a unified architecture that maintains consistent representations of drift and diffusion functions across various datasets.\nOverall, it advances the state of data-driven SDE discovery by combining neural network amortization, in-context learning, and distributional inference techniques.", "questions": "1. What exactly does $\\Delta y^2$ represent? If $\\Delta y$ already carries the full information, what is the purpose of introducing $\\Delta y^2$?\n\n2. The input consists of K(L−1) tuples. Is it provided as a sequence, similar to the standard Transformer input format? If only fewer tuples are available, would the model still work?\n\n3. The paper focuses only on diagonal diffusion. Is there a specific reason or limitation behind this choice? Additionally, are there any other constraints on the diffusion term, such as requiring it to be invertible? If invertibility is not enforced, does the SDE formulation fully encompass ODEs as a special case? And how well does the model perform on tasks governed by deterministic dynamics?\n\n4. Is there any restriction on the step size? What is the upper bound, if any? In general, the allowable step size is often related to the stiffness of the equation. For this type of SDE foundation model, how are different scales—both large and small—handled in practice?\n\n5. Building on the previous question, suppose an SDE from the training set has its drift term scaled by a large coefficient. Would the resulting SDE, potentially exhibiting significantly different dynamics, still be accurately learned by the model?\n\n6. How does the proposed method compare to neural network approaches that learn a specific SDE when evaluated on a given single test case? What is the magnitude of the performance difference？", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 1, "originality": 4, "strengths_and_weaknesses": "Strengths:\n\n1. The proposed novel FIM-SDE model leverages in-context learning, enabling it to infer drift and diffusion functions directly from raw data without retraining for each new system. This approach significantly enhances flexibility and efficiency in SDE discovery.\n\n2. FIM-SDE demonstrates robustness to noisy, irregular, and sparsely sampled data, which are common in real-world scenarios. \n\n3. The model can perform inference rapidly in a single forward pass. \n\nWeakness:\nAs with many neural approaches, formal guarantees on approximation accuracy and convergence are limited. That said, the work still lacks a clear explanation of why only a few demonstrations suffice to train the SDE foundation model.", "was_revised": true, "final_justification": "The detailed response, as well as the additional experiments provided, has addressed my concerns.", "rebuttal": "We thank the reviewer for their detailed review and insightful comments, which have helped us improve the presentation of our manuscript. Below, we address each of the points raised.\n\n**@W1:** To answer this point precisely we would appreciate a brief clarification: does the reviewer mean by demonstrations \n\n(i) the number of context tuples per SDE instance, or \n\n(ii) the number of distinct SDEs used during pre‑training?\n\n**@Q1:** Both Gaussian‑process regression and symbolic regression recover drift and diffusion by exploiting Eqs. 2 and 3: \n\n- the expectation of $\\\\Delta y$ is proportional to the local drift,\n- the expectation of $\\\\Delta y^2$ is proportional to the local diffusion.\n\nWe therefore explicitly provide both $\\\\Delta y$ and $\\\\Delta y^2$ as inputs to FIM‑SDE, allowing the network to learn separate feature channels for drift and diffusion estimation.\n\n**@Q2**: As noted in lines 189-193, we *assume* the sequential nature of the time series does not provide additional information for the *local* estimation of drift and diffusion, beyond that encoded into the one-step transitions ($\\\\Delta y$, $\\\\Delta y^2$). Consequently, FIM‑SDE treats each of the $K(L-1)$ tuples as iid variables (i.e. no positional encodings are used). \n\n*Number of tuples required by FIM-SDE:* During pre‑training, we randomly subsample between $128$ and $12\\\\ 800$ tuples per SDE, so the model is exposed to widely varying context sizes. In our experiments, the $12$ target data sets range from $5\\\\ 000$ tuples (canonical systems) to $20\\\\ 000$ tuples (stock data). FIM-SDE nicely generalizes to every case.\n\n**[PROPOSED MODIFICATION:]** We thank the reviewer for making us realize that this information was missing from the Appendix. We will include an appendix table that summarises context‑size statistics for all training and target sets.\n\nFinally, we have additionally revisited the performance of FIM-SDE on the double well (Eq. 24) and damped linear (E. 25-26) problems when varying the context size (in tuples). We considered the clean and dense observation case ($\\\\rho=0, \\\\Delta \\\\tau = 0.002$). The table below shows that FIM-SDE can correctly infer the damped linear oscillator with only $2\\\\ 000$ tuples. The double well problem, albeit 1D, is more complex because it requires characterising a state-dependent diffusion. FIM-SDE requires at least $4\\\\ 000$ tuples to correctly capture this structure. \n\n| Context Size | Double Well (MMD)   | Damped Linear (MMD) |\n|-------------:|:-------------------:|:-------------------:|\n| $500     $   | $0.155 \\\\pm 0.090 $  | $ 0.128 \\\\pm 0.068 $ |\n| $1\\\\ 000  $   | $0.115 \\\\pm 0.094 $  | $ 0.014 \\\\pm 0.005 $ |\n| $2\\\\ 000  $   | $0.053 \\\\pm 0.073 $  | $ 0.006 \\\\pm 0.005 $ |\n| $3\\\\ 000  $   | $0.042 \\\\pm 0.086 $  | $ 0.002 \\\\pm 0.003 $ |\n| $4\\\\ 000  $   | $0.004 \\\\pm 0.003 $  | $ 0.000 \\\\pm 0.000 $ |\n| $5\\\\ 000  $   | $0.004 \\\\pm 0.003 $  | $ 0.000 \\\\pm 0.000 $ |\n| $50\\\\ 000 $   | $0.006 \\\\pm 0.006 $  | $ 0.000 \\\\pm 0.000 $ |\n\n\n\nThese results showcase the FIM-SDE’s ability to handle a wide range of context sizes, and will also be included into the Appendix.\n\n**@Q3.a (Is there a specific reason or limitation behind this choice?)**: For this first study we opted for diagonal diffusion to limit the number of functions to infer. Nothing in the architecture prevents non‑diagonal diffusion, and extending pre‑training to that setting is straightforward. We leave this extension for future work. Let us remark, nevertheless, that diagonal and state-dependent diffusion matrices represent a rich family of SDEs with wide application in the natural sciences. What is more, this assumption is consistent with all baseline models.\n\n**@Q3.b (are there any other constraints on the diffusion term?)**: No additional constraints are imposed beyond the polynomial diffusion of degree $\\\\leq 2$ prior used in pre-training (Section 3.1, lines 157-163).\n\n**@Q3.c (does the SDE formulation fully encompass ODEs as a special case?)**: Yes. The training corpus includes SDEs with **constant or zero diffusion**, so FIM‑SDE has seen purely deterministic trajectories. Nonetheless, because most training cases have non‑zero diffusion, the model is biased towards inferring non‑constant diffusion. The results we report in **@Q6** below illustrate this bias. They also demonstrate that the bias can be corrected via fine-tuning.\n\n**@Q4**: Each target data set has its own spatial and temporal scales, so all inputs are normalised as described in Section 3.2 (lines 187–188) and Appendix E.1. During data generation we simulated our sampled SDEs on a fine grid $\\\\Delta t = 10^{-3}$, but recorded observations at three spacings: $\\\\Delta \\\\tau \\in \\\\{10^{-1}, 10^{-4}, 10^{-3}\\\\}$. Pretraining FIM-SDE on such paths teaches it to handle both sparse and dense observations. \n\nNotably, in our target datasets $\\\\Delta \\\\tau$ ranges from $0.2$ (wind speeds) to $2 \\\\times 10^{-5}$ (stock data) and FIM-SDE is able to handle all systems across that span.\n\n**[PROPOSED MODIFICATION:]** We again thank the reviewer for pointing out that this information was missing. We will include an appendix table that reports the inter-observation time scales for training and target systems.\n\n**@Q5**:  Scaling the drift can be viewed (e.g. through an Itô change of variables) as scaling the diffusion downwards. As explained in our reply **@Q3.c**, the pre-training distribution already contains vanishing‑diffusion cases, so the model remains capable of fitting such systems. Fine‑tuning on a handful of trajectories further removes any residual bias toward non‑zero diffusion (see **@Q6** below).\n\n**@Q6**: Comparison with neural network approaches.\n\n*Baseline.* LatentSDE [1] is a widely adopted neural model for SDE inference, so we use it as the task‑specific baseline.\n\n*Experimental setup.* We follow [1] exactly and consider the stochastic Lorenz attractor with parameters $(\\\\sigma, \\\\rho, \\\\beta)=(10, 28, 8/3)$ and a diagonal diffusion matrix with scale $0.15$. LatentSDE is trained on $1\\\\ 024$ paths $\\\\times\\\\ 41$ time points, sampled from the initial condition $\\\\mathcal{N}(0, 1)$. FIM‑SDE sees the same data as context. We evaluate it (i) zero‑shot and (ii) after fine‑tuning on those paths.\n\n*Evaluation setup.* For each of three initial‑condition sets: $\\\\{\\\\mathcal{N}(0, 1), \\\\mathcal{N}(0, 2), (1, 1, 1)\\\\}$, we generate ground‑truth trajectories and compute the Maximum Mean Discrepancy (MMD) to samples from each model.\n\n*Results*: \n\n| Model         | $\\\\mathcal{N}(0,1)$ | $(1, 1, 1)$ | $\\\\mathcal{N}(0, 2)$ |\n|:--------------|:------------------:|:-----------:|:-------------------:|\n| LatentSDE     |  $0.051708$        |  $0.135003$ |  $0.054335$         |\n| FIM-SDE (zero-shot) |  $0.403933$        |  $1.151080$ |  $0.378515$         |\n| FIM-SDE (fine-tuned) |  $0.003838$        |  $0.037159$ |  $0.003335$         |\n\nLatentSDE clearly outperforms FIM‑SDE when the latter is evaluated zero‑shot. After fine‑tuning for a few hundred epochs, however, FIM‑SDE surpasses LatentSDE on all three tests.\n\nLet us remark that FIM‑SDE is only fine-tuned on the training set of LatentSDE — which was sampled from the initial condition $\\\\mathcal{N}(0, 1)$. Paths sampled from the other two initial conditions are only used during model evaluations.\n\n**Why is the zero‑shot performance weaker?** \n\nThe data are drift‑dominated. We invite the reviewer to quickly check Fig 6 in the LatentSDE paper [1]. Because our pre‑training set contains mostly non‑zero diffusion functions, FIM‑SDE is biased toward inferring finite diffusion, which increases the variance of its generated paths that lead to larger MMD scores. We will add a diagnostic plot of the zero-shot inferred diffusion illustrating this bias to the Appendix.\nFine‑tuning on the task data quickly removes the bias and restores accuracy.\n\n**Convergence speed** \n\nA recent benchmark, SDEmatching [2], which studied the same stochastic Lorentz attractor setup,\nshows that LatentSDE typically converges only after ${\\\\approx}\\\\ 5\\\\ 000$ iterations (we have verified this), whereas SDEmatching needs ${\\\\approx}\\\\ 500$.\n\nNext, we compare those convergence speeds against (i) fine‑tuning FIM‑SDE and (ii) training FIM‑SDE from scratch:\n\n|    FIM Training Mode   | It.  50     | It.  100    | It.  200    | It.  500    | It.  1 000 | It.  2 000 | It.  5 000 |\n|:-------------------|:------------:|:------------:|:------------:|:------------:|:------------:|:------------:|:------------:|\n|      Fine-tune      | $0.012888$   | $0.003838$   | $0.003247$   | $0.000000$   | $0.0000000$  | $0.0000000$  | $0.0000000$  |\n| Train from scratch | $0.749462$   | $0.491508$   | $0.397677$   | $0.203864$   | $0.0982177$  | $0.0154879$  | $0.0125062$  |\n\n\n*Key take‑aways*\n- Fine‑tuning FIM‑SDE reaches near‑optimal performance within ${\\\\leq}\\\\ 100$ epochs — much faster than both LatentSDE and SDEmatching.\n- Pre‑training provides an excellent initialisation; training from scratch requires two orders of magnitude more iterations to match the same error.\n\nWe also invite the reviewer to check our response **@Q3** for reviewer Kq3M, where we present a comparison between FIM-SDE and LatentSDE on the real-world datasets.\n\n**[PROPOSED MODIFICATION:]** All these results will be inserted into the manuscript.\n\n**NOTE ON TABLE 2B (paper) vs TABLE ABOVE**\n\nThe results in the first table above differ from those in Table 2b of the main text, as the latter were computed using an incorrect normalization factor. However, the corrected results presented in this rebuttal can be easily reproduced using the pretrained FIM-SDE model provided with our submission.\n\n*References*:\n\n[1] X. Li, et al. Scalable Gradients for Stochastic Differential Equations. AISTATS 2020.\n\n[2] Bartosh et al. SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations. ICML 2025"}, {"review_id": "LR2RI2QUYg", "reviewer": "Reviewer_np9z", "summary": "In this paper, the authors introduce FIM-SDE for estimating the underlying dynamics of SDEs from noisy time series data. This is a pretrained recognition model that performs in-context (zero-shot) inference of both drift and diffusion functions without requiring fine-tuning or prior knowledge of the system. A notable strength of the method is it can learn from many simulated examples and works across a wide range of systems, from physical models to real-world data like stock prices and wind speeds. It's a flexible and general tool for understanding complex, noisy dynamics.", "questions": "- Regarding the data generation process described in Lines 152–153, could the authors clarify whether the use of randomly sampled dynamics truly reflects real-world scenarios? From my perspective, many sampled dynamics may lack physical interpretability or relevance.\n- Regarding the proposed FIM-SDE, I am curious about the validity of treating x as the spatial location. This implicitly reformulates the SDE as a PDE, but x here is not location of fixed grid as required by SDEs.", "limitations": "- The related work on Neural SDEs appears to be outdated and insufficient. The authors should incorporate discussions of more recent work in Neural SDEs and clarify how FIM-SDE differs from these approaches.\n- The test data are limited to quite low-dimensional systems (at most 3), leaving the scalability of the method on higher-dimensional dynamics underexplored.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "- This work investigates the problem of identifying latent SDEs from noisy observations. This is a very appealing and practical challenge. The contributions in both dataset construction and model design are noteworthy and appreciated.\n- The proposed FIM-SDE model introduces a somewhat novel approach to SDE inference by leveraging amortized learning and neural operator principles.\n- The experiments are sufficient and convincing, verifying the superiority of the proposed method.", "was_revised": true, "final_justification": "All of my questions have been addressed by the authors. I think this work is solid and well-executed, and thus I recommend accept.", "rebuttal": "We thank the reviewer for their encouraging feedback, as well as for highlighting potential limitations and raising thoughtful questions. Below, we address each of the points.\n\n**@Q1**: One of our core assumptions is that the class of SDEs with polynomial drift, of degree at most three, and polynomial diffusion, of degree at most two, already cover many models used in the natural and social sciences. All $8$ canonical systems in Section 4.1 fall inside this class. \n\nWhether such randomly sampled SDEs reflect real-world scenarios can be judged from our real‑world evaluations. Below we reproduce Table 2.a and add\n\n(i) LatentSDE [1] (a strong neural network baseline), and \n\n(ii) FIM-SDE after fine-tuning on each dataset\n\n| Model           | Wind               | Oil                | Facebook           | Tesla                |\n|:----------------|:------------------:|:------------------:|:------------------:|:--------------------:|\n| BISDE           | $0.000 \\\\pm 0.000$ | $0.137 \\\\pm 0.097$ | $0.053 \\\\pm 0.060$ | $0.013 \\\\pm 0.015$   |\n| LatentSDE       | $0.350 \\\\pm 0.212$ | $0.235 \\\\pm 0.167$ | $0.014 \\\\pm 0.018$ | $0.016 \\\\pm 0.023$   |\n| FIM-SDE (zero-shot) | $0.330 \\\\pm 0.114$ | $0.199 \\\\pm 0.137$ | $0.012 \\\\pm 0.026$ | $0.000 \\\\pm 0.000$   |\n| FIM-SDE (fine-tuned)   | $0.029 \\\\pm 0.034$ | $0.068 \\\\pm 0.042$ | $0.006 \\\\pm 0.007$ | $0.004 \\\\pm 0.008$   |\n\n*Key take‑aways*:\n\n- *In zero‑shot mode*, FIM‑SDE is already competitive with both symbolic (BISDE) and neural (LatentSDE) methods, suggesting that pre‑training on our randomly sampled SDEs captures useful structure.\n- The structure provides a good basis for inference via finetuning. A brief fine‑tune via Eq. (12) (App. D) lets FIM‑SDE outperform both baselines on all four real‑world systems.\n\n**[PROPOSED MODIFICATION:]** We will include the new results into Section 4.\n\n**@Q2**: We may be misunderstanding the concern, so we kindly ask for clarification. The variable $\\\\mathbf{x}$ in Eq. 1 denotes *the state of the system*. FIM-SDE implements a map from: \n\n(i) (the space of) noisy observations on sampled paths drawn from SDEs like Eq. 1; to \n\n(ii) (the space of) *state-dependent* drift and diffusion functions that govern those paths. \n\nNo PDE reformulation is involved.\n\n**@Limitation1**: We thank the reviewer for this observation. Appendix A provides a more detailed discussion of related work than the overview in the Introduction. There, we included six recent neural SDE models which, to the best of our knowledge, cover the most recent advances in neural SDE modeling.\n\n**[PROPOSED MODIFICATION:]** To make this discussion more visible we will\n\n1. create a dedicated Related Work section in the main paper, and\n\n2. move the discussion of the latest neural‑SDE papers from the appendix into that section, explicitly contrasting their objectives and assumptions with ours.\n\n**@Limitation2**: We acknowledge this limitation. As stated in both our abstract and introduction (line 79), the present work **targets low‑dimensional SDEs (1 to 3 dimensions) only**. Section 4 (lines 240–241) and the conclusion (lines 327–332) both reiterate this scope and outline the challenges of extending to higher‑D systems.\n\nIn the reply **@W2** to Reviewer Kq3M we explain one of the main reasons why we, in this first study, restrict our attention to low-dimensional SDEs. We plan to address this in future work.\n\n*References*:\n\n[1] X. Li, et al. Scalable Gradients for Stochastic Differential Equations. AISTATS 2020."}, {"review_id": "QnEPR0Qcu3", "reviewer": "Reviewer_P9Vv", "summary": "FIM-SDE introduces a pre-trained neural network capable of performing zero-shot estimations of drift and diffusion functions of SDEs from noisy, irregularly sampled time series data without requiring parameter fine-tuning for each new system. The approach leverages amortized inference and neural operator concepts by training a Transformer-based architecture on a large synthetic dataset of 600K low-dimensional SDEs with polynomial drift functions (up to degree 3) and diagonal diffusion matrices, where the model learns to map corrupted SDE observations to their underlying governing functions through a context matrix encoding and functional attention mechanism. FIM-SDE addresses key limitations of existing methods—heavy reliance on prior knowledge, sensitivity to noise, and computational challenges of neural SDE training—by providing a simulation-free solution that can generalize across diverse dynamical systems. Extensive evaluation of eight canonical SDE systems and four real-world datasets (including stock prices, oil prices, and wind speed fluctuations) demonstrates that FIM-SDE consistently outperforms traditional Gaussian process regression and symbolic methods, particularly under noisy conditions while offering the practical advantage of immediate deployment without system-specific optimization, positioning it as a foundation model for automated scientific discovery in stochastic dynamical systems.", "questions": "Can you clearly demonstrate some usecases for the paper? For instance, can you compare the training and regression computing tradeoffs between this approach and the on-the-go approaches?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. This idea of a GPT-alike SDE fitting model is novel and interesting. While the performance against existing methods is not always stellar, this is a good first attempt.\n2. The paper is well-written, very clear motivation for the paper.\n\nWeakness:\n1. Like mentioned above, the experimental results are not stellar always.\n2. The computational cost seems really high. The idea of such a model is that at inference time this is faster and cheaper than other on the go model, but it is not so clear.", "was_revised": true, "final_justification": "My concerns are addressed by the author's detailed explanations and new experimental results, and will remain my positive feedback.", "rebuttal": "We sincerely thank the reviewer for their time, for their comments and questions, and for their positive assessment of our submission. Below, we address each of the points raised.\n\n**@W1:** Let us begin by summarising the experiments we have prepared for the rebuttal into two tables:\n\n1. *Real‑world data (Table 2a, extended)*. We added (i) FIM‑SDE fine‑tuned on each data set and (ii) LatentSDE [1] as an additional, neural network baseline.\n\n| Model           | Wind               | Oil                | Facebook           | Tesla                |\n|:----------------|:------------------:|:------------------:|:------------------:|:--------------------:|\n| BISDE           | $0.000 \\\\pm 0.000$ | $0.137 \\\\pm 0.097$ | $0.053 \\\\pm 0.060$ | $0.013 \\\\pm 0.015$   |\n| LatentSDE       | $0.350 \\\\pm 0.212$ | $0.235 \\\\pm 0.167$ | $0.014 \\\\pm 0.018$ | $0.016 \\\\pm 0.023$   |\n| FIM-SDE (zero-shot) | $0.330 \\\\pm 0.114$ | $0.199 \\\\pm 0.137$ | $0.012 \\\\pm 0.026$ | $0.000 \\\\pm 0.000$   |\n| FIM-SDE (fine-tuned)   | $0.029 \\\\pm 0.034$ | $0.068 \\\\pm 0.042$ | $0.006 \\\\pm 0.007$ | $0.004 \\\\pm 0.008$   |\n\n\n2. *Stochastic Lorenz attractor (Table 2b, corrected)*. The original numbers were unfortunately computed using an incorrect normalization factor. The corrected scores below can be easily reproduced using the pretrained FIM-SDE model provided with our submission.\n\n| Model         | $\\\\mathcal{N}(0,1)$ | $(1, 1, 1)$ | $\\\\mathcal{N}(0, 2)$ |\n|:--------------|:--------------------|:------------|:---------------------|\n| LatentSDE     |  $0.051708$         |  $0.135003$ |  $0.054335$          |\n| FIM-SDE (zero-shot) |  $0.403933$         |  $1.151080$ |  $0.378515$          |\n| FIM-SDE (fine-tuned) |  $0.003838$         |  $0.037159$ |  $0.003335$          |\n\nKey take‑aways:\n\n- The **zero-shot** performance of FIM-SDE on real-world data is not only competitive with symbolic regression methods (BISDE)  but also with neural models (LatentSDE).\n- A short fine‑tune (${\\\\leq}\\\\ 100$ iterations) improves FIM‑SDE dramatically, surpassing LatentSDE, which needs ${\\\\approx}\\\\ 5\\\\ 000$ iterations to reach its reported scores (see **@W2** & **@Q1** below).\n-  The stochastic Lorenz attractor highlights a bias in our pre‑training distribution, as discussed in **@Q6** to reviewer WGed. Nevertheless, the same quick fine‑tune eliminates the gap and again outperforms LatentSDE.\n\nThese results demonstrate that FIM-SDE provides a competitive zero-shot solution to a wide variety of dynamical systems. They also demonstrate that FIM-SDE can easily be fine-tuned to synthetic and real-world datasets and that the result outperforms neural SDE solutions. \n\n**@W2** & **@Q1**: Since the second weakness and the main question raised are related, we address them together in what follows. \n\nFIM‑SDE is intended as an *off‑the‑shelf* tool for scientists who have noisy trajectories but little ML expertise. Like all foundation models, **pre‑training is a one‑time, centralised expense**. End‑users download the checkpoint and either (i) use it zero‑shot, or (ii) fine‑tune briefly on their data. The figures below therefore *exclude* pre‑training cost and focus on the compute a user would spend per new data set.\n\n*Use-case Scenario*: You receive the four real‑world data sets of Section 4.2 and have the following methods available\n\n- BISDE [3]: symbolic regression requiring a hand‑crafted library;\n- LatentSDE [1]: neural model trained from scratch per data set;\n- FIM‑SDE: pre‑trained checkpoint, optionally fine‑tuned.\n\nThe task is to infer the SDE that best characterizes the datasets.\n\n*Trade-offs*:\n\n| Method               | Prior knowledge required            | Training iterations | MMD quality |\n|----------------------|-------------------------------------|---------------------|-------------|\n| BISDE                | Hand‑crafted basis for each system  | $0$ (effectively)   | good        |\n| LatentSDE            | none                                | $5\\\\ 000$ (typical) | good        |\n| FIM-SDE (zero‑shot)  | none                                | $0$                 | good        |\n| FIM-SDE (fine‑tuned) | none                                | $\\\\mathbf{100}$     | **best**    |\n\nBy not relying on prior knowledge and already achieving strong zero-shot performance, FIM‑SDE is the most suitable choice when compute time is limited or unavailable. When fine-tuning is permitted, FIM‑SDE further improves and consistently outperforms both baselines across all systems, requiring only $100$ epochs to reach state-of-the-art quality.\n\n*Details on convergence:  LatentSDE vs FIM-SDE fine-tuning*: The table below compares the performance trajectories of LatentSDE and FIM‑SDE over the course of training. LatentSDE requires thousands of iterations to converge, while FIM‑SDE reaches comparable or superior performance after only $100$ fine-tuning epochs—representing a $\\\\, > 50 \\times$ reduction in compute. This highlights the efficiency of leveraging a pre-trained FIM‑SDE model.\n\n\n\n| Model                            | Wind               | Oil                | Facebook           | Tesla                |\n|:---------------------------------|:------------------:|:------------------:|:------------------:|:--------------------:|\n| LatentSDE (Epoch 500)            | $1.019 \\\\pm 0.078$ | $0.382 \\\\pm 0.236$ | $0.091 \\\\pm 0.043$ | $0.066 \\\\pm 0.059$   |\n| LatentSDE (Epoch 1000)           | $0.874 \\\\pm 0.127$ | $0.224 \\\\pm 0.185$ | $0.071 \\\\pm 0.037$ | $0.052 \\\\pm 0.056$   |\n| LatentSDE (Epoch 2000)           | $0.318 \\\\pm 0.164$ | $0.205 \\\\pm 0.172$ | $0.034 \\\\pm 0.035$ | $0.023 \\\\pm 0.028$   |\n| LatentSDE (Epoch 5000)           | $0.350 \\\\pm 0.212$ | $0.235 \\\\pm 0.167$ | $0.014 \\\\pm 0.018$ | $0.016 \\\\pm 0.023$   |\n|                                  |                    |                    |                    |                      |\n| FIM-SDE (zero-shot)              | $0.330 \\\\pm 0.11$4 | $0.199 \\\\pm 0.137$ | $0.012 \\\\pm 0.026$ | $0.000 \\\\pm 0.000$   |\n| FIM-SDE (fine-tuned) (Epoch 20)  | $0.114 \\\\pm 0.06$9 | $0.071 \\\\pm 0.094$ | $0.020 \\\\pm 0.019$ | $0.021 \\\\pm 0.014$   |\n| FIM-SDE (fine-tuned) (Epoch 40)  | $0.075 \\\\pm 0.05$2 | $0.048 \\\\pm 0.057$ | $0.019 \\\\pm 0.026$ | $0.020 \\\\pm 0.028$   |\n| FIM-SDE (fine-tuned) (Epoch 60)  | $0.042 \\\\pm 0.04$4 | $0.049 \\\\pm 0.034$ | $0.003 \\\\pm 0.007$ | $0.013 \\\\pm 0.018$   |\n| FIM-SDE (fine-tuned) (Epoch 80)  | $0.043 \\\\pm 0.07$7 | $0.046 \\\\pm 0.028$ | $0.002 \\\\pm 0.005$ | $0.011 \\\\pm 0.016$   |\n| FIM-SDE (fine-tuned) (Epoch 100) | $0.029 \\\\pm 0.03$4 | $0.068 \\\\pm 0.042$ | $0.006 \\\\pm 0.007$ | $0.004 \\\\pm 0.008$   |\n\n\nThese results illustrate a clear use-case for FIM‑SDE: when facing a new real‑world system, one can immediately deploy the zero‑shot model without any training, obtaining results competitive with fully trained baselines. If better accuracy is needed and minimal compute is available, fine-tuning FIM‑SDE for just $100$ iterations suffices to outperform both symbolic and neural methods trained from scratch. This flexibility—in balancing training cost and regression quality—sets FIM‑SDE apart from on-the-go approaches, and makes it a practical foundation model for scientific time series.\n\n**[PROPOSED MODIFICATION:]** We will add all this new data to the Appendix.\n\n*References:*\n\n[1] X. Li, et al. Scalable Gradients for Stochastic Differential Equations. AISTATS 2020.\n\n[2] Bartosh et al. SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations. ICML 2025\n\n[3] Wang et al. Data-driven discovery of stochastic differential equations. Engineering, 2022."}]

Author rebuttal:

Meta-review: This paper introduces FIM-SDE, a framework for estimating the dynamics of stochastic differential equations from noisy time series data. The authors successfully pre-train a model to map SDE realizations to their drift and diffusion coefficients, allowing for zero-shot generalization to unseen SDEs. They generate a dataset of 600k low-dimensional SDEs and demonstrate the tool's effectiveness compared to existing methods.

All reviewers note that the authors have effectively addressed previous concerns. The strengths of the paper include its novel approach to in-context learning, robustness to noisy data, and rapid inference capabilities. While some limitations remain, such as the focus on low-dimensional systems and the need for clearer explanations of certain equations, the overall contributions are significant. The reviewers unanimously recommend acceptance, acknowledging the paper's potential to advance the field of SDE discovery.
Decision: Accept (poster)