LSH for Enc-Dec attention about trax

[lukaszkaiser]

@py4 - @Enumaris was definitely right: it's not trivial. We found a trick to go around it very recently though: run the encoder, then concatenate with target embeddings and run decoder on everything, then split for output. It has similar parameter count and complexity to enc-dec Transformer while allowing Q=K. Early experiments look ok, we're running more and it's all very new, but you can see the code here: https://github.com/google/trax/blob/master/trax/models/reformer/reformer.py#L709 ; Give it a run if it looks good to you (or let us know, reopen the issue if you have problems).

[Enumaris]

It does not seem trivial, given the Q=K requirement of LSH implemented here, to implement it also for Enc-Dec attention where Q and K come from different sources and can't be equated to one another - even often dimension-ally. Without the Q=K requirement, it seems batching and queries not having keys to attend to become issues. This is based off of my reading of the paper. It would be great to have the authors comment on this.