I've read the paper quickly, the main idea is simple and interesting, but maybe a little dubious (it's kind of an accuracy for memory trade-off).

In the transformer architecture one has to compute QKT.

QKT=(hd * Wq * WkT)heT (equation (2) page 3 in the paper).

Where hd is the hidden state of the decoder, and he is the hidden state of the encoder, and Wq and Wd are some parameters matrices, and T denotes the transposition operation.

By grouping the calculation this way, in a transformer encoder-decoder architecture, they can build and use only a single index (you index the he vectors using a vector database) for all the decoder layers queries. Instead of having to build 2 * L * H indices (with L the number of layers of the decoder and H the number of head in the decoder).

But what makes it a little dubious, is that this transformation mean you make your near neighbor queries in a space of dimension "dimension of the hidden state", instead of "dimension of a head" that is H times smaller.

So if you had to build 2 * L * H indices each index would be H times smaller.

So you only gain a factor 2 * L. But the trade-off is that you are doing a near neighbor search in higher dimension where you are then subjected to the curse of dimensionality (the higher the dimension the more similar all points are to each other). Whereas the whole point of projections in transformer is to lower the dimension so that the knn search make more sense. So to get the same accuracy, your near-neighbor search engine will have to work a lot harder.

Also as an approximation of the transformer, because it's using some knn search, it comes with the problems associated with it (for example harder to train because more sparse, and a tendency to hyperfocus), but it can be complemented with low-rank linearization of the attention to also have the neural net act on the gist rather than the closest neighbors.