search

tensorly / torch

TensorLy-Torch: Deep Tensor Learning with TensorLy and PyTorch
http://tensorly.org/torch/
BSD 3-Clause "New" or "Revised" License
70 stars 18 forks source link

Tensorized Embedding Layers #6

Closed colehawkins closed 2 years ago

colehawkins commented 2 years ago

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.

JeanKossaifi commented 2 years ago

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.

colehawkins commented 2 years ago

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.

JeanKossaifi commented 2 years ago

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 .

colehawkins commented 2 years ago

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.

JeanKossaifi commented 2 years ago

Sounds great! Looking forward to your PR! :)

JeanKossaifi commented 2 years ago

Adding in #10