Zero-shot Imputation with Foundation Inference Models for Dynamical Systems

Authors: Patrick Seifner, Kostadin Cvejoski, Antonia Körner, Ramses J Sanchez
Primary area: learning on time series and dynamical systems
Keywords: Zero-shot imputation, foundation models, time series imputation, dynamical systems, amortized inference, zero-shot interpolation, foundation models for time series

Abstract:
Dynamical systems governed by ordinary differential equations (ODEs) serve as models for a vast number of natural and social phenomena. In this work, we offer a fresh perspective on the classical problem of imputing missing time series data, whose underlying dynamics are assumed to be determined by ODEs. Specifically, we revisit ideas from amortized inference and neural operators, and propose a novel supervised learning framework for *zero-shot time series imputation*, through parametric functions satisfying some (hidden) ODEs. Our proposal consists of two components. First, a broad probability distribution over the space of ODE solutions, observation times and noise mechanisms, with which we generate a large, synthetic dataset of (hidden) ODE solutions, along with their noisy and sparse observations. Second, a neural recognition model that is trained *offline*, to map the generated time series onto the spaces of initial conditions and time derivatives of the (hidden) ODE solutions, which we then integrate to impute the missing data. We empirically demonstrate that *one and the same* (pretrained) recognition model can perform zero-shot imputation across 63 distinct time series with missing values, each sampled from widely different dynamical systems. Likewise, we demonstrate that it can perform zero-shot imputation of missing high-dimensional data in 10 vastly different settings, spanning human motion, air quality, traffic and electricity studies, as well as Navier-Stokes simulations — *without requiring any fine-tuning*. What is more, our proposal often outperforms state-of-the-art methods, which are trained on the target datasets.

Our pretrained model, repository and tutorials are available online.

Decision: accept (poster)
Reviewer scores: 5, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730584660413

thread_type:
official_review

<thread_object>
creator: reviewer_6fr3

<content>
summary:
This paper addresses the challenges of imputing missing time-series data. A framework for zero-shot time-series imputation is proposed with 1) a synthetic data generation model for sampling a set of ODE solutions and 2) a neural recognition model mapping the time series data onto parametric functions. Empirical results present improved performance of imputation compared to baseline models.

strengths:
1. The problem of imputing time-series data is important and the idea of linking amortized inference and neural operator is interesting.
2. The methodology of the foundation inference model is well-structured.
3. Experiments were performed on the imputation of point-wise and temporal missing patterns to present the effectiveness of the proposed method.

weaknesses:
1. I think the idea of zero-shot amortized inference framework is somewhat incremental. There have been works on amortized inference for few-shot time-series forecasting (e.g. [1] [2] [3][4]) in terms of lacking sufficient observation and learning single dynamics, where the idea of amortized inference learning the prior knowledge of dynamics is similar to the proposed method. Could the authors add a discussion about these works and the benefits of the proposed method? Also, it would be good if the authors could make a comparison of some of these methods in experiments.
2. In Equation 3, 4, and 5, do $\phi^{\theta}_i$ and $\psi^{\theta}$ share the same parameters $\theta$? Similar problem to the parameter $\psi$ in FIM model in Equation 7 and 8.

questions:
Please find the questions in the Weaknesses section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675762603

thread_type:
official_review

<thread_object>
creator: reviewer_8y1g

<content>
summary:
The paper develops a ODE-based foundation model for time series imputation. The idea is to simulate a large collection of incomplete time series including both point-wise and temporal missing patterns. A recognition model is trained to map such time series to the initial condition, and time derivatives at each time step, with both mean and log-variance prediction. The design of the recognition model is inspired by deep ONet. There is no need to fine tune the model on application-specific data. The experiments show the improvement over many sota methods specifically trained on application data.

strengths:
1. very interesting work on time series imputation, potentially pointing out a new direction
2. design of neural recognition model is novel and interesting, albeit inspired by deepOnet. Such "borrowing" is still interesting. 
3. zero-shot performance is surprisingly good

weaknesses:
1. as admitted by the authors, one limitation is that when the actual application does not well match the assumption of the simulation data, the performance can deteriorate. Though the paper focuses on zero-shot prediction, it will be great to discuss  possible methods for fine tuning or adaptation to specific applications. Since your training assumes the ground-truth of the initial and time derivatives of each sampled trajectories are known, and the neural recognition models directly learn to fit them, such training cannot directly apply to real dataset where these ground-truth is unknown. 
2.  The design of the neural recognition model seems not well justified, especially for the branch-net component. why do you use RNN? Why not using attention mechanism instead? Given RNN/LSTM's are replaced by attention  nearly every where, the authors should explain their rationale of such choices, either theoretically or empirically. 
3. How are the hyperparameters selected? Why using 1024 dimensional embeddings? Are this done by a rigorous validation process? If so, how?  It will be good for the authors to specify the details of hyperparmeter selection and validation process.

questions:
see above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730684726800

thread_type:
official_review

<thread_object>
creator: reviewer_nbhn

<content>
summary:
This paper introduces foundation model for time-series imputation, using several parametric functions of ODEs for training the model. The authors generate many synthetic noisy, irregular observations to train recognition model in offline manner.  They use combination of objective function with respect to initial value and data points to train this model. This trained model can conduct zero-shot imputation on unseen time series dataset with different dimensions. They validate their method on various different benchmarks in several domains.

