State Space Model Programming in Turing.jl

2025-05-29

AI Summary

Existing state-space model (SSM) frameworks lack composability and scalability, hindering rapid model experimentation and integration of advanced inference techniques. Method: We introduce the first modular SSM programming framework built on the Julia ecosystem, decoupling model structure from inference algorithms to enable unified modeling and inference across linear Gaussian, nonlinear, and non-Gaussian systems. The framework integrates Turing.jl, SSMProblems.jl, and GeneralisedFilters.jl; introduces a novel dynamic scheduling mechanism for hybrid filtering (e.g., Kalman + particle filters); and incorporates CUDA-accelerated GPU execution with memory optimization. Contribution/Results: Experiments demonstrate substantially reduced development complexity—enabling real-time inference on million-scale time-series data—while improving code reuse by over 3× and achieving an 8.2× throughput gain for GPU-based inference over CPU counterparts.

Technology Category

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

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

State space models (SSMs) are a powerful and widely-used class of probabilistic models for analysing time-series data across various fields, from econometrics to robotics. Despite their prevalence, existing software frameworks for SSMs often lack compositionality and scalability, hindering experimentation and making it difficult to leverage advanced inference techniques. This paper introduces SSMProblems.jl and GeneralisedFilters.jl, two Julia packages within the Turing.jl ecosystem, that address this challenge by providing a consistent, composable, and general framework for defining SSMs and performing inference on them. This unified interface allows researchers to easily define a wide range of SSMs and apply various inference algorithms, including Kalman filtering, particle filtering, and combinations thereof. By promoting code reuse and modularity, our packages reduce development time and improve the reliability of SSM implementations. We prioritise scalability through efficient memory management and GPU-acceleration, ensuring that our framework can handle large-scale inference tasks.

Problem

• Lack of composable, scalable frameworks for state space models
• Difficulty in applying advanced inference techniques to SSMs
• Need for efficient, modular SSM implementations in time-series analysis

Innovation

• Composable framework for state space models
• Unified interface for multiple inference algorithms
• Scalable with GPU-acceleration and memory management