Compositional simulation-based inference for time series

We thank the reviewer for the constructive feedback and for recognizing the novelty and importance of our work. We address each of your concerns in detail below and outline the steps we will take to improve clarity and practical applicability in response to your review. We also thank the reviewer for pointing out smaller issues within the manuscript, which we will address in the updated version of the manuscript.

On Weakness:

“My main criticism is that the paper should include a clearer description of what might constitute a "good" proposal distribution $p(x_t)$, especially when it is first introduced… Should we aim for p~(xt) to resemble a prior predictive simulation for a randomly chosen t?”

We thank the reviewer for their valuable feedback. The performance of the proposed method is to some degree, sensitive to the selection of proposals, as shown in Fig. 2b and Fig. 8 (where we ablate over a large number of different Gaussian proposals). We agree that the current manuscript does not sufficiently discuss their construction. In response to your feedback, we have provided a practical guide to proposal construction (Sec. 3 in here). This will also be included in the revised manuscript, and we aim to improve our presentation in Sec. 3.

On Questions:

“Is xt in the SIR example the three states St,It,Rt? Correct me if I am misunderstanding, but if this is the case, the proposal is p~(xt)=Unif(0,5), presumably independent for each state S, I and R. At any given time t, I would expect the states to be correlated …”

Indeed, as the reviewer recognizes, the proposal we use is non-optimal (but by choice). Proposals that align with task correlation structure are expected to perform better. We decided to use a “simple” proposal across all benchmark tasks to demonstrate that our method also performs well in these cases, given that good proposals are not necessarily trivial to obtain (without using simulations themselves!). Although, it can be worth spending some simulations to construct a good proposal (see here, Sec. 3).

"Further, I presume the method is only valid if we observe all three states?"

Yes, we require Markovness. Marginalizing the state can lead to non-Markovian simulators.

“Could the proposal p~(xt) be sequentially improved as the posterior is learned?...”

Yes. Sequential methods in SBI usually target only a posterior for a specific observation $x_o^{0:T}$, whereas we here aim for an amortized approximation. Notably, a good proposal for a specific observation (or multiple ones) is to simply sample $x_o^t$ for random t (as at inference time, we will evaluate this observation). The proposal $p(x^t)$ can be freely updated over time in the sequential setting. However, the best way to improve it is not totally clear. An interesting direction for future work could be to update the proposal based on the Fisher information provided by $x_o^t$ (per time), which would avoid wasting simulations in regions that are uninformative for the $\theta$ (and tractable given that in FNSE, the networks estimate the score).

We thank the reviewer for your thoughtful review and constructive feedback. We are glad that you found our approach to simulation-based inference (SBI) for Markov chains conceptually appealing and well-aligned with the aim of increasing efficiency through local transition factor learning. We address each of your points below and outline the steps we will take to improve clarity and expand upon the technical and practical details in the revised manuscript.

On Weaknesses:

“The presentation of the method needs some improvements…. should remain self-contained ... .To save space, the FNPE approximation method could be placed in the appendix, as it is not used in the experiments…”

Thank you, for pointing out this issue in clarity. Our goal was to provide a concise overview, but we understand that this made it difficult to follow without referring to prior works. We will follow your suggestions and adapt the manuscript accordingly:

• We will move parts from FNPE to the appendix and expand on GAUSS/JAC
• We will adapt the corresponding sections to be self-contained.

“The presentation of the Kolmogorov flow experiments is missing details”

We thank the reviewer for pointing this out. We will add missing details and move important details from the appendix into the main section.

“The main weakness of the method right now is that the number of timesteps investigated in the paper remains smaller (up to 100),...”

We want to point out that previous work in the i.i.d. case only used up to 30 i.i.d. samples in their empirical evaluation (Geffner et al. 2022, Linhart et al. 2024). We acknowledge that for time series, it is even more relevant to support larger lengths. We chose the limit of 100 when comparing all the tasks, as we require ground-truth reference samplers for each task within the main evaluation, and not all samplers scale efficiently and robustly beyond. For the tasks where analytic or efficient sampling algorithms are available, we have also extended the analysis to 500 or 1000 transitions (Fig 2ab, Appendix Fig. 7). We aimed to address the performance degradation for FNSE over more and more timesteps (till 500) and diffusion solvers in Appendix Fig. 7 within the original manuscript. The performance did not degrade strongly over time in these tasks using GAUSS or JAC (in contrast to FNPE).

“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”

The increase is naturally attributed to the natural divergence of the time series over time (due to noise and different parameters). We will add this explanation to the manuscript.

On Question:

