Foundation Inference Models for Markov Jump Processes

Summary

The paper explores the possibility of using transformers to directly infer the parameters of a Markov jump process (MJPs) from a noisy time data set obtained from a time-dependent process. They first train the machine to correctly predict the model parameters by using synthetic datasets containing data obtained with different realizations of MJPs defined with different parameter sets. Later, they show that the models trained in this way are nevertheless able to derive the parameters of a real simulated or experimental process, even though they were not trained on this data. They also show that one-shot inference is as good as state-of-the-art methods when trained on the target datasets.

Strengths

The authors present a very simple yet powerful and practical idea of using neural networks to extract information from a time-dependent process. In contrast to traditional unsupervised approaches, whose goal is to learn the model parameters to generate samples that are as similar as possible to the original, their goal is a fully supervised one to correctly predict the model parameters in inverse synthetic experiments. They show that the neural network is not only powerful enough to infer all the parameters occurring in the master equation in controlled experiments, but also that they can later successfully apply these trained models to unseen real-world datasets.

Weaknesses

Despite the originality of the work, there are some aspects that remain unclear to me. I find that important details about how to apply the trained models to real data in practice are missing. I also think the limitations are not properly explored.

Questions

Do they make their predictions on the basis of coarse-grained data? If so, how are these representations obtained? If not, how is the data standardized so that they can transfer the models? I did not understand if the adjacency matrix is learned or not, and I am a bit confused about the information given to the machine in real datasets if not. I find that the limitations of this approach are not properly explored or discussed in the text. Are the processes used to generate the synthetic datasets realistic in a general setup? Are they able to describe complex processes like the data generated with the trap model? or just Gaussian processes such as a Brownian motion?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

I discused my doubts about the limitations above.

Summary

The authors propose a foundation model for Markov jump processes (MJPs) that is trained on sequences drawn from synthetic MJPs to predict the corresponding rate matrix and initial state distribution. High-level arguments are presented for the justification of model trained in such a way to be able to generalize to unseen sequences drawn from different processes. Extensive experimental results are shown demonstrating the model's performance after initial training on synthetic data and without any further fine-tuning.

Strengths

The paper is very clearly written and the approach is motivated well. The lack of finetuning in the experiments showcase well the generalizability of the proposed foundation model. Many experiments are showcased that do a good job of demonstrating the intended strengths of the model.

Weaknesses

As pointed out in the paper, the model naturally has weaknesses generalizing to data that lies outside of the initially exposed distribution that it was exposed to in training, such as on sequences with very high rates or with states that lie outside the initial support. Outside of this, I could not find any other weaknesses of the paper; I found it to be very good in general.

Questions

I personally believe the success demonstrated by the foundation model with how well it generalizes being due to the simplicity (not a bad thing!) of the data, i.e., single dimensional across time and is Markov. Would you agree with this or am I missing something? Additionally, I was wondering if "zero-shot" is really appropriate to use here. Typically, zero-shot is used when trying to produce predictions when there is no labeled data to learn from, i.e., predicting "dogs" vs. "cats" with an image model that has seen neither during training. This is commonly achieved by providing some side-information, such as a description of a class. Here the model does indeed not see any "labeled" rate matrices; however, it does see many different sequences realized from this matrix. I see it as less zero-shot and moreso that the inference problem has been amortized. By this, I mean that other methods directly learn a MJP whereas this predicts the parameters to one given a batch of sequences. I am not saying any of this to detract from the paper, but rather just to dial down how to actually think about what it is doing and the language used to describe it. Any thoughts on this would be great to hear.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors adequately describe the limitations.

Summary

This study describes a foundation model for a specific stochastic process, called Markov Jump Process (MJP). The foundation model, called FIM, is trained with a large number of synthetically generated MJP data set. It is shown that the pre-trained FIM can make a zero-shot inference. The capability of FIM is demonstrated by using a few data set.

Strengths

The paper is well written. It clearly formulates the problem setup and the proposed approach, which makes it read well. While the proposed method is exploratory, it demonstrates some interesting capability. The numerical experiments are limited, but well documented.

Weaknesses

The numerical experiments are not strong enough. Numerical experiments should achieve either 1) comparing the performance of the proposed model against a cohort of the state-of-the-art methods, or 2) to systematically demonstrate the capability and limitations of the proposed model and investigate the behavior of the model under a range of important conditions. However, the manuscript does not achieve both. Based on the Appendix, I believe that the authors would have performed an extensive number of studies. A more detailed analysis of the model behaviors with respect to the range of parameters will be appreciated.

Questions

1. It is not clearly explained how the data are scaled. Since the training data is generated from equations, it should be fine. But for an inference the input data should be correctly scaled. How is the inference data scaled? Are there any effects of the scaling on the inference result? 2. Similar to question 1, how the time scale is determined. The model predicts "rate", which depends on the time scale of the data. What method do you use for the inference time? Is the time scaling method general enough to be applied for any real-life MJP? 3. A more detailed explanation about the number of states is required. The model is trained with the maximum number of states of six. What happens when the number of states is less than 6? What do $A$ and $F$ look like? Do the redundant elements of the matrices become exactly zero? If not, how do you normalize the transition probability? 4. Similar to 3, it is interesting to see how the model behaves if the number of states is larger than the maximum number of states in the training, as in real-life problems it is difficult to know or restrict the number of states. 5. During the training, when the number of states is less than six, how do you compute $Var\hat{f}_{ij}$? Note that $Var\hat{f}_{ij}$ is used in the denominator and in a log. 6. How accurate is the prediction of the initial state, $\pi_0$? 7. The prediction requires a prior distribution for $A$. What are the effects of $A$ on the prediction? What happens if the data is generated from a MJP with a different distribution of $A$? 8. The inference of FIM requires a fairly large amount of input data. A conventional approach is to fit a state-space model with the same amount of data to learn the MJP parameters and make an inference. How does FIM compare against the standard ML approach?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

This work presents a framework for amortizing inference on Markov Jump Processes by learning a foundational model in a supervised fashion from synthetic data. Once learned, the "foundation model" is shown to be successful at zero-shot inference in MJP across a range of domains, out-performing SOTA models that are fine-tuned to the target datasets.

Strengths

Clear exposition of why problem is important and the motivation behind the approach. Clear exposition of the method and results.

Weaknesses

Discussion of prior literature could be improved. Some relevant work is on amortized inference in other types of Bayesian models. E.g. "Inference Networks for Sequential Monte Carlo in Graphical Models" learns a supervised model on synthetic data for zero shot inference on Bayesian regression coefficients.

Questions

You mention a main limitation is that generalization is poor outside of the family of models assumed in the synthetic training data. What happens to generalization as you scale up the NN model and widen the family of MJPs in the training data (like for LLMs)?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are adequately discussed.

Response to Rebuttal

I thank the authors for their thoughtful response to my review and am glad to keep my score as accept. I look forward to future work in this area building upon this paper.