Compositional simulation-based inference for time series

## Abstract

abstract Amortized simulation-based inference (SBI) methods train neural networks on simulated data to perform Bayesian inference. While this approach avoids the need for tractable likelihoods, it often requires a large number of simulations and has been challenging to scale to time-series data. Scientific simulators frequently emulate real-world dynamics through thousands of single-state transitions over time. We propose an SBI framework that can exploit such Markovian simulators by locally identifying parameters consistent with individual state transitions. We then compose these local results to obtain a posterior over parameters that align with the entire time series observation. We focus on applying this approach to neural posterior score estimation but also show how it can be applied, e.g., to neural likelihood (ratio) estimation. We demonstrate that our approach is more simulation-efficient than directly estimating the global posterior on several synthetic benchmark tasks and simulators used in ecology and epidemiology. Finally, we validate scalability and simulation efficiency of our approach by applying it to a high-dimensional Kolmogorov flow simulator with around one million dimensions in the data domain.

## 1 Introduction

Numerical simulations are a central approach for tackling problems in a wide range of scientific and engineering disciplines, including physics (Brehmer & Cranmer, ; Dax et al., ), molecular dynamics (Hollingsworth & Dror, ), neuroscience (Gonçalves et al., ) and climate science (Watson-Parris et al., ).
Simulators often include at least some parameters that cannot be measured experimentally. Inferring such parameters from observed data is a fundamental challenge.
Bayesian inference provides a principled approach to identifying parameters that align with empirical observations. Standard algorithms for Bayesian inference, such as Markov Chain Monte Carlo (MCMC) (Gilks et al., ) and variational inference (Beal, ), generally require access to the likelihoods $p({\bm{x}}|{\bm{\theta}})$. However, for many simulators, directly *evaluating* the likelihood remains intractable, rendering conventional Bayesian approaches inapplicable. However, *generating* synthetic data ${\bm{x}}\sim p({\bm{x}}|{\bm{\theta}})$ is feasible for numerical simulators. Simulation-based inference (SBI) methods offer a powerful alternative to perform Bayesian inference for such simulator models with intractable likelihoods (Cranmer et al., ).

Figure: Fig 1: Illustrative figure of the methodology: Factorized Neural Score Estimation (FNSE) method. Here, $a$ denotes the diffusion time, and ${\bm{\theta}}_{a}$ is the associated noisy parameter. The training process involves single-state transitions starting from arbitrary proposal $\tilde{p}({\bm{x}}^{t})$, with different parameters sampled from a prior distribution. During inference, the time series is divided into single state transitions, and each state transition is evaluated by the neural network to estimate local posterior scores. These local estimates are then aggregated to form a global approximation, which is subsequently used to sample from the overall posterior distribution.
Refer to caption: x1.png

Classical SBI methods, such as Approximate Bayesian Computation (ABC) (Beaumont et al., ) and synthetic likelihoods (Wood, ), struggle to effectively scale to high-dimensional simulations. To address this,
SBI methods using neural networks have been developed, which train networks to represent likelihoods (Papamakarios et al., ; Glöckler et al., ; Boelts et al., ), likelihood ratios (Durkan et al., ; Hermans et al., ; ; Miller et al., ), posteriors (Papamakarios & Murray, ; Lueckmann et al., ; Greenberg et al., ; Deistler et al., ; Geffner et al., ; Wildberger et al., ; Sharrock et al., ) or target several properties at once (Radev et al., ; Gloeckler et al., ). These methods allow for parameter inference without requiring additional simulations after training, making them more efficient than traditional approaches (Lueckmann et al., ): they effectively amortize the cost of the simulation, and/or the full inference approach.

However, applying these neural amortized methods to time series simulations can be challenging due to the high computational cost of repeated simulator calls. Running numerous time series simulations—possibly with varying sequence lengths—can be computationally prohibitive or at least wasteful. Instead, it seems advantageous to exploit the temporal structure of these simulators. For example, many scientific simulators are based on solving (stochastic) differential equations that model real-world time series dynamics through state transitions over time. Consequently, many time series simulators have an inherently Markovian structure, which can be leveraged for efficient inference. In particular, we can tackle the problem on the level of single-state transitions, a simpler task, that can be tackled with fewer simulations.

In this work, we propose a framework for efficient simulation-based inference for Markovian time series simulators. Unlike other neural SBI methods that require long sequences for training, our approach
locally estimates parameters consistent with a single state transition. By aggregating these local solutions, we recover the global posterior for time series of arbitrary length, significantly reducing simulation costs (Fig. 1). In this work, we focus on applying this framework to posterior score-based approaches and additionally apply it to likelihood(-ratio) estimation methods. A related challenge has been addressed in prior work  (Geffner et al., ; Linhart et al., ; Boelts et al., ) which has focused on efficient inference on i.i.d. (independently and identically distributed) observations. We empirically evaluate our approach on a series of synthetic benchmark tasks, as well as simulators from ecology and epidemiology. Finally, we demonstrate the superior performance of this approach on high-dimensional tasks, such as the Kolmogorov flow problem, with improved simulation efficiency and scalability.

## 2 Preliminaries

### 2.1 Problem setting

We target simulators $p({\bm{x}}^{0:T}|{\bm{\theta}})$ that output time series ${\bm{x}}^{0:T}:=({\bm{x}}^{0},...,{\bm{x}}^{T})\in\mathbb{R}^{T\times d}$. Many simulators in scientific fields are implemented using (stochastic) differential equations, which inherently satisfy the Markov property. We focus on this class of simulators, where the likelihood takes the following form

$$ $p({\bm{x}}^{0:T}|{\bm{\theta}})=p({\bm{x}}^{0}|{\bm{\theta}})\prod_{t=0}^{T-1} p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}}).$ (1) $$

To simplify the exposition, we constrain the initial condition to be independent of the parameter, $p({\bm{x}}^{0}|{\bm{\theta}})=p({\bm{x}}^{0})$, and assume that the simulation transition $p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$ does not depend on time $t$. We relax both of these assumptions in Appendix A.2, and also address the general case of higher-order Markov chains in Appendix A.2.3. We assume one can sample from the state-transition function ${\bm{x}}^{t+1}\sim\mathcal{T}({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$, but not to its probability density.

### 2.2 Amortized methods for simulation-based inference

Neural Posterior Estimation (NPE): Neural Posterior Estimation (NPE; (Papamakarios & Murray, ; Lueckmann et al., ; Greenberg et al., ; Wildberger et al., )) methods typically use a conditional neural density estimator, e.g., a normalizing flow (Rezende & Mohamed, ; Papamakarios et al., ; Durkan et al., ) to approximate the target posterior distribution. The model $q_{\phi}({\bm{\theta}}|{\bm{x}}^{0:T})$ is trained via maximum likelihood to estimate $p({\bm{\theta}}|{\bm{x}}^{0:T})$ from a dataset of parameter-data pairs $\{({\bm{\theta}}_{i},{\bm{x}}_{i}^{0:T})\}_{i=1}$, where $({\bm{\theta}}_{i},{\bm{x}}_{i}^{0:T})\sim p({\bm{\theta}})p({\bm{x}}^{0:T}|{
\bm{\theta}})$.
For observations of inhomogeneous size, NPE often incorporates an embedding network.
Embedding networks can be designed as recurrent networks (Lueckmann et al., ), permutation-invariant architectures (Chan et al., ; Radev et al., ), or attention-based models (Schmitt et al., ) to respect known invariances in the data.
At inference time, NPE yields an amortized approximation of the posterior for any observation ${\bm{x}}_{o}^{0:T}$ therefore offering rapid evaluations or samples of the posterior.

Neural likelihood(-ratio) estimation (NLE, NRE):
Neural Likelihood(-Ratio) Estimation (NLE (Papamakarios et al., ), NRE (Hermans et al., ; ; Miller et al., )) methods do not directly approximate the posterior distribution but instead train a surrogate $q_{\psi}({\bm{x}}_{i}^{0:T}|{\bm{\theta}})$ for the likelihood $p({\bm{x}}_{i}^{0:T}|{\bm{\theta}})$ or likelihood ratio $p({\bm{x}}_{i}^{0:T}|{\bm{\theta}})/p({\bm{x}}_{i}^{0:T})$. Once the likelihood is accessible via the surrogate, standard techniques such as MCMC or variational inference can be applied to perform inference. These approacehs circumvent the problem of not having access to the likelihood, but still leaves us with the drawbacks of standard inference methods, i.e., slow sampling and potential failure modes such as robustly handling multimodality (Glöckler et al., ). Nonetheless, they also inherit the flexibility of standard inference techniques, such as the ability to handle multiple independent (or conditionally independent) observations by simply multiplying the likelihood terms (Boelts et al., ).

Neural score estimation (NSE): Score-based diffusion models are a powerful tool for generating samples from a target distribution. Common diffusion models (Song & Ermon, ; Song et al., ; Ho et al., ) are based on stochastic differential equations (SDEs) that can be expressed as $d{\bm{\theta}}_{a}=f(a){\bm{\theta}}_{a}da+g(a)d\bm{w}$(^1^11The diffusion time is denoted by $a$, against convention, to distinguish it from $t$ in the time series.) with $\bm{w}$ being a standard Wiener process. The drift and diffusion coefficients $f$ and $g$ are chosen such that the solution to this SDE defines a diffusion process that transforms any initial distribution into a simple noise distribution $p_{A}({\bm{\theta}}_{A})=\mathcal{N}({\bm{\theta}}_{A},\mu_{A},\sigma_{A}^{2}I)$. Samples from any target, such as the posterior $p({\bm{\theta}}|{\bm{x}}^{0:T})$, can then be obtained by simulating the reverse diffusion process (Anderson, )

$$ $d{\bm{\theta}}_{a}=\left\{f(a)-g^{2}(a)\cdot\nabla_{{\bm{\theta}}_{a}}\log p({ \bm{\theta}}_{a}|{\bm{x}}^{0:T})\right\}da+g(a)d\bm{\tilde{w}},$ (2) $$

where $\bm{\tilde{w}}$ is a backward-in-time Wiener process. In practice, we can not access the analytic form of the score $s({\bm{\theta}}_{a}|{\bm{x}}^{0:T})=\nabla_{{\bm{\theta}}_{a}}\log p({\bm{
\theta}}_{a}|{\bm{x}}^{0:T})$, but we can estimate it from samples using conditional denoising score-matching  (Hyvärinen & Dayan, ; Song et al., )

$$ $\mathcal{L}(\phi)=\mathbb{E}_{a,{\bm{\theta}},{\bm{\theta}}_{a},{\bm{x}}^{0:T} }\left[\lambda(a)||s_{\phi}({\bm{\theta}}_{a}|{\bm{x}}^{0:T})-\nabla_{{\bm{ \theta}}_{a}}\log p({\bm{\theta}}_{a}|{\bm{\theta}})||_{2}^{2}\right]$ $$

where $\lambda$ denotes a positive weighting function and $p({\bm{\theta}}_{a}|{\bm{\theta}})=\mathcal{N}({\bm{\theta}}_{a};s(a){\bm{
\theta}},\sigma(a)^{2}{\bm{I}})$.

This recently proposed approach has been highly successful across various tasks (Geffner et al., ; Wildberger et al., ; Sharrock et al., ; Gloeckler et al., ). It offers a trade-off between the efficiency of the static NPE method and the flexibility of slower but more flexible NL(R)E method. By enabling feasible post-hoc modifications and composability through appropriate adjustments to the backward diffusion process, it bridges the gap between these two approaches (Geffner et al., ; Gloeckler et al., ).

Figure: Algorithm 1 Training

## 3 Methods

### 3.1 Global solution from single-step transitions

Direct estimation of the global target distribution ($p({\bm{\theta}}|{\bm{x}}^{0:T})$ in NPE or $p({\bm{x}}^{0:T}|{\bm{\theta}})$ in NLE; Sec. 2.2) often requires a large number of simulations of entire time series, leading to high computational costs. We mitigate the problem of expensive simulation calls by leveraging the Markov factorization in Equation 1.

Since the transition probability $p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$ encapsulates all relevant information about the parameters, we expect that, given sufficiently many single-step transitions, it should be possible to infer the global target from observations of arbitrary time lengths. Building on this insight, we propose a general framework that focuses on solving the local inverse problem. In this framework, we aim to estimate a local target $s_{\text{local}}$ using single-step transition simulation data $\mathcal{D}=\{(\theta_{i},\mathbf{x}_{i}^{t:t+1})\}_{i=1}$ (Alg. 1). Afterward, we estimate a global target by aggregating these local solutions (Alg. 2), using a composition rule (compose, Sec. 3.2.2). In the following, we apply the framework to NSE, NLE, and NRE.

### 3.2 Factorized neural score estimation (FNSE)

#### 3.2.1 Local score estimation

Assuming that $p({\bm{x}}^{0:T}|{\bm{\theta}})$ satisfies Eq. 1 and $p({\bm{x}}^{0}|{\bm{\theta}})=p({\bm{x}}^{0})$, the global target $\nabla_{\bm{\theta}}\log p({\bm{\theta}}|{\bm{x}}^{0:T})$ in NSE can be factorized as

$$ $\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T}_{o})=\sum_{t=0}^{T-1 }\nabla_{{\bm{\theta}}}\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t}_{o},{\bm{x}}^{ t+1}_{o})-(T-1)\cdot\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}).$ (3) $$

