Zero-shot Imputation with Foundation Inference Models for Dynamical Systems

We would like to thank the reviewer for the helpful comments and questions, which we know will improve the presentation of our work. Below we address each of them. **@W1**: Could the reviewer please specify the references they refer to? As soon as we know what references the reviewer means, we will address all the proposed questions. **@W2**: We use $\theta$ to denote the complete set of trainable parameters inside FIM-l (see line 259). Specifically, the set of MLPs $\phi_1$, …, $\phi_5$ and the sequence processing network $\psi_1$ all have different trainable parameters. We denote the set of all these parameters with $\theta$. Similarly, we use $\varphi$ to denote the complete set of trainable parameters inside FIM **that do not belong to** FIM-l (see lines 345, 346). We make the distinction between the trainable parameters in these two models because FIM leverages the *pretrained* components of FIM-l. Please note that this distinction is necessary, because the *pretrained* parameters of FIM-l (i.e. $\theta$) are **not** updated during the (pre)training of FIM (see lines 352, 353). **Proposed modification:** 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$*. Does this modification address the weakness pointed out by the reviewer?

Thank you for pointing out the missing reference and the clarification of the notation. Please find the detailed work below. Given the limited time, the authors could focus on the discussion of the few-shot learning method vs your proposed method. If time allows, the authors could include some comparison experiments. [1] Hewitt, Luke B., et al. "The variational homoencoder: Learning to learn high capacity generative models from few examples." arXiv preprint arXiv:1807.08919 (2018). [2] Singh, Gautam, et al. "Sequential neural processes." Advances in Neural Information Processing Systems 32 (2019). [3] Jiang, Xiajun, et al. "Sequential latent variable models for few-shot high-dimensional time-series forecasting." The Eleventh International Conference on Learning Representations. 2023. [4] Wang, Rui, Robin Walters, and Rose Yu. "Meta-learning dynamics forecasting using task inference." Advances in Neural Information Processing Systems 35 (2022): 21640-21653.

We thank the reviewer for engaging with us in the discussion. We hope the reviewer will now allow us to further comment on their reply, for we think there *might be* a fundamental point that we have not yet succeeded in conveying to the reviewer, namely that *we use one and the same pre-trained model to impute the data* **in all experiments**, *without any parameter finetuning*. We have mentioned this throughout the paper and, in fact, the weights of the pre-trained model are attached to our submission (we will also make them accessible on Hugging Face). Yet the last sentence of the reviewer’s comment seems to indicate that this feature was overlooked. *We apologize if we are misreading this comment*. If not, let us continue: **@** *I would encourage the authors to conduct experiments from Van der Pol oscillator to Air quality*: The prepositions “from” and “to” indicate transfer. Keeping the concept of transfer in mind — within the context of our work — one could say that our methodology transfers the knowledge encoded into our synthetic dataset (that is, our two assumptions in lines 165-178) from the pretrained weights onto the Van der Pol oscillator and Air quality imputation tasks. By this we mean that *after (pre)training on the synthetic dataset, we do not modify the (numerical values of the) parameters of the neural networks building FIM-l nor FIM*. We directly apply our pretrained models — without modification — **throughout all our experiments, in all experimental sections**. In particular, we use the same pre-trained FIM-l model when processing both the Van der Pol oscillator (Section 4.1) *and* the Air quality dataset (Section 4.2) , without modification. We hope this clarifies why we referred to our proposal as "not problem specific", as compared to the four references cited by the reviewer. Let us continue by adding that we never said meta-learning is problem specific. In fact, all our remarks and discussion focused only on the works referenced by the reviewer. Quoting *word-by-word* our remark above: *“They [The four works referenced by the reviewer] need to be trained on datasets from their target domains, which makes both their inferred representations and optimized weights problem specific.”* By this we meant that — keeping with the example of Jiang et al. (2023) – the parameters of the neural network building of the model of Jiang et al. (2023) are **not the same** for the bouncing ball experiment and the turbulent flow experiment. Similarly, the semantic content of the shared representations inferred by the model of Jiang et al. (2023) is **not the same** for the bouncing ball experiment and the turbulent flow experiment. In the former case it corresponds to gravity, whereas in the latter it corresponds to buoyant force. *In contrast*, our method **maintains the same network parameters and representation semantics throughout all experiments**. The representations consistently correspond to the time derivative and initial conditions of the hidden interpolating function, regardless of the target dataset. All that being said, let us conclude by acknowledging that “problem specific” and “not problem specific” were perhaps not the best adjectives to make our points.

