Optimizing LSH attention runtime about reformer-pytorch HOT 4 CLOSED

I was thinking more of implementation than the research problem of LSH itself. Well, I'll just leave it for now then.

from reformer-pytorch.

@erickrf yes, I think this was noted in this chart in the paper

The payoffs only come at sequences ~2048 or above, depending on the hash rate. I introduced an extra keyword full_attn_thres, which you can set at your desired threshold to auto-switch between full and LSH attention, so that there is not wasted time and computation at smaller sequences. The other thing to note is, you could set the n_hashes to be smaller. Per my conversations with some of the authors, 4 was enough in most of their runs.

from reformer-pytorch.

@erickrf there is also perhaps room to improve on LSH. https://science.sciencemag.org/content/358/6364/793/tab-figures-data But that is more a topic for longer term research. I would be happy to make the framework more flexible to experimentation if that is something you would be interested in

from reformer-pytorch.

@erickrf if you have any ideas for improving on performance, welcoming PRs :)

from reformer-pytorch.