Foundation Inference Models for Markov Jump Processes
Authors: David Berghaus
Kostadin Cvejoski
Patrick Seifner
Cesar Ojeda
Ramses J Sanchez
Venue: NeurIPS 2024 poster (2024)
Keywords: Zero-shot inference
Markov jump process
Inference of Markov processes
Foundation models
Foundation models for time series
time series
Primary area: probabilistic_methods
TL;DR: We introduce a framework for zero-shot inference of Markov jump processes from time series data. Our foundation model performs on par with state-of-the-art models which are finetuned on the target datasets.

Abstract: Markov jump processes are continuous-time stochastic processes which describe dynamical systems evolving in discrete state spaces. These processes find wide application in the natural sciences and machine learning, but their inference is known to be far from trivial. In this work we introduce a methodology for *zero-shot inference* of Markov jump processes (MJPs), on bounded state spaces, from noisy and sparse observations, which consists of two components. First, a broad probability distribution over families of MJPs, as well as over possible observation times and noise mechanisms, with which we simulate a synthetic dataset of hidden MJPs and their noisy observations. Second, a neural recognition model that processes subsets of the simulated observations, and that is trained to output the initial condition and rate matrix of the target MJP in a supervised way. We empirically demonstrate that *one and the same* (pretrained) recognition model can infer, *in a zero-shot fashion*, hidden MJPs evolving in state spaces of different dimensionalities. Specifically, we infer MJPs which describe (i) discrete flashing ratchet systems, which are a type of Brownian motors, and the conformational dynamics in (ii) molecular simulations, (iii) experimental ion channel data and (iv) simple protein folding models. What is more, we show that our model performs on par with state-of-the-art models which are trained on the target datasets.

Our pretrained model is available online.

