Exact M-Step for Certain Emissions Classes

[bantin]

Currently, fitting an LDS always does an approximate update of the emissions parameters during the m-step (see line 719 of lds.py). We should instead check if the emissions update has a closed form (as in the linear-gaussian case) and do an exact update if possible.

[pabloi]

Is this on its way? I can volunteer to try to implement it (at least the linear-gaussian case).

Would there be a benefit of implementing the exact E-step for linear-gaussian dynamics too (i.e. Kalman smoother)? I'm not familiar with variational inference, but there may be a performance (computation time) advantage in using the closed-form solution given by the Kalman smoother when possible. Either way, it may be interesting to implement the Kalman filter and smoother just as as nice-to-have functionality.

[bantin]

Hi @pabloi ! We haven't started on this yet, but would love to have you implement it. We recently added similar functionality which does an exact M-step for the Linear-Gaussian Dynamics. In that case, we use the output of the Kalman smoother as you mentioned. For reference, we have standard and information-form versions of the Kalman Filter and Kalman smoother in messages.py

I think a good starting point would be to look through the M-step code for Linear-Gaussian Dynamics (This is the AutoRegressiveObservations class in observations.py), since this will likely have a similar flavor.

[pabloi]

Ok, I'll take a look at the existing m-step for the dynamics and try to replicate that.

Regarding the E-step: if K-smoother is already implemented, it should be straightforward to use it instead of Laplace EM with structured mean-field for the linear-gaussian LDS (as the comments on the code say, it would be equivalent, so using it may only improve performance for this special case).

I also noticed the LDS does not have a log-probability computation implemented. If the kalman smoother is implemented, that should be straightforward too, or am I missing something?

[bantin]

Great! Let me know if you have questions. As you mentioned, using Laplace-EM with a Linear Gaussian LDS is equivalent to doing exact EM (assuming Newton's method for the continuous state update). There's a balance here between having code that is general enough to work with all models, versus a small performance boost, and so far we have been leaning towards a more general approach. I think this generality is necessary to reduce the overall complexity of the package (which is already fairly complicated). That said, we might revisit this in the future!

Yes, the LDS log-probability computation is currently not implemented, but it should be straightforward to do so using the Kalman smoother. We would love help on that! We would just need to cache the previous values of the Log-normalizer (these are needed as inputs to the information-form version of the Kalman smoother)