Here, $\tilde{p}({\bm{\theta}}|{\bm{x}}^{t},{\bm{x}}^{t+1})$ is a posterior with ${\bm{x}}^{t}$ following any proposal distribution ${\bm{x}}^{t}\sim\tilde{p}({\bm{x}}^{t})$ and ${\bm{x}}^{t+1}$ the state transition ${\bm{x}}^{t+1}\sim p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$ respectively (Appendix Sec. A.1).
The factorization (3) implies that global posterior is fully characterized by $s({\bm{\theta}}|{\bm{x}}^{t},{\bm{x}}^{t+1})=\nabla_{\theta}\log\tilde{p}({\bm
{\theta}}|{\bm{x}}^{t},{\bm{x}}^{t+1})$. We can aim to estimate this quantity using only single-state transitions by minimizing the following loss

$$ $\mathcal{L}(\phi)=\mathbb{E}_{a,{\bm{\theta}},{\bm{\theta}}_{a},{\bm{x}}^{t},{ \bm{x}}^{t+1}}\left[\lambda(a)||s_{\phi}({\bm{\theta}}_{a}|{\bm{x}}^{t},{\bm{x }}^{t+1})-\nabla_{{\bm{\theta}}_{a}}\log p({\bm{\theta}}_{a}|{\bm{\theta}})||_ {2}^{2}\right]$ (4) $$

given any proposal distribution ${\bm{x}}^{t}\sim\tilde{p}({\bm{x}}^{t})$.
Consequently, we can learn a local score estimator, as outlined in Alg. 1, by empirically minimizing this loss given a dataset of single-step simulations. This result generalizes to higher-order Markov chains, which require proposing more than one past observation (Appendix A.2.3 for details).

There is lots of flexibility in the choice of the proposal distribution $\tilde{p}({\bm{x}}^{t})$ in theory. The only requirement is that the proposed state must be independent of the parameters involved in the current state transition being emulated (see Alg. 1 line 4, Sec. A.1).
However, this choice can have significant practical implications given a finite simulation budget, as further investigated in Sec. 4. In particular, this choice makes it possible to allocate simulation resources to inference questions of interest. For instance, if the goal is to perform inference on a predefined set of observations, the proposal can be tailored to predominantly sample values from this set. In general, however, designing a good proposal can be challenging.

Geffner et al. () also proposed a score factorization in the i.i.d. setting to handle data-sets with multiple *independent* observations.
The factorization (3)
can be regarded as an extension of this factorization from the i.i.d. case to the Markov setting. Similarly, we can extend our results to partially factorized approaches by amortizing over multiple transitions instead of just a single one (see Appendix Sec. A.2.4).

#### 3.2.2 Local score composition

At inference time, we have to approximate $s_{\text{glob}}({\bm{\theta}}_{a},{\bm{x}}_{o}^{0:T})=\nabla_{{\bm{\theta}}_{a
}}\log p({\bm{\theta}}_{a}|{\bm{x}}_{o}^{0:T})$, by composing all local scores estimated by the network (Algorithm 2 the compose function). In the end, $s_{\text{glob}}$ will be passed to a diffusion sampler to obtain samples from the global posterior $p({\bm{\theta}}|{\bm{x}}_{o}^{0:T})$.

Unfortunately, $p({\bm{x}}^{0:T}|{\bm{\theta}}_{a})=\int p({\bm{x}}^{0:T}|{\bm{\theta}})p({\bm
{\theta}}|{\bm{\theta}}_{a})d{\bm{\theta}}$ no longer satisfies the Markov property, even when $p({\bm{x}}^{0:T}|{\bm{\theta}})$ does (Weilbach et al., ; Gloeckler et al., ; Geffner et al., ; Rozet & Louppe, ). This renders
Eq. 3 invalid for any $a>0$, requiring some corrections.
This issue has been addressed in the i.i.d. setting by Geffner et al. () (referred to as FNPE) and Linhart et al. () (referred to as GAUSS/JAC), and we will adapt their approach to the Markov setting.

FNPE: Geffner et al. () avoids the issue by constructing an alternative sequence of distributions $q_{a}~{}(0\leq a\leq A)$ that satisfies $q_{0}=p({\bm{\theta}}|{\bm{x}}^{0:T})$ and $q_{A}=\mathcal{N}({\bm{\theta}};\mu_{A},\frac{\sigma_{A}^{2}}{T}I)$. Adapting their proposal to the Markovian case, we obtain

$$ $\nabla_{{\bm{\theta}}_{a}}\log q({\bm{\theta}}_{a}|{\bm{x}}^{0:T})=\sum_{t=0}^ {T-1}s_{\phi}({\bm{\theta}}_{a}|{\bm{x}}^{t},{\bm{x}}^{t+1})+\frac{(1-T)(A-a)} {A}\nabla_{\theta}\log p({\bm{\theta}}).$ $$

While this approach is generally applicable, it renders the backward SDE (Eq. 2) invalid as the marginals no longer follow the corresponding forward diffusion process. This requires the use of MCMC-based sampling corrections, such as unadjusted Langevin dynamics (Geffner et al., ) or similar approaches (Sjöberg et al., ).

GAUSS/JAC:
Linhart et al. () have recently proposed an alternative approach by approximating the marginal scores derived from the forward diffusion process. This enables the utilization of the reverse SDE and consequently allows for the application of standard diffusion-based sampling techniques. A second-order approximation (i.e. using Gaussian assumptions), suggests the approximation

$$ $\nabla_{{\bm{\theta}}_{a}}\log p({\bm{\theta}}_{a}|{\bm{x}}^{0:T})\approx \mathbf{\Lambda}({\bm{\theta}}_{a})^{-1}\left(\sum_{t=0}^{T-1}\mathbf{\Sigma}_ {a,t,t+1}^{-1}s_{\phi}({\bm{\theta}}_{a}|{\bm{x}}^{t},{\bm{x}}^{t+1})+(1-T) \mathbf{\Sigma}_{a}^{-1}\nabla_{{\bm{\theta}}_{a}}\log p({\bm{\theta}}_{a}),\right)$ $$

where $\mathbf{\Lambda}({\bm{\theta}}_{a})=\sum_{t=0}^{T-1}\mathbf{\Sigma}_{a,t,t+1}^
{-1}+(1-T)\mathbf{\Sigma}_{a}^{-1}$.
This approach requires specifying two ”hyperparameters”: the denoising prior and posterior precision matrices. The prior precision matrix is typically known analytically, along with the marginal prior score. In contrast, the denoising posterior precision matrix must be estimated. To address this, GAUSS assumes a Gaussian posterior and analytically computes the denoising posterior covariances. Alternatively, JAC iteratively estimates it using Tweedie’s Moment projection (Boys et al., ), leveraging the Jacobian of the score network. Furthermore, the above approximation is only valid if $\mathbf{\Lambda}$ is positive definite. Unlike Linhart et al. (), who clip the diffusion process to a specific prior range to avoid numerical issues, we adapt all estimated posterior precision matrices to ensure positive definiteness of $\mathbf{\Lambda}$. This adaptation resolves the problem even for long time series (i.e., $T>100$; see Appendix C.2 for details). Unless otherwise specified, we use the GAUSS approximation as the default.

### 3.3 Factorized likelihood(-ratio) estimation (FNLE, FNRE)

For likelihood(-ratio) estimation, applying the approach outlined above is straightforward: by assumption, the likelihood factorizes as shown in Eq. 1. Consequently, we can instead directly learn the transition density $p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$, similarly to the method described in Alg. 1, by replacing the score matching loss with the appropriate likelihood (or likelihood-ratio) loss (Papamakarios et al., ; Hermans et al., ; ; Miller et al., ).

Once the transition density is obtained, the global log-likelihood approximation $\ell_{\text{glob}}$ can be computed directly from Eq. 1 (i.e., by simply summing up the local approximations). Classical MCMC techniques can be employed for sampling. We use reference implementations of NLE and NRE as implemented in the sbi package  (Tejero-Cantero et al., ), adapted to the Markovian setting.

## 4 Empirical results

### 4.1 Evaluation approach

We empirically assess the accuracy and computational efficiency of the proposed approach, comparing it against non-factorized NPE with an appropriately chosen embedding net as a baseline. We note that the simulation budget is not determined by the number of simulator calls $N$, but rather by the number of calls to the state-transition function, i.e., $T\cdot N$. We, therefore, configure the NPE baseline with a $T_{\text{max}}=10$ steps, i.e., if the simulation budget is $10k$, the NPE baseline used $1000$ simulations, each of which is simulated for $10$ steps.We additionally use data augmentation by duplicating shortened variants ($T<T_{\text{max}}$) of these simulations. This is to help the RNN embedding net to generalize to different sequence lengths. To also amortize over different initial conditions, the initial condition for each ten-step simulation was drawn from the proposal also used in the factorized methods.

As main metrics for comparison, we use sliced Wasserstein distance (sW1) and Classifier two sample test accuracy (C2ST) on reference posterior samples (Bischoff et al., ; Lopez-Paz & Oquab, ; Kolouri et al., ). We average the estimated value over a total of $10$ randomly drawn observations. The whole process, i.e., training, sampling, and evaluation, was repeated five times starting from different random seeds.

We begin by evaluating the methods on a set of newly designed benchmark tasks for Markovian simulators (Sec. 4.2). Next, we apply the approach to classical models from ecology and epidemiology, including the stochastic Lotka-Volterra and SIR models (Sec. 4.3). Finally, we demonstrate the scalability of the method on a large-scale Kolmogorov flow task, where we perform inference on very high-dimensional data using only 200k simulator steps (Sec. 4.4)