Foundation Inference Models for Markov Jump Processes
David Berghaus Kostadin Cvejoski', Patrick Seifner
César Ojeda' & Ramsés J. Sánchez!
Lamarr Institute's Fraunhofer IAIS?, University of Bonn' & University of Potsdam"
{david. berghaus, kostadin sei reifnerccs.. uni-bonn. de, ojedamarin pww.diiippoammmmmmmm de, sanchezêcs. wnw.bonn.... de
fraunhofer.c. aiaas..aaooeee.
Abstract
Markov jump processes continuous-ting titeitmuus--teeeeeo stochastic processes which describe dynamical systems evolving in discrete state spaces. These processes find wide application in the natural seiences and machine learning, but their inference is known to be far from trivial. In this work we introduce a methodology for zero- shot inference of Markov jump processes (MJPs,, on bounded state spaces, from noisy and sparse observations, which consists of two components. First, a broad probability distribution over families of MJPs, well over possible observation times and noise mechanisms, with which we simulate synthetic dataset of hidden MJPs and their noisy observations. Second, neural recognition model that processes subsets of the simulated observations, and that is trained to output the initial condition and rate matrix of the target MJP in supervised way. We empirically demonstrate that one and the same (pretrained) recognition model can infer, zero-shot fashion, hidden MJPs evolving in state spaces of different dimensionalities. Specifically, we infer MJPs which describe (i) discrete flashing ratchet systems, which are a type of Brownian motors, and the conformational dynamics in (ii) molecular simulations, (iii) experimental ion channel data and (iv) simple protein folding models. What more, we show that our model performs on par with state-of-the art models which are trained the target datasets. Our pretrained model, repository and tutorials are available online!.
Introduction
Very often one encounters dynamic phenomena wildly different nature, that display features which can be reasonably described in terms of macroscopic variable that jumps among finite set of long-lived, metastable discrete states. Think, for example, of the changes in economic activity of country, which exhibit jumps between recession and expansion states (Hamilton, 1989), or the internal motion in proteins or enzymes, which feature jumps between different conformational states (Elber and Karplus, 1987). The states in these phenomena every jump event among them is rare, at least as compared to every other event (or subprocess, or fluctuation) that composes the phenomenon and that occurs, by construction, within the metastable states. Such a description in terms of macroscopic variables effectively decouples the fast, intra-state events from the slow, inter-state ones, and allows for a simple probabílistic treatment of the jumping sequences as Markov stochastic processes: the Markov Jump Processes (MJPs). In this work we are interested in the general problem of inferring the MJPs that best describe empirical (time series) data, recorded from dynamic phenomena of very different kinds.
said to be long-lived, inasmuch as
'nttpe://Eiimaciienn thub... 10/ St/DpeeFFM/immmmmmmmaa www.iitr./iiiii.eeiii html
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Cuurent [pA]
State
0
4
8
10
0.15
0.20
0.25
01:30
0.35
0.40
Time [..u.
Time
Figure 1:
Processes of very different nature (seem to) feature similar jump processes. Left: State values (blue circles) recorded from the discrete flashing ratchet process (black line). Right: Current signal (blue line) recorded from the viral potassium channel KCV together with one possible coarse-gra coarse-grained grained representation (black line).
To set the stage, let us assume that we want to study some D-dimensional empirical process z(t) : R+ R.D, which features long- livellvee dynamic modes, trapped in some discrete set of metastable states. Let us call this set X. Let also assume that we can obtain macroscopic, coarse-grained representation from z(t) - say, with a clustering algorithm - in which the fast, intra-state events have been integrated out (i.e. marginalized). Let us call this macroscopic variable (1): : R+ + X. If we now make the Markov assumption and define the quantity f( (2))z))) At as the infinitesimal probability of observing one jump from state z' (at some time t). into different state (at time t + Åt), we can immediately write down, following standard arguments (Gardiner, 2009), differential equation that describes the probability distribution pMo (I, t), over the discrete set of metastable states X, which encapsulates the state of the process X (t) as time evolves, that
:
dpaar(s, () dt
(f(r)a)) (f(r)al)Paarr(d 1) #'
Σ
(zr, t)).
Equation is the so-called master equation the MJP whose solutions are completely characterized by an initial condition (zur(z, t 0) and the transition rates X' x X - Rt.
With these preliminaries in mind, we shall say that infer an MJP from set of (noisy) observations z(r))..... ((T) on the process (t), recorded at some observation times .i..... means to infer both the transition rates and the initial condition determining the hidden MJP ( that best explains the observations. In practice, statisticians typically assume that they directly observe the coarse- grained process X (1). That is, they assume they have access to the (possibly noisy) values Tı, T,, taken by (L) at the observation times ......: (see Section 2). We shall start from the same assumptions. Statisticians then tackle the inference problem by (i) defining some (typically complex) model that encodes, in one way or the other, equation I above; (ii) parameterizing the model with some trainable parameter set 0; and (iii updating to fit the empirical dataset.
One issue with this approach is that it turns the inference of hidden MJPs into an instance of an unsupervised learning problem, which, as history shows, far from trivial (see Section 2). Another major issue is that, if one happens to succeed in training said model, the trained parameter set d* will usually be overly specific the training set { TI..... (zt, T)), which means will likely struggle to handle second empirical process, even if the latter can be described by a similar MJP. Figure contains snapshots from two empirical processes of very different nature. The figure on the left shows set of observations (blue circles) recorded from the discrete flashing ratchet process (black line). The figure on the right shows the ion flow across a cell membrane, which jumps between different activity levels (blue line). Despite the vast differences between the physical mechanisms underlying each of these processes, the coarse-grained representations of the second one (black line) is abstract enough to be strikingly similar to the first one. Now, we expect that at this level of representation - one could train a single inference model to fit each process (separately). Unfortunately, also expect that an inference model trained to fit only one of these coorse-grainee processes, will have a hard time describing the second one.
In this paper we will argue that the notion of an MJP description (in coarse-grained space) is simple enough, that it can be encoded into the weights of single neural network model. Indeed, instead of training, in an unsupervised manner, complex model (which somehow encodes the

---

log Var F
4/
K
Wi
22
Figure 2: Foundation Inference Model (FIM) for MJP. Left: Graphical model of the FIM (synthetic) data generation mechanism. Filled (empty) circles represent observed (unobserved) random variables. The light-blue rectangle represents the continuous-time MJP trajectory, which is observed discretely in time. See main text for details regarding notation. Right: Inference model. The network 1/1 is called K times to process K different time series. Their outputs is first processed by the attention network 21 and then by the FNNs Φ1, 62 and Φ3 to obtain the estimates F, log Var F and a.. respectively.
master equation) on a single empirical process; we will train, in supervised manner, a simple neural network model synthetic dataset that is composed of many different MJPs, and hence implicitly encodes the master equation. This procedure can be understood as an amortization of the probabilistic inference process through single recognition model, 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, these previous works do, our (pretrained) recognition model as auxiliary to Monte Carlo or expectation propagation methods, we employ it to directly infer hidden MJPs from various synthetic, simulation and experimental datasets, without any parameter fine-tuning. We thus adopt the "zero-shot" terminology introduced by Larochelle et al. (2008), by which we mean that our procedure aims recognize objects (i.e. MJPs) whose instances (i.e. noisy and sparse series of observations on them) may have not been seen during training. We have recently shown that such an amortization can be used to train a recognition model to perform zero-shot imputation of time series data (Seifner et al., 2024). Below we demonstrate that it can also be used to train model of minimal inductive biases, to perform zero-shot inference of hidden MJPs from empirical processes of very different kinds, which take values in state spaces of different sizes. We shall call this recognition model Foundation Inference Model (FIM) for Markov jump processes.
In what follows, we first review both classical and recent solutions to the MJP inference problem in Section 2. We then introduce the FIM methodology in Section 3, which consists of a synthetic data generation model and a neural recognition model. In Section 4 we empirically demonstrate that our methodology is able to infer MJPs from discrete flashing ratchet process, as well as from molecular dynamics simulations and experimental ion channel data, all in zero-shot fashion, while performing on par with seate-of-the-eee models which are trained on the target datasets. Finally, Section 5 closes the paper with some concluding remarks about future work, while Section 6 comments on the main limitations of our methodology.
2 Related Work
The inference of MJP from noisy and sparse observations (in coarse-grained space) is by now classical problem machine learning. There are three main lines of research. The first (and earliest) one attempts to directly optimize the MJP transition rates, to maximize the likelihood of the discretely observed MJP via expectation maximization (Asmussen et al., 1996; Bladt and Sørensen, 2005; Metzner et al., 2007). Thus, these works encode the MJP inductive bias directly into their architecture. The second line of research leverages Bayesian framework to infer the posterior distribution over the transition rates, through various Markov chain Monte Carlo (MCMC) algorithms (Boys et al., 2008; Fearnhead and Sherlock, 2006; Rao and Teg, 2013; Hajiaghayi et al., 2014). Accordingly, these
We name our model foundation model because it aligns with the definition proposed by Bommasani et al. (2021). Indeed, they define foundation models as any model that is trained on broad data (generally using self supersuper scale) that can be adapted to wide range of downstream tasks.
3

---

ioaullation-bseen approaches encode the MJP inductive bias directly into their trainable sampling distributions. The third one, also Bayesian in character, involves variational inference. Within it, one finds again MCMC (Zhang et al., 2017), as well as expectation maximization (Opper and Sanguinetti, 2007) and moment-based (Wildner and Koeppl, 2019) approaches. More recently, Seifner and Sánchez (2023) used neural variational inference (Kingma and Welling, 2013) and neural ODEs (Chen et al., 2018) to infer an implicit distribution over the MJP transition rates. All these variational methods encode the MJP inductive bias into their training objective and, in some cases, into their architecture too.
Besides the model of Seifner and Sánchez (2023), which automatically infers the coarse-graing representation X (t) from D-dimensional, countinuous signals, all the solutions above tackle the MIP inference problem directly in coarse- -aanse-g space. Yet below, we also investigate the conformational dynamies of physical systems for which the recorded data lies continuous space. To approach such type of problems, we will first need to define a goarse-grained representation of the state space large body of works, within the molecular simulation
of interest. Fortunately for us, there community, precisely dealing with different methods to obtain such representations, and we refer the reader to e.g. Noé et al. (2020) for a review. McGibbon and Pande (2015), for example, leveraged one such method to infer the MJP transition rates describing a molecular dynamics simulation via maximum likelihood. Alternatively, researchers have also treated the conformational states in these systems as core sets, and inferred phenomenological MJP rates from them (Schütte et al., 2011 1), or modelled the fast intra-state events as diffusion processes, indexed by a hidden MJP, and inferred the latter either via MCMC (Kilic et al., 2021; Köhs et al., 2022) or variational (Horenko et al., 2006; Köhs et al., 2021) methods.
In this work we tackle the classical MJP inference problem on coarse-grained space and present, to the best of our knowledge, its first zero-shot solution.
3 Foundation Inference Models
In this section we introduce a novel methodology for zero-shot inference of Markov jump processes which frames the inference task as supervised learning problem. Our main assumption is that the space of realizable MJPs', which take values on bounded state spaces that are not too large, is simple enough to be covered by heuristically constructed synthetic distribution over noisy and discretely observed MJPs. If this assumption were to hold, a model trained to infer the hidden MJPs within a synthetic dataset sampled from this distribution would automatically perform zero-shot inference on any unseen sequence of empirical observations. We do not intend to formally prove this assumption. Rather, we will empirically demonstrate that a model trained in such a way can indeed perform zero-shot inference of MJPs in variety of cases.
Our methodology has two components. First, a data generation model that encodes our believes about the class of realizable MJPs we aim to model. Second, a neural recognition model that maps subsets of the simulated MJP observations onto the initial condition and rate matrix of their target MJPs. We will explore the details of these two components in the following sections.
3.1 Synthetic Data Generation Model
In this subsection we define broad distribution over possible MJPs, observation times and noise mechanisms, with which we simulate an ensemble of noisy, discretely observed MJPs. Before we start, let us remark that will slightly abuse notation and denote both probability distributions and their densities with the same symbols. Similarly, we will also denote both random variables and their values with the same symbols.
Let us denote the size of the largest state space we include in our ensemble with C, and arrange all transition rates, for every MJPs within the ensemble, into C x C rate matrices. Let us label these matrices with F. We define the probability of recording the noisy sequence 2'....., xf E X, at the
By realizable MJPs we mean here MJPs that can be inferred from physical processes, given the typical experimental constraints, like egg. temporal spatial resolution.

---

observation times 0 <<= < ... < 77 < T, with T the observation time horizon, as follows
II Poose
Pross (a/) (T's Ps: Proue (z Ta To) Po ( 71, T..... (llp))
(2)
Next, we specify the different components of Eq. 2, starting from the right.
Distribution over initial conditions. The distribution p( p(Toolo)). with hyperparameter Po. is defined over the C-simplex, and encodes our beliefs about the initial state (i.e. the preparation) of the system. It enters the master equation as the class probabilities of the categorical distribution over the states of the system, at the start of the process. That is Pur( = 0) = Cat(TTo). We either choose wo to be the class probabilities of the stationary distribution of the process, sample it from Dirichlet distribution. Appendix provides the specifics.
Distribution over rate matrices. The distribution Prates (F) A, pp) over the rate matrices encodes our beliefs about the class of MJPs we expect find in practice. We define it to cover MJPs with state spaces whose sizes range from 2 until C, because we want our FIM to be able to handle processes taking values in all those spaces. The distribution conditioned on the adjacency matrix A, which encodes only connected state spaces (iee. irreducible embedded Markov chains only), and a hyperparameter Py which encodes the range of rate values within the ensemble. Specifically, we define the transition rates Fij anj where dij is the corresponding entry of A and fúy is sampled from a set of Beta distributions, with different hyperparameters py. Note that these choices restrict the values of the transition rates within the ensemble to the interval (0, 1) and hence, they restrict the number of resolvable transitions within the time horizon T of the simulation. We refer the reader to Appendix B, where we specify the prior p(A. (())p(pp)) and its consequences, as well as give details about the sampling procedure. We also discuss the main limitations of choosing a Beta prior over the transition rates in Section 6.
Distribution over observation grids. The distribution Peia (Ti. Ti n(pr), with hyperparameter fr. gives the probability of observing the MJP at the times .., and thus encodes our uncertainty about the recording process. Given that we do not know a priori whether the data will be recorded regularly irregularly in time, nor we know its recording frequency, we define this distribution to cover both regular and irregular cases, as well as various recording frequencies. Note that the number of observation points on the grid is variable. Please see Appendix B for details.
Distribution over noise process. Just as the (instantaneous) solution of the master equation Paur (zlt, F, To). the noise distribution Paoise (z' |x, Pz). with hyperparameter is defined over the set of metastable states X. Recall that FIM solves the MJP inference problem directly coarse-grained space. The noise distributions then encodes both, possible measurement errors that propagate through the coarse-graines representation, or noise in the coarse-grained representation itself. We provide details of its implementation in Appendix B.
We use the generative model, Eq. above, to generate N MJPs, taking values on state spaces with sizes ranging from 2 to C. We then sample K paths per MJP, with probability p( (K), on the interval [0, 27. The jth instance of the dataset thus consists of K paths and is given by
F2 ~ Prace (F (FAA,,P))) Prs); and Toz ~ p(Toolo)), with (A,, ~ ~
so that {x,((1) K (K.,) and { (a/ses Pmuise (x' (e'|X,n [X (T)He))
/ ~ GillespieeFF Toj); A=1
(3)
(k,1)=(1,,)) with Tyki, Jykl J==1 Pria(TI..... nlpr),
where Gillespie denotes the Gillespie algorithm we use to sample the MJP paths (see Algorithm 1). Note that we make the number of paths (K above) per MJP random, because we do not know a priori how many realizations (i.e. experiments), from the empirical process of interest, will be available at the inference time. We refer the reader to Appendix B for additional details.
Figure illustrates the complete data generation process.
3.2 Supervised Recognition Model
In this subsection we introduce a neural recognition model that processes a set of K time series of the form (<2/11 Tki)...., (2/g, 7ha) [x)]K k= as generated by the procedure in Eq. 3 above, and estimates
5

---

the intensity rate matrix F and initial distribution To of the hidden MJP. Practically speaking, we would like the model be able to infer MJPs from time series with observation times on any scale. To ensure this, we first normalize all observation times to lie on the unit interval, by dividing them by the maximum observation time Timx = max (Tkl,..., The T and then rescale the output of the model accordingly (see Appendix C for details).
Let us use i and I denote fted-forward, sequence processing networks, and attention networks, respectively. Thus v can denote e.g. LSTM or Transformer networks, while Q can denote e.g. a self-attention mechanism. Let us also denote the networks' parameters with e.
We first process each time series with a network 1 to get a set of K embeddings, which we then summarize into global representation h& through the attention network ?2 In equations, we write
b = Q (hie, bkee 0) with hRg- = v/ (1/ Thi...., L'a. Tkl. 0) and k = I,..., K. (4)
Next we use the global representation to get estimate of the intensity rate matrix, which artificially model as a Gaussian variable with positive mean, and the initial distribution of the hidden MJP follows
F exp(éi 0)), Var F = exp(02(ho. 0)) and to = c (bo. 0),
(5)
where the exponential function ensures the positivity of our estimates, and the variance is used to represent the model's unceriainty in the estimation of the rates (Seifner et al., 2024). The right panel of Figure summarizes the recognition model, and Appendix provides additional information about the inputs to, outputs of and rescalings done by the model.
Training objective. We train the model maximize the likelihood its predictions, taking care of the exact zeros (i.e. the missing links) in the data. To wit
C =
E
Σ
aij
(fg
fu)) Var Paj
+ 2
log Varfu] - A(1 -as) and + Var fu]}
1j ml
C
E Tamp
Σ
To0 log 7,
(6)
where the second term is nothing but the mean-squared error of the predicted rates fuy (and its standard deviation) when the corresponding link is missing, and can be understood as a regularizer with weight A. The latter is a hyperparameter.
FIM context number. During training, FIM processes a variable number K of time series, which lies the interval [Kain ] Similarly, each one of these time series has variable number of observation points, which lies on the interval [imin, /mmx]. We shall say that FIM needs a bare minimum of Lii input data points to function. Perhaps unsurprisingly, perform bests when processing Kuachmax data points. Going significantly beyond this number seems nevertheless to decrease the performance of FIM. We invite the reader check Appendix D for details.
have empirically seen that FIM
Let us define then, for the sake of convenience, the FIM context number c(K,1) Kl as the number of input points* FIM makes use of to estimate and πe.
Experiments
In this section we test methodology five datasets of varying complexity, and corrupted by noise signals of very different nature, whose hidden MJPs are known to take values in state spaces of different sizes. In what follows use one and the same (pretrained) FIM infer hidden MJPs from all these datasets, without any parameter fine-tuning. Our FIM was (pre)trained on a dataset of 45K MJPs, defined over state spaces whose sizes range from 2 to 6. A maximum of (K =)300 realizations (paths) per MJP were observed during training, everyone of which spanned a time-horizon T 10, recorded a maximum of 100 time points, 1% of which were mislabeled. Given these specifications, FIM is expected to perform best for the context number c(300, 100) during evaluation. Additional
We can think about as the context length in large language models.

---

2V
V
b
GROUND TRUTH
1.00
1.00
1.00
NEURALMJP
1.06
1.17
1.14
FIM 1.11(7)
0.99
0.98(5)
Figure Illustration of the six-state dis- crete flashing ratchet model. The potential V switched on and off at rate r. The transition rates fūj J allow the particle to propagate through the ring.
Table 1: Inference of the discrete flashing ratchet process. The FIM results correspond to FIM evalu- ations with context number c(300, 50), averaged over 15 batches.
information regarding model architecture, hyperparameter selection and other training details can be found in Appendix D.
Depending on the dataset, we compare our findings against the NeuralMJP model of Seifner and Sánchez (2023), the switching diffusion model (SDiff) of Köhs et al. (2021), and the discrete-time Markov model (VampNets) of Mardt et al. (2017).
Baselines:
All these baselines are trained the target datasets.
4.1 The Discrete Flashing Ratchet (DFR): A Proof of Concept
In statistical physics, the ratchet effect refers the rectification of thermal fluctuations into directed motion to produce work, and goes all the way back to Feynman (Feynman et al., 1965). Here we consider a simple example thereof, in which Brownian particle, immersed in a thermal bath unit temperature, moves on a one the-dimensional lattice. The particle is subject to a linear, periodic and asymmetric potential of maximum height 2V that is switched on and off at constant rate r. The potential has three possible values when is switched on, which correspond to three of the states of the system. The particle jumps among them with rate fiy. When the potential is switched off, the particle jumps freely with rate f.j We can therefore think the system as a six-state system, as illustrated in Figure 3. Similar to Roldán and Parrondo (2010), we define the transition rates
V (-56-1)) for i,j 10,1,2); f = b, for i,j E (3, 5).
J3 = exp (
(7)
Given these specifics, we consider the parameter set (V, B) = 1, together with the dataset simulated by Seifner and Sánchez (2023), which consists of 5000 paths (in coarse-grained space) recorded irregular grid of 50 time points. The task is to infer (V,r, B) from these time series. NeuralMJP infers a global distribution over the rate matrices and hence relies on their entire train set, which amounts to about 4500 time series. We therefore report FIM evaluations with context number c(300, 50) on that same train set, averaged over 15 (non-overlapping) batches in Table
The results show that FIM performs on par with (or even better than) NeuralMJP, despite not having been trained on the data. Note in particular that our results are sharply peaked around their mean, indicating that context of c(300, 50) points only contains enough information to describe the data well. What is more, Table 16 in the Appendix demonstrates that FIM can infer vanishing transition rates well (see Eq. 6). Now, being able to infer the rate matrix in zero-shot mode allows us to immediately estimate a number of observables of interest without any training. Stationary distributions, relaxation times and mean first-passaee times (see Appendix A for their definition), as well as dime-dependent moments, can all be computed zero-shot via FIM. For example, we report on the left block of Figure 4 the time-dependent class probabilities (i.e. the master eq. solutions) computed with the FIM-inferred rate matrix (black), against the ground-truth solution (blue). The agreement is very good.
Zero-shot estimation of entropy production. The DFR model is interesting because the random switching combined with the asymmetry in the potential make it more likely for the particle to jump towards the right (see Figure 4). Indeed, that is the ratchet effect. As a consequence, the system features stationary distribution with net current the so-called non non-equilibrium steady state
7

---

0.30
1.2
Prediction
0.25
Entropy Production
1.0-
Ground Truth
Prohsdility
Ps
Po
0.8
0.20-
0.6
0.15
0.4
0.10
0.2
0.0
0.5
1.0
2.0
3
Time
Figure 4: Zero-shot inference of DFR process. Left: master eq. solution Pxur t) as time evolves, wrt. the (averaged) FIM-inferree rate matrix is shown in black. The ground-truth solution is shown in blue. Right: Total entropy production computed from FIM (over time-horizon T 2.5 [a.u.]). The model works remarkably well for continuous range of potential values.
(Ajdari and Prost, 1992), which is characterized by non-vanishing (stochastic) entropy production. The development of (neural) estimators of entropy production very active topic of current research (see egg. Kim et al. (2020) and Otsubo et al. (2022)). Given that the entropy production can be written down in closed form function of both the rate matrix and the master eq. solution (see e.g. Seifert (2012)), we can readily use FIM to estimate it.
Figure 4 displays the total entropy production computed with FIM for set of different potentials. The results are averaged over 15 FIM evaluations with c(300, 50) and are again in very good agreement with the ground truth. It is noteworthy that FIM, trained captures well contimous set of MJPs. That is, we evaluate one and the same FIM over different datasets, each sampled from DFR model with different potential value. In sharp contrast, state-of- the-art models need to be retrained for every new potential value (Kim et al., 2020).
heuristically constructed dataset,
Zero-shot simulation of the DFR process. Inferring the rate matrix and initial condition of a MJP process entails that one can also sample from it. Our FIM can thus be used as a zero-shot generative model for MJPs. However, to test the quality of said MJP realizations wrt. some target MJP, need a distance between the two. Here we propose to use the Hellinger distance (Le Cam and Yang, 2000) to first estimate the divergence between a sequence of (local) histogram pairs, recorded at given set of observation times, and then average the local estimates along time. Appendix F. empirically demonstrates that this pragmatically defined MJP distance is sensible.
Table reports the time-averaged Hellinger distance between 1000 (ground-truth) DFR paths and 1000 paths sampled from (the MJPs inferred by) NeuralMJP and FIM. We repeat this calculation 100 times, for 1000 newly sampled paths from NeuralM.IP and FIM, but the same 1000 target paths, to compute the mean values and error bars in the Table. The results show that the zero-shot DFR simulation obtained through FIM is on par with the NeuralMJP-baae simulation, wrt. the ground truth.
4.2 Switching Ion Channel (lonCh): Zero-Shot Inference of Three-State MJP
In this section we study the conformational dynamics of the viral ion channel KCv which exhibits three metastable states (Gazzarrini et al., 2006). Specifically, analyse the ion flow across the membrane as the system jumps between its metastable configurations. This ion flow was recorded at frequency of 5kHz over second. Figure I shows one snapshot of these recordings, which were made available to us via private communication (see the (ednnwwledgements)) Our goal is to infer physical observables like the stationary distribution and mean first- pass-passage times of the conformational dynamics, and to compare our findings against the SDiff model Köhs et al. (2021) and NeuralMJP.
The recordings live in real space, which means that we first need obtain a coarse-grained represen- tation (CGR) from them, before we can apply FIM. Here consider two CGRs: the CGR inferred

---

Dataset
NEURALMUP
FIM
BOTTOM
MIDDLE
TOP
DFR 0.30(0.06)
0.27(0.06)
SDIFF
0.17961
0.14987
0.67052
TONCH
0.48(0.02)
(.40().02)
NEURALMJJ
0.17672
0.09472
0.72856
ADP
1.38((.52)
1.39(0.47)
FIM-NMJP
0.18224
0.10156
0.71621
PFOLD
0.015(0.015) 0.014(0.014)
FIM-GMM
0.19330
0.08124
0.725416
Table 2: Time-averaged Hellinger distances between empirical processes and samples from either NeuralMJJ or FIM [in a le-2 scale] (lower is better). Mean and std. are computed from set of 100 histograms
Table 3: Stationary distribution inferred from the
switching ion channel experiment. FIM-NMJP and FIM-GMAN correspond to our inference from differ- ent coarse-grained representations. The results agree well.
by NcuralMJP and naive CGR obtained with a Gaussian Mixture Model (GMM), Given that we only have 5000 observations available, make use of single FIM evaluation with context number c(50, 100). We infer two FIM rate matrices, one per each CGR, which we label as FIM-NMJP and FIM-GMM.
Table 3 contains the inferred stationary distributions from all models and evidences that a single FIM evaluation is enough unveil the long-time asymptotics of the process. Similarly, Table 15 in the Appendix, which contains the inferred mean-first passage times, demonstrates that FIM makes the same inference about the short-term dynamics of the process do SDiff and NeuralMJP. See Appendix F for additional results.
Zero-shot simulation of switching ion channel process. Just as we did with the DFR process, can use FIM to simulate the switching ion channel process in coarse- coarse-grained space. Since only paths evaluate NeuralMJP against FIM-NMJP. To construct
on the same CG space the target distribution, we leverage another 30 seconds of measurements, which amount to 150K observations that have not been seen by any of the models. The results in Table indicate that zero-shot simulations is statistically closer to the ground-truth process than the NeuralMJP simulation.
be compared,
4.3 Alanine Dipeptide (ADP): Zero-Shot Inference of Six-State MJP
Alanine dipeptide is 22-atom molecule widely used as benchmark in molecular dynamics simulation studies. Its popularity stems from the fact that the heavy-atom dynamics, which jumps between six metastable states, can be fully described in terms of the dihedral (torsional) angles & and ø (see e.g. Mironov et al. (2019) for details).
We examine all-atom ADP simulation of microsecond, which made available us via private communication (see the Acknowledgements below), and compare against hoth, the VampNets model of Mardt et al. (2017) and MJU. The data consists of the values taken by the dihedral angles as time evolves and thus needs to be mapped onto some coarse-grained space. We again make use of NeuralMJP obtain a CGR. We then use FIM with context number c(300, 100) to process 32 100-point time windows of the simulation and compute an average rate matrix. Note that this is the optimal context number of our pretrained model. Table 4 (and Appendix F.2) confirms that, once again, FIM can infer the same physical properties from the ADP simulation as the baselines.
Zero-shot simulation of the alanine dipeptide. Simulations in coarse- gaanse-rained space for molecular dynamics is a high-interest research direction (Husic et al., 2020). Here we demonstrate that FIM can be used simulate the ADP process in zero-shot mode. Indeed, Table 2 reports the distance from both MJP and FIM to a target ADP process, computed from 200 paths with 100 observations each. Once more, FIM performs comparable to NeuralMJP.
4.4 Zero-Shot Inference of Two-State MJJPs
Finally, we consider two additional systems that feature jumps between two metastable states: simple protein folding model and two-mode switching system. We invite the reader to check out Appendix F.5 and E.6 for the details. That being said, Table 8 reports the distance of both N.uralMJJP and FIM wrt. the empirical protein folding process (PFold). The high variance indicates that the distance cannot resolve any difference between the processes given the available number of samples.

---

PROBABILITY PER STATE
RELANATIION TIME SCALES (IN 118)
III
IV
V
VI
VAMPNETS
0.30
0.24
0.20
0.15
0.11
0.01
0.008
0,009
0.055
0.065
1.920
NEURALMUP
0.30
0.31
0.23
0.10
0.05
0.01
0.009
0.009
0.043
0.069
0.774
FIM
0.28
0.28
0.24
0.07
0.10
0.03
0.008
0.009
0.079
0.118
0.611
Table Left: stationary distribution of the ADP process. The states are ordered in such a way that the ADP conformations associated with given state are comparable between the VampNets and NeuralMJP CGRs. Right: relaxation time scales to stationarity. FIM agrees well with both baselines.
Conclusions
In this work we introduced a novel methodology for zero-shot inference Markov jump processes and its Foundation Inference Model (FIM). We empirically demonstrated that one and the same FIM can used to estimate stationary distributions, relaxation times, mean first-passage times, time-dependent moments and thermodynamic quantities (i.e. the entropy production) from noisy and discretely observed MJPs, taking values state spaces of different dimensionalities, all zero-shot mode. To the best of our knowledge, FIM is also the first zero-shot generative model for MJPs.
Future work shall involve extending our methodology to Birth and Death processes. as well as considering more complex (prior) transition rate distributions. See our discussion on Limitations in the next section, for details.
Limitations
The main limitations of our methodology clearly involve our synthetic distribution. Evaluating FIM on empirical datasets whose distribution significantly deviates from our synthetic distribution will, inevitably, yield poor estimates. Consider Figure 4 (right), for example. The performance of FIM quickly deteriorates for 2 3, for which ratio between the largest and smallest rates gets larger than about three orders of magnitude. These cases are unlikely under our prior Beta distributions, and hence effectively lie outside of our synthetic distribution.
More generally, the MJP dynamics underlying phenomena that feature long-lived, metastable states, ultimately depends on the shape the energy landscape characterizing the X' inasmuch as the transition rates between metastable states and, j (faj in our notation) are characterized by the depth of the energy traps (that is, the height of the barrier between them).
In equations, we write
ES fuj = exp
where Ej is the jth trap depth, and T is the temperature of the system. Therefore, the distribution over energy traps determines the distribution over transition rates.
Just to give an example, if we studied systems with exponentially distributed energy traps as egg. in the classical Trap model of glassy systems of Bouchaud (1992)- - we would immediately find p(f) x T fT-1 Transition rates sampled from such power-law distributions clearly lie outside our ensemble Beta distributions, even we use our rescaling trick. Future work shall explore training FIM on synthetic MJPs featuring power-law-distributed transition rates.
Acknowledgements
This research has been funded by the Federal Ministry of Education and Research of Germany and the state of North-Rhine Westphalia as part of the Lamarr Institute for Machine Learning and Artificial Intelligence. Additionally, César Ojeda was supported by Deutsche Forschungsgemeinschat (DFG) - Project-ID 318763901 SFBB 1294.
We would like thank Lukas Köhs for sharing the experimental ion channel data with us. The actual experiment was carried out by Kerri Kukovetz and Oliver Rauh while working in the lab of Gerhard
10

---

Thiel of TU Darmstadt. Similarly, we would like to thank Nick Charron and Cecilia Clementi, from the Theoretical and Computational Biophysics group of the Freie Universität Berlin, for sharing the all-atom alanine dipeptide simulation data with us. The simulation was carried out by Christoph Wehmeyer while working the research group of Frank Noé of the Freie Universität Berlin.
References
Armand Ajdari and Jaxques Prost. Mouvement induit par un potentiel périodique de basse symétrie: diélectrophorese pulsée. Comptes rendus de I' l'Académie des sciences. Série 2, Mécanique, Physique, Chimie, Sciences de l'univers, Sciences de la Terre, 315(13):1635-163) 1992.
Søren Asmussen, Olle Nerman. and Marita Olsson. Fitting phase-type distributions via the em algorithm. Scandinavian Journal of Statistics, pages 419-441, 1996.
Mogens Bladt and Michael Sørensen. Statistical inference for discretely observed markov jump processes. Journal of the Royal Statistical Society: Series (Statistical Methodology), 67(3): 395-410, 2005.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint wrXiv:2108.01 072507, 2021.
Jean-Philippe Bouchaud. Weak ergodicity breaking and aging in disordered systems. Journal de Physique I, 2(9): 705-17133 1992.
Richard J Boys, Darren J Wilkinson, and Thomas BL Kirkwood. Bayesian inference for discretely observed stochastic kinetic model. Statistics and Computing. 18(2):225-33, 2008.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Kristjanson Duvenaud. Neural ordinary differential equations. In Neural Information Processing Systems, 2018
R Elber and Martin Karplus. Multiple conformational states of proteins: molecular dynamics analysis of myoglobin. Science, 235(4786):318-321, 1987.
Erdös and A Rényi. On random graphs i. Publ. math. debrecen, 6(290-297):18, 1959.
Paul Fearnhead and Chris Sherlock. An exact gibbs sampler for the markov-modulated poisson process. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 68(5): 767-784, 2006.
Richard P Feynman, Robert B Leighton, Matthew Sands, and Everett M Hafner. The feynman lectures on physics; vol. i. American Journal of Physics, 33(9):750-755 1965.
Crispin Gardiner. Stochastic methods: handbook for the natural and social sciences. 2009.
Sabrina Gazzarrini, Ming Kang, Svetlana Epimashko, James Van Etten, Jack Dainty, Gerhard Thiel, and Anna Moroni. Chlorella virus mt325 encodes water and potassium channels that interact synergistically. Proceedings of the National Academy of Sciences, 103(14):5355-5366 2006.
Daniel T. Gillespie. Exact stochastic simulation of coupled chemical reactions. The Journal of Physical Chemistry, 81:2340-220 2261, 1977.
Monir Hajiaghayi, Bonnie Kirkpatrick, Liangliang Wang, and Alexandre Bouchard-Cöre Efficient continuous-time markov chain estimation. In International Conference on Machine Learning, pages 638-646. PMLR, 2014.
James Hamilton. A new approach to the economic analysis of nonstationary time series and the business cycle. Econometrica: Journal of the econometric society, pages 357-384, 1989.
Nicolas Heess, Daniel Tarlow, and John Winn. Learning to pass expectation propagation messages. Advances in Neural Information Processing Systems, 26, 2013.
Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory. Neural Computation, 9()): 7735-1780. 1997.

---

Illia Horenko, Evelyn Dittmer, Alexander Fischer, and Christof Schütte. Automated model reduction for complex systems exhibiting metastability. Multiscale Modeling & Simulation, 5(3):802-827, 2006.
Brooke E. Husic, Nicholas E. Charron, Dominik Lemm, Jiang Wang, Adrià Pérez, Maciej Majewski, Andreas Krämer, Yaoyi Chen, Simon Olsson, Gianni de Fabritiis, Frank Noé, and Cecilia Clementi. Coarse graining molecular dynamics with graph neural networks. The Journal of Chemical Physics, 153(11):194101, 2020.
Zeliha Kilic, Ioannis Sgouralis, and Steve Pressé. Generalizing hmms to continuous time for fast kinetics: Hidden markov jump processes. Biophysical journal, 120(3):409-44 423, 2021.
Dong-Ky Kyum Kim, Youngkyoung Bae, Sangyun Lee, and Hawoong Jeong. Learning entropy produc- tion via neural networks. Physical Review Letters, 125(14): 2020.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Lukas Köhs, Bastian Alt, and Heinz Koeppl. Variational inference for continuous-time switching dynamical systems. In Advances in Neural Information Processing Systems, volume 34, pages 20545-20557, 2021.
Lukas Köhs. Bastian Alt, and Heinz Koeppl. Markov chain monte carlo for continuous-time switching dynamical systems. In International Conference on Machine Learning, pages 1450-11454. PMLR, 2022.
Hugo Larochelle, Dumitru Erhan, and Yoshua Bengio. Zero-data learning of new tasks. In AAAI, volume 1, page 3, 2008.
Lucien Marie Le Cam and Grace Lo Yang. Asymptotics in statistics: some basic concepts. Springer Science & Business Media, 2000.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 10000 2017.
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. Vampnets for deep learning of molecular kinetics. Nature Communications, 9, 2017.
Robert McGibbon and Vijay S Pande. Efficient maximum likelihood parameterization of continuous- time markov processes. The Journal of chemical physics, 143(3):034109, 2015.
Philipp Metzner, Illia Horenko, and Christof Schütte. Generator estimation of markov jump processes based on incomplete observations nonequidistant in time. Phys. Rev. E, 76:066702, Dec 2007.
Vladimir Mironov, Yuri Alexeev, Vikram Khipple Mulligan, and Dmitri G. Fedorov. A systematic study of minima in alanine dipeptide. Journal of Computational Chemistry, 40(2):297-309, 2019.
Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annual review of physical chemistry, 71:361-390, 2020.
M. Opper and G. Sanguinetti. Variational inference for markov jump processes. In NIPS, 2007.
Shun Otsubo, Sreekanth Manikandan, Takahiro Sagawa, and Supriya Krishnamurthy. Estimating time-dependent entropy production from non-equilibrium trajectories. Communications Physics, (1):1, 2022.
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.
Vinayak Rao and Yee Whte Teg. Fast meme sampling for markov jump processes and extensions. Journal of Machine Learning Research, 14(11), 2013.
É. Roldán and J. M. R. Parrondo. Estimating dissipation from single stationary trajectories. Physical review letters, 105 15:150607, 2010.
12

---

Christof Schütte, Frank Noé, Jianfeng Lu, Marco Sarich, and Eric Vanden Mijnenn. Markov state models based on milestoning. The Journal of chemical physics, 134(20), 2011.
Udo Seifert. Stochastic thermodynamics, fluctuation theorems and molecular machines. Reports on Progress in Physics, 75(12): nov 2012. doi: 10. 110880034-4835///222000 URL https://dx.co doi. org/ / 10.
Patrick Seifner and Ramsés J Sánchez. Neural markov jump processes. In International Conference Machine Learning, pages 30523-30552 PMLR, 2023.
Patrick Seifner, Kostadin Cvejoski, and Ramses Sanchez. Foundational inference models for dynamical systems. arXiv preprint arXiv:2402.0777 07594, 2024.
Andreas Stuhlmüller, Jacob Taylor, and Noah Goodman. Learning stochastic inverses. Advances in neural information processing systems, 26, 2013.
Benjamin Trendelkamp-Schnee Schrrorr and Frank Noé. Efficient estimation rare-event kinetics. arXiv: Chemical Physics, 2014.
Yasemin Bozkurt Varolgünes, T. Bereau, and Joseph F. Rudzinski. Interpretable embeddings from molecular simulations using gaussian mixture variational autoencoders. Machine Learning: Science and Technology. 1, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6088 6008, 2017.
Christian Wildner and Heinz Koeppl. nosent-based variational inference for markov jump processes. In International Conference on Machine Learning, pages 6766-6775 PMLR, 2019.
Boqian Zhang, Jiangwei Pan, and Vinayak A Rao. Collapsed variational bayes for markov jump processes. Advances in Neural Information Processing Systems, 30, 2017.
13

---

A Background on MJPs
In this section we provide some brief background on MJPs and describe how physical quantities such as the stationary distributions, relaxation times and mean first passage times can be computed from the intensity matrix. Additionally, we mention how trajectories for MJPs can be sampled using the Gillespie algorithm.
A.1 Background on Markov Jump Processes in Continuous Time
Markov jump processes stochastic models used describe systems that transition between states at random times. These processes are characterized by the Markov property where the future state depends only on the current state, on the sequence of events that preceded
continuous- time MJP (t) has right-continuous, seecewise-constann paths and takes values in a countable state space X' over time interval [0, 77. The instantaneous probability rate of transitioning from state to x is defined as
f(z|=',t)
lim 0 At P Paup At At
(9)
where Prup (x, 0/', t') denotes the transition probability.
The evolution of the state probabilities (az t) is governed by the master equation
dpaa (x, 1)
Σ t2
((rlaa))srr(*) ()) f(x' Pra (z. (z.(t) 1)).
(10)
dt
For homogeneous MJPs with time-independent transition rates, the master equation in matrix form is
dpsor (I, .)
(t) PMr(()) F,
(11)
dt
with the solution given by the matrix exponential
Paur (t) PMur exp((F))
(12)
A.2 Stationary Distribution
The stationary distribution Pup of a homogeneous MJP a probability distribution over the state space X' that satisfies the condition Par F = O. This implies that the stationary distribution is left eigenvector of the rate matrix corresponding to the eigenvalue O.
A.3 Relaxation Times
The relaxation time of a homogeneous MJP is determined by its non-zero eigenvalues λ2 ......, daxk- ...
These eigenvalues define the time scales of the process: [Re(A2) |Re(As))=1.... ..... |Re(> xx) These time scales are indicative of the exponential rates of decay toward the stationary distribution. The relaxation time, which is the longest of these time scales, dominates the long-term convergence behavior. the eigenvalue corresponding to the relaxation time has a non-zero imaginary part, then this means that the system does not converge into fixed stationary distribution but that instead ends in a periodic oscillation.
A.4 Mean First-Passage Times (MFPT)
For MJP starting in state i X, the first pastaaa time to another state j E X' is defined the carliest time t at which the MJP reaches state j. given it started in state i. The mean first-passage time (MFPT) Tıy is the expected value of this time. For finite state, tmme-homgenncous MJP, the MFPT's can be determined by solving a series of linear equations for each state j, distinct from i, with the initial condition that Tu 0
To 0 1 Σ F FATTK 0, i
(13)
14

---

A.5 The Gillespie Algorithm for Continuous-Time Time Markov Jump Processes
The Gillespie algorithm (Gillespie, 1977) is stochastic simulation algorithm used to generate trajectories of Markov jump processes in continuous time. The algorithm proceeds as follows:
Algorithm 1 Gillespie Algorithm for Markov Jump Processes
INPUT: The intensity matrix the initial state distribution No, the starting time to and the end
time fend
2: Initialize the time t the starting time to
3: Initialize the system's state to initial state TT
4: While < fend do
5: Calculate the intensity À 11/F sa from state
6: Sample the time to the next event from an exponential distribution with rate
7: Update the time t <- +
8:
If 2 fend then exit loop Calculate transition probabilities
9:
-F F.,. for each possible next state j
10:
Set ps to zero because we allow for no self jumps
11:
Sample the next state s' from the distribution defined by p
12:
Update the system's s state $ f-
Record the state and time
14: End while
OUTPUT: The trajectory of states and times
B Synthetic Dataset Generation: Statistics and other Details
This section is a continuation of section 3.1 and provides more details on the generation of our synthetic training dataset. Additionally, we provide some statistics about the dataset distribution.
B.1 Prior Distributions and their Implementation
In this subsection we give additional details about our data generation mechanism.
Distribution over rate matrices. Our data generation procedure starts by sampling the entries fúg of the intensity matrix from the following beta distributions
((ffglp)) = (eta(py = (a, B)), with p(a) = Uniform( {1,2})
and p(3) Uniform((1,3,33 103). (14)
Both these discrete uniform distribution define the prior p(p,) = p(a)p(3).
The choices for and were made heuristically, obtain reasonable (i.e. varied) distributions the number of jumps (see e.g. Figure 5). We remark that we fixed this set of training distributions before evaluating the model on the evaluation sets, in order to prevent from introducing unwanted biases into the distribution hyperparameters by optimizing on the evaluation set.
Next we define the prior over the adjacency matrix as
5(A-) P(A) = - JJ Perces Reayl (. 0.5)
(15)
where d-)) labels the Dirac delta distribution and J denotes the matrix for which all off-diagonal entries are and the diagonal ones are O. Furthermore PentusRimyi labels the Erdös-Rényl model (Erdös and Rényi, 1959), for which each link is defined via independent Bernoulli variable, with some fixed, global probability p. here set to 1 state networks fully connected, whether the other 50 percent not.
2- Equation 15 indicates that (in average) 50 percent of our
Our motivation for this prior is that often happens real world processes that the intensity matrices are not fully connected. Let remark, however, that we only accept the Frdös-Rény' sample if the corresponding graph is connected that is, the system cannot get stuck into single state, Both these distributions implicitly define Pracs (F)A, pJ), for F.j agg fij.
15

---

40000
55000
20000
000e+
5090
20
20
20
30
Saater of Jammm
Naniles jurpee
Nexiter jurge
(a) 2D
(b) 3D
(c) 4D
20000
15400
10000
5000
5406
Nurcher praps
Narther of yraqs.
(d) SD
(c) 6D
Figure 5: Distributions of the number of jumps per trajectory. We used the same distributions as the training set and sampled up to time 10. The figures are based on 1000 processes with 300 paths per process.
Remark on generalization beyond prior rate distribution. We remark that while all entries of the intensity matrix seen during training lie on the interval [0, 1], the model can still predict intensities outside this interval. We empirically demonstrated that this in indeed the case on the widely different target sets of the experimental section, in the main text. The reason behind this that we normalize the maximum time among all input paths to be 1, and rescale the predicted intensities accordingly. Ultimately, what matters the difference among the rates (and therefore among the observation times) within the target time series. Our approach for sampling intensity matrices resulted in a vast variety of different processes.
The distribution of the number of jumps per trajectory is shown in figure 5 and that of relaxation times is shown figure 6.
Distribution over initial conditions. We choose half of our initial distributions in our synthetic ensemble to be the stationary distribution of the MJP Plue- The motivation for this is that it often happens that real life experiments produce very long observations of system in equilibrium. The second half of our initial distributions To are randomly sampled from Dirichlet distribution Dir(po). where we heuristically choose Pa = 50. In equations, we write
((TO) = 2 PMUP F)
2 Dir(po 50).
(16)
Distribution observation grids. In practice, the exact jump (i.e. transition) times not known. We therefore first generate observations of the state of the system on regular grid with a maximum of I. 100 points. We then randomly mask out some observations from this fixed regular grid, in order to make the model grid independent. Half of our (subsampled) observation grids are chosen to be regular, Lee they are strided with strides E {1, 2, 3, 4). The other half are chosen be irregular, through a Bernoulli filter (or mask) with Psurvival E {1/4, 1/2] applied to the base (L = 100) grid.
Distribution over noise process. Because real world data is often noisy also add noise to the labels. If a state observation is selected to be mislabeled, the new label is randomly chosen from a uniform distribution over all states. We investigate two different configurations in this project, one with 1% label noise (p, = 0.01) and one with 10% label noise (p,, 0.1).
MJP simulation. We sample the jumps between different states with an algorithm due to (Gillespie, 1977) (see A.5). We sample jumps between times 0 and 10 because almost all of our processes are in cquilibrium by then (see figure 6).
16

---

700
150
20C
Vregeney
130
832
200
20
20
20
20
20
40
20
30
50
Reassalis
Tiaee
Bedicsiio Tite
Didication Tixe
(a) 2D OP: 0% NCP: 2.6%
(b) 3D OP: 13.9% NCP: 4.1%
(c) 4D OP: 19.9% NCP: 4.2%
100
I0C-
300
120-
200-
100
50
Psleestion Time
Esducation Time
(d) 5D OP: 19.3% NCP: 4.5%
(c) 6D - OP: 18.8% NCP: 4.8%
Figure Distributions of the relaxation times. We also report the percentage processes that converge into an oscillating distribution (OP) and the percentage of processes that have a relaxation time which is larger than the maximum sampling time (NCP) of our training data (given by fend 10). The figures are based on 1000 processes.
Training Dataset Size The synthetic dataset on which our models were trained consists of 25k six-state processes, and processes of 2-5 states, resulting in a total size of 45k processes. For each of these processes we sampled 300 paths.
Distribution over the number of MJP paths (K). While we generate the data with 300 paths per process, we want to ensure that the model is able to handle datasets with less than 300 paths. For this reason, we shuffle the training data the beginning of every epoch and distribute it into batches with path counts 1, 11, 21, ,,, 300. We found that such a static selection of the path counts is better than random selection, because random selection leads to oscillating loss functions (because the model obviously gets a larger loss for samples with fewer paths), and thus training instabilities. Since we do not always select all paths per process but instead select a random subset of them, the data that the model processes changes during every epoch, which helps in reducing overfitting.
C How to use the Model: Inputs, Outputs and Rescalings
In this section, give details about the inputs and outputs of pretrained recognition model. We also comment on the internal rescalings done by the model, in order to be able to infer MJPs from time series with observation times of any scale.
C.1 Input
The model takes as input three parameters:
1. The observation grids (shape: [num_paths K, grid_size LJ): The observation times [Tki,..., Th) R
k= padded to the maximum length I..
2. The observation values (shape: houm_atths K, grid_size LJ): The noisy observa- tion values { = padded to the maximum length Note that these values are integers lying on the discrete set {0, 1....,,C--11.
3. The dimension e: The (a priori known) dimension of the process as an integer between
2 and C. If this dimension is unknown, the model returns a C' C' rate matrix whose rank might (approximately) be smaller than C, which indicates a hidden state-space of size smaller than C.
17

---

We recommend users to use the model only within its training range. That is, with up to a maximum of K 300 paths, and grids up a maximum of L 100 points.
C.2 Internal Rescaling
Internally, the model does the following:
1. It computes the maximum observation time:
max TRa Tha Taak-z
(17)
2. It normalizes the observation times between and 1:
(TRI,-..,TT)E {TaI.-... Thl ) k=1/ Tmxx-
(18)
3. It computes the inter-event times ÂTk = Th,++1 - Tkis for I,...,KK
4. It transforms the observation values to one-hot-encodings.
5. It predicts the (normalized) off-diagonal elements the intensity matrix and variance matrix well the initial distribution (here we working with the maximum supported dimension, that is C)..
6. It rescales back the estimates the original time scale:
intensity matrix F intensity matrix P //Taxx:
(19)
Var F < Var FarF//r F/Tmax TTaxx
(20)
Note that, as we empirically demonstrated in the paper, this rescaling procedure allows us to work with real-world MJPs of arbitrary time scales. For example, the time scales for the switching ion channel dataset were more than 500 times smaller than the time scales in our training datasct.
C.3 Support for varying State Space Sizes
We now elaborate how the model can deal with processes whose state spaces have sizes C. We arranged all the target rate matrices F within our training dataset, for MJPs with state spaces of size < C, to be the leftmost block diagonal e x matrix within C' x C' matrix of zeros, so that the redundant matrix elements are always zero. As can read from equation (6) of the main text, we train FIM to predict zeros for those redundant matrix elements.
In practice, however, our trained FIM does not exactly predict zeros for those redundant matrix elements. In our experiments, the user knows a priori the number of states c of the hidden process, so we explicitly set the redundant matrix elements to zero, and only then compute the corrected diagonal (i.e. the normalization) of the output rate matrix.
We afterwards select the entries of the predicted intensity matrix and variance matrix well as the first e entries of the predicted initial distribution.
We refer the reader to our library for additional details.
C.4 Output
The output the model consists three parameters:
1. The intensity matrix F (shape: [c, cl).
2. The variance matrix Var F (shape: [c,e]).
3. The initial distribution mo (shape [c]).
18

---

D Model Architecture and Experimental Setup
In this section we provide more details about the architecture of our models and the hyperparameters.
D.1 Model Architecture
Path encoder 1/1 We evaluated two different approaches for the path encoder th The first approach utilizes a bidirectional LSTM (Hochreiter and Schmidhuber, 1997) as 1/1, while the second approach employs a transformer (Vaswani et al., 2017) for 1/. The time series embeddings are denoted by h&0 (see Equation 4). The input the encoder th1 is (x41= Tkl, X Tkt), where Tki = [Tkt, 5kt),
5yd Fkl
T(k- 1yr. and X E {0, 1] is the one-hot encoding of the system's 's state.
Path attention network Qz. We tested two approaches. The first approach uses classical self- attention Vaswani et al. (2017) and selects the last embedding. For the second approach used an approach we denote as learnable query attention which is equivalent to classical multi-head attention with the exception that we do not compute the query based on the input, but instead make it a learnable parameter, i.e.,
MultiHead (Q, K, V) = Concat (hancctthhad........ head, Attention(Q:; HI.K W H..kWY),
(21)
(22)
where HiKK E RXXoodd denotes concatenation of bi....,hk, W.K. W. Redmmedet d and Qi E Rxxdx is the learnable query matrix. The output dimension of the learnable query attention is therefore independent of the number of input tokens.
D.2 Experimental Setup
Hyperparameters were tuned using a grid search method. The optimizer utilized was AdamW W (Loshchilov and Hutter, 2017), with learning rate and weight decay both set 1e-4. A batch size of 128 was used. During the grid search, we experimented with the hidden size of the path encoder (164, 128, 256, 512)), the hidden size of the path attention network ([128, 256]), and various MI.P architectures for Φ1, 02, and db ([[32, 32], [128, 128]]).
Hyperparameter tuning:
Training procedure: All models were trained on two A 100 80Gb GPUs for approximately 500 epochs approximately 2.5 days on average per model. Early stopping employed as the stopping criterion. The models were trained by maximizing the likelihood.
Final model parameters: The final models (FIM-MJP 1 1% Noise and FIM-MJP 10% Noise) have the following hyperparameters: Path encoder hidden (atnn_siiee_______ = 256 (the final models used a BiLSTM); Path attention network Ωii 16, dę 128 (the final models used the learnable query approach); d1, 62, Φy = [128, 128].
Pretrained models: Our pretrained models are also available online
E Ablation Studies
In this section, we study the performance of the models with different architectures. Additionally, we study the behavior of the performance of the models with respect varying numbers of states and varying number of paths.
E.1 General Remarks about the Error Bars and Context Number
If the evaluation set is larger than the optimal context number "aKmassmaaa we split the evaluation set into batches and give these to the model independently (because the model does not work well to give the model more paths than during training, see table 8). Afterwards, we compute the mean of the predictions among the batches and report the mean RMSE of the intensity entries (if the ground-truth available). This makes easier compare our model against previous works which have also used the full dataset to make predictions. Interestingly, we find that the RMSE of this averaged prediction
hhttps://gitt httpg/twww.ccoocceoeee////aaaaaa
19

---

is often sigmificantly better than the mean RMSE among the batches. For example for the DFR dataset the RMSE of the averaged prediction 0.0617, while the average RMSE of the batches is 0.122. If the dataset has been split into multiple batches, we report the RMSE together with the standard deviation of the RMSE among the batches. The reported confidence is the mean predicted variance of the model (recall that we are using Gaussian log-likelihood during training).
E.2 Performance of the Model by varying its Architecture
The ablation study presented in Table 5 evaluates the impact of different model features on the performance by comparing various combinations architectures and attention mechanisms with varying numbers of paths, and their corresponding RMSE values. The study examines models using a BiLSTM or Transformer, with and without self-attention and learnable query attention, across I, 100, and 300 paths. The results indicate that increasing the number of paths consistently reduces RMSE (see section E.4 for more details), demonstrating the benefit considering more paths during training. Specifically, using BiLSTM with learnable query attention achieves an RMSE of 0.193 + 0.031 with a single path, significantly improving to 0.048 = 0.011 with 100 paths, and further to 0.0457 + 0.0 with 300 paths. Similarly, Transformer with learnable query attention shows an RMSE of 0.196 + 0.031 for single path, 0.049 + 0.011 for 100 paths, and 0.0458 + 0.0 for 300 paths. The inclusion of self-attention in the Transformer models slightly improves performance, with the best RMSE of 0.0459 + 0.0 achieved when both self-attention and learnable query attention are used with 300 paths. In this case since many of the processes contain one path it beneficial to use the learnable query attention over the standard self-attention mechanism.
# Paths
BiLSTM
Transformer
Self Attention
Learnable Query Attention
RMSE
1
0.193 = 0.031
1
0.196 = 0.031
1
0.197 + 0.015
100
0.048==0000
100
0.049 1049:00000 0.011
100
0.054
300
0.04571:..0
300
0.0458 + 0.0
300
0.0459 0.0
Table 5: Comparison of model features with different number of paths and their RMSE. This table presents ablation study comparing the performance of models using BiLSTM and Transformer architectures, with and without self-attention and learnable query attention, across different numbers of paths (1, 100, and 300). The performance is measured by the Root Mean Square Error (RMSE), with lower values indicating better model accuracy. The study highlights that both the architectural choices and the number of paths significantly impact model performance, with the best results achieved using a combination of attention mechanisms and a higher number of paths.
Figure presents a series of line plots illustrating the impact of different hyperparameter settings on the RMSE of the model. The first subplot shows the RMSE 4h path encoder, with hidden sizes 64, 128, 256, and 512. The RMSE increases as the hidden size increases, with the lowest RMSE observed RMSE as a function of the architecture size of Φi, comparing two architectures: [2x32] and [2x128) 128]. The RMSE decreases as the architecture size increases, indicating better performance with larger architecture size for d1. The third subplot examines the RMSE based on the architecture size of 62, with two architectures tested: [2x32] and [28]. There no significant difference in RMSE between the two sizes, suggesting that the choice of architecture size for 62 does not markedly affect model performance. The fourth subplot investigates the RMSE as a function of the hidden size the Ω component, with hidden sizes 128 and 256 tested, and results shown for different a hidden sizes (64, 128, 256, and 512). The RMSE remains relatively stable across different hidden sizes of Ω, with slight variations observed depending some components, such as 1/1 and Φ1, are more sensitive to changes in hyperparameters, emphasizing the importance of selecting appropriate hyperparameters to optimize model performance.
function of the hidden size of the
hidden size of 256. The second subplot displays the
the hidden size of v/, Overall, the plots highlight that
20

---

0.036
start 1-1
0.054
6 16
sizel eh Jm22
R0USE
0,002
sizel el Jm012
0.050
0.046
0'
512
2602
2x:28
2x125
Stze
Oy
Figure 7: Impact of Hyperparameters on RMSE. The figure shows four line plots illustrating the effect of hyperparameters on model RMSE. The first plot shows RMSE increases with larger 10 hidden sizes, being lowest at 256. The second plot indicates lower RMSE with larger Ø1 architecture size ([2x 12 The third plot shows minimal RMSE impact from Φ architecture size. The fourth plot shows RMSE stability across different Si hidden sizes, with slight variations based on th. This highlights the importance of tuning 1 and Φ1 for optimal performance.
1% Noise Data 10% Noise Data
FIM-MJP 1%
0.046
0.199
FIM-MJP 10%
0.096
0.087
Table 6: Performance of FIM-MJP 1% and FIM-MJP 10% synthetic datasets with different noise levels. We use weighted average among the datasets with different numbers of states to compute a final RMSE.
Table 6 compares the performance of two models, FIM-MJP 1% and FIM-MJP 10%, on synthetic datasets with noise levels of 1 %% and 10%, measured in terms of RMSE. For datasets with 1 1% noise, the FIM-MJP 1% model achieves an RMSE of 0.046, indicating good performance, but its RMSE 10% noise data, showing decreased performance with higher noise. Conversely, the FIM-MJP 10% model, trained with 10% noise data, has an RMSE of 0.096 on 1 %% noise data, higher than the FIM-MJP 1% model the same data, but achieves lower RMSE of 0.087 on 10% noise data, demonstrating better performance under high noise conditions. This indicates that the FIM-MJP 10% model is more robust noise, maintaining consistent performance across varying noise levels, while the FIM-MJP 1% model excels low noise environments but struggles with higher noise. The results highlight the importance of training with appropriate noise levels to ensure robust model performance across different noise conditions.
increases significantly to 0.199
E.3 Performance of the Model with varying Number of States
We compare the performance of our models on processes with varying number of states. Note that our model always outputs 6 dimensional intensity matrix. However, in these experiments we only use the rows and columns that correspond the lower-dimensional process. This improves the comparability between different dimensions as lower-dimensional processes obviously have many zero-entries in their intensity matrix which would make it easier for the model achieve good RMSE score.
It can be seen in Table 7 that the moleil performs well among all different dimensions. As expected, dimensiinai processes seem to be easier for the model. Additionally, Table 7 shows the performance of model which has only been trained on six-state processes. The performance of this native six-state-men model for six number of states is very similar to the multi-state-modde which shows that having more states during training does not reduce the single-state performance. As expected, the performance the six-state model on processes with lower numbers of states is significantly worse, but still better than random.
E.4 Performance of the Model with varying Number of Paths during Evaluation
One of the advantages of our model architecture that can handle arbitrary number of paths. We therefore use model that trained at maximum 300 paths and assess its performance
21

---

States
Multi-State RMSE
Multi-Staan Confidence
6-State RMSE
6-State Confidence
0.026
0,028
0.129
0.056
3
0.037
0,030
0,113
0,049
4
0,046
0,037
0.087
0.046
5
0.054
0.040
0.066
0,041
6
0.059
0.044
0.059
0.044
Table 7: Performance of the multi-state and six-state models (which has only been trained processes with six states) on synthetic test sets with varying number of states
with varying number of paths during evaluation. The results are presented in Table 8. When being inside the training range, the performance and the confidence of the model goes down as the model is given fewer paths per evaluation, which is to be expected. Interestingly, the performance of the learnable-query (LQ) model peaks at 500 paths instead at 300, which was the maximum training range. One possible explanation for this might be that we still close enough to the training range while being able to use the full data (note that the dataset contains 5000 paths which not divisible by 300, we have leave some of the data out). Going too far beyond the training range does however not work well, for example processing all 5000 paths at once leads to very poor performance, although the model (falsely) become very confident. Another insight from this experiment is that the self-attention (SA) architecture behaves significantly worse when going beyond the maximum number of paths that was during training. This is another reason why chose the (LQ) architecture over the (SA) architecture for the final version of our model.
#Paths during Evaluation
RMSE (LQ)
Confidence (LQ)
RMSE (SA)
Confidence (SA)
0.548 = 0.067
0.838
0.579 + 0.074
0.898
30
0.074 = 0.081
0.263
0.075 + 0.070
0.264
100
0.061 = 0.039
0.143
0.060 = 0.035
0.142
300
0.056 + 0.023
0.089
0.059 + 0.024
0.085
500
0.053 + 0.014
0.069
0.074 0.021
0.061
1000
0.067 + 0.012
0.037
0.229 + 0.025
0.029
5000
0.818 + 0.000
0.000
2.135 + 0.000
0.000
Table 8: Performance of FIM-MJP 1% given varying number of paths during the evaluation on the DFR dataset with regular grid. (LQ) denotes seumable-query-attentio (see section D.1), (SA) denotes self-attention.
F
Additional Results
This section contains more of results which did not fit into the main text. We begin this section by providing more details on the Hellinger distance which we used as a metric to assess the performance of our models. Afterwards, provide more results and background DFR datasets. Additionally, we introduce two two-state MJPs, given by the protein folding datasets (F.5) and the two-mode switching system (F6), which we use evaluate our models and compare it against previous works.
the ADP, ion channel and
F1 Hellinger Distance
Real-world empirical datasets of MJPs provide no knowledge of ground truth solution. For this reason we present a new metric that can be used to compare the performance of the inference of various models based only the empirical data. Our metric of choice is the Hellinger distance which a measure of the dissimilarity between two probability distributions. Given two discrete probability distributions P (P1.....Ps. and (g1,..., 9k), the Hellinger distance is defined as
1 H(P
Σ (VD: 24.)2.
(23)
/2
22

---

For our empirical cases, the class probabilities of the discrete probability distributions are not known explicitly. We therefore approximate them by using the empirical distributions, given by the (normalized) histograms of the observed states at the observation grids.
We test this approach on the DFR process by first sampling specified number of paths for the potential V = 1 using the Gillespie algorithm, which we then consider as the target distribution. Counting states among the different paths then yields histograms of the states for every time step. We repeat the same procedure for different choices of V. Afterwards we compute the Hellinger distance between the newly sampled histogram and the target distribution for every time step. Figure 8 shows that the distance indeed goes down as we approach the target distribution, which provides heuristic evidence of the effectiveness of our metric. The Hellinger distances for various models are shown in Table 2 and Table 9.
As see, FIM-MJP performs as well (and sometimes better) the current otate-of-the-art model NJU
Averagy Helliggr Distaner
0.05
Muster Paths
0.04
1000 Paths
0.03
0.02
0.01
0.00
0.0
0.2
0.4
0.6
0.8
1.0
Figure 8: Tver-Average Hellinger distance for varying potentials on the DFR. The plot shows the Hellinger distance target dataset that was sampled from DFR with V = 1 on a grid of 50 points between O and 2.5. The means and standard deviations were computed by sampling 100 histograms per dataset. As expected, the distance decreases the voltage gets closer to the voltage of the target dataset. We also remark that the scale the distances gets smaller as one takes more paths into account and converge the distance of the solutions of the master equation.
Dataset
NeuralMJP
FIM-MJP 1% Noise
FIM-MJP 10% Noise
ADP
1.38 + 0.52
1.39 + 0.47
1.35 = 0.42
Ion Channel
0.48 0.02
0.41 + 0.02
1.78 + 0.03
Protein Folding
0.015 + 0.015
0.014 + 0.014
0.024 + 0.026
DFR
0.30 0.06
0.27 + 0.06
0.28 = 0.06
Table 9: Comparison of the time-average Hellinger distances for various models. We used the same labels as NeuralMJP to make the results comparable. The errors are the standard deviation among 100 sampled histograms. The target datasets contain 200 paths for ADP, 1500 paths for Ion Channel, 2000 paths for Protein Folding and 1000 paths for the DFR. The distances are reported in a scale le-2. We remark that the high variance of the distances the Protein Folding dataset is caused by the models performing basically perfect predictions, which causes the oscillations to be noise. We verified this claim by confirming that the distances of the predictions of the models are small as the distance of the target dataset to additional simulated data.
F.2 Alanine Dipeptide
We use the dataset of Husic et al. (2020), which models the conformal dynamics of ADP, for evaluating our model. This dataset was provided to us via private communication. The dataset consists of 9800 paths on grids of size 100 and has the sines and cosines of the Ramachandran angles as features:
23

---

sin db cos y, sin 6 and cos Φ. We use KMeans to classify the data into states. The reason why we did choose GMM as for the other datasets is that we could initialize KMeans with hand-selected values to try to achieve similar classification like those learned by NeuralMJP (Seifner and Sánchez, 2023), see Figure 9. Still, the classification is very different and thus also leads to very different results (see Table 10). We use 9600 paths to evaluate our models. Our results are shown in Table 10. Table reports the stationary distributions and compares them to previous works, while Table 12 reports the ordered time scales.
II
III
III
VI
VI
Figure Comparison of the classifications between KMeans (left) and NeuralMJP (right).
Model
Intensity Matrix
61.32
53.15
0.19
7.89
0.06
0.02
47.29
-50.37
0.05
11.97
0.04
0.01
NeuralMIP
16.81
0.16
0.22
0.31
0.2
-3.86
2:96
1.13
1.73
0.46
0.66
18.78
22.76
-59.35 + 2.11
48.72 + 1.90
0.33 + 0.08
10.14 + 1.47
0.09 + 0.08
0.07 à 0.07
50.54 1: 3.25
-57.62
2.99
0.44 1: 0.04
6.44 1.18
0.09 0.09
0.10::0.0
FIM MJP 1% Nuise (NeuralMJP Label))
38.31 4.63
0.40 0.08
33.71
0.10
4.97
49.14
4.80
13.16
7.36 122.66
0.21 4 0.19
0.30
0.32
0.25 1: 0.27
0.43 1: 0.60
0.20 1 0.24
0.30 1: 0.33
-2.40 3.06
1.23 : 1.69
0.44 0.46
1.12 1.56
4.79
4.58
40.51
3.82
0.3
7.96 1.09
0.15
0.22 0.11
FIM-MIP 10% Noise (NeurilMIP Labels)
39.18 5.42
0.27 - 0.04
28.24 à 4.14
0.44
58.36 à 8.72
13.05 1.66
129.02 10.51
11.35 1.81
0.32 0.007
1.1:0.11
0.68 0.28
1
9.61 1:
9.32 1: 6.63
5.53 1: 3.97
4.36 3.009
48.11 29.51
14.3 9.01
6.72
5.90
115.442 8.87
8.73
0.23 0.24
157.16 13.99
165.37 13.64
6.67 = 1.78
L.17 0.24
0.22 0.16
0.14
22.26 + 3.68
9.84 3.100
375.77 = 20.00
342 13 19.80
0.71 0.71 0.66
0.84 28420.68 0.65
FIM-MP 1% Nosse (KMeans Labels)
0.93 0.15
1.37 D.16
305.86 20.47
308.48 088..48. +: 20.30
0.25 0.19
0.07 = 0.09
0.28 0.33
1.14
0.28 0.38
0.18 0.23
4.81 = 27:33
-6.44 9.08.
94.45
91.38 +
+ 1.98 + 0.49
0.84 + 0.15
0.32
0.29
184.86 = 20.63
190.00 00 19:41
0.49 + 0.23
0.84 0.32
1.93 0.93
FIM-MUP Nosse (KMeans Labels)
5.93 1.57
0.74
13.71
2.48
266.40 +: 18.43
241.54 11.4442 17.98
0.85 0.18
4.48
0.52
3.45 1.82
17.28
11.78
42.3
6.83
2.43 2 0.50
7.14 + 3.09
6.11 + 2.37
6.62 - 2.24
16.38
7.84
38.66
5.39
Table 10: Comparison of intensity matrices for the ADP dataset. The time scales are in nanoseconds.

---

PROBABILITY PER STATE
I
I1
III
IV
V
VI
VAMPNETS
0.30
0.24
0.20
0.15
0.11
0.01
NEURALMJP
0.30
0.31
0.23
0.10
0.05
0.01
FIM-MJP 1% NOISE
0.28
0.28
0.24
0.07
0.10
0.03
FIM-MJP 10% NOISE
0.30
0.30
0.31
0.06
0.01
0.02
Table 11: Comparison of the stationary distribution on the ADP dataset of FIM-MJP, VAMPnets Mardt al. (2017) and NeuralMJJ (Seifner and Sánchez, 2023). The states are ordered such that the protein conformations associated to given state are comparable in both models. We use the labels of MJuralMJP to evaluate FIM-MJP.
RELAXATION
SCALES ns)
VAMPNETS
0.008
0.009
0.055
0.065
1.920
GMVAE
0.003
0.003
0.033
0.065
1,430
MSM
1.490
NEURALMJP
0.009
0.009
0.043
0.069
0.774
FIM-MJP 1% NOISE (NEURALMUP LABELS)
0.008
0.009
0.079
0.118
0.611
FIM-MJP 10% NOISE (NEURALMJP LABELS)
0.007
0.011
0.019
0.038
0.091
FIM-MUP % NOISE (KMEANS LABELS)
0.001
0.003
0.046
0.142
0.455
FIM-MJP 10% NOISE (KMEANS LABELS)
0.002
0.004
0.018
0.034
0.070
Table 12: Relaxation time scales for six-state Markov models of ADP. The time scales are ordered by size and reported nanoseconds. VAMPnet results are taken from Mardt et (2017), GMVAE from Varolgünes et al. (2019), MSM from Trendelkamp-Schroer and Noé (2014) and NeuralMJP from (Seifner and Sánchez, 2023).
F.3 Ion Channel
We consider the is observation window that has been used in (Köhs al., 2021) and (Seifner and Sánchez, 2023) and split it into 50 paths of 100 points. This dataset was provided to via private communication. We then apply a Gaussian Mixture Model (GMM) classify the experimental data into discrete states as shown in figure 10.
6
Current PAA]
Observations
GMM Classification
2-
0-
0.15
0.20
0.25
0.30
0.35
0.40
Time [s]
Figure 10: Classification of the ion channel dataset into states.
The predictions our models and NJuralMJP are shown table 13. Table 14 reports the stationary distributions and Table 15 reports the mean fist-passage times.
25

---

Model
Intensity Matrix
-57.73
55.81
1.93
NeuralMIP
102.13
-306.93
204.81
26.
FIM-MIP 1% Noise (NeuraiMIP Labels)
110.65
334.06
223. 00
0.78
31.53
- -3.30
-92.63
85.80
6.79
(NeuralMIP Labels)
2.86
32.72
35.58
-11.337
114.65
1.73
FIM-MP 1% Noise (GMM Labels)
271,88
716.32
444.64
0.56
49.69
50.25
104.01
FIM MIP 10% Noise (GMM Labels)
82.72
2215588
132.86
2.89
40.29
-43.18
Table 13: Comparison of intensity matrices for the ion channel dataset. We cannot report error bars here because the dataset is small that it gets processed single batch.
BOTTOM
MIDDLE
ТOP
KÓHS ET AL. (2021)
0.17961
0.14987
0.67052
NEURALMMIP (1 SEC)
0.17672
0.09472
0.72856
FIM-MUP 1% NOISE (NEURALMUP LABELS)
0.18224
0.10156
0.71621
FIM-MJP 10% NOISE (NEURALMUP LABELS)
0.14229
0.23090
0.62682
FIM-MJP 1% NOISE (GMM LABELS)
0.19330
0.08124
0.72546
FIM- MJP 10% NOISE (GMM LABELS)
0.17348
0.19610
0.63042
Table 14: Stationary distribution for the switching ion channel process when trained the one-second window.
KOHS ET AL. (2021)
NEURALMJP
FIM-MJP 1% NOISE
(NEURALMJP LABELS)
/
BOTTOM
MIDDLE
TOP
BOTTOM
MIDDLE
TOP
BOTTOM
MIDDLE
TOP
BOTTOM
0.
0.068
0.054
0.
0.019
0.031
0
0.017
0.027
MIDDLE
0.133
0.
0.033
0.083
0.
0.014
0.068
0
0.012
TOP
0.181
0.092
0.119
0.038
0.
0.098
0.031
FIM-MIP 10% NOISE
FIM-MJP 1 NOISE
FIM-MJP 10% NOISE
(NEURALMJP LABELS)
(GMM LABELS)
(GMM LABELS)
Tis
BOrTOM MIDDLE
TOP
BOTTOM
MIDDLE
TOP
BOTTOM
MIDDLE
TOP
BOTTOM
0.013
0.026
0
0.009
0.016
0
0.011
0.022
MIDDLE
0.063
0
0.016
0.036
0
0.007
0.045
0.013
ToP
0.086
0.029
0.055
0.02
0
0.065
0.024
0
Table 15: Mean farst-passage times of the predictions of various models on the Switching Ion Channel dataset. We compare against (Köhs et al., 2021) and NeuralMJP (Seifner and Sánchez, 2023). Entry j in row is mean first-passage time of transition - j of the corresponding model.
E.4 Discrete Flashing Ratchet
We use the same datasets that were used Seifner and Sánchez (2023), which contains 5000 paths on grids of size 50 that lie between times 0 and 2.5. This dataset was provided to us via private communication. We used 4500 paths to evaluate our model. The predicted intensity matrices for the DFR and the ground truth shown in table 16.
26

---

Mode
Intencity Matrix
1..97
0.61
0.37
1.65
-3.26
0.61
Ground Truth
2.72
0.87
-3
3J
-1.88 + 0.09
0.32 0.06
0,31 + 0.05
0.99 ds 0.09
0.03
0.00
0.03 + 0.01
1: 0.12
-3 0.113
0.57 1 0.10
0.06
:: 0.01
1.04 - 0.14
0.05 0.0.1
FIM-MIP 18 Noise
0.97 +: 0.10
0.00 0.01
1.60 0.19
0,04 + 0.01
-5.00
3300 + 0.17
0.12
0.99 0.10
0.97 4 0.09
0.05 +: 0.01
0.96 0.12
0.05 1 0.01
0.95 1 0.15
-3.05 + 0.18
1.01 1 0.111
0.00 0:01
+
0.10
1,03
33,00
-1.61 0.10
0.46 0.07
4 0.03
0.88 à: 0.10
0.02 0.00
0.02 4:
HIM-MIP 10% Noise
0.87 4: 0.12
2.68 0.34
0.03 4: 0.01
1.47 0.17
-4.93 0:49
0.03 + 0.00
-2.53 0.20
0.06 = 0.01
0.80 4.00
0.06 606:::00
0.80 4: 0.10
0.65 0.25
0.04 4: 0.01
0.84 4: 0.12
0.03 0.00
0.84 0.11
-2.61
0,87 1 0.10
0.93
Table 16: Comparison of intensity matrices for the DFR dataset on the irregular grid.
F,5
Modeling Protein Folding through Bistable Dynamics
The work of Mardt et al. (2017) introduces simple protein folding model via a 10th step trajectory simulation in a 5-dimensional Brownian dynamics framework, governed by:
dx((t) VU(((()) (2)(())
with the potential (xx) being dependent solely on the norm r(x) |x| as follows:
U { -25((((z) 3]2
if((() <3
0.5/r(2) [r()) 3/2 if rff(() 2
This model exhibits bistability in the norm r(r), encapsulating two states akin to the folded and unfolded conformations protein.
We use the dataset Seifner and Sánchez (2023) and apply a Caussian-Mixture--e the dataset into two states. The decision boundary of the classifier seems to be based the absolute absolute value of the radius, namely the classifier seems to classify all states with a radius smaller than approximately 2 into the lower state (see figure 11).
to classify
Seifner and Sánchez (2023) generated 1000 trajectories, each with 100 steps after 1000-step burn-in period. We used 900 paths to evaluate our model. The results are shown in table 17. Table 18 compares the stationary distributions obtained from our models to the ones from Mardt et (2017) and Seifner and Sánchez (2023).
27

---

Low STD
HIGH STD
MARDT ET AL. (2017)
0.73
0.27
NEURALMUP
0.74
0.26
FIM-MJP 1% NOISE (NEURALMJP LABELS)
0.73
0.27
FIM-MJP 10% NOISE (NEURALMJP LABELS)
0.62
0.38
MJP 1% NOISE (GMM LABELS)
0.70
0.30
FIM-MJP 10% NOISE (GMM LABELS)
0.65
0.35
Table 18: Stationary distribution of the model predictions the protein folding dataset
6
5-
4-
Radius
3-
2
0
0
100
200
300
400
500
Time
Figure 11: Classification of the protein folding dataset into Low and a High state. The GMM- Classifier has learned decision boundary close to the radius 2.
Low STD
HIGH STD
HIGH STD - Low STD
NEURALMJP
0.028
0.085
FIM-MJP 1% NOISE (NEURALMJP LABELS)
0.019 + 0.003
0.054 + 0.011
FIM-MJP 10% NOISE (NEURALMJP LABELS)
0.034 + 0.005
0.055 + 0.008
FIM-MJP 1% NOISE (GMM LABELS)
0.054 1 0.005
0.154 1: 0.018
FJM-MJP 10% NOISE (GMM LABELS)
0.050 = 0.006
0.093 1 0.011
Table Predicted transition rates on the protein folding dataset
F66 A Toy Two-Mode Switching System
In their study, Köhs et al. (2021) produced time series derived from the trajectory of a switching stochastic differential equation
dy(t) a=(0) v(t)) + 0.5dW (t).
with parameters (i = 02 = 1.5, 31 -1 and 832 = For a concise overview the generation process, the reader is directed to (Köhs et al., 2021) for comprehensive details. We use the same dataset that was generated (Seifner and Sánchez, 2023) using the code of (Köhs et al., 2021) which contains 256 paths of length 67 to evaluate model. Our results are shown in table 20.
28

---

BOTTOM TOP
ToP - BOTTOM
GROUND TRUTH
0.2
0.2
KÖHS AL. (2021)
0.64
0.63
NEURALMUP
0.19
0.36
FIM-MJP 1% NOISE
0.43
0,25
FIM-MJP 10% NOISE
0.23
0.15
Table 19: Two-Mode Switching System transition rates. We do not report error bars here because the dataset is small that it runs single batch.
F.7 Initial Distributions
For completeness, we report in this section the initial distributions predicted by FIM-MJP on various datasets well as the heuristic initial distribution (which is computed simply by counting the numher of state occurrences at the first observation). We observe that FIM-MJP typically captures the initial distribution quite well. An exception is the Two-Mode Switching System for which FIM-MJP falsely predicts a non-zero probability of the first state. This might happen because we did not capture this case in our training distribution which could be improvement for future work.
DATASET
PREDICTED To
HEURISTIC TO
DER =1
[0.22, 0.15, 0.1, 0.16, 0.17]
[0.3, 0.14, 0.06, 0.2, 0.17 0.14] [0.24, 0.08, 0.68
IONCH
[0.14, 0.75
ADP
[0.34, 0.3, 0.23, 0.110 0.02, 0.0]
[0.33, 0.29, 0. 0.25, 0.11, 0.02, 0.0]
TWO-MODE SYSTEM
00.32, 0.68
[0.0, 1.0] 00.73, 0.27]
PROTEIN FOLDING
00775, 0.25
Table 20: Comparison of the predicted initial distribution of the model versus the heuristic initial distribution of various datasets.
29

---

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: Yes, the claims of the abstract and in the introduction are shown in the contributions of sections and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 be 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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, Section 6 is devoted the limitations of our approach.
Guidelines:
The answer NA means that the paper has no 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 a separate "Limitations" section 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 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 on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be 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 images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged 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.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
30

---

Justification:
We do not present any theoretical results in this work.
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 the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but 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
Question: 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 the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We share our code and trained models, as well as the synthetic data used to evaluate our models The synthetic training data is however too large be published, but can be regenerated with our code. The relevant hyperparameters are stated in Appendix B. Lastly, we cannot share all evaluation data because we do not own it. We provide references in the Acknowledgment to the data owners.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be 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 to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility 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, or other means that are appropriate to the research performed.
While NeurIPS 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 to reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(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-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 particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
31

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: Our code and models openly available. For the availability of the data, please refer to the above point.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips..
Saaiides/CodSSumiiie
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for 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. wttp:ccccpccccccc..ecece..ccccpa gww.d/gg///e/eeeeeee CCodeeubeesseeeeeeeeee issiissiionn iessionnoiccyy for more details.
The authors should provide instructions on 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 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
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: All the training details are described in section D.2 in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: Our results are reported with error bars possible.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
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 with given experimental conditions).
32

---

The method for calculating the error bars should be explained (closed form formula, call 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 the standard error of the mean.
It 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 figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The resources that are used for computation are described in section D.2 the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or 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 in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: httpp//////////./////// cc/ cc/public//eieecccccccccc Coublic/Ethiin htte//hhiisiiiieeeeeeeeeee csGuidelines!
Answer: [Yes]
Justification:
Our conducted research does not clash with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS 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 socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is fundamental research that has no impact on society.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
33

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 not 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, it 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 as 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 to 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: Our paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 usage guidelines 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
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?
Question:
Answer: [Yes]
Justification: We have credited the owners of the evaluation data and referenced the related work on which this project has been built on.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
34

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithheeeee thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The new asset of this paper are the code and the models which well documented.
Justification:
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
This paper did not involve crowdsourcing or reasearch with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor 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 the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper did not involve crowdsourcing or reasearch with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending 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.
35

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
36


Reviews:
[{"review_id": "fbhbuJNJH6", "reviewer": "Reviewer_EGs8", "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.", "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)?", "limitations": "Limitations are adequately discussed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First of all, we would like to thank the reviewer for taking the time to review our work and for the helpful comments and questions.\n\n**@W1**: We thank the reviewer for pointing this reference out. We will include it into our related work section, together with other representative references on amortized inference and offline learning. \n\n**@Q1**: Indeed, we expect that if one leverages FIM to infer the rate matrix from an empirical process whose path distribution lies well outside our synthetic training distribution, one will obtain poor rate estimates.  As we commented in the limitation section, Figure 4 represents such an example. \n\nAs implied by the reviewer, there are two paths one can follow to improve the performance of FIM in such cases: \n\n(i) scaling the parameter count of FIM, and \n\n(ii) widening the MJP training distribution. \n\n*Regarding parameter scalings*, we have empirically observed that, for our current 45K MJP synthetic dataset, increasing the parameter count does not necessarily improve the performance of FIM. This can be seen, for example, in Figure 7 of Appendix D.2, in which, broadly speaking, the performance of different FIM architectures, with respect to our synthetic distribution, is statistically similar. We invite the reviewer to read our discussion in lines 592-606 for details.  \n\n*Regarding the MJP training distribution*, we do expect that widening our synthetic training distribution will improve the performance of FIM in e.g. the cases reported in Fig. 4. Future work will explore how to define wider synthetic distributions that contain rare or exceptional processes, and how scaling of the FIM parameter count in such cases affects their encoding.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VO8MZJcUDJ", "reviewer": "Reviewer_sVKK", "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.", "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?\n\n2. 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?\n\n3. 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?\n\n4. 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.\n\n5. 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. \n\n6. How accurate is the prediction of the initial state, $\\pi_0$? \n\n7. 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$?\n\n8. 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?", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Let us thank the reviewer for the detailed review, as well as for the proposed questions and weaknesses. However, let us remark here that some of the proposed weaknesses are somewhat vague. Accordingly, below we ask the reviewer to be more specific in (some of) their remarks. \n\n**@W1** (On comparisons against state-of-the-art): We tackle the problem of MJP inference from noisy data for processes defined on bounded state spaces. A prominent application of these processes is the modeling of conformational dynamics in proteins and molecules, and in our paper we considered two such examples: the conformational switching of ion channels and the molecular dynamics of a 22-atom molecule. We compared FIM against\n\n- The VampNets model of Mardt et al. (2017): the go-to, state-of-the-art deep-learning approach for modeling molecular dynamics via discrete-time Markov chains;\n- The SDiff model of Koehns et al. (2021): a recent variational model that has been applied to the modeling of switching ion channels; and \n- The NeuralMJP model of Seifner & Sanchez (2023): to the best of our knowledge, the most recent, state-of-the-art neural network model available online, which leverages neural variational inference and neural ODEs to infer hidden MJPs.\n\nWe therefore believe that we have compared our methodology against a set of state-of-the-art models. Nevertheless, we would be very happy if the reviewer could point out any state-of-the-art model we did not include into our related work.\n\n**@W2a** (On the systematic demonstration of the capabilities of FIM): We have empirically demonstrated that one and the same pretrained FIM can be used to estimate stationary distributions, relaxation times, mean first-passage times, time-dependent moments and thermodynamic quantities (i.e. the entropy production) from five very different, noisy and discretely observed MJPs, taking values in state spaces of different dimensionalities, without the need of any fine-tuning. What is more, FIM was shown to perform on par with all the baselines.\n\nWe believe that these tests cover well the capabilities of FIM. We therefore would very much appreciate if the reviewer could be more specific as to what other demonstrations they are referring to. \n\n**@W2b** (On the systematic demonstration of the limitations of FIM): As we explained in the limitations section of our paper, the main limitations of FIM are related to the synthetic training distribution.  We expect that if one leverages FIM to infer the rate matrix from an empirical process whose path distribution lies well outside our synthetic training distribution, one will obtain poor rate estimates. Figure 4 represents such an example.\n\nBeyond this, we have also explored how other features of our synthetic training distribution affect FIM. Indeed, as we explained in lines 199-207 and lines 212-217 of the main text, FIM is expected to perform best for the context number $c(300, 100)$, a number which is specified by the training distribution. We studied how FIM handles context numbers that are different from the optimum in both Appendix D.1 and D.4, and we refer the reviewer to them for details. \n\nSimilarly, we have also studied the effect of training FIM on a synthetic dataset which contains only six-state MJPs in Appendix D.3. \n\nAgain, we believe all these experiments cover well the limitations and features of FIM, especially as regards the synthetic training distribution, and we would be glad if the reviewer could be more specific as to which type of experiments they refer to. \n\n**@Q1, Q2, Q3**: We kindly refer the reviewer to our general comment above.\n\n**@Q4**: As explained in the general comment above, FIM always returns a $C\\times C$ matrix. It is therefore not possible for it to predict intensity matrices of dimension higher than C. \n\n**@Q5**: As can be read from equation 6, during training we ask FIM to predict zeros for those redundant matrix elements. Indeed, the first two terms in equation 6 are both multiplied by $a_{ij}$, which is zero when the target rate is zero. Similarly, the second two terms are multiplied by $(1-a_{ij})$, so that only these terms are active when the target rate is zero.\n\n**@Q6**: Let us answer this question by reporting our predictions for the DFR dataset with $V=1$. The ground-truth initial condition is given by the distribution\n\n$\\pi_0 = [0.301, 0.137, 0.062, 0.200, 0.159, 0.141]$.\n\nFIM predicts\n\n$\\hat \\pi_0 = [0.224, 0.147, 0.112, 0.189, 0.155, 0.173]$.\n\nOur estimates are therefore reasonably accurate. For completeness we will report both our estimated $\\hat \\pi_0$ and their empirical counterpart for all target datasets in the Appendix.\n\n**@Q7**: The adjacency matrix $\\textbf{A}$ only appears as part of our training distribution, and it simply defines the level of sparsity of the embedded Markov chain of the MJPs within the training dataset. Thus, no information about the adjacency matrix is needed during inference. \n\nAs explained in Appendix B.1, lines 479-484, the adjacency matrix is sampled from an Erdős–Rényi model, with edge probability 0.5, and is rejected if the corresponding graph is not connected. This prior on $A$ is, just as all other priors in Eq. 2, part of our definition of the FIM synthetic training distribution. Our experiments corroborate that these heuristic priors, which we define in Appendix B.1, are expressive enough for our trained FIM model to perform well in our 5 target datasets. \n\n**@Q8**: In our experiments we make use of as many data points as the baselines in all experiments, for fairness in comparison. We give the sizes of each of these datasets in each subsection of section 4 in the paper. \n\nWe however do not fully understand what the reviewer means by the last question: *How does FIM compare against the standard ML approach?* All our baselines are different instances of the “standard ML approach”. Indeed, the goal of our paper is to show that FIM provides a good alternative to the standard paradigm.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RDJt7Iv5KC", "reviewer": "Reviewer_5km2", "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.", "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?\n\nAdditionally, 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.", "limitations": "The authors adequately describe the limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n\nOutside of this, I could not find any other weaknesses of the paper; I found it to be very good in general.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer  for both the detailed review and the kind words about our work. Below we address each of the comments and suggestions.\n\n**@Q1** (FIM success due to simplicity of the data): We completely agree with the reviewer. Indeed, we argued in Section 3, lines 119-127 of our paper, that the main assumption underlying our FIM methodology is that *the space of realizable MJPs, which take values on bounded state spaces that are not too large, is simple enough to be covered by a heuristically constructed synthetic distribution over noisy and discretely observed MJPs*.\n\nThat being said, we are currently working on extending our methodology to other families of Markov processes,  like e.g. single MJPs on unbounded state spaces, or multidimensional (i.e. coupled) MJPs, for which, despite their significantly larger state space, we think we can define, via information theoretic arguments, synthetic distributions that cover well their realizable trajectories.\n\n**@Q2** (On the zero-shot terminology): We understand the point raised by the reviewer. We use the zero-shot terminology in the spirit of the work in Larochelle et al. (2008). That is, for us, zero-shot learning aims to recognize objects (e.g. stochastic processes) whose instances (e.g. samples or paths) may not have been seen during training. \n\nAs the reviewer writes, FIM does indeed not see any labels (i.e. rate matrices) during inference. That is one of our first motivations to use the zero-shot terminology. We understand that what we propose is, perhaps, an extension of this terminology to the case of inference of stochastic processes, which we intended to go hand in hand with our foundation (inference) model terminology (which is also an extension of the nlp terminology). \n\nNevertheless, we do not want to create any confusion as to how FIM works. We completely agree with the reviewer in that our methodology can be fully understood as an amortized inference procedure. In fact, reviewer EGs8 also drew our attention to the work of Paige&Wood (2018), who described a procedure to train a recognition model offline via amortized inference. \n\nWhat we therefore propose is to\n- first, add to the related work section a paragraph on amortized inference and offline learning, where we reference Paige&Wood (2018), together with other representative works in this direction, and connect them with our approach; and\n- second, explain in Section 3 what our motivation for the zero-shot terminology is, and how we understand it as an amortized inference procedure. \n\nDoes the reviewer think that these modifications will help avoid misunderstandings with our terminology? \n\n*Links to references:*\n\n- Larochelle et al. (2008): https://cdn.aaai.org/AAAI/2008/AAAI08-103.pdf\n- Paige&Wood (2018): https://arxiv.org/pdf/1602.06701", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "okxPtHLcUI", "reviewer": "Reviewer_pfZc", "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.", "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?\n\nI 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.\n\nI 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?", "limitations": "I discused my doubts about the limitations above.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**@W1** (On how to apply the model): We kindly refer the reader to our general response above. We hope it addresses the reviewer’s comment on the lack of details on how to apply FIM to any dataset. It will be included into our manuscript (see summary above).\n\n**@W2** (On limitations): Please see our response to your Q3 below.\n\n**@Q1** (Do they make their predictions on the basis of coarse-grained (CGR) data): Yes, FIM takes coarse-grained data as input. \n\nAs we explained in our introduction (see e.g. lines 49-54 and lines 116-117), our goal is to tackle the classical MJP inference problem on *coarse-grained (i.e. discrete) space*.\n\nIn practice, one can obtain these CGRs by e.g. leveraging pretrained models specifically tailored to performing CGR inference in specific datasets or, more naively, by simply using clustering algorithms. \n\nIn our experiments we used the CGR inferred by NeuralMJP, for the sake of fairness when comparing against it. We also used naive CGRs obtained via Gaussian mixture models (see e.g. Section 4.2, lines 283-287) and the KMeans clustering algorithm (see e.g. Section 4.3, lines 305-307 and Appendix E.2, lines 674-683). Table 3 in the main text and Tables 12, 15, 17 and 18 in the Appendix contain our estimated observables wrt. the different CGRs we employed for each of our datasets.\n\nLet us finish this response with the following remark. The inference of CGR should be understood as a different problem from that of the MJP inference. Accordingly, it is to be considered outside the scope of the present work (see e.g. our discussion in lines 102-117 in the related work section). To include such an inference into the FIM methodology, would amount to including an emission probability model into the generative model (Eq. 2) of our synthetic training distribution. Effectively, this would be equivalent to replacing our noise distributions with more complex models. We do not pursue this avenue, but the present work serves as a basis for such an extension.\n\n**@Q2** (On the adjacency matrix): The adjacency matrix $\\mathbf{A}$ only appears in our model as part of our synthetic training distribution, and it simply defines the level of sparsity of the embedded Markov chain of the MJPs within the training dataset. Accordingly, the adjacency matrix $\\textbf{A}$ is *not needed* as input to the model. One only requires the time series of the observation times and the corresponding values (states) in CG space (see also the general reply above for details).\n\nFinally, note that since FIM returns the $C \\times C$ inferred rate matrix $\\mathbf{\\hat F}$ (see e.g. Eq. 5), the target adjacency matrix can be understood as being implicitly inferred by FIM (that is, it can be extracted from the inferred $\\mathbf{\\hat F}$ matrix).\n\n**@Q3**. (On limitations): The question is twofold. Let us start by answering whether our synthetic datasets are realistic enough to be useful in a general setup. If we interpret this question as whether our synthetic distributions are wide enough to cover empirical path distributions of real world problems, our answer is *yes*. Indeed, we leveraged one and the same pretrained FIM to infer hidden MJPs from 5 very different target datasets. In particular, the switching ion dataset is a real-world, experimental dataset recorded from a viral ion channel, and the ADP dataset is an all-atom simulation of the alanine dipeptide molecule. Both these datasets correspond to very complex dynamical systems. Yet FIM, trained on our heuristically constructed synthetic dataset, is able to make the same inference as the baselines, without the need of any fine-tuning. \n\nThe second question concerns the applicability of our pretrained FIM to diffusion processes and to trajectories sampled from the trap model. Traditionally, MJPs and diffusion processes are treated (or better, defined) as different types of (continuous-time) stochastic process, as they have different supports. Indeed, MJPs are defined as processes which take values on discrete sets (like e.g. the integers), whereas diffusion processes are defined as processes which take values in continuous sets (like e.g. the real numbers). Our methodology is designed to tackle the inference of MJP only, and hence cannot be used to infer diffusion processes. \n\nNevertheless, let us remark, for the sake of completeness, that there are certain limits in which diffusion process can be asymptotically described in terms of MJPs, as first studied by van Kampen and Kubo (see e.g. Gardiner, 2009, chapter 11, or the very recent work of Winkler et al. 2024). These limits involve state spaces that are typically unbounded, and hence are out of the scope of the present work, which focuses, as we stressed in our abstract, on the inference of MJP on bounded state spaces.\n\nSimilarly, trap-like models, understood as phenomenological models for relaxation times (as in the mean-field ferromagnetic model of Griffiths et al. 1966, or the model for aging in disordered systems of Bouchaud 1992), can be understood as MJPs on unbounded state spaces (i.e. space of infinitely many metastable states). Our methodology targets the inference of MJP on bounded state spaces and hence cannot be used to model trajectories from trap-like models.  \n\nLet us finish this section by kindly referring the reviewer to our response to weakness W2b of Reviewer sVKK, where we discuss the limitations of our methodology, especially with respect to the synthetic training distribution.\n\n*Links to references*:\n- Gardiner (2009): https://link.springer.com/book/9783540707127\n- Winkler et al. (2024): https://arxiv.org/pdf/2405.03549\n- Griffiths et al. (1966): https://journals.aps.org/pr/abstract/10.1103/PhysRev.149.301\n- Bouchaud (1992): https://jp1.journaldephysique.org/articles/jp1/abs/1992/09/jp1v2p1705/jp1v2p1705.html", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
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 have however noticed a couple of common questions among (some of) the reviewers, namely 

1. How is the input data of FIM scaled, so that FIM can be applied to real-world data?
2. How does FIM handle MJPs whose state space size is smaller than the maximum size $C$ in the training ensemble?

Let us address each of these, and summarise the main updates we will do to the manuscript, which we believe answer all questions/concerns of the reviewers, here.

**1. On the scaling of the input data**: As we explained in Appendix C.2, lines 549-556, FIM takes as input: 

- the normalized sequence of observation times, which lie on the interval $[0, 1]$;
- the values of the process at those observation times, which are sequences of integers taken from the set $\{0, 1, …, C-1\}$, where $C$ is the size of the largest state space in our training dataset;
- the number of observations per simulation; and 
- the time scale $T^*$, which we define as the observation time of the last event in the (unnormalized) input time series.

(*On the scaling of the observation values*).  The observation values do not need to be scaled since they are nothing but sequences of integers from the set $\{0, 1, 2, … C-1 \}$. Indeed, as we explained in our introduction (see e.g. lines 49-54 and lines 116-117), our aim is to tackle the classical MJP inference problem on coarse-grained (i.e. discrete) space. 

(*On the scaling of the observation times*). In order to be able to process arbitrary input data, FIM requires the observations times to lie on the interval $[0, 1]$. Indeed, as we explained in Appendix B.1, lines 485-490 and lines 505-507, our target rate matrices take, by construction, values between 0 and 1 only. We use these matrices to simulate MJPs on the time interval [0, 10], and then map our random observation times to the interval [0, 1]. We do this mapping by normalizing all observation times wrt. the maximum observation time in the set of paths we simulate per MJP. That is, we define the *time scale* of the simulated paths to be the time $T^*< 10$ of the last event in the path set of a given MJP. Naturally, this normalization affects the target rates (which have units of one over time) only by the factor $T^*$. *FIM is trained to predict the normalized rates*. 

The true (unnormalized) rates can then simply be computed by multiplying the matrix $\mathbf{\hat F}$ returned by FIM with $1/T^*$. In practice, this normalization happens inside FIM, which is why the model requires $T^*$ as input.

Note that, as we empirically demonstrated in the paper, this rescaling procedure allows us to work with real-world MJPs of arbitrary time scales. For example, the time scales for the switching ion channel dataset were more than 500 times smaller than the time scales in our training dataset.

**2. FIM and MJP with different state space sizes**. As we mentioned above, FIM takes as input integer sequences (i.e. the observations) on the discrete set $\{0, 1, 2, … C-1 \}$. FIM always returns the $C(C-1)$ off-diagonal elements of the $C \times C$ inferred rate matrix $\mathbf{\hat F}$, whose diagonal elements are then computed as $\hat F_{ii} = \sum_{j \neq i} \hat F_{ij}$. 

In our experiments we trained FIM on a synthetic dataset with $C=6$. 

Now, we arrange all the target rate matrices $\mathbf{F}$ within our training dataset for MJPs with state spaces of size $c < C=6$, to be the leftmost block diagonal $c \times c$ matrix within a $6\times 6$ matrix of zeros, so that the redundant matrix elements are always zero. As can be read from equation 6, and as we explained in lines 196-198, we train FIM to predict zeros for those redundant matrix elements. 

In practice, however, our trained FIM does not exactly predict zeros for those redundant matrix elements. In our experiments the user knows a priori the number of states of the hidden process, so we explicitly set the redundant matrix elements to zero, and only then compute the corrected diagonal (i.e. the normalization) of the output rate matrix. Note that this assumption is typically made by the baselines (see e.g. Seifner & Sanchez, 2023). 

To illustrate that FIM nevertheless predicts very small values for the entries of the rate matrix which are supposed to be zero, we considered again the 3-state switching ion dataset, and computed again the stationary distribution from the complete $6\times 6$ output matrix, without zeroing out the elements outside the $3\times 3$ leftmost block diagonal matrix. We obtained the following:

$
\hat p^* = (0.154, 0.226, 0.590, 0.010, 0.008, 0.011).
$

Note that, as expected, the last three entries have very small probabilities. Also note that the first three entries agree well with the results reported in Table 3. 

**SUMMARY OF MAIN MODIFICATIONS**:

- We will add to the related work section a paragraph on amortized inference and offline learning, where we reference the work of Paige&Wood (2018), together with other representative works in this direction (Weakness 1, Reviewer EGs8);
- We will explain in Section 3 what our motivation for the zero-shot terminology is, and how we understand it as an amortized inference procedure (Question 2, Reviewer 5km2);
- We will move the subsection “How to use the model” of Appendix C.2 into a new Appendix section. This new Appendix will gather and extend what we wrote above **on the scaling of the input data** (Question 1 and 2, Reviewer sVKK and Weakness 1, Reviewer pfZc);
- We will include a new Appendix which will contain our notes **on FIM and MJP with different state space sizes** which we wrote above (Question 3, Reviewer sVKK).
- We will report both our estimated initial conditions ($\hat \pi_0$) and their empirical counterpart for all datasets in the Appendix (Question 6, Reviewer sVKK).

Meta-review: The reviewers agree that this is an interesting paper that deserves to be published. A limitation is that this work uses empirical rather than theoretical investigations to demonstrate the usefulness of the proposed idea, and the paper would be strengthened if this experiments more clearly demonstrated the limitations of this methodology. The authors might want to think about this (as well as other reviewer comments) when they prepare the final version of this paper.
Decision: Accept (poster)