“"MCMC-based sampling corrections, such as unadjusted Langevin dynamics". The ULA doesn't include a correction step (it is "unadjusted"). What do the authors mean by that?”

We do mean ULA. In the context of diffusion models, samplers improved using MCMC methods like ULA, which are often referred to as predictor-corrector methods (where ULA would be the corrector). We will clarify this.

“Is there a particular motivation behind using Sliced Wasserstein Distance for evaluation … Otherwise, I'd suggest using the Sinkhorn Divergence instead.”

The main motivation was that sliced Wasserstein distances were also used in related prior work (Linhart et al. 2024). Furthermore, it is “parameter-free” in contrast to MMD (kernel) and Sinkhorn (regularization strength) and can be computed very efficiently (making it, in our eyes, a good metric). C2ST is a more common metric within related literature, is saturated, and is very sensitive for narrow distributions, making it ill-suited to show relative performance differences. Slicing preserves the metric axioms and the weak continuity of the divergence, implying that the sliced divergence will share similar topological properties (Kimia Nadjahi, 2022). At the same time, we agree that Sinkhorn divergence would be a good metric and might be more interpretable for people familiar with optimal transport. Yet, we think there is no large enough benefit to providing this additionally (or instead).

On Questions

“Is there any reason why NPE performs so badly in the C2ST statistic in Figs 3c and 3d? “

NPE results in Fig 3c seem to align with Lueckman et al. results (C2ST ~ 1.). However, while the tasks are inspired by Lueckman et al. 2022, there are some key differences that explain the discrepancy:

• The considered tasks are different. As Markov models, they are implemented through a stochastic differential equation (SIR additionally has a stochastic initial condition). Both involve a state-dependent diffusion term (state-correlated noise) which makes the task more challenging. Observations also differ (details can be found in Appendix B.1).
• Our simulation budget is reported based on the number of “transition steps.” Depending on how many “transition steps” per full simulation are required, the simulation budget for other work must be adjusted accordingly.

“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?”

Within our evaluation, the training budget (equaling the number of transition steps) is fixed across methods to, e.g., 10k. This means:

• Factorized methods obtain 10k pairs $(\theta, x^t, x^{t+1})$ as training data.
• NPE using 10-step simulation, in contrast, receives 1k paris of $(\theta, x^{0,10})$. Totaling to exactly the same number of required transitions. In addition, this dataset is augmented with $(\theta, x^{0,i})$ for $i<10$ to promote generalization across sequence lengths (without additional simulations).

In Fig. 2a and Fig2c, we evaluate the corresponding methods on sequences generated using $T=1,10,100$ transitions while the training budget of steps is fixed to 10k or 100k transition evaluations.

We thank the reviewer for the detailed review. We are glad that the reviewer found our approach intuitive, timely, and well-presented. In the following, we will address the reviewer's concerns and clarify aspects of our work. Below, we respond to each of your points individually.

On Weaknesses

"The proposed method is very incremental".

We recognize that our approach draws upon recent advances in score composition (Geffner et al., 2023; Linhart et al., 2024), and is based on an simple insight (which, furthermore, we maintain is novel!) -- with this conceptually simple change, we can substantially increase the scope of applicability of this approach. Almost all mechanistic simulators of time-series models are Markovian. Simulation-based inference for time-series models has been highly challenging, and our insight provides a general and effective way of addressing this issue! We do not agree that the effect of this insight, and thus the contribution of the study, is incremental! Our contribution lies in demonstrating how, based on our observation, these techniques can be adapted and effectively utilized for simulation-based inference in Markovian simulators – a “timely and significant task”. We demonstrate that posteriors can be efficiently approximated using only single (or few) step simulations. To our knowledge, this approach is novel and relevant, and addresses an important problem.

“The evaluations are in my opinion, not fully convincing”

The reviewer points out that we only compare our method to "exactly one baseline," which they consider "a bit insufficient given the wealth of SBI methods." In response, we have expanded our experimental evaluation to include additional SBI methods. Specifically, we have incorporated Neural Score Estimation (NSE) as a baseline (as well as approaches based on neural summary statistics, see below). Furthermore, we thank the reviewer for highlighting related literature, and we have included several of these methods in our evaluation to provide a more comprehensive comparison.

The reviewer pointed out that we could have “just computed summary statistics and used this as conditioning variables”. We want to clarify that, for our experimental evaluation using NPE, we require summary statistics that receive a sequence of variable size, compressing it into a static vector (in a likelihood-free setting, [1] would itself require estimating the likelihood ahead of time). This statistic should not only be nearly sufficient but also generalized to a “long” time series (larger than observed within training). We want to point out that this is a general and open problem in sequence models (e.g. [E1]). We do provide a sensible, commonly used baseline using an RNN (see, e.g., [E2]). Finding ``the best’’ approach lies outside the scope of the study. But we agree that an additional embedding approach would be useful to comparison and included it, which performed similar as our original baseline (see here, NPE’ Sec 2.).

