State Space Model Programming in Turing.jl

TL;DR

This paper introduces Julia packages SSMProblems.jl and GeneralisedFilters.jl within Turing.jl, providing a scalable, composable framework for defining and performing inference on state space models, enhancing flexibility and efficiency.

Contribution

It presents a unified, modular framework for state space models in Julia, enabling flexible inference and scalability with GPU support, addressing limitations of existing tools.

Findings

Supports multiple inference algorithms including Kalman and particle filtering

Enhances scalability with memory management and GPU acceleration

Facilitates easy definition and experimentation with SSMs

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.