### 4.2 Benchmarks

To investigate several properties of the proposed approach, we develop several synthetic tasks of first-order Markovian time series with associated reference posterior samplers (details in Appendix B.1).

Gaussian RW: A Gaussian Random Walk of form ${\bm{x}}^{t+1}=\alpha\cdot{\bm{x}}^{t}+{\bm{\theta}}+\epsilon$ for $\epsilon\sim\mathcal{N}(0,{\bm{I}}),\alpha<1.$ with ${\bm{x}}^{t}\in\mathbb{R}^{d}$ and ${\bm{\theta}}\in\mathbb{R}^{d}$. This simple task offers an analytic Gaussian posterior.

Mixture RW: A Mixture of Gaussian Random Walk of form ${\bm{x}}^{t+1}={\bm{x}}^{t}+u\cdot\theta+\epsilon$ for $\epsilon\sim\mathcal{N}(0,{\bm{I}}),u\sim\text{Unif}(\{-1.,1.\})$ and ${\bm{x}}^{t},{\bm{\theta}}\in\mathbb{R}^{d}$. By design, this task has a mixture of Gaussian transition density and a non-Gaussian bimodal posterior.

Periodic/Linear SDE: Both tasks can be formulated as linear stochastic differential equation $d{\bm{x}}^{t}={\bm{A}}({\bm{\theta}}){\bm{x}}^{t}+{\bm{B}}({\bm{\theta}})d{\bm
{w}}^{t}$. The periodic SDE is nonlinearly parameterized ${\bm{A}}$ to be oscillatory with ${\bm{\theta}},{\bm{x}}^{t}\in\mathbb{R}^{2}$. This task has a posterior with four modes symmetric around the origin. The linear SDE task s every entry of ${\bm{A}}$ and ${\bm{B}}$ with ${\bm{x}}^{t}\in\mathbb{R}^{3}$ and ${\bm{\theta}}\in\mathbb{R}^{18}$.

Double well: A nonlinear SDE $d{\bm{x}}^{t}=\theta_{1}{\bm{x}}^{t}+\theta_{2}\left({\bm{x}}^{t}\right)^{3}+
\sigma d{\bm{w}}^{t}$, which samples from a double-well potential with modes position depending on $\theta_{1},\theta_{2}$ (Singer, ; Cai et al., ).

We first examine the scalability of all evaluated methods as a function of the length and dimensionality of the time series, while maintaining a fixed simulation budget. Specifically, we assess each method on the Gaussian random walk (RW) example in 1, 2, and 10 dimensions, using a total of 10k simulations (Fig. 2a). Both FNLE and FNSE demonstrate superior simulation efficiency compared to the NPE baselines. FNRE, however, appears to be more sensitive to local errors, with these errors accumulating over time, resulting in a consistent decline in performance for long time series. A similar observation and hypothesis was made by Geffner et al. () for i.i.d. data. Interestingly, this effect is largely reduced in the FNLE approach, which was not part of their empirical analysis (but successfully applied by Boelts et al. ()). On the other hand, the computational cost escalates significantly on long time series, as evaluating the global likelihood requires $T$ forward passes for each iteration of MCMC sampling, making FNLE and FNRE relatively slow (especially NLE, as a forward pass through the normalizing flow in FNLE is more costly than the classifier used in FNRE, Fig. 2a). In contrast, the sampling cost for FSNE remains moderate due to its more efficient sampling method, while performance remains similar to or better than FNLE.

Figure: Fig 2: Benchmarks: We validate our method on a Gaussian Random Walk with different dimensions for different lengths (i.e. Transitions), also tracking sampling times (a). We assess FNSE score accumulation over Gaussian RW and Periodic SDE tasks using a fixed Euler–Maruyama sampler (b). We compare methods across tasks and transition steps (c). Finally, we examine the effect of the proposal on NFSE trained with 10k simulations from a normal proposal of varying standard deviation (d).
Refer to caption: x2.png

In contrast to FNLE and FNRE, FNSE can introduce errors when composing the local scores. Therefore, we investigate the scaling behavior of different score-composition techniques for long time series (Fig. 2b, Appendix Fig. 5b). Previous work within the i.i.d case did not investigate performance beyond thirty samples (Geffner et al., ; Linhart et al., ). Specifically, for time series exceeding a length of $T>100$, the FNPE method starts to become numerically unstable due to diverging trajectories. Although the assumptions of the GAUSS/JAC approximations are violated (as the Mixture RW has a non-Gaussian posterior), these methods still tend to perform well in practice in our implementation. Even when Langevin corrections (Song et al., ; Geffner et al., ) are introduced, the GAUSS and JAC method remain superior, especially for large $T$ (Fig. A7, Appendix Sec. C.2).

Next, we conducted benchmarks across several tasks. Overall, FNSE and FNLE show higher accuracy compared to FNRE and NPE in most tasks (Fig. 2c, Appendix Fig. 5c). An exception is the double-well task, where NPE performs relatively well due to the SDE’s rapid convergence to a stationary distribution, enabling the RNN to learn a generalizable summary statistic for inference. FNSE strongly outperforms other methods in cases of multimodal transition densities (e.g., Mixture RW) and performs comparably to likelihood-based approaches in simpler cases (e.g., Periodic/Linear SDE) even if the posterior is strongly multimodal (Periodic SDE). Learning from single-step transitions is an extreme end of a spectrum of possibilities. We therefore also evaluate this benchmark on *partially* factorized variants using five transitions each (Appendix Fig. 6). Overall, this could improve (e.g., Double Well), but it could also hurt (e.g., Mixture RW) performance relative to the NPE baseline.

Finally, we investigate the impact of the proposal distribution, specifically for the FNSE method trained over 10k simulation steps (Fig. 2d, extended results in Fig. A8). We use a Normal distribution centered at zero, adjusting its standard deviation ranging from $\sigma=0.1$ (too narrow) to $\sigma=1000$ (too wide). Our results indicate a relatively broad range of standard deviations that result in good performance. If the distribution is too narrow, the time series observation may reach values outside the training domain, while if it is too wide, it the model is trained on values not observed during evaluation.

Figure: Fig 3: Lotka Volterra and SIR experiments : The FNSE approximate posterior (predictive) of the best performing FNSE model using 100k transition steps to train, visualized on subsequences from a fixed observation and associated true parameter. (a, b). We then show the quantitative performance in terms of C2ST and sW1 for each task on ten randomly selected observations, each run repeated five times (c, d).
Refer to caption: x3.png

### 4.3 Lotka Volterra and SIR

To further evaluate our method, we tested it on two famous models from ecology and epidemiology: the Lotka-Volterra and Susceptible-Infected-Recovered (SIR) simulators. The Lotka-Volterra simulator models predator-prey dynamics through four parameters that govern prey and predator growth, hunting rates, and mortality. Populations evolve over time according to interconnected stochastic differential equations. The SIR (Susceptible-Infected-Recovered) simulator is a fundamental model for understanding infectious disease spread. Although commonly used for benchmarking SBI on fixed-size observations Lueckmann et al. (), we adapted these models for our analysis (details in Appendix Sec. B.1).

We begin by demonstrating how the posteriors evolve as more time points are observed using the NFSE method (Fig. 3ab). In both tasks, we find that the posteriors converge to the true parameter values, and the posterior predictive simulations increasingly align with the time series observation. Notably, unlike in previous i.i.d. scenarios (Geffner et al., ; Linhart et al., ), the posterior uncertainty about the parameters does not necessarily decrease as additional time points are included. This is particularly evident in the SIR task, where the posterior remains largely unchanged between $T=10$ and $T=100$ (Fig. 3b). This can be explained by the fact that the initial dynamics contain significant information about the parameters, while the later dynamics do not, i.e., the infected population consistently declines to zero, and the susceptible and recovered populations converge to steady-state values that are independent of the parameterization.

We trained all methods on this task using $100k$ transition simulations. We then evaluate the accuracy of all models for performing inference on sequences of length $T=1,10,100$. For both tasks we see that the factorized methods (except FNRE; in LV) generally outperform the NPE baseline (Fig. 3cd). NFSE, in particular, significantly outperforms any other approaches on the SIR task.

### 4.4 Kolmogorov flow

Finally, we consider a task with very high dimensional observations, previously considered in the context of data-assimilation (Rozet & Louppe, ). The simulator models incompressible fluid dynamics governed by the Navier-Stokes equations:

$$ $\displaystyle\frac{\partial\mathbf{u}}{\partial t}$ $\displaystyle=-\mathbf{u}\cdot\nabla\mathbf{u}+\frac{1}{\text{Re}}\nabla^{2} \mathbf{u}-\frac{1}{\rho}\nabla p+\mathbf{f}\qquad\qquad\nabla\cdot\mathbf{u}=0,$ $$

where $\mathbf{u}$ is the velocity field, Re is the Reynolds number, $\rho$ is the fluid density, $p$ is the pressure field, and $\mathbf{f}$ is the external forcing. Following Kochkov et al. (); Rozet & Louppe (), we consider a two-dimensional domain $[0,2\pi]^{2}$ with periodic boundary conditions and an external forcing $\mathbf{f}$ that corresponds to Kolmogorov forcing with linear damping Chandler & Kerswell (); Boffetta & Ecke (). We take the Reynolds number and density as free parameters ${\bm{\theta}}=(\text{Re},\rho)$. We utilize the jax-cfd library (Kochkov et al., ) to solve the Navier-Stokes equations on a $64\times 64$ grid.

In contrast to Rozet & Louppe (), we consider the problem of parameter inference on ${\bm{\theta}}$.
We apply FNSE to this problem, using only 200k transition evaluations, which successfully recovers an amortized posterior estimator that generalizes to long-term observations (Fig. 4). Notably, in such a high-dimensional state space, a simple proposal distribution would be ineffective. But, we can use both the initial distribution and simulator to propose a variety of feasible states (details in Appendix Sec. C.4). The posterior distributions for observations with both $T=10$ and $T=100$ are well-concentrated around the true parameter values that generated the specific observations (Fig. 4a,b). Posterior predictive samples are, on average, significantly closer to the observed data compared to prior predictives (in mean absolute error calculated over 1000 predictive samples from 50 different observations). Although the difference is less pronounced for $T=100$, this is primarily due to the divergence for slight parameter modification on long simulations. To quantitatively further investigate the performance, we perform a simulation-based calibration analysis (Talts et al., ) (Fig. 9). Overall, given the constraints on the simulation budget, the posterior calibration is good but deteriorates for larger time series.

Figure: Fig 4: Kolmogorov flow experiment: Top row shows a single selected observation. We visualize the posterior distribution considering (a, left top) only using the first ten transitions of the selected observation, along with a quantitative comparison of the mean absolute error (MAE) between posterior predictive samples and prior predictive samples of ten different observations (left, bottom), including the one shown above. The vorticity of two selected predictive samples is visualized on the right. This analysis is then repeated for the entire observation (b).
Refer to caption: x4.png

## 5 Discussion

### 5.1 Related Work

We build upon previous work in the i.i.d. observation setting (Geffner et al., ; Linhart et al., ; Boelts et al., ), extending these methods to Markovian time series. In the context of data assimilation, which aims to estimate $p({\bm{x}}^{0:T})$, Rozet & Louppe () introduced a technique for learning a local score network that estimates the score of $p({\bm{x}}^{t})$ based on a limited neighborhood ${\bm{x}}^{t-k:t+k}$ (i.e., its Markov blanket). In contrast, our approach emphasizes identifying a local inverse problem that can be composed to produce a global solution. Recently, a simulation-based conditional kernel density approximation for inference in SDEs was proposed (Cai et al., ), which can be viewed as a special case of FNLE, substituting normalizing flows with kernel density estimation.