strengths:
- This paper is well written and organized. I enjoy reading this paper.
- To the best of my knowledge, this paper is novel in that this firstly introduces the zero-shot imputation of time-series with dimension free inference. 
- I think the training objective and the way they give supervision to the model is also novel.
- The experimental results are impressive, covering a wide range of datasets and scenarios enough to prove the zero-shot capability of this method. Additionally, the paper faithfully includes details for reproducing the experiments which is helpful

weaknesses:
- Although this model impute the data in zero-shot manner, experimental results are not that impressive in Table 2. Also, the authors use reported value from the other papers which can question whether comparison is really fair.
-  I think it would be better to include ODE based methods for the baseline considering concept of this paper.

questions:
- I think some ablation studies for measuring the generalization capability of this model will be beneficial.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707250653

thread_type:
official_review

<thread_object>
creator: reviewer_6udz

<content>
summary:
This paper proposes a supervised learning framework for zero-shot time series imputation in dynamical systems of any dimensionality. The framework uses a synthetic data generation model based on two key assumptions: 
(1) time series with point-wise missing data have simple ODE-based interpolation solutions, and 
(2) time series with temporal missing patterns are locally simple. 
The authors introduce two neural interpolation models, FIM and FIM-l, designed for point-wise and temporal missing patterns, respectively.
The approach is evaulated over 8 datasets

strengths:
— The proposed FIM has great zero-shot imputation performance and achieves SOTA results on multiple datasets.
— Unlike prior work, this paper proposes a zero-shot approach that can be used on processes of any dimensionality.
— The synthetic data generation model offers a general method for creating meaningful synthetic time series, which could support future research.
- The authors are honest about the limitation of the approach, especially those related to the synthetic distributions used in the analysis and setting (ODEs) as opposed to real-world data which can exhibit any (or no ) distrubution.

weaknesses:
- Fig 2 is difficult to read, the colors are too  similar
- The paper is too "buzzwordy", I would have  preferred more of a proper technical discussion
- The results of the proposed method are not always close to SOTA or better than SOTA

questions:
— The FIM-l model only works for time series that follow a ‘simple’ distribution. Since existing imputation models already perform well on real-world and synthetic data, the practical value of the proposed model might be limited.
— The authors did not explain clearly how the FIM could handle processes of any dimensionality, and how the performance of the model varies with the dimensionality in practice.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732618458756

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply to reviewers

comment:
We would like to thank all reviewers for their valuable comments and questions. We have carefully read and addressed all of them, for each reviewer separately, in the discussions below. 

We are happy to provide any additional information the reviewers might need during the discussion period, in order to make an informed decision. For now, let us summarise the main updates we propose to do to the manuscript, which we believe address all questions and concerns of the reviewers.

**Summary of proposed modifications**

1. To make the feature of our methodology that allows us to process time series of any length and large complexity more visible to readers, we will modify the title *“Processing data of any length and dimensionality with FIM-l”* of the subsection in lines 295-310 into *“Beyond “simple” functions: Processing data of any length and dimensionality with FIM-l”.* This modification addresses **Question 1 of reviewer 6Udz**.

2. To address **weakness 2 from reviewer nBHN** we will specify, near lines 373-377, which of our baselines are ODE-based and to which subsection (i.e. 4.1, 4.2 or 4.3) they are pertinent.

3. We will create a new appendix titled *“On fine-tuning FIM”*, which will include and expand upon our answer to **weakness 1 from reviewer 8Y1g**. Our answer demonstrates that fine-tuning FIM is plausible, interesting and of practical value. It also puts forward the way we plan to pursue this idea in future work.

4. To address **weakness 2 and 3 from reviewer 8Y1g**, we will include the missing architecture details pertinent to our ablation studies of Appendix B.5. These details justify the specifics of the model architecture of our final model.

5. We will create a new appendix titled *“Related Work on Meta-learning”*, which will contain and expand upon our answer to **weakness 1 proposed by reviewer 6Fr3**. It will also make reference to other representative references on meta-learning and time series modeling. Our answer explains the key differences between our zero-shot approach and recent meta-learning proposals for dynamical systems, thereby elucidating the novelty of our methodology.

6. We will change the sentence *"Let us also denote the trainable network parameters with $\theta$"* in line 259 into *“Let us also denote the set of all trainable parameters in these networks with $\theta$”* to address **weakness 2 from reviewer 6Fr3**.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper addresses the problem of imputing missing values in time series whose dynamics are assumed to follow unknown ODEs. The objective is zero-shot imputation, i.e. performing imputation without re-training from a series of new observations. The authors propose a “foundation models” that is pretrained on a large number of synthetic series to generate ODE features. The methodology proceeds in two steps. First generate a large set of noisy ODE solutions corresponding to sparse and noisy time series, second learn to map the time series data to parametric functions consisting of the initial condition and the evolution function of the ODE that generated the data. The later is the model that best interpolates the data. At inference, using this interpolation property, the pretrained model  is used for imputation in two settings: pointwise and temporal imputation. The proposed approach is evaluated on diverse time series datasets.

The reviewers acknowledge the originality of the approach targeting zero shot imputation and of the methodology which implements a novel method for amortized inference. They all stress that the proposed method is limited to relatively simple time series defined by ODEs. However since these limitations are clearly outlined in the paper and considering the novelty of the approach that opens new perspectives and the SOTA performance in the considered evaluation setting, I propose an acceptation.