Foundation Inference Models for Markov Jump Processes

2024-06-10

Neural Information Processing Systems

AI Summary

This work addresses the zero-shot inference challenge for continuous-time, discrete-state Markov jump processes (MJPs) under high noise and sparse observations. We propose the first general-purpose foundational inference model, built upon a generalized prior-driven synthetic data generation framework coupled with neural supervised learning. The model jointly infers the rate matrix and initial state distribution, and natively accommodates diverse observation noise models and irregular sampling patterns. Its key innovation is cross-dimensional and cross-physical-scenario zero-shot generalization: a single pre-trained model achieves accurate inference across MJPs with varying numbers of states and distinct dynamical mechanisms—without fine-tuning. Evaluated on four real-world tasks—discrete flashing ratchet systems, molecular conformational dynamics, ion channel recordings, and simplified protein folding—the model matches state-of-the-art fine-tuned methods in performance, significantly advancing the universality and practical applicability of MJP modeling.

Technology Category

• Machine Learning: Probabilistic Circuits and Graphical Models
• Reasoning under Uncertainty: Relational Probabilistic Models
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

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

Problem

• Inferring Markov jump processes from noisy, sparse observations.
• Developing a zero-shot inference method for MJPs in bounded state spaces.
• Applying the model to diverse systems like molecular simulations and protein folding.

Innovation

• Zero-shot inference for Markov jump processes
• Neural network predicts MJP parameters
• Generalizes across diverse state space dimensions