Moreover, any Kalman filtering method (Kalman, ; Wan & Van Der Merwe, ; Arasaratnam & Haykin, ; Julier & Uhlmann, ) is inherently related, as these methods can approximate the marginal log-likelihood in an online (non-amortized) manner, but do require the availability of likelihoods.
Finally, parameter inference for (stochastic) differential equations has been extensively studied across various fields, including numerical analysis (Plessix, ) and probability theory (Brouste et al., ), both of which rely on the explicit likelihood functions.

### 5.2 Limitations and Future Work

In this study, our primary focus is on time-invariant Markovian stochastic processes, with an extension to encompass time-varying cases and processes featuring parameterized initial distributions. We note that Hidden Markov Models, although widely applied in practical scenarios, do not fall within the scope of our method.
Various approaches are viable, particularly in combination with data assimilation techniques (Rozet & Louppe, ) the Markovian hidden state is recovered to which our technique can be directly applied. Importantly, our approach applies to a broad spectrum of probabilistic models: first, identify a smaller inference problem from which the global target can be recovered; second, locally solve it. However, addressing these aspects remains a direction for future research.

The score corrections caused by score aggregation described in Subsection 3.2.2 could be further improved. In their methods, Linhart et al. () approximates the data distribution by a single Gaussian distribution, which can be violated in real-world scenarios. Our results demonstrate that this approximation empirically still performs well even if these assumptions are strongly violated.

### 5.3 Conclusions

We introduced a simulation-efficient framework for amortized inference in Markovian simulators. Although flexible embedding networks are commonly used for high-dimensional time series (Lueckmann et al., ; Radev et al., ; Schmitt et al., ), they often demand extensive amounts of simulations and can be fragile when faced with data perturbations (Cannon et al., ; Gloeckler et al., ). Success in these methods hinges on the embedding network’s ability to capture robust and generalizable representations. In contrast, our approach decomposes the inference task into smaller, locally solvable problems, reducing computational costs and enhancing scalability for large-scale, complex simulations.

#### Acknowledgments

We thank all members of the Mackelab, as well as Keisuke Yano and Yoshinobu Kawahara for discussions. MG and JHM are supported by the German Research Foundation (DFG) through Germany’s Excellence Strategy (EXC-Number 2064/1, Project number 390727645), the German Federal Ministry of Education and Research (Tubingen AI Center), the ”Certification and Foundations of Safe Machine Learning Systems in Healthcare” project funded by the Carl Zeiss Foundation, and the European Union (ERC, DeepCoMechTome, 101089288). MG
is member of the International Max Planck Research School for Intelligent Systems (IMPRS-IS). ST and KF are supported by JST CREST JPMJCR2015. ST is supported by Grant-in-Aid for Research Activity Start-up 23K19966 and Grant-in-Aid for Early-Career Scientists 24K20750. KF is partially supported by JSPS Grant-in-Aid for Transformative Research Areas (A) 22H05106.

### 5.4 Software and data

We use JAX (Bradbury et al., ) as the computational backbone and hydra (Yadan, ) to track configuration. For reference implementation of baselines, we use sbi (Tejero-Cantero et al., ). Code to reproduce the experiments can be found in [https://github.com/mackelab/markovsbi/](https://github.com/mackelab/markovsbi/).

#### Author contributions

Conceptualization: MG, ST, KF, JHM. Methodology: MG, ST. Software: MG. Investigation and Analysis (Benchmarks): MG, ST. Investigation and Analysis (Lotka Volterra, SIR and Kolmogorov Flow): MG. Visualization: MG. Writing: MG, ST. Writing (Review $\&$ Editing): KF, JHM. Funding acquisition: ST, KF, JHM. Supervision: KF, JHM.

## Appendix A Methods

### A.1 Derivation of Score Factorization ( 3 )

For ${\bm{x}}^{0:T}$ following the simulator $p({\bm{x}}^{0:T}|{\bm{\theta}})$, the global score $\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T})$ is given by

$$ $\displaystyle\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T})$ $\displaystyle=\nabla_{\bm{\theta}}\log\left\{p({\bm{\theta}}|{\bm{x}}^{0:T}) \cdot p({\bm{x}}^{0:T})\right\}=\nabla_{\bm{\theta}}\log\left\{p({\bm{\theta}} )\cdot p({\bm{x}}^{0:T}|{\bm{\theta}})\right\}$ $\displaystyle=\nabla_{\bm{\theta}}\log\left\{p({\bm{\theta}})\cdot p({\bm{x}}^ {0}|{\bm{\theta}})\prod_{t=0}^{T-1}p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}} )\right\}$ $\displaystyle=\nabla_{\bm{\theta}}\log\left\{p({\bm{\theta}})\cdot p({\bm{x}}^ {0}|{\bm{\theta}})\right\}+\nabla_{\bm{\theta}}\log\prod_{t=0}^{T-1}p({\bm{x}} ^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$ $\displaystyle=\nabla_{\bm{\theta}}\log\left\{p({\bm{\theta}}|{\bm{x}}^{0}) \cdot p({\bm{x}}^{0})\right\}+\nabla_{\bm{\theta}}\log\prod_{t=0}^{T-1}p({\bm{ x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$ $\displaystyle=\nabla_{\bm{\theta}}\log p({\bm{\theta}}|{\bm{x}}^{0})+\sum_{t=0 }^{T-1}\nabla_{\bm{\theta}}\log p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}}).$ (5) $$

Let ${{\bm{x}}}^{t}$ follow any proposal distribution $\tilde{p}({{\bm{x}}}^{t})$, and ${\bm{x}}^{t+1}$ follow the simulation transition $p({\bm{x}}^{t+1}|{{\bm{x}}}^{t},{\bm{\theta}})$ given ${{\bm{x}}}^{t}$.
Noting that the replacement does not affect the conditional distribution $p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$, we may assume ${{\bm{x}}}^{t}\sim\tilde{p}({{\bm{x}}}^{t})$ in (5).

By Chain Rule, we have

$$ $\displaystyle\tilde{p}({\bm{x}}^{t+1},{\bm{\theta}}|{{\bm{x}}}^{t})$ $\displaystyle=p({\bm{x}}^{t+1}|{{\bm{x}}}^{t},{\bm{\theta}})\cdot\tilde{p}({ \bm{\theta}}|{{\bm{x}}}^{t})$ (6) $\displaystyle=\tilde{p}({\bm{\theta}}|{\bm{x}}^{t+1},{{\bm{x}}}^{t})\cdot \tilde{p}({\bm{x}}^{t+1}|{{\bm{x}}}^{t}),$ $$

where $\tilde{p}$ means ${\bm{x}}^{t}$ follows the proposal $\tilde{p}({\bm{x}}^{t})$.
This leads us to the equality

$$ $p({\bm{x}}^{t+1}|{{\bm{x}}}^{t},{\bm{\theta}})=\frac{\tilde{p}({\bm{\theta}}|{ \bm{x}}^{t+1},{{\bm{x}}}^{t})\cdot\tilde{p}({\bm{x}}^{t+1}|{{\bm{x}}}^{t})}{ \tilde{p}({\bm{\theta}}|{{\bm{x}}}^{t})}=\frac{\tilde{p}({\bm{\theta}}|{\bm{x} }^{t+1},{{\bm{x}}}^{t})\cdot\tilde{p}({\bm{x}}^{t+1}|{{\bm{x}}}^{t})}{{p}({\bm {\theta}})}.$ $$

Here, the second equality is derived from the equality $\tilde{p}({\bm{\theta}}|{{\bm{x}}}^{t})=p({\bm{\theta}})$, to which the fact that $\tilde{p}({{\bm{x}}}^{t})$ does not involve a parameter leads.
Substituting this equality into Eq. (5), we have

$$ $\displaystyle(\ref{eq: Score_Decomposition})=\nabla_{\bm{\theta}}\log p({\bm{ \theta}}|{\bm{x}}^{0})+\sum_{t=0}^{T-1}\nabla_{\bm{\theta}}\log\frac{\tilde{p} ({\bm{\theta}}|{\bm{x}}^{t+1},{{\bm{x}}}^{t})\cdot\tilde{p}({\bm{x}}^{t+1}|{{ \bm{x}}}^{t})}{p({\bm{\theta}})}$ $\displaystyle=\nabla_{\bm{\theta}}\log p({\bm{\theta}}|{\bm{x}}^{0})+\sum_{t=0 }^{T-1}\nabla_{\bm{\theta}}\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t+1},{{\bm{x} }}^{t})-T\cdot\nabla_{\bm{\theta}}\log p({\bm{\theta}})$ (7) $\displaystyle=\sum_{t=0}^{T-1}\nabla_{\bm{\theta}}\log\tilde{p}({\bm{\theta}}| {\bm{x}}^{t+1},{{\bm{x}}}^{t})-(T-1)\cdot\nabla_{\bm{\theta}}\log p({\bm{ \theta}}).$ $$

It concludes the proof. Here, the final equality is derived from the the assumption $p({\bm{x}}^{0}|{\bm{\theta}})=p({\bm{x}}^{0})$.

### A.2 Extensions

#### A.2.1 Time Variant Simulation

The assumption of time-invariant transition where $p({\bm{x}}^{t+1}|{\bm{x}}^{t+1},{\bm{\theta}})$ does not depend on $t$ in the main manuscript can be relaxed. Regarding FSNE, the score decomposition (3) also holds for a time-variant simulator, but the score $\nabla_{\bm{\theta}}\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t},{\bm{x}}^{t+1})$ depends on time $t$.
To address the issue, we train a time-variant score $s_{\phi}({\bm{\theta}}_{a}|{\bm{x}}^{t},{\bm{x}}^{t+1},t)$ using training data ${\mathcal{D}}=\left\{({\bm{\theta}}_{i},{\bm{x}}_{i}^{t_{i}},{\bm{x}}_{i}^{t_{
i}},t_{i})\right\}$, where $t_{i}$ is drawn from a random number generator, ${\bm{x}}_{i}^{t_{i}}$ is drawn from a proposal distribution $\tilde{p}({\bm{x}})$, and ${\bm{x}}^{t_{i}+1}_{i}\sim\mathcal{T}({\bm{x}}^{t_{i}+1}|{\bm{x}}^{t_{i}}_{i},
{\bm{\theta}}_{i})$. FNPE and FNRE can be handled in the same way as FNLE by adding $t$ as a variable. Experimental evaluations are presented in Appendix C.1.

#### A.2.2 Initial Parameter Estimation

The condition on initial state $p({\bm{x}}^{0}|{\bm{\theta}})$ can also be relaxed. Regarding FNSE, the global score $\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T}_{o})$ can be factorized by

$$ $\displaystyle\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T}_{o})= \nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}_{o}^{0})$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{ }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}+\sum_{t=0}^{T-1}\nabla_{{\bm{\theta}} }\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t}_{o},{\bm{x}}^{t+1}_{o})-T\cdot\nabla _{{\bm{\theta}}}\log p({\bm{\theta}})$ $$