Let us start by thanking the reviewer for both the detailed review and the kind words about our work. Below we address each of the weaknesses, comments and questions. **@W1**: we thank the reviewer for this very interesting suggestion. Although outside the scope of the present work, we agree that fine-tuning FIM to specific applications would be of great interest and value. As pointed out by the reviewer, in practice one does not have access to any ground-truth initial condition, nor to any time derivative (of the hidden interpolating function). However, the model can still be optimised to **reconstruct** the available data. To illustrate the plausibility of this approach in a controlled setting, we have fine-tuned FIM-l to reconstruct the Lorentz system data we studied in Appendix E. The details of this dataset can be found in lines 1821-1829 of that Appendix. **Setting up the fine-tuning**: 1. FIM-l takes as input the noisy observations along one path and returns (its best guess of) the interpolation function, which can be evaluated *at any time*. We fine-tuned FIM-l in two settings, namely (a) to reconstruct the noisy input data, and (b) to reconstruct the ground-truth, clean values of the Lorentz system’s solution, at the input observation times. Clearly, case (a) is the realistic setting. We consider case (b) as a control case, because FIM-l was pretrained to reconstruct the clean path (for details, check out Appendix B.3 where we introduced the training objective). 2. As we described in the Subsection titled *Processing Data of Any Length and Dimensionality with FIM-l* in lines 295-310 of the main text, FIM-l can be used to process either sets of up to $L_{max}$ observations simultaneously, or subsets (or windows) with fewer observations, followed by a combination of the local window estimates. This trick is what allows us to interpolate time series of any length and handle interpolation functions which lie outside our distribution of “simple” functions. The **number of windows** into which we split the target dataset **is a hyperparameter** of the model, as it is clearly problem dependent. Indeed, Figure 8 of Appendix E demonstrates the effect of modifying it on the Lorentz system data. To fine-tune FIM-l, we therefore need to specify the number of windows hyperparameter. That is, we need to specify *the temporal scale at which we want to fine-tune the model*. In our simple experiment below we chose to fine-tune FIM-l on **four (4) windows**. 3. Other details: We fine-tuned FIM-l with a learning rate of 1e-6, which was the learning rate at which we stopped pretraining. We fine-tuned it for two epochs only, on all components of the Lorenz systems (i.e. channel independence strategy) **Results**: The Table below (which expands on Fig 8) contains our results. | Window count | Pre-trained | Fine-tuned on noisy obs | Fine-tuned on clean obs | |:--------------:|:--------------------:|:-------------------------:|:-------------------------:| | 1 | 10.32 $\pm$ 1.478 | 8.983 $\pm$ 0.932 | 9.001 $\pm$ 0.945 | | 4 | 7.136 $\pm$ 0.697 | 5.936 $\pm$ 0.432 | 5.947 $\pm$ 0.433 | | 8 | 2.860 $\pm$ 0.237 | 2.708 $\pm$ 0.176 | 2.704 $\pm$ 0.175 | | 12 | 1.956 $\pm$ 0.161 | 1.913 $\pm$ 0.144 | 1.911 $\pm$ 0.143 | | 16 | 1.652 $\pm$ 0.126 | 1.702 $\pm$ 0.127 | 1.703 $\pm$ 0.126 | | 20 | 1.526 $\pm$ 0.125 | 1.586 $\pm$ 0.112 | 1.588 $\pm$ 0.112 | | 24 | 1.504 $\pm$ 0.138 | 1.532 $\pm$ 0.108 | 1.534 $\pm$ 0.108 | | 28 | 1.516 $\pm$ 0.133 | 1.505 $\pm$ 0.098 | 1.506 $\pm$ 0.098 | | 32 | 1.561 $\pm$ 0.142 | 1.500 $\pm$ 0.118 | 1.501 $\pm$ 0.118 | - The first thing we notice is that the performance of FIM-l clearly improves at the time scale (i.e. at the window count) on which it was fine-tuned. Indeed, we fine-tuned it on 4 windows. Very importantly, the fine-tuning process does not negatively perturb the interpolation performance at different time scales (i.e. at fewer or larger window counts). Therefore, *there is no catastrophic forgetting*. This observation alone demonstrates the plausibility of fine-tuning with FIM-l. - The second thing we notice is that fine-tuning on the noisy and clean data had a similar effect. The continuous nature of the interpolation function appears to make the fine-tuning robust to noise in the target data. However, we have also observed that if one fine-tunes FIM-l on the noisy data for many more epochs, or simply with a larger learning rate, its performance at different time scales does deteriorate. This occurs because the model starts fitting the noise. Therefore, to properly address fine-tuning within our framework, we would need to define an *emission model* that handles the noise.

