Related packages

[gdalle]

Hi there @briandepasquale and @rsenne.
Recently @THargreaves reached out to me to chat about SSMProblems.jl and HiddenMarkovModels.jl. Both seem closely related to StateSpaceDynamics.jl, so I was wondering if collaboration / coordination efforts could be worth exploring?
We could also extend the list to StateSpaceModels.jl by @guilhermebodin.

[rsenne]

Hey @gdalle! This has certainly been something we have talked about, and we were planning in the future to talk to you about likely interfacing HiddenMarkovModels.jl with StateSpaceDynamics.jl. From our end, we would love to collaborate/coordinate as @briandepasquale and I are very excited about getting Computational Neuroscientists (and neuroscientists in general) into the Julia space, and having others help our efforts is really exciting. I think as @briandepasquale mentioned in a prior issue, our eventual development target is models like the rSLDS. Our goal is mainly to develop a Julia equivalent to packages like the SSM and Dynamax. We have just now successfully implemented the Poisson LDS and have started writing a JOSS submission as we get to optimal performance. Would love to hear your thoughts!

[rsenne]

I took a look at the issue @THargreaves opened on HiddenMarkovModels.jl and think there’s a strong case for collaboration on this front in the future. For context on the inference/learning methods we’re using in StateSpaceDynamics.jl, the approach centers on direct optimization of the complete-data log-likelihood, paired with a Laplace approximation of the posterior distribution for the continuous hidden states. In the canonical Kalman Smoother case, this approach yields identical results, but it’s powerful because it generalizes easily to other observation distributions (e.g., Poisson, Bernoulli), allowing us to construct a variety of models quickly. This approach and its theoretical foundation are well outlined in A New Look at State-Space Models for Neural Data.

This framework also scales to hierarchical structures, like in the recurrent Switching Linear Dynamical System (rSLDS), where we can do "Variational Laplace EM" as described in this paper. The idea here is a CAVI-style parameter update where we apply a variational update to the discrete hidden states, followed by a Laplace EM step on the continuous states.

This general approach might complement future work on HiddenMarkovModels.jl and SSMProblems.jl well.

[THargreaves]

This all sounds very exciting. We would certainly be interested in looking into the idea of pooling our efforts together. The vast majority of our work is currently in a private repository which we hope to be sharing in a week. We're currently putting the final touches on a workshop paper for POPL 2025 with a deadline on Wednesday so are a bit pre-occupied with that. Once that is done, I'll take a proper look through your work and have a think about what the overlap might me.

At a high-level though we've been focusing on implementing "advanced" filtering and smoothing methods in composable and flexible ways. This includes Rao-Blackwellised particle filters/smoothers, PMCMC methods, etc. with a focus on GPU implementations using CUDA.jl (with the plan to later generalise to KernelAbstractions.jl). Perhaps of a bit of a different direction to your work but there could likely still be some common ground worth exploring.

A quick question for now: the two motivating Python packages you mentioned are (will be for SSM) built on Jax. Would you similarly be looking for GPU-acceleration, either through CUDA.jl, KernelAbstractions.jl or some wrapper for Jax?

[rsenne]

Sounds good! The main branch is pretty out of date as of right now so make sure to poke your head through the other branches (right now the most up to date for LDS type models is Poisson_speedup).

As for your other point, in the future I would love to add GPU-acceleration via CUDA.jl, but admittedly this has not been a major focus yet as right now I've been working a lot on getting the most basic model offerings reliable/efficient. But, ideally this would be available in the future.