as proven in Eq. (7). This implies that the global score is obtained by estimating the two scores $\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}_{a}|{\bm{x}}_{o}^{0})$ and $\nabla_{{\bm{\theta}}}\log\tilde{p}({\bm{\theta}}_{a}|{\bm{x}}^{t}_{o},{\bm{x}
}^{t+1}_{o})$ separately. The estimation of these two scores can be combined by modeling a score function that outputs different scores depending on the number of input variables. FNPE and FNRE can be extended by estimating the two likelihoods $p({\bm{x}}_{o}^{0}|{\bm{\theta}})$ and $p({\bm{x}}_{o}^{t+1}|{\bm{x}}_{o}^{t},{\bm{\theta}})$, similar to FNSE.

#### A.2.3 Higher order Markov chains

Our proposal can be generalized to simulators with high order markov chain

$$ $p({\bm{x}}^{0:T}|{\bm{\theta}})=p({\bm{x}}^{0:m-1}|{\bm{\theta}})\prod_{t=m-1} ^{T-1}p({\bm{x}}^{t+1}|{\bm{x}}^{t-m+1:t+1},{\bm{\theta}})$ (8) $$

with its degree $m$. Under the setting, the score $\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T})$ has the factorization

$$ $\displaystyle\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:T}_{o})= \nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}_{o}^{0:m-1})$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{ }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}+\sum_{t=m-1}^{T-m-1}\nabla_{{\bm{ \theta}}}\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t-m+1:t+1}_{0})-(T-m+1)\cdot \nabla_{{\bm{\theta}}}\log p({\bm{\theta}}),$ $$

which implies that the global score can be recovered by the two scores $\nabla_{{\bm{\theta}}}\log p({\bm{\theta}}|{\bm{x}}^{0:m-1})$ and $\nabla_{{\bm{\theta}}}\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t-m+1:t+1})$. Here, ${\bm{x}}^{t-m:t+1}$ follows a proposal $\tilde{p}({\bm{x}}^{t-m:t+1})$ and ${\bm{x}}^{t-m+1}$ follows the simulation transition $p({\bm{x}}^{t-m+1}|{\bm{x}}^{t-m:t+1},{\bm{\theta}})$ in the score $\nabla_{{\bm{\theta}}}\log\tilde{p}({\bm{\theta}}|{\bm{x}}^{t-m+1:t+1})$.
The estimation of these two scores can be merged same as the initial parameter estimation case in the last section.

#### A.2.4 Partial Factorization

Instead of considering just a single transition, the methodology naturally extends to multiple transitions, as also explored by Geffner et al. (). In PFNSE with $M$, we target $\nabla_{\bm{\theta}}\log p({\bm{\theta}}|{\bm{x}}^{t},\dots,{\bm{x}}^{t+M})$, while in PFNLE or PFNRE, we target $p({\bm{x}}^{t},\dots,{\bm{x}}^{t+M}|{\bm{\theta}})$ (up to constants), using a dataset of $M$-step simulations. A key advantage of partially factorized methods is that they require fewer network evaluations for inference on a fixed-size observation, which can accelerate sampling and reduce the accumulation of local errors over time. However, they may require more simulations for effective training.

## Appendix B Experiment Detail

### B.1 Tasks

In this section, we elaborate on all tasks used in the analysis in detail. For all tasks, we used a standard normal prior but eventually reparameterized them to certain constraints if necessary for the task (see individual tasks for specification).

Note that the training aims to also amortize over the initial distribution $p({\bm{x}}^{0})$. We evaluate the performance always on a specific and fixed initial distribution (see individual tasks for specification), which is also different from the proposal $\tilde{p}({\bm{x}}^{t})$.

##### Gaussian RW:

This task serves as a simple baseline. It serves as an extension of the Gaussian Linear task introduced by Lueckmann et al. () to the Markovian case. It is defined by a Gaussian Random Walk of form

$$ ${\bm{x}}^{t+1}=\alpha\cdot{\bm{x}}^{t}+{\bm{\theta}}+\epsilon\quad\text{for} \quad\epsilon\sim\mathcal{N}(0,{\bm{I}}),\quad\alpha=0.9,$ $$

with ${\bm{x}}^{t},{\bm{\theta}}\in\mathbb{R}^{d}$ and ${\bm{\theta}}\in\mathbb{R}^{d}$. This task offers an analytic solution for the posterior, which is Gaussian. Dimension can be set as wanted, we choose $d=1,2,10$. As a proposal, we choose $\tilde{p}({\bm{x}}^{t})=\mathcal{N}({\bm{x}}^{t};\bm{0},\sqrt{10}{\bm{I}})$, which is motivated by the fact that it is the variance of the corresponding stationary distribution. For evaluation, we fixed $p({\bm{x}}^{0})=\delta({\bm{x}}^{0})$ (i.e., a point mass at zero).

##### Mixture RW:

A Mixture of Gaussian Random Walk of form

$$ ${\bm{x}}^{t+1}={\bm{x}}^{t}+u\cdot{\bm{\theta}}+\epsilon\quad\text{for}\quad \epsilon\sim\mathcal{N}(0,{\bm{I}}),\quad u\sim\text{Unif}(\{-1,1\}),$ $$

with ${\bm{x}}^{t},{\bm{\theta}}\in\mathbb{R}^{d}$. We choose $d=2,5$. As proposal, we choose $\tilde{p}({\bm{x}}^{t})=\mathcal{N}({\bm{x}}^{t};\bm{0},\sqrt{10}{\bm{I}})$. For evaluation, we fixed $p({\bm{x}}^{0})=\delta({\bm{x}}^{0})$ (i.e., a point mass at zero).

By design, this task has a tractable transition density, which is a mixture of Gaussians, hence allowing exact computation of the marginal log-likelihood. In contrast to the previous task, this offers a non-Gaussian bimodal posterior. We generate reference posterior samples using Hamiltonian Monte Carlo (HMC, (Betancourt, )) with 5 integration steps. To avoid mode collapse, we run 100 parallel chains, each initialized on different samples obtained by importance resampling of 50 initial samples from the prior. We run the sampler for 600 burn-in iterations, during which the integration step size is adapted to accept around $80\%$ of the proposed steps. We additionally thin the accepted samples in the chain by a factor of six to avoid autocorrelation.

##### Periodic SDE:

The periodic SDE is a linear SDE parameterized as follows

$$ $d{\bm{x}}^{t}=\begin{pmatrix}0&\theta_{2}^{2}\\ \theta_{1}^{2}&0\\ \end{pmatrix}{\bm{x}}^{t}\,dt+0.1{\bm{I}}\,d{\bm{w}}^{t},$ $$

which gives us a stochastic oscillator. Each transition corresponds to simulating this SDE for $0.1$ ms, using 20 Euler Maruyama steps. As proposal, we choose $\tilde{p}({\bm{x}}^{t})=\mathcal{N}({\bm{x}}^{t};\bm{0},{\bm{I}})$. For evaluation, we fixed the initial distribution to $p({\bm{x}}^{0})=(-0.5,0.5)^{T}$, to quickly start with oscillations.

The corresponding marginal log-likelihood is computed using a Kalman filter Kalman (). Reference posterior samples are obtained in combination with HMC, using a procedure similar to the one discussed above. This task offers a multimodal posterior (4 modes) with all modes being point symmetric around the origin.

##### Linear SDE:

The linear SDE task is given by

$$ ${\bm{x}}^{t}=({\bm{A}}({\bm{\theta}})-2{\bm{I}}){\bm{x}}^{t}dt+(0.5{\bm{B}}({ \bm{\theta}})+0.5{\bm{I}})d\bm{w}^{t}$ $$

where every entry of ${\bm{A}}$ and ${\bm{B}}$ is directly parameterized by ${\bm{\theta}}\in\mathbb{R}^{18}$ and ${\bm{x}}^{t}\in\mathbb{R}^{3}$. Reference posterior samples are obtained in the same way as for the previous task.

##### Double well:

This is a nonlinear SDE

$$ $d{\bm{x}}^{t}=\theta_{1}{\bm{x}}^{t}\,dt+\theta_{2}\left({\bm{x}}^{t}\right)^{ 3}\,dt+\sigma\,d{\bm{w}}^{t},$ $$

which samples from a double-well potential with modes position depending on $\theta_{1},\theta_{2}$ (Singer, ). We use a combination of particle filter (Doucet et al., ) and pseudo-marginal-like Metropolis-Hastings MCMC (Andrieu & Roberts, ) sampler to generate reference posterior samples.

##### Lotka-Volterra:

We use a stochastic Lotka-Volterra model, a classic system of differential equations used to describe predator-prey interactions. We use a stochastic variant, where the amount of noise also depends on the population of each species. The model dynamics are given by the following equations:

$$ $\displaystyle\frac{dx^{t}}{dt}$ $\displaystyle=\alpha\cdot x^{t}-\beta\cdot x^{t}\cdot y^{t}dt+\sigma x^{t}y^{t }dw_{1}^{t},$ $\displaystyle\frac{dy^{t}}{dt}$ $\displaystyle=-\gamma\cdot y^{t}+\delta\cdot x^{t}\cdot y^{t}dt+\sigma x^{t}y^ {t}dw_{2}^{t},$ $$

where $x$ represents the prey population and $y$ represents the predator population. The parameters $\alpha,\beta,\gamma,$ and $\delta$ control the interaction between the species. The noise hyperparameter $\sigma$ was set to $0.05$. The dynamic are constrained to remain positive. As a proposal we use $\tilde{p}({\bm{x}}^{t})=\text{Unif}(0,10)$. As evaluation time the initial distribution was fixed to $p({\bm{x}}^{0})=\delta({\bm{x}}^{0}-(1.,0.5)^{T})$.

To sample from the posterior by traditional means, we use a Particle Filter to obtain a stochastic estimator of the marginal likelihood. We use this stochastic estimator in a pseudo-marginal-like Metropolis-Hastings MCMC algorithm to obtain reference samples from the posterior distribution (Doucet et al., ; Andrieu & Roberts, ).

##### SIR:

We use a stochastic SIR method defined through the following equations:

$$ $\displaystyle\frac{dS^{t}}{dt}$ $\displaystyle=-\beta S\cdot I+0.01dw_{1}^{t},$ $\displaystyle\frac{dI^{t}}{dt}$ $\displaystyle=\beta S\cdot I-\gamma I+0.1\cdot Idw_{2}^{t},$ $\displaystyle\frac{dR^{t}}{dt}$ $\displaystyle=\gamma I+0.01dw_{3}^{t}.$ $$

here $S$, $I$, and $R$ represent the susceptible, infected, and recovered populations. We use as a proposal $\tilde{p}({\bm{x}}^{t})=\text{Unif}(0,5)$. As an initial distribution, we sample $S^{0}\sim\text{Unif}(2.5,5.)$, $I^{0}\sim\text{Unif}(0.,2.5)$ and $R^{0}\sim\delta(R^{0})$.

To sample from the posterior by traditional means, we use a Particle Filter to obtain a stochastic estimator of the marginal likelihood. We use this stochastic estimator in a pseudo-marginal-like Metropolis-Hastings MCMC algorithm to obtain reference samples from the posterior distribution.

### B.2 Training and evaluation

For the NPE baseline, we utilize a 5-layer neural spline flow (Durkan et al., ), with each layer parameterized by a 2-layer MLP having a hidden dimension of 50. Additionally, we employ a Gated Recurrent Unit (GRU) network Cho et al. () as the embedding network, also with a
hidden dimension of 50. To enhance the RNN embedding network’s ability to generalize across different sequence lengths, we apply data augmentation. Specifically, for each subsequence $T<T_{\text{max}}=10$, we duplicate $10\%$ of randomly selected parameter-data pairs and shorten the corresponding data time series to length $T$. This does not add any more simulation calls but directly trains the neural net for sequences $T<T_{\text{max}}$.