*(Continuation from above)*: One could picture an emission model, (pre-)trained on our synthetic dataset, which would take as input both the interpolation function $x(t)$ — represented on some sensor grid, like in DeepONets — and the noisy data $y_1, \dots, y_l$. The emission model would then output the probability of observing the noisy value $y_i$ *at any desire time $\tau_i$*, that is $p(y_i|x(\tau_i))$. Coupling FIM together with this emission model would allow to finetune FIM (or the emission model, or both) to any target dataset. We plan to pursue this general direction in the near future. **Proposed modification**: We will include the discussion above as a new Appendix titled: *On fine-tuning FIM* into the manuscript. **@W2**: We purposely tried to make the definition of FIM-l in Section 3.2.1 agnostic to the architecture of the sequence processing network $\psi_1$. That being said, in our ablation studies of Appendix B.5 (lines 1069-1112) we did experiment with the architecture and found that both BiLSTM and Transformer networks yielded similar performance (Table 5). There is no specific reason why we reported our results with BiLSTM instead of with Transformer in the main text. Finally, we noticed that we did not include the specifics for these two networks in our manuscript. These follow below: - Transformer: MLP: 4x256, N-of-layers: 4, QKV-dim: 256, N-of-heads 8, output_dim: 256, total-parameter-count: ~20M - BiLSTM: MLP: 4x1024, hidden-dim: 512 (i.e. 256 per each directions), total-parameter-count: ~20M We will include these details back into our manuscript. **@W3**: In Table 4 of Appendix B.5 (lines 1069-1112), we report our studies on the effect of changing the size of the training dataset and the parameter count of the BiLSTM-based FIM-l. As can be seen in this Table, going beyond the 20M parameter count did not improve the performance of the model. Unfortunately, here again we did not include the specifics of the architectures. These follow: - 2M: MLP: 4x256, hidden-dim: 256 - 20M: MLP: 4x1024, hidden-dim: 512 - 50M: MLP: 3x2048, hidden-dim: 1024 These results justify our hyperparameter selection. We hope the comments and arguments above answer the questions and weaknesses proposed by the reviewer.

