Title: In-Context Learning of Stochastic Differential Equations with Foundation Inference Models
Venue: NeurIPS 2025
Accepted: yes
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
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFIM-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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\nCan 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS AND WEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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？\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n*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)\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}