For FNLE and FNRE we use adapted reference implementation as implemented in the sbi package (Tejero-Cantero et al., ). In FNLE, we use a 5-layer Masked Autoregressive Flow (Papamakarios et al., ), each parameterized by a 2-layer MLP with a hidden dimension of 50. In FNRE we use resnet classifier with two blocks each considering 2 layers with hidden dimension of 50. As MCMC sampling algorithm we use a per-axis slice sampling algorithm. To avoid mode-collapse, we run 100 parallel chains.

Both approaches are trained with a training batch size of 1000 until convergence, as determined by the default early stopping routine.

For FNSE we use a custom implementation in JAX (Bradbury et al., ). In all experiments, we use the Variance Preserving SDE (Song et al., ) using

$$ $f(t)=-0.5\cdot(\beta_{\text{min}}+t\cdot(\beta_{\text{max}}-\beta_{\text{min}} )),\qquad g(t)=\sqrt{\beta_{\text{min}}+t\cdot(\beta_{\text{max}}-\beta_{\text {min}})}$ $$

We set $\beta_{\text{min}}=0.1$, and $\beta_{\text{max}}=10$ for all experiments. Both for the time interval $[1e-2,1.]$. The associated conditional means and variances can be derived from this SDE (Song et al., ).

For the score estimation network, we use a 5-layer MLP with a hidden dimension of 50 and GELU activations. The diffusion time is embedded using a random Fourier embedding. Precondition to the scoring network by performing time-dependent z-scoring (Karras et al., ). We use the denoising score matching loss with weighting function as in Song et al. ().

We use an AdamW optimizer with a learning rate of $5e-4$ with a cosine schedule and a training batch size of $1000$. Similar to the SBI routine, we use early stopping, but with a maximum number of $5000$ epochs.

## Appendix C Additional Experiments

### C.1 Time-Variant Gaussian Random Work

To evaluate the performance of the time-variant extension described in Appendix A.2.1, we consider the Time-Variant Gaussian Random Walk (TV Gaussian RW)

$$ ${\bm{x}}_{t+1}=\left(\alpha+\frac{1}{t+1}\right)\cdot{\bm{x}}_{t}+{\bm{\theta} }+\epsilon\quad\text{for}\quad\epsilon\sim\mathcal{N}(0,{\bm{I}}),\,\alpha=0.9 ,{\bm{x}}_{t},{\bm{\theta}}\in{\mathbb{R}}^{d}$ $$

with ${\bm{x}}_{t}$ multiplied by the time-variant coefficient $\left(\alpha+\frac{1}{t+1}\right)$. We demonstrate the scalability of our methods with dimension $d$ ranging from $1,2$, to $10$, and the number of observation transitions ranging from $1$, $10$, to $100$. As a comparison metric, we use C2ST. The simulation budget is fixed at $10k$ and $100k$, consistent with the problem settings in the main section. Note that the training data include the time variable $t$, and therefore, the training data for estimating the transition distribution $p({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}})$ are reduced to $\frac{1}{100}$ compared to the time-invariant setting in the main section. The maximum time length $T_{\text{max}}$ in training is set to $100$.

The experimental results are presented in Fig. 10. We can confirm that the proposed methods outperform the NPE baseline overall. In particular, it has been confirmed that FNSE scales better for longer observation times compared to other methods.

### C.2 Score composition rules

We utilize the GAUSS composition rule as our default method:

$$ $\nabla_{{\bm{\theta}}_{a}}\log p({\bm{\theta}}_{a}|{\bm{x}}^{0:T})\approx \mathbf{\Lambda}({\bm{\theta}}_{a})^{-1}\left(\sum_{t=0}^{T-1}\mathbf{\Sigma}_ {a,t,t+1}^{-1}s_{\phi}({\bm{\theta}}_{a}|{\bm{x}}^{t},{\bm{x}}^{t+1})+(1-T) \mathbf{\Sigma}_{a}^{-1}\nabla_{{\bm{\theta}}_{a}}\log p({\bm{\theta}}_{a})\right)$ $$

with

$$ $\mathbf{\Lambda}({\bm{\theta}}_{a})=\sum_{t=0}^{T-1}\mathbf{\Sigma}_{a,t,t+1}^ {-1}+(1-T)\mathbf{\Sigma}_{a}^{-1}.$ $$

We estimate the posterior precision matrices $\mathbf{\Sigma}_{a,t,t+1}^{-1}$ as recommended by Linhart et al. (). For each posterior $p({\bm{\theta}}|{\bm{x}}^{t},{\bm{x}}^{t+1})$, we produce $500\cdot\dim_{\bm{\theta}}$ samples in parallel from the diffusion model. These samples are then used to estimate the associated covariance matrix. Although this approach introduces a slight constant overhead prior to each sampling run, it does not require the iterative estimation of Jacobians, as is necessary with the JAC method.

Importantly, since the GAUSS method is derived under Gaussian assumptions, it presupposes that the posterior covariance matrices (and their corresponding scores) are smaller than those of the prior due to observational constraints. This assumption may not hold true, particularly in multi-modal tasks. For instance, in the Periodic SDE task, four modes are symmetrically positioned around the origin. While the variances of each mode may decrease, the total variance often still increases, especially when the modes are distant from the origin. This situation can render $\Lambda({\bm{\theta}}_{a})$ non-positive definite, thus invalidating the approximation.

To address these challenges, we make minimal adjustments to the initially estimated $\mathbf{\Sigma}_{a,t,t+1}^{-1}$ to ensure that $\Lambda({\bm{\theta}}_{a})$ remains positive definite. We achieve this by considering its eigenvalue decomposition $\Lambda=VAV^{-1}$, allowing us to identify the minimal adjustment required to enforce positive definiteness: $\Lambda_{-}=-V\min(A,0)V^{-1}$. Consequently, we update $\hat{\mathbf{\Sigma}}_{a,t,t+1}^{-1}=\mathbf{\Sigma}_{a,t,t+1}^{-1}+\frac{
\Lambda_{-}}{T}+\epsilon I$ for a small nugget $\epsilon$. The $\Lambda$ computed with these adjusted covariance matrices is guaranteed to be positive definite. The JAC method encounters similar issues, and we apply an analogous adjustment to its initial estimates.

Overall, this strategy significantly enhances the performance of both GAUSS and JAC, even when their underlying assumptions are not fully validated (see Fig. 7 for benchmarks). This improvement is particularly critical for scaling to long time series, especially when the number of function evaluations is reduced (see Fig. 7b).

It is noteworthy that the GAUSS and JAC methods yield theoretically equivalent results in the Gaussian RW task, with no violations of assumptions (the only difference lies in their approach to estimating posterior precision matrices). As such, the performance of both methods is quite similar (Fig. 7, first column). However, when assumptions are violated, the differences between the two approaches become apparent. The GAUSS method presumes a Gaussian posterior and analytically computes the associated denoising posteriors. In contrast, the JAC method iteratively derives the best Gaussian approximation for the denoising posterior, as discussed by Boys et al. (). Notably, for samplers with fewer steps on non-Gaussian tasks, JAC tends to perform slightly better than GAUSS on long time series (Fig. 7b, columns two onwards). The imposed regularity conditions appear to resolve the numerical challenges associated with this approach, as observed by Linhart et al. ().

### C.3 Proposal distribution

The choice of proposal distribution can either be an advantage or a challenge. In our benchmark tasks, we primarily selected simple yet reasonable distributions based on our prior knowledge of the dynamics involved. However, we examine the impact of this choice more closely in the Gaussian RW and Mixture RW tasks (see Fig. 8).

We set the proposal distribution to a Gaussian with a zero mean and variable variance, ranging from very narrow to quite wide. Importantly, at evaluation, we initialize the distribution as a point mass at zero. For a very narrow proposal, we anticipate strong performance at $T=1$. However, as $T$ increases, performance may decline because the time series dynamics could extend beyond the high-probability support of the proposal. Conversely, excessively wide proposals may also hinder performance if the values at the observation time fall outside the training set. It is crucial to note that this performance evaluation is relative to our specific evaluation pipeline. In both tasks, evaluations are conducted on random walks that consistently start from zero, thus constraining the range of attainable values given a maximum number of steps.

Overall, there is indeed an optimal proposal for all tasks, which is close to, but not identical to, our intuitively chosen options. As expected, the significance of the proposal also increases with dimensionality.

It is important to emphasize that the proposals we utilized are generally not optimal. Other tasks, such as the Periodic/Linear SDE, Double Well, SIR, and Lotka-Volterra models, exhibit a high correlation among different variables, which the proposal we chose does not capture. We believe that identifying a more suitable proposal could greatly enhance the efficiency of our methods, although doing so may prove nontrivial (and might, itself, require many simulation calls). Nevertheless, to ensure a fair comparison with the NPE baseline, we opted for relatively simple proposals.

### C.4 Kolmogorov flow

For the Kolmogorov flow, we modify the score architecture. We utilize a convolutional embedding network composed of four blocks, each containing a Conv2D layer, GroupNorm, and GeLU activation. The output is then processed by a two-layer MLP to produce a 100-dimensional embedding. This embedding is generated from the two frames passed to the score network and subsequently concatenated. On this embedding, we employ the same five-layer MLP that we use for all other tasks, but we increase the hidden dimension to 400.

It is important to note that in the NPE approach, we would need to construct an embedding network not only for individual ”images” but also for a complete video of arbitrary length, which poses significant challenges.

As an initial distribution, we utilize a filtered velocity field provided by jax_cfg, with a maximum velocity and peak wave number of three. Our approach closely follows the methodology outlined by Rozet & Louppe (). The simulator performs 100 solve steps per transition, with a step size of $10^{-3}$ seconds; each transition thus emulates 10 milliseconds. To introduce stochasticity into the PDE, we add small amounts of Gaussian noise with a standard deviation of $5\cdot 10^{-3}$ after each transition. The simulations are conducted on a regular $64\times 64$ grid.

Designing an effective proposal distribution in a high-dimensional setting is challenging. Simply using the initial distribution is unlikely to yield satisfactory performance, as the dynamics evolve over time. Ideally, we would like to obtain samples from $p({\bm{x}}^{t})$ for a range of different $t$; however, this approach would require a significant simulation budget to step to each $t$.

To address this challenge while still generating a diverse set of simulations, we employ the following scheme:

- (i)
Sample an initial value ${\bm{x}}^{0}\sim p({\bm{x}}^{0})$
- (ii)
For $t=0,\dots,T$:
–
Sample ${\bm{\theta}}_{i}\sim p({\bm{\theta}})$
–
Perform a transition ${\bm{x}}_{t+1}\sim\mathcal{T}({\bm{x}}^{t+1}|{\bm{x}}^{t},{\bm{\theta}}_{i})$

Notably, this procedure does not violate the requirement that $\tilde{p}({\bm{x}}^{T})$ is independent of ${\bm{\theta}}$, as each transition is performed with a different, independent parameter.

For 1,000 initializations, we run this procedure for 100 steps, and for 10,000 initializations, we run it for 10 steps, totaling exactly 200,000 transition emulations.