We thank the reviewer for pointing out the work on “neural sufficient statistics”, and agree that it would be a valuable addition to our experimental validation. Specifically, we included Neural sufficient statistics [3], based on the Infomax criterium using the distance correlation approximation and Sliced Neural sufficient statistics [2] (see here, N(L/P)E+(+) Sec 2.). These baselines could sometimes improve upon our original NPE baseline, but overall, they performed similarly and distinctively from the factorized versions. In addition, we will extend our related work section to include [1-4]. We will incorporate these new extended results in Appendix C within the next day.

[E1] Zhou, Y., Alon, U., Chen, X., Wang, X., Agarwal, R., & Zhou, D. (2024). Transformers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371.

[E2] Dyer, J., Cannon, P. W., & Schmon, S. M. (2021, June). Deep signature statistics for likelihood-free time-series models. In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models.

Thank you very much for addressing my concerns and questions.

After having read the revision to the paper and the other reviews, I am willing to raise my score to 6. The technical contribution, which I still find to be fairly incremental, prohibits a higher score for me.

We thank the reviewer for the detailed and extensive review. We appreciate that the reviewer finds our paper “self-contained” and “easy to follow,” addressing a “relevant” problem with an “original” idea.

On Weakness

“the proposed method is a straightforward extension of FNPE”

Our approach draws upon recent advances in score composition (Geffner et al., 2023; Linhart et al., 2024), but we do not regard this as a weakness-- rather, our contribution lies in noticing how these techniques can be adapted (in a conceptually simple manner) to substantially improve simulation-based inference for Markovian simulators of time series. We show that posteriors can be efficiently approximated using only single (or few) step simulations. This approach is novel and effective. We do agree that the wording of “general SBI framework” should be tempered, and have modified this in the manuscript (we now call it “approach to SBI in time series”).

“I find it a bit unconvincing that FNPE can robustly approximate the correct global posteriors with sparse training. … suppose that most information about θ 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?”

We want to clarify that this study focuses specifically on Markov simulators (Hidden Markov Models (HMM), which the reviewer mentions are beyond the scope of this work, as noted in “Limitations” (Section 5.2.)). To clarify the scope further and highlight the inclusion of significant simulators, we will expand Section 1 with additional explanations and illustrative examples. However, for Markovian simulators, inference can be performed using just pairwise transitions by composing scores as we describe!

For example, let's assume that we have a Markov process in which most information about the parameters is contained in long-term simulation from a specific initial distribution ($T >> 0$), as in the example by the reviewer. Our approach can still leverage this information if the proposal shares sufficient support with the marginal distributions at $T$. We agree with the reviewer that having stationarity is beneficial for our methodology; for example, the stationary distribution would serve as a good proposal in the above scenario. For an unstable process, any ML-based approach will inevitably struggle given that, at some point, trajectories will simply fall out of training distribution (for too large T). Our approach is not excluded from this problem; however, for any fixed maximum $T$, the proposal can be designed to avoid this problem (similar to the maximum $T$ in training and other methods). We will aim to clarify this in the revised manuscript.

“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!)”.

We carefully checked the reviewer's concern by reevaluating parts of the benchmark using 100 test simulations, which did not lead to qualitative changes but confirmed and strengthened our initial results (see here, Sec. 1). We also want to clarify that we not only visualize the performance over 10 test simulations, but we plot the performance over five independent runs (training + evaluation), each considering different test simulations. This approach is in line with related studies (Lueckmann et al. 2021, Dax et al. 2023, Sharrock et al. 2024) that use 10 test samples but only one run (e.g. Geffner et al. 2023 use 6 test samples for five independent runs, Linhart et al. 2024 uses 25 test samples for one independent run).

“...at least some practically relevant metrics with an absolute interpretation”

We used metrics that are commonly used in the literature, and we do maintain that our current suite of metrics is sufficient to demonstrate our claims and, indeed, has an absolute interpretation (0 in sW and 0.5 in C2ST correspond to perfect point-wise approximations).

We do, however, agree with the reviewer that a calibration analysis would strengthen our manuscript, and we hence performed an additional analysis on the Lotka Volterra and SIR task (see here Sec. 4). The results are in line with our original benchmark, verifying that our methods are better calibrated than the NPE baseline (although not perfectly across time steps as also indicated by imperfect sW/C2ST due to limited training budget). We will incorporate this in Appendix C [now E.2], within the next few days.

