Estimating SSMs based on LANs

[DominiqueMakowski]

The team behind the paper showcasing the usage of LANs to SSMs is actively working on a Python package (https://github.com/lnccbrown/HSSM) that implements this "fast" approach to compute SSM parameters.

And so I was curious about what do you think of this approach? Is this something that could be implemented in Julia? Are other alternatives interesting as well? Is it going to be less relevant once Turing works with efficient ADs like Enzyme?

[itsdfish]

This is something I have played around with in Julia using Flux.jl. Here is the repo: https://github.com/itsdfish/LANSandBox.jl

I don't remember what state I left it in. I might need to clean it up.

In general, I think machine learning approaches and standard approaches to Bayesian parameter estimation will both be relevant for a while. LANs are most useful the logpdf does not have a closed form solution, in which case it is possible to use a LAN to approximate the logpdf. The main challenge is generating sufficient training data and training the LAN. Once the LAN is trained, it is generally quite fast to evaluate the logpdf relative to Monte Carlo simulation. However, in Turing, NUTS is a bit slow with LANs. So having a faster AD like Enzyme would be helpful.

There are other neural network approaches, such as BayesFlow, which combines a normalizing flow with a summary network to perform amortized posterior inference, meaning you don't need to perform a separate MCMC sampling procedure because the posterior distribution is learned directly. In this regard, I think it's better than LAN. The main downsides are that with BayesFlow the prior distribution is fixed. With LANS, you can use any valid prior, but there is always a cost for inference.

Developers at Turing are working on a normalizing flow implementation, which is one of the main pieces of BayesFlow. A while back I mentioned the idea of implementing BayesFlow in Turing, but didn't really receive a response. Given that they are developing a normalizing flow package, it might be worth filing an issue and linking to the paper.

[itsdfish]

This is a good question. I think Julia would greatly benefit from a package similar to BayesFlow. Given that it is a general inference method, I think it is relevant for SSMs but outside the scope of this package. For this reason, I will close this issue. However, if something like BayesFlow is developed, I will consider either linking to the package or providing a worked example if it is not slow to run.