How to use past_key_values in RAG model?
Author: agi-templar
Pull request: no

Hi the team,

Great work on this amazing project! Currently I'm trying to test RAG's performance on some simple QA tasks and I read the document of [RAGforSequenceGeneration](https://huggingface.co/transformers/model_doc/rag.html?highlight=rag#ragsequenceforgeneration). I wonder how exactly to use `past_key_values` to speed up the generation?

I run the following code which is from your example:

```python
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", use_dummy_dataset=True)
# 1. Encode
question_hidden_states = model.question_encoder(input_ids)[0]
# 2. Retrieve
docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt")
doc_scores = torch.bmm(question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)).squeeze(1)
# 3. Forward to generator
outputs = model(context_input_ids=docs_dict["context_input_ids"], context_attention_mask=docs_dict["context_attention_mask"], doc_scores=doc_scores, decoder_input_ids=labels)
```

and get `outputs` that has five keys:

```python
['logits', 'doc_scores', 'past_key_values', 'context_input_ids', 'generator_enc_last_hidden_state']
```

I then tried to use the model's `generate()` function to get the actual sequence. For the `past_key_values`, what should I plug in here? The document said: "Tuple consists of two elements: `encoder_outputs` of the RAG model (see `encoder_outputs`) and`past_key_values` of the underlying generator. " Thus, I composed a tuple with the `outputs['generator_enc_last_hidden_state']` and the `outputs['past_key_values']` from the latest step generation. It still does not work. 

During debugging, I checked the tensor size of the `outputs['past_key_values']`. It is a list of 12 elements (n_layers) and each element is a tuple of four tensors: two of them are of `torch.Size([5, 16, 11, 64])` and two of them are of `torch.Size([5, 16, 300, 64])`. It seems like the later ones are key-value pairs (same size, 300 as hidden), but what are the previous ones (with the size of `torch.Size([5, 16, 11, 64])`)?

I would really appreciate it if you can point me out more references (docs, code) about how to use the RAG model properly!

Many Thanks!



(I think @patrickvonplaten, @lhoestq are the right people who can help.)


Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

Hey @DapangLiu,

I'm very sorry to reply so late here! When doing generation with `RegSequenceForGeneration`, the `past_key_values` are automatically used to speed up generation - so you don't have to do anything ;-) 

RAG uses Bart as the decoder and RAG as well as Bart have set `use_cache=True` by default. This then means that `past_key_values` are automatically returned after the first forward pass,
see: https://github.com/huggingface/transformers/blob/1fec32adc6a4840123d5ec5ff5cf419c02342b5a/src/transformers/models/bart/modeling_bart.py#L442
and:
https://github.com/huggingface/transformers/blob/1fec32adc6a4840123d5ec5ff5cf419c02342b5a/src/transformers/models/bart/modeling_bart.py#L1478

Those `past_key_values` are then continuously forwarded during generation :-)

Hi @patrickvonplaten,

Thanks a lot for your response! The reason why I asked this is because I want to add some noise to the `past_key_values`.

As I mentioned before, I found there were two parts that have different dimensions in the `past_key_values` --- my guess is  one is for the cross-attention from the encoder and the other is for the decoder-only key value paris --- but I'm not sure whether it is correct (any docs I can read?).

Or they are both decoder-only key value pairs? Why are they in different dimensions?

Sorry for so many questions. If you could refer me to some docs I can study by myself. Thanks!

Exactly one are the saved `cross-attention` projections and the second ones are the decoder past key and value projections of all previous steps. You should see that during generation only the tensors of one of the two parts change -> those are the decoder-only past key value pairs :-)

Thank you! Very helpful!