We thank the reviewer for their encouraging feedback, as well as for the proposed questions and weaknesses. We’re very grateful for their recognition of our work. **@W1**: Regarding fairness, our personal experience suggests that one rarely gets baselines to perform as well as their authors do, simply because the authors of said baselines are intimately familiar with the nuances of their proposal. We thus think that by reporting scores from other papers, we are being as fair as possible to the baselines. That said, we are also sure we evaluate our models on *exactly the same setting as the authors do*, for we obtained all target datasets we used in Table 2, as preprocessed by the authors, from their public repositories. We thus hope their published datasets are indeed the same ones they used for computing their scores in their publications. If this is the case, we believe our evaluations are fair with respect to the baselines. Does the reviewer disagree? Or are we misunderstanding the reviewer’s comment? **@W2**: We did include ODE baselines. Indeed, every baseline we compare against in Section 4.1 is ODE-based. To wit: ODEFormer, LatentODE (details in Appendix E.4, lines 1810-1921) and NeuralODEProcesses (details in Appendix E.5, lines 1921 - 1985) are all ODE-based. In Section 4.2, we compared against BRITS, which leverages a linear ODE to model the values of the times series between observations. Finally, in Section 4.3, we compared against both LatentODE and npODE (details in Appendix G.1, lines 2226-2305), which are ODE-based. We however noticed we did not make explicit reference to the comparison against npODE in the main text. We will add this reference to the manuscript. **@Q1**: We did include ablation studies in Appendix B.5 (lines 1069-1112). Specifically, we studied changing: - the parameter count in the models; - the training dataset size; - the architecture of the sequence processing network $\psi_1$ (i.e. LSTM vs Transformer); and - the maximum and minimum number of context points seen by the model during training. While our experiments in Section 4 of the main text demonstrated that our models perform well on many different synthetic and real-world applications and therefore that they *generalise* beyond the (synthetic) training dataset, we are unsure about how to set up an ablation study measuring their generalisation capabilities. Could the reviewer please be more precise about their expectations of the setup of such a study?

Before we start, let us thank the reviewer for their review, as well as for the proposed questions and weaknesses. We attempt to answer them in what follows. **@W1**: We will modify the colour to make the difference between the ground-truth trajectory and our estimates more apparent. **@W3**: Let us simply comment here that all baseline’s scores in Table 2 correspond to the scores of SOTA models that are trained on the target datasets. In contrast, FIM-l is neither trained nor finetuned on the target datasets. It is not a priori obvious that a pre-trained model can outperform, or even perform on par, with models that are actually trained on the target datasets. Indeed, most foundation models for time series, which mostly focus on forecasting, often perform *at best* on par with models that are trained or fine-tuned on the target datasets. See e.g. [1, 2, 3]. Our findings of Table 2 are in agreement with this. **@Q1**: We thank the reviewer for this question, because it highlights a very important feature of our methodology, namely that it can perform zero-shot imputation of point-wise missing values in complex time series, *via their decomposition into simpler time series*. It seems however that we did not stress this feature enough. FIM-l is indeed trained on what we define as “simple” functions. However, in the subsection titled *“Processing data of any length and dimensionality with FIM-l”* of the main text, in lines 295-310, we explain how to deal with data whose interpolating function is **not** represented by our ensemble of “simple” functions. In short, our key idea is that the domain of any complex function can be decomposed into subdomains or regions on which the function is locally “simple”. We therefore first make use of FIM-l to process the data on each of these regions (or subdomains), and then combine the local estimates into a global one. We provided the specifics on how we implemented these operations in Appendix B.6.2 (lines 1138-1195). Note that this simple approach is what allowed us to impute point-wise missing data sampled from complex analytic functions, like those in Figure 2. It also allowed us to impute point-wise missing data in real-world scenarios, as shown on the left and centre panel of Figure 3, and as demonstrated in Table 2. We invite the reader to also check Figure 6 in the Appendix, which contains the reconstructed data for all 63 ODE solutions in ODEBench, many of which feature complex functions. In sum, all these cases feature data whose interpolating functions do not belong to our ensemble of “simple” functions, and yet we are able to impute them with our windowing approach of section B.6.2 and FIM-l. **Proposed modification:** In order to make this feature of our methodology more visible to readers, we will modify the title *“Processing data of any length and dimensionality with FIM-l”* to *“Beyond “simple” functions: Processing data of any length and dimensionality with FIM-l”*. Does the reviewer think this modification makes our point clearer? We hope the clarifications above have helped convey to the reviewer how we use FIM-l to process data that extends beyond our ensemble of “simple” functions. **We also hope they elucidate the practical value of FIM-l**. To wit, as demonstrated in all our experiments, the practical value of FIM-l lies in its ability to be used off the shelf across very different cases, without the need for any fine-tuning. It also works as the backbone of FIM which, as we demonstrated in 4.3, can perform zero-shot imputation of the significantly harder temporal missing patterns, and outperform models trained on the target datasets. **@Q2**: In lines 189-193 we first explain that our methodology takes a channel independent strategy. That is, it treats each component of every target dataset as independent, and therefore processes them independently with FIM-l. Later, in the aforementioned subsection titled *“Processing data of any length and dimensionality with FIM-l”* of lines 295-310, we write: we adopt a channel independent strategy and process each component of any target, D-dimensional process independently with FIM-l. There, we also refer the reader to Appendix B.6 (lines 1115-1195), where we further expatiate on this treatment. We believe these sentences explain how FIM handles processes of any dimensionality. If these sentences are insufficient, could the reviewer please specify which aspects remain unclear?

