Entropic Matching for Expectation Propagation of Markov Jump Processes

TL;DR

The paper tackles latent-state inference for continuous-time Markov jump processes by marrying entropic matching with expectation propagation, enabling tractable, closed-form updates for a practical variational family. By applying this framework to chemical reaction networks with mass-action kinetics and a linear-Gaussian observation model, the authors derive FFBS-like dynamics for filtering and smoothing and integrate EP site updates to tighten posterior approximations, along with an EM-style scheme for parameter learning. Empirical results on Lotka–Volterra, motility, gene transcription, and enzyme-kinetics CRNs show improved accuracy of posterior means over several baselines and demonstrate scalability relative to particle-based methods. The approach provides a principled, fast alternative for complex continuous-time Bayesian inference in high-dimensional CRN models, while acknowledging limitations of the product-Poisson approximation and suggesting avenues for more expressive variational families and broader MJPs in future work.

Abstract

We propose a novel, tractable latent state inference scheme for Markov jump processes, for which exact inference is often intractable. Our approach is based on an entropic matching framework that can be embedded into the well-known expectation propagation algorithm. We demonstrate the effectiveness of our method by providing closed-form results for a simple family of approximate distributions and apply it to the general class of chemical reaction networks, which are a crucial tool for modeling in systems biology. Moreover, we derive closed-form expressions for point estimation of the underlying parameters using an approximate expectation maximization procedure. We evaluate our method across various chemical reaction networks and compare it to multiple baseline approaches, demonstrating superior performance in approximating the mean of the posterior process. Finally, we discuss the limitations of our method and potential avenues for future improvement, highlighting its promising direction for addressing complex continuous-time Bayesian inference problems.

Figures (6)

• Figure 1: Probabilistic graphical model and approximate inference scheme for mjp. The continuous-time Markov process $X_{[0,T]}=\{X(t) \mid t \in [0,T]\}$ emits the observations $\{Y_1,\dots, Y_N\}$. The approximate inference scheme consists of a continuous-time message passing algorithm, depicted by the dashed lines.
• Figure 2: Simulation and latent state inference of a Lotka-Volterra model. Solid lines represent the ground truth trajectory, while crosses indicate the observations. Approximate inference results of our method are visualized with dashed lines for the posterior mean using and the background indicates the inferred marginal state probabilities.
• Figure 3: Three selected species from the motility model adapted from wilkinson2010parameter, where only the species $\text{SigD}$ is noisily observed (crosses). The plots compare the results of our method and smc approach by indicating their resulting mean and standard deviation.
• Figure 4: Additional results of the motility model. Here, the posterior mean tracks roughly the time-average of the ground-truth signal. However, since for this species the effective realizations are very low, random fluctuations play a more significant role, making accurate tracking more challenging. The results of our method closely aligns with the smc results.
• Figure 5: Simulation of the Gene Expression model. The protein species is visualized in orange. Dashed lines denote variational posterior mean, solid lines denote ground truth trajectory, the background indicates the inferred marginal state probabilities and the crosses indicate observations.