To assess the average performance across different observations and observation lengths, we conducted a simulation-based calibration analysis  (Talts et al., ) (see Fig. 9). This analysis evaluates whether the true parameter ranks across posterior samples are uniformly distributed, as would be expected under the true posterior distribution. In other words, it examines whether the estimated posteriors adequately cover the true parameters. The results indicate a generally good, though not perfect, calibration of the estimated posteriors. However, the calibration tends to deteriorate with longer observation lengths. In this scenario, the posterior becomes very narrow, making the metric sensitive to such narrow distributions (which increases the likelihood of missing the true posterior). Furthermore, this suggests that for certain observations, the approximate posterior may be biased. Given the limited amount of training data, such outcomes are to be expected.

Figure: Fig 5: Benchmark with C2ST metric: We validate our method on a Gaussian Random Walk with different dimensions for different lengths (i.e. Transitions), also tracking sampling times (a). We assess FNSE score accumulation over Gaussian RW and Periodic SDE tasks using a fixed Euler–Maruyama sampler (b). We compare methods across tasks and transition steps (c). Finally, we examine the effect of the proposal on NFSE trained with 10k simulations from a normal proposal of varying standard deviation (d).
Refer to caption: x5.png

Figure: Fig 6: Benchmark partially factorized methods: We show the benchmark performance using partially factorized methods PFNLE, PFNRE, and PFNSE. Each is trained with 5 steps. The performance is shown with respect to C2ST (a) and sliced Wasserstein distance (b).
Refer to caption: x6.png

Figure: Fig 7: Benchmark score composition methods: Performance of score estimators trained with 10k simulations using different diffusion samplers (default and fewer steps) and score composition methods. Each column corresponds to a different task, and each row represents either a standard Euler-Maruyama diffusion sampler or one equipped with a Langevin corrector, both performing 5 steps. The analysis is conducted for two time discretizations: (a) 500 steps and (b) 100 steps.
Refer to caption: x7.png

Figure: Fig 8: Benchmark proposal sensitivity: Each column represents a different task. The top row displays results using the C2ST metric, while the bottom row shows the sliced Wasserstein distance. FNSE models are trained with a fixed budget of 10k simulations, with each using different proposal distributions characterized by a mean of zero and varying standard deviations (x-axis).
Refer to caption: x8.png

Figure: Fig 9: Simulation-based calibration Kolmogorov flow: The simulation-based calibration results present the empirical cumulative density functions (CDF) of the ground-truth parameters, ranked according to the inferred posteriors derived from 100 different observations. A well-calibrated posterior should exhibit uniformly distributed ranks, as highlighted by the shaded gray area. Repeated for data emulated for T=1,10,20,50,100 steps).
Refer to caption: x9.png

Figure: Fig 10: Time Variant Gaussian random walk: We train all methods on 10k (top row) and 100k (bottom row) simulation on the time-variant Gaussian RW talk. All methods have also been adapted to input the current time (as given by integers ranging from zero to 100).
Refer to caption: x10.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: * Is there a reason why comparison with NSE is not done? 
* Can FNSE be also made sequential in a similar manner as other simulation-efficient SBI methods such as SNPE/SNLE/SNRE? 
* Which score composition method is used in the experiments when implementing FNSE? (I might have missed if this is mentioned in the appendix)
soundness: 2
strengths: The paper addresses the relevant problem of computational complexity in SBI which arises when the simulator is costly to sample from. 

The experimental evaluation seems thorough, and the results are significant for FNSE.
summary: This paper proposes a simulation-based inference (SBI) method for state-space models where the transition dynamics is Markovian. The core idea is to simulate many single-state transitions and then aggregate them instead of simulating entire trajectories of time-series, so as to reduce the total number of simulator calls.
weaknesses: My main concern is related to the novelty of the paper. The main technical contribution lies in deriving the factorized NSE method, which as the authors mention, is an extension of the setting of the Geffner et al. (2023) paper to the Markov setting. An in-depth analysis/discussion around selecting the proposal $\tilde{p}(\mathbf x^t)$ and providing recommendations on its choice would have strengthened the paper significantly and added to the technical contribution. The paragraph discussing the proposal is kind of confusing. The authors say that "there is lots of flexibility in the choice of the proposal...", but also that "...designing a good proposal can be challenging". I appreciate the experiments evaluating the sensitivity of the results to the choice of proposal, but it is not clear how the proposals are chosen in the first place.

The clarity of writing can be improved by providing examples, especially when it comes to motivating the setting, the core idea, and the assumptions. For instance, the authors say "...given sufficiently many single-step transitions, it should be possible to infer the global target...". A discussion around the cases in which this does (and does not) hold would help the reader understand the scope of this work. Another example is the requirement that "...the proposed state must be independent of the parameters involved in the current state transition...", where it is not clear how restrictive this requirement is, what kind of cases does it cover, etc. Similarly, some examples of expensive real-world scientific simulators in para 3 of the intro would be nice to have.

On the topic of clarity, I found Section 3.2.2 difficult to follow. In particular, it is not clear when the authors are talking about background of previous work (Geffner 2023 and Linhart 2024), and which part is their contribution. Perhaps some of the discussion regarding the implementation details can be moved elsewhere, as they are not easy to follow without reading those two works anyway.  

The paper is presented as a sample-efficient method for Markovian time-series simulators. However, the experiments measure the performance as a function of the number of transitions. To conclude that the proposed method is "simulation-efficient", I would have expected performance to be plotted as a function the number of calls to the simulator (or computational cost) for a fixed number of transitions.  

Would be nice to have a small discussion on other sample-efficient SBI methods in the related works (e.g. based on sequential sampling, Bayesian optimization, etc.).

Minor comment: Typo in line 125-126 (the word "approaches")

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: Why does sW seem to, counterintuitively, increase for increasing number of transitions (Figure 2)?
Will the method be robust to missing or corrupted data?
soundness: 2
strengths: - The paper is self-contained, easy to follow, and it attempts to address a challenging problem that is relevant to the field as a whole. 
- The method is applicable to different classes of SBI methods (e.g., likelihood approximation, likelihood ratio approximation, direct posterior estimation).
- The idea is original and can stimulate further research into efficient SBI on dynamic models, especially when simulation budgets are scarce.
summary: The paper applies compositional score-matching, previously introduced for SBI on exchangeable models, to non-exchangeable, Markovian models. The appealing aspect of the methodology is that a score estimator can provide good posterior estimates without simulating the full process for each batch instance, but instead simulating single transitions per batch instance and learning to aggregate across time. The hope is that one can achieve the same posterior accuracy with less than $N * T$ evaluations, where $N$ is the simulation budget and $T$ is the maximum time horizon of the simulator.
weaknesses: - As far as I understand it, the proposed method is a straightforward extension of FNPE with an additional input to the score estimator (Eq.6 in [1]). As such, the claim that a new “general SBI framework” is proposed requires some calibration. In contrast, the authors could highlight and extend the empirical aspects of the work.

- While the basic idea is rather appealing, I find it a bit unconvincing that FNPE can robustly approximate the correct global posteriors with sparse training (the same goes for FNLE). There is an aspect of the methodology that strikes me as magical: suppose that most information about $\theta$ in a time-varying signal is contained in later time segments (e.g., as in certain non-stationary signals) or subject to latent transitions (e.g., as in HMMs). It would be incredibly hard to get to that information if the simulation is not evolved for longer $T$, but the factorized approach is supposed to somehow get enough training signal from very sparse simulations in all relevant time windows. How is it possible for the score network to properly learn the correct composition without any assumptions on signal stationarity or smoothness? I assume the reasonable performance on the toy examples can be attributed to the rather short signal lengths and the use of simple, low-dimensional models (Fig.9 reveals striking miscalibration for longer time series on the only challenging model, confirming my fears). 

- Overall, the evaluation lacks robustness and I am concerned that it is subject to randomness given the extremely limited number of test simulations used (10!). Since the work is explicitly situated in an amortized inference context and pitched as such, the evaluation could have been much more comprehensive, featuring hundreds or even thousands of test simulations and at least some practically relevant metrics with an absolute interpretation, such as calibration error or other measures used for evaluating computational faithfulness in Bayesian analysis ([2]). In addition, there are no ablation studies (e.g., one such study could vary the simulation budget or use an adaptive solver for better-than-uniform sampling, another study can consider extremely long time horizons $T$), all models only have a few parameters (SIR and LV are toy models from the 60s, but presented as distinct from the toy model section), and it is hard to tell if the performance on the only non-trivial model (4.4) is acceptable given the lack of ground-truth and the poor calibration for not not even that long time horizons ($T=100$).

- It may be helpful to clarify some points in Section 2.2 regarding related work on amortized inference. It might strengthen the section to highlight that amortized Bayesian inference is explicitly introduced in [3] and extensively discussed and expanded upon in inference compilation methods [4, 5, along with additional references therein]. Including this line of research could offer a more comprehensive overview of the field. Additionally, [6, 7] don't seem to introduce, discuss, or validate amortized methods; instead, they focus on sequential likelihood-free techniques, such as SNPE [7], without suggesting or claiming to use amortization. In fact, [6] cursorily note that learning over the prior predictive is possible but ultimately dismiss it as “grossly inefficient” (p.3). 

[1] Geffner, T., Papamakarios, G., & Mnih, A. (2023, July). Compositional score modeling for simulation-based inference. In International Conference on Machine Learning (pp. 11098-11116). PMLR.

[2] Gelman, A., Vehtari, A., Simpson, D., Margossian, C. C., Carpenter, B., Yao, Y., ... & Modrák, M. (2020). Bayesian workflow. arXiv preprint arXiv:2011.01808.

[3] Gershman, S., & Goodman, N. (2014). Amortized inference in probabilistic reasoning. In Proceedings of the Annual Meeting of the Cognitive Science Society (Vol. 36, No. 36).

[4] Le, T. A., Baydin, A. G., & Wood, F. (2017). Inference compilation and universal probabilistic programming. In Artificial Intelligence and Statistics (pp. 1338-1348). PMLR.

[5] Wu, M., Choi, K., Goodman, N., & Ermon, S. (2020). Meta-amortized variational inference and learning. In Proceedings of the AAAI Conference on Artificial Intelligence.

[6] Papamakarios, G., & Murray, I. (2016). Fast ε-free inference of simulation models with bayesian conditional density estimation. Advances in Neural Information Processing Systems, 29.

[7] Lueckmann, J. M., Goncalves, P. J., Bassetto, G., Öcal, K., Nonnenmacher, M., & Macke, J. H. (2017). Flexible statistical inference for mechanistic models of neural dynamics. Advances in Neural Information Processing Systems, 30.


