probml / dynamax

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

Add spectral methods for HMM learning #283

Open slinderman opened 2 years ago

slinderman commented 2 years 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 2 years 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 2 years 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.