Implement optimized closed-form gradients for Kalman Filter

[jessegrabowski]

I'm quite interested in the results of this paper. The authors derive closed-form gradients for backprop through Kalman Filters. Specifically equations 28-31.

They report a 38x speedup over autodiff gradients from PyTorch. I suspect (with no evidence) that the gradient computations are where the default PyMC sampler really fall down, so this might even make non-JAX sampling of SS models palatable.

[victorgarcia98]

Hey @jessegrabowski,

Cool paper!
I would like to contribute towards implementing (or at least trying) this gradients. To start, I can try to reproduce the paper and then I'll need some help to know where to integrate these gradients.

What do you think? 😄

[jessegrabowski]

Yeah that's how I'd proceed. You have closed form gradients from the current implementation, so basically you can just try to implement their equations 28-31 and get it to match with what we currently have. Hooking everything together should be a second step.

For a test model use a local level model. I recommend this because it doesn't have any parameters (except for the innovation covariance, and measurement error if included)

A big contribution to just getting the ball rolling would also just be to work out the notation differences between what this paper uses and what we use (we follow Durbin and Koopman Time Series Analysis by Statespace Methods).

[victorgarcia98]

Awesome! Sounds like I plan :D

I'll create a table comparing the notations!