probml / dynamax

State Space Models library in JAX
https://probml.github.io/dynamax/
MIT License
636 stars 69 forks source link

Add spectral methods for HMM learning #283

Open slinderman opened 1 year ago

slinderman commented 1 year ago

In addition to SGD and EM, we should support spectral learning methods. See for example,

Anandkumar, Animashree, et al. "Tensor decompositions for learning latent variable models." Journal of machine learning research 15 (2014): 2773-2832. link

@JeanKossaifi, it sounds like you may have some code for this in Tensorly?

Thanks @Anima-Lab for suggesting this!

JeanKossaifi commented 1 year ago

Thanks @slinderman!

TensorLy does provide all tools for tensor learning (including tensor decomposition) and is backend agnostic (and tested with JAX) so will also work transparently with dynamax. We will also be adding Tensor LDA soon.

Which methods exactly are you looking to add?

slinderman commented 1 year ago

I didn't realize TensorLy worked on JAX — cool! I am most familiar with the tensor decomposition approach described in the paper above, but to be honest I need to dig into the details. It's been many years, but I remember liking this short note too: https://arxiv.org/pdf/1204.2477.pdf. Any suggestions welcome.