“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)”

We do perform a number of ablation studies: Specifically, we already ablate on different relevant diffusion solvers and composition rules on a larger up of transitions (Fig. 7, [edit now Fig. 10]). Including additional adaptive solvers is out of the scope of this study as we do not aim or claim to improve standard diffusion samplers. We additionally ablate on proposals (Fig. 8 [edit now Fig. 11]), and partially factorized approaches (Fig. 6 [edit now Fig. 9]).

“It may be helpful to clarify some points in Section 2.2 regarding related work on amortized inference. … . Additionally, [6, 7] don't seem to introduce, discuss, or validate amortized methods…”

We will incorporate the suggestions [2,3,4,5] into our manuscript and thereby expand our literature review to also include amortized inference in a likelihood-based setting. In contrast, our work is situated within the context of SBI, where we consider "likelihood-free" or "black-box" scenarios.

On Questions

"Why does sW seem to, counterintuitively, increase for increasing the number of transitions (Figure 2)?"

The x-axis corresponds to the number of transitions used to generate the test observations on a fixed training budget (the length of the test time series). The training budget is fixed for all methods to achieve an overall 10k (or 100k) transition simulation. It is expected that the sW distance on all methods increases to some degree. For factorized methods, some local errors might accumulate. For NPE, the embedding network might not be able to generalize to larger sequences than observed in training (e.g., to 100).

"Will the method be robust to missing or corrupted data?"

Thank you for highlighting this important point. Our method can indeed handle missing or corrupted data. It relies on the Markov properties of time-series simulators at observation points, which are preserved even when data points are missing. One approach to address this is by amortizing over the time step size, $\Delta t$, allowing the method to "skip" missing observations. This will be discussed in Appendix A.

We thank the reviewer for the positive evaluation of our manuscript and the constructive feedback. The reviewer found the paper “exceptionally well-written” and that it provided valuable insights into the field. We appreciate your recognition of the strengths of our work, especially the practical relevance and the generalization capabilities of our method. We address your comments and questions below.

On Weakness:

“ Specifically, the proposal distribution seems like a critical design choice”

We agree that the current manuscript does not sufficiently discuss the proposal's construction. In response to your variable comment, we have provided the two practical and general approaches to construct the proposal (Sec. 3 in here), and we will revise our submission to include them.

the score composition rules prove surprisingly robust even when their assumptions are violated.

We can confirm that the Gaussian correction consistently improves performance, a conclusion also supported by Linhart et al. (2024). However, this finding only demonstrates that the correction enhances performance relative to the version without it; it does not imply that the Gaussian approximation perfectly represents the true correction term. As discussed in Section 5.2, some room for improvement remains, particularly in relaxing the Gaussian assumption, which represents an important direction for future research.

On Questions:

“How sensitive is the more complex Kolmogorov flow example to the choice of proposal distribution?”

The selection of the proposal distribution is important, especially in such structured high-dimensional simulations. It should be designed to ensure that its support encompasses possible simulation outputs over time (see answer on the proposal), and our construction in Appendix C.4 aims to address this in a simulation-efficient manner.

“Beyond the empirical results demonstrated, do the authors have a deeper understanding of the robustness of score composition rules?”

The Gaussian approximative assumptions in both GAUSS and JAC are only used within approximating the “correction term”, which involves the “denoised posterior” $p(\theta|\theta_a,x)$. While its influence is less significant at the end of the diffusion process (near noise), it becomes important once we approach the data distribution (Appendix C, Linhart et al. 2024). Intuitively, in this region, if now e.g., the posterior is multimodal, the distribution $p(\theta|\theta_a,x)$ is approximately not (as $\theta_a$ is $\theta$ + low amounts of noise so that the modes will collapse around the location of $\theta_a$) and can thus be to some degree approximate by a Gaussian. Of course, this cannot be claimed for larger noise levels.

Anyway, a big reason for preferring the Gaussian approximation is its simplicity and computational tractability. The method by Linhart et al. 2024 can, for example, be analogously extended with Mixtures of Gaussians (under similar assumptions). Yet, at least a naive application will run into computational challenges that must be addressed approximately (products of Gaussian will stay a Mixture of Gaussian but with exponentially increasing components, which in our case would be time steps).

“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)?”

Longer training simulations are expected to enable the training of score networks based on extended states, potentially leading to superior performance in long-term predictions. But instead, it incurs significant simulation costs during training. I will include this trade-off between prediction performance and simulation cost in Appendix of the revision.

