Tensorized Embedding Layers
Closed this issue · 6 comments
Recent work has shown promising results on tensorized embedding layers for memory reduction.
I have a tensorly-based implementation of an embedding layer (lookup table) that supports the CP/Tucker/TTM/TT format. If that's of interest, I would be happy to submit a PR to this repo. This implementation constructs the appropriate rows without forming the full embedding table.
Hi @colehawkins, indeed! I have been working on general indexing FactorizedMatrices and am almost done with TTM and CP. That allows to access rows of a TTM/CP table as e.g. TTM[[1, 2, 3], :].
Would be happy to join forces! I'll be unable to work on this for a few days but after this will finalize and push what I already have.
I'm happy to start working off of whatever you push.
I was doing exactly the same. It seems like the natural thing would be to finish the indexing, then wrap it in the embedding layer. I think I have a decent indexing function sorted for all four tensor formats, and hopefully it covers the Tucker/TT after you're done.
Hi @colehawkins - sorry for the delay, I haven't had much time. I wanted to unify tensor-factorizations and tensorized tensors but not yet sure whether it's a good idea - let me know what you think :)
I created a PR (#7) with most of the features ready though: the tensorized tensors are in tltorch/factorized_tensors/tensorized_matrices.py
I generalized TT-matrices to any orders and renamed it BlockTT to differentiate from TT. It supports generalized indexing too. I also wrote indexing for CP, for both the regular and tensorized version. Only Tucker indexing is missing .
No worries at all.
Based on your PR I can have a go at the tensorized embedding table, and Tucker should be trivial to add in later if I've got everything else working.
Sounds great! Looking forward to your PR! :)
Adding in #10