Implementation of Memorizing Transformers (ICLR 2022), attention net augmented with indexing and retrieval of memories using approximate nearest neighbors, in Pytorch

Hi lucidrains,

This looks like excellent work! I have gone through the original paper and your repo, and am now trying to reproduce the model from the paper as closely as possible. Of course, the modifications you made such as hybrid attention instead of sigmoid gate are fine.

Specifically, I would like to be able to try some of the variations in Table 4:

Suppose I'm interested in the 4th to last row with Context 512 Memory 8192 XL cache 512. Can you help me the model arguments to do that? Here is my initial attempt, with reference to Section 4.2:

model = MemorizingTransformer(
    num_tokens = 32000, # vocab 32k
    dim = 1024, 
    depth = 12,
    memorizing_layers = 9,
    max_knn_memories = 8192, # Memory column
    num_retrieved_memories = 32,
    clear_memories_on_sos_token_id = 1,
    xl_memory_layers = (6, 7, 8, 9),  # not sure about this?
    xl_max_memories = 512, # XL cache column
    shift_knn_memories_down = 1, 
    shift_xl_memories_down = 1,
    # which argument corresponds to Context column?
).cuda()

A second question is what are the model arguments to reproduce to first row of Table 4, with no memory nor XL cache? Thanks in advance.

Thanks for the nice implementation. I've adapted this code for my own use, so I don't have the whole stack that would reproduce this bug. However, you can check for yourself.

The following code ought to update the KNN objects in the KNNMemory class:

@delayed
def knn_add(knn, key, db_offset):
    knn.add(key, ids = knn_insert_ids + db_offset)

Parallel(n_jobs = self.n_jobs)(knn_add(*args) for args in zip(knns, keys, db_offsets))

[link to that code here]

However, even after repeated calls to add to the memory, calling KNNMemory.search() results in empty values. If you view self.knns at this point, self.is_trained remains False.

When I modify the code as follows, this fixes the issue.

@delayed
def knn_add(knn, key, db_offset):
    knn.add(key, ids = knn_insert_ids + db_offset)
    return knn

updated_knns = Parallel(n_jobs = self.n_jobs)(knn_add(*args) for args in zip(knns, keys, db_offsets))
self.knns = updated_knns

This will allow searches to return actual values.

Hello and thanks for this implementation!

Do you know of any solutions to efficiently solve the "hard reset" problem in FAISS? I know that one could use IndexFlatL2 but that's not really efficient.

Thank you!

when I run train.py, error like this ,"index out of range: Tried to access index 10218 out of table with 255 rows. at /pytorch/aten/src/TH/generic/THTensorEvenMoreMath.cpp:418"happens

Thank you so much for the great implementation. I would like to ask whether your implementation for Memorizing Transformer could support multi-card distributed training like original paper. If you distribute the memorizingtrransformer model you created to each GPU, then every GPU would hold a memory with a retrieval faiss index. Therefore, each model on different GPU holds different memory database and retrieval index, which is different from the original paper. I regard that each model on different GPU should share the same retrieval context. This problem confuses me a lot.

Thank you so much for your time. Looking forward to your response!

I have two questions about the key and value calculation in Attention (and similarly for KNNAttention).

The relevant line is: https://github.com/lucidrains/memorizing-transformers-pytorch/blob/83fa1479d6f7881dd977fbff55681e709e3b250e/memorizing_transformers_pytorch/memorizing_transformers_pytorch.py#L135

1. Why is there only one Linear layer to_kv, instead of 2 linear layers to_k and to_v?
2. Why is the last dimension dim_head*2? I get that *2 is for both k and v, but what about dim_head? I thought q, k, v should all have the same final dimension (i.e. inner_dim==dim_head*heads). My understanding is that this means that either a) there is only 1 attention head, or for b) all heads, k and v are shared. Is there a reason this is done, or am I misunderstanding?

In your Attention class for Performer, q, k, v all have the same dimensions.

Thanks in advance!

https://github.com/lucidrains/memorizing-transformers-pytorch/blob/83fa1479d6f7881dd977fbff55681e709e3b250e/memorizing_transformers_pytorch/memorizing_transformers_pytorch.py#L237

Shouldn't this be (1-scale)?