We thank the reviewer for the detailed review and valuable feedback on our manuscript. We are happy that the reviewer recognizes the relevance of addressing computational complexity in simulation-based inference (SBI) for costly simulators and appreciates the “thoroughness” and “significance” of our experimental results for FNSE. We acknowledge the concerns and suggestions, and we aim to improve our paper accordingly. Below, we address each of your points in detail.

On Weaknesses

"An in-depth analysis/discussion around selecting the proposal and providing recommendations on its choice would have strengthened the paper"

We thank the reviewer for their valuable feedback. We agree that the current manuscript does not sufficiently discuss the construction of the proposal. The proposal is task-specific. It can, e.g., be chosen by simple domain knowledge (or observations available beforehand). Within the main manuscript, we use rather “naive” proposals, e.g., given a SIR model, we know that dynamics will be bounded between population size and zero, so we use a Uniform before that interval. This might not be applicable; if such knowledge is unavailable, then we might have to use the simulator within construction. This would require an upfront design cost but leads to better proposals. In response to your comments, we have provided a more general guideline to construct the proposal (see here, Sec. 3). And investigated this approach in the Lotka Volterra model, which demonstrated that the upfront cost of designing a proposal via simulations could improve performance even if this cost is subtracted from training budget. This will also be included in the revised manuscript, and we aim to improve our presentation within Sec. 3.

"The clarity of writing can be improved"

Thank you for your insightful suggestion. We fully agree that including examples will help readers better understand our paper. In response, we will add examples to Section 1 to clarify the scope and assumptions of the paper and aim to incorporate your suggestions.

"I found Section 3.2.2 difficult to follow."

We acknowledge that the current manuscript can be difficult to follow without referencing prior works. In response to your suggestion, we will address the following points in the revised paper:

• Ensure the explanations and notations are self-contained.
• Move certain technical details, which are not essential for understanding the main content, to the Appendix.

“Would be nice to have a small discussion on other sample-efficient SBI methods in the related works”

We agree that this should be mentioned and will include it in the revised version.

“ I would have expected performance to be plotted as a function the number of calls to the simulator”

We perform evaluations on different simulation budgets (number of total training transition steps), and within all panels, the budget across methods is the same, allowing discussion about simulation efficiency. We decided to plot it in terms of transition because it is important that the performance translates to a variable-sized time series. Yet, we agree that plotting by simulation budget is more common in related literature. We thus will provide such a figure in the appendix, in addition to our current approach (See here, Sec. 4, Fig.4).

On Questions

“Is there a reason why comparison with NSE is not done?”

No, we included it as an additional baseline within the Appendix (see here, Sec 2). This performed similarly to NPE, as expected, with small improvements in some tasks.

”Can FNSE be also made sequential in a similar manner as other simulation-efficient SBI methods such as SNPE/SNLE/SNRE?”

Yes. Sequential methods in SBI usually target only a posterior for a specific observation $x_o^{0:T}$, whereas we here aim for an amortized approximation. Notably, a rather good proposal for a specific observation (or a set of observations) is to propose $x_o^t$ for random $t$ (optionally with noise). The proposal $p(x_t)$ can also be updated over time in “rounds” without requiring corrections. In addition, one can, for example, use truncated sequential methods for FNSE, which was the only method found effective in the context of score matching (Sharrok et al. 2024). This method does truncate the prior based on the “posterior support” (usually some 99% highest density regions), which allows training without additional correction. Yet, within our framework, this truncation must be based on the “joint” support of the “local” posteriors instead of the global posterior support (as this is our training target). We leave this discussion to future work.

”Which score composition method is used in the experiments when implementing FNSE?“

We used the GAUSS method if not explicitly stated otherwise.

Multi-fidelity approach: Warne, D. J., Prescott, T. P., Baker, R. E., and Simpson, M. J. (2022). Multifidelity multilevel Monte Carlo to accelerate approximate Bayesian parameter inference for partially observed stochastic processes. Journal of Computational Physics, 469:111543.

Using quasi-Monte Carlo sampling: Niu, Z., Meier, J., and Briol, F.-X. (2023). Discrepancy-based inference for intractable generative models using quasi-Monte Carlo. Electronic Journal of Statistics, 17(1):1411–1456.

Using cost function of the simulator: Bharti, A., Huang, D., Kaski, S., and Briol, F.-X. (2024). Cost-aware simulation-based inference. arxiv:2410.07930.

Method based on early stopping: Prangle, D. (2016). Lazy ABC. Statistics and Computing, 26:171–185.