Loglikelihood parallelization

[rsenne]

Hey! Currently working a lot with the package for future integration with StateSpaceDynamics.jl. Really great package! One quick thing that could be of use is parallelizing the log-likelihood calculation for single sequences of data? If I'm understanding the code correctly you parallelize for multiple sequences of data using an @threads call, but not for single sequences, instead opting for @simd. I think this makes a lot of sense for most observation models, but one issue I've been running into is for models where the likelihood calculation is quite numerically intensive, like in this paper here. I think an optional way to integrate alternative parallelization could be nice for these use cases, but would love to hear your thoughts/make sure I'm not misunderstanding your code anywhere.

[gdalle]

What kind of parallelization do you have in mind? The same kind they use in dynamax? Since data has to be processed sequentially, it's not as trivial as parallelizing across sequences.

[gdalle]

Also, have you profiled your StateSpaceDynamics code yet? If it is performance you're looking for, the bottlenecks are not always in the places you'd expect. I tried to do it myself but the README example didn't run.

I'm not saying that parallelization is a bad idea, but maybe there is much lower hanging fruit (like type stability) for whatever experiments you need to run right now.

[rsenne]

Thank you for the response! I think I may not have been clear enough about my suggestion. I'm not proposing to parallelize the sequential processing of the forward algorithm itself (which would indeed be challenging due to the dependencies between time steps).

Instead I'm proposing that for models where the observation log-likelihood is numerically intensive, and we assume the observations for all t are independent through the latent states (e.g., a switching LDS model would not qualify here), we can parallelize just the observation loglikelihood. Then from here one could move onto the forward/backward pass. Right now, I think your code works perfectly for most models (it is faster than what we are doing!) But, for the situation I'm describing I found that the speedup if this approach was quite helpful in the model I linked in my previous comment.

Also, yes! I've profiled a good bit, still much to be done....and I have quite a few ideas for further speeding up (one of which is to just use HiddenMarkovModels.jl). I'll also make look into why that example wasn't working and push an update to the README.md. Thanks for the heads up.

[gdalle]

I'm not proposing to parallelize the sequential processing of the forward algorithm itself (which would indeed be challenging due to the dependencies between time steps).

Note that it is actually possible, as I discovered while reviewing the dynamax paper.

Instead I'm proposing that for models where the observation log-likelihood is numerically intensive, and we assume the observations for all t are independent through the latent states (e.g., a switching LDS model would not qualify here), we can parallelize just the observation loglikelihood.

I think this can be achieved with the AbstractHMM interface, by expanding it to include obs_logdensities instead of just obs_distributions. Then the user is free to parallelize it if they want. What do you think?

[rsenne]

Note that it is actually possible, as I discovered while reviewing the dynamax paper.

Is this what you are referring to? Seems like an interesting future direction...

I think this can be achieved with the AbstractHMM interface, by expanding it to include obs_logdensities instead of just obs_distributions. Then the user is free to parallelize it if they want. What do you think?

I think this makes perfect sense and seems like the most painless solution.

Off topic:

I tried to do it myself but the README example didn't run.

Just looked into this. I can get this to run, but it seems to be a bit flaky with random init. I have a hunch about why this is, but setting trials=1, seems to be a reasonable temporary fix if you want to start profiling.