**On performance vs dimensionality**: The performance of FIM — or of any other method that takes a channel independent strategy — does not directly nor necessarily depend on the dimensionality of the target process. Rather, it depends on the complexity of the dynamics characteristic of each of the components in the target process, and whether this complexity is represented by the training distribution of the model in question. Indeed, a channel independent strategy simply does not rely on the correlations present among the components of the target process. Interestingly enough, it has recently been argued (see [4]) that models trained using this strategy outperform models which are trained to exploit correlations across the components. Regarding FIM-l specifically, let us kindly direct the attention of the reviewer to e.g. Table 11 of the Appendix, where we report the MAEs of FIM evaluated on ODEBench, separated by the dimensionality of the target process. There we see that it’s easier for FIM-l to fit the four dimensional processes than, say, the one dimensional ones. The three dimensional processes, which include many chaotic systems, are in contrast the hardest to fit. Figure 6 in the Appendix paints a similar picture. We hope these arguments answer the questions of the reviewer. **References:** 1. Liu et al. (KDD 2024): Generative Pretrained Hierarchical Transformer for Time Series Forecasting; 2. Darlow et al. (ICLR 2024): DAM: Towards a Foundation Model for Time Series Forecasting; 3. Das et al. (ICML 2024): A decoder-only foundation model for time-series forecasting. 4. Han et al. (2024): The Capacity and Robustness Trade-off: Revisiting the Channel Independent Strategy for Multivariate Time Series Forecasting

The discussion period is about to end. We therefore would like to thank once again the reviewers for their valuable input and summarize the output of the discussions. As we described in our first general reply to all reviewers, many of the questions and concerns raised by the reviewers were resolved via small modifications to either the appendix or the main text. Other questions opened up interesting discussions on the possibility of fine-tuning within our framework, and the connection of our framework with recent meta-learning approaches for dynamical systems. These discussions will be incorporated and expanded upon in our manuscript as two new appendices. Beyond these points, Reviewer nBHN raised concerns about fairness in our comparison with the baselines in Table 2. We clarified that the baseline scores, the code for their computation, and the preprocessed datasets were obtained from a recent public benchmark. To further ensure fairness, we propose cloning the benchmark's public repository, running it, and evaluating our pretrained models using their evaluation scripts. Details of these additional experiments will be included into a new appendix Finally, Reviewer 6Udz considered the results in Table 2, where our proposal does not consistently achieve state-of-the-art performance, to be a weakness. In response, we note that all baseline scores correspond to state-of-the-art models trained directly on the target datasets, whereas our proposal is neither trained nor fine-tuned on these datasets. It is not immediately evident that a zero-shot model should outperform or match models specifically trained on the target datasets. Indeed, most foundation models for time series, which primarily focus on forecasting, typically perform at best on par with models trained or fine-tuned on the target datasets (see, e.g., [1, 2, 3]). Our findings in Table 2 align with this trend. **References** 1. Liu et al. (KDD 2024): Generative Pretrained Hierarchical Transformer for Time Series Forecasting; 2. Darlow et al. (ICLR 2024): DAM: Towards a Foundation Model for Time Series Forecasting; 3. Das et al. (ICML 2024): A decoder-only foundation model for time-series forecasting.