*Minor*
Some attention to typos and neologisms is needed (e.g., P2L107 -> one can sample, …, to its probability density?, P7L377, etc.).

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: - How sensitive is the more complex Kolmogorov flow example to the choice of proposal distribution?
- Beyond the empirical results demonstrated, do the authors have a deeper understanding of the robustness of score composition rules?
- The generalization beyond the number of timesteps seen is one of the most useful properties of the method -- how do the generalization capabilities depend on the number of timesteps used during training (e.g., max 10 vs 100 timesteps during training, then 1000 during evaluations)?
soundness: 4
strengths: - The paper is exceptionally well-written. Secs. 2 and 3 set up the problem well, and describe the contributions in the context of existing literature. I feel like I learned a lot about the surrounding field beyond the specific contributions of the paper by reading these sections. Limitations are discussed upfront, and some extension settings explored in an initial way.
- The experiments in the paper are well-motivated. The paper looks at toy models, standard benchmarks, as well as more challenging simulators that highlight some of the specific advantages of the method (e.g., generalization beyond the number of steps used during training).
- The paper explores a timely question of broad scientific as well practical real-world interest. The generalization capabilities of the method beyond the number of timesteps trained for seem especially practical (and maybe should be highlighted more) -- many applications are bottlenecked by the inability to practically train on the longer-horizon timescales needed during deployment.
summary: The paper proposes a class of methods to perform neural simulation-based inference on time-series data when the forward model has Markovian structure, i.e. when the next step only depends on the current one. This is the case in many physical systems defined by PDEs, for example. The method exploits this structure of the simulator / forward model to learn locally (in timestep) amortized estimators for global parameters of interest that describe the system in question. The paper exploits many recent advances in amortized simulation-based inference and adapts them to the Markovian time-series setting.
weaknesses: While the core methods are described nicely, some practical aspects receive less thorough treatment. Specifically, the proposal distribution seems like a critical design choice (especially for complex problems like Kolmogorov flow), and the score composition rules prove surprisingly robust even when their assumptions are violated (e.g. GAUSS in non-Gaussian setting). There is limited discussion of why, and limited principled guidance on these implementation choices that practitioners might need to apply the method to new domains. A more thorough discussion here could significantly strengthen the paper in particular for practitioners across domains.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - Is there any reason why NPE performs so badly in the C2ST statistic in Figs 3c and 3d? Previously literature, e.g., [1-2], seems to report different results.
- Could you kindly explain the exact experimental setup in Fig2a and Fig2c? Does, e.g., $100$ transitions mean that NPE received $100$ data samples of length $100$ while for $10$ transitions it received $1000$ samples?

#### References
- [1] Flow Matching for Scalable Simulation-Based Inference, Figure 4, https://arxiv.org/pdf/2305.17161
- [2] Benchmarking Simulation-Based Inference, Figure 2, https://arxiv.org/abs/2101.04653
soundness: 3
strengths: - The paper proposes an intuitive approach for SBI for (very) high-dimensional time series.
- The topic (SBI for high-dimensional time series data) is very timely and relevant given the (seemingly) increased interested of applied researchers.
- The paper is well written and easy to follow.
summary: The authors introduce a novel approach for simulation-based inference (SBI) for time series models of which the joint probability density is Markovian. The work builds on the recently introduced framework of score-based generative models for SBI which aims to approximate the score of the posterior distribution when multiple data points (in this case time points) are conditioned on. They show that their approach is beneficial over a state-of-the-art method in time series benchmarks.
weaknesses: - Despite the intuitive idea, the proposed method is very incremental. As far as I can see, instead of modelling the score $s(\theta_a|x^t)$ the paper proposes to model $s(\theta_a|x^t, x^{t-1})$, i.e., adding one variable to the score network input, in order to make the method amendable for time series? 
- The evaluations are in my opinion not fully convincing. 
    - The authors chose to compare their approach to exactly one other baseline which I find a bit insufficient given the wealth of SBI methods and the fact that the authors propose approaches for local FNLE, FNRE and FNSE.
    -  The baseline NPE uses an RNN as an embedding network. Given the authors assume Markovianity, they could have, e.g., just computed summary statistics and used this as conditioning variables, instead of running a costly RNN (which also needs to be trained in addition to the flow layers). Alternatively approaches that reduce the dimensionality more efficiently like neural sufficient statistics could have been benchmarked for a more thorough evaluation [1,3].
  - Relatedly, there is work on high-dimensional SBI that the authors could potentially have included in their evaluations: [1-4]

#### References
- [1] Generalized massive optimal data compression, 2017, https://arxiv.org/abs/1712.00012
- [2] Neural Approximate Sufficient Statistics for Implicit Models, 2021, https://arxiv.org/abs/2010.10079
- [3] Is Learning Summary Statistics Necessary for Likelihood-free Inference?, 2023, https://proceedings.mlr.press/v202/chen23h.html
- [4] Simulation-based inference using surjective sequential neural likelihood estimation, 2023, https://arxiv.org/abs/2308.01054

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: Is there a particular motivation behind using Sliced Wasserstein Distance for evaluation? If using it as the main criterion, it would be nice to provide an intuition as of what features of differences in distribution this distance is capturing. Otherwise, I'd suggest using the Sinkhorn Divergence instead -- with a small regularization parameter, it is essentially a direct approximation of the Wasserstein distance.

**Minor**

- The use of "time-variant" and "time invariant" should be replaced by the more standard "(time) homogenous/unhomogenous".
- "MCMC-based sampling corrections, such as unadjusted Langevin dynamics". The ULA doesn't include a correction step (it is "undajusted"). What do the authors mean by that?
- To produce symmetric quotation marks, consider using ``text'' instead of "text" in latex.
soundness: 3
strengths: The paper proposes a conceptually simple, attractive adaptation of SBI methods to Markov chains. For a total simulation budget, the shift towards learning the transition factor allows to increase the number of available observations. Moreover, learning the transition probabilities may be a simpler problem than learning the parameter-to-whole-timeseries relationship, as the observation is lower dimensional in the former case.
summary: This work proposes a set of Simulation-Based Inference (SBI) methods to perform approxiamte Bayesian (parameter) inference given an
observation that takes the from of a specific class of timeseries (finite-length homogenous Markov chains).
These methods are adaptations of existing score, likelihood, and ratio-based estimation methods, whose original formulation
assume i.i.d observations, and not observations arising from a Markov Chain.

The methods do so by estimating ratio/likelihood/scores "locally", e.g. to the individual factors of the Markov chain.
They then produce a "global" posterior estimate (e.g. an estimate of the posterior of the parameters given the entire Markov chain sample)
by relying on the factorization structure of the Markov chain probability. The composition is straightforward in the case of ratio and likelihood methods,
while when using scores, additional approximations are necessary due to the non-Markovian structure of the blurred posteriors estimated in the previous step.

These approximations are direct adaptations of the one already used in prior work:
- one performs iterative Langevin-based sampling of an annealing path of distributions bridging from a standard Gaussian distribution to the true posterior, and whose intermediate distributions are precisely the SDE solutions whose scores were estimated.
- one produces an approximation of the scores of the SDE initialized at the posterior of the parameter given the entire Markov trajectory, which can then be used to simulate trajectories from the time-reversed SDE, whose final iterate is marginally distributed according to the true posterior.

The performance of the method is investigated on a set of experiments, which include standard benchmark models and a Kolmogorov flow model with high dimensional observations. The simulators are adapted (using noise injection) to fit the formalism of stochastic inference methods.
weaknesses: **Regarding Presentation**

The presentation of the method needs some improvements. 

First, regarding the the sampling part, the presentation of FNPE and Gauss do not define precisely key details of the approximations.
- In FNPE, where is $q_a$ used for instance?
- In GAUSS/JAC, What is $\Sigma_{a,t, t+1}$?

I had to look at each of the separate papers that introduced such approximations in the iid case to understand what was going on.
The paper should be updated to define all terms properly -- most of will constitute background. Even though the adaptations made by the method are almost independent from the these terms, the paper should remain self contained as far as the core method is concerned. To save space, the FNPE approximation method could be placed in the appendix, as it is not used in the experiments.

Second, the local score estimation section could also be made clearer: for instance, this section
starts by breaking down the true (unblurred) posterior (Equation (3)), highlighting how knowing the score of each
factor is enough to get the score of the resulting posterior. However, what is actually needed for sampling is the blurred posterior,
which cannot be factorized as Equation (3). Equation (3) thus ends up not being very relevant to the method.


**Regarding experiments**: the presentation of the Kolmogorov flow experiments is missing details: for instance, it does not include any mention of stochasticity, which is confusing given that the method is about estimating probability distributions.


**Regarding performance**
- The main weakness of the method right now is that the number of timesteps investigated in the paper remains smaller (up to 100)
compared to the regime of some application of interest, like neuroscience.
- In the Kolmogorov flow experiments, the MAE of the posterior predictive significantly increases for when using 100 instead of 10. It would be nice if the authors commented on this point in the paper.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: Is $x^t$ in the SIR example the three states $\{S^t, I^t, R^t\}$? Correct me if I am misunderstanding, but if this is the case, the proposal is  $\tilde{p}(x^t) = \text{Unif}(0, 5)$, presumably independent for for each state S, I and R. At any given time $t$, I would expect the states to be correlated (e.g. more recovered likely implies less infected and susceptible individuals). The use of an independent uniform does not capture this. Further, I presume the method is only valid  if we observe all three states?

As mentioned above: Should we aim for $\tilde{p}(x^t)$ to resemble a prior predictive simulation for a randomly chosen $t$? Could the proposal $\tilde{p}(x^t)$ be sequentially improved as the posterior is learned?
soundness: 3
strengths: The method is to the best of my knowledge novel and provides a solution to utilizing the known dependency structure of Markovian simulators; knowledge which is often neglected in simulation-based inference.  The application area is an important area and handling sequence data in simulation based inference is known to be often challenging (in terms of performance and computational cost). The experiments are convincing and consider a set of familiar but interesting models.
summary: The authors consider performing simulation-based inference for Markovian simulators. They propose two primary methods, based on score matching, and likelihood (ratio) estimation. In both cases, they utilize the Markovian structure to increase the efficiency of inference, by performing inference using data from single-state transitions, rather than requiring whole sequence simulation. In the likelihood case, the global likelihood is the product of the transition functions $p(x^{t+1}|x^t, \theta)$ (eq. 1), which once estimated can be used with MCMC to infer the posterior. Similarly, for estimating the score, they again use the Markovian structure to factorize the score (eq. 3), and use a "local" score estimator to approximate the global score. In both cases, a proposal distribution is used, $\tilde{p}(x^t)$, from which a sample is drawn, and a single transition is performed, which is used for learning.
weaknesses: My main criticism is that the paper should include a clearer description of what might constitute a "good" proposal distribution $\tilde{p}(x^t)$, especially when it is first introduced. I further feel that the requirement for choosing this proposal is a major problem for the practical utility of the introduced methods, and the authors should better address this problem. E.g. should we aim for it to resemble a prior predictive simulation for a randomly chosen $t$? Could the proposal $\tilde{p}(x^t)$ be sequentially improved as the posterior is learned? The use of hand-picked simple distributions is somewhat concerning for broader applicability.

Some smaller issues:
- The names of the contributed methods, at least FNSE, should be introduced earlier, for example, in the final paragraph of the introduction (FNSE is in Figure 1, but it's easy to miss the name). 
- "except FNRE; in LV": I don't believe LV abreviation was introduced, presumably Lotka-Volterra.
- Text size in figures is often a little too small.
- There is a reasonable degree of entangling of previous work in section 3.2.2, which made it hard to read. For example, the abbreviation FNPE is introduced without definition, initially implied to refer to the method by Geffner. It is then introduced and thereafter used as an adaptation of the Geffner method. It also refers to a score-based method, so the use of NPE is confusing as it implies a neural posterior estimation method, using e.g. flows.

Meta-review: The reviewers recommend acceptance (6-5-8-6-6-6). The paper presents a simulation-based inference approach for Markovian state-space models, leveraging the structure of the forward model to reduce the computational cost of the inference. The approach is well-motivated and the results are convincing. The author-reviewer discussion has been constructive and has led to a number of clarifications and improvements, with a better discussion of the related work and the addition of new results. The main concern raised by the reviewers is the incremental contribution of the paper, as the key idea appears to be a straightforward (but novel) extension of previous works. Nevertheless, the reviewers agree that the paper is well-executed and that the results are convincing. For these reasons, I recommend acceptance. I encourage the authors to address the remaining concerns and to implement the modifications discussed with the reviewers in the final version of the paper.

Decision: Accept (Poster)