RAG model implementation seems different from the paper
Author: XinyuHua
Pull request: no

Hi folks,

Thanks for open-sourcing RAG! After reading the model description in the paper and the actual code, I noticed a few discrepancies:

1. The marginalization for `RagSequenceForGeneration` seems a bit strange. From line 998 to line 1001 ([link](https://github.com/huggingface/transformers/blob/master/src/transformers/models/rag/modeling_rag.py#L998)), only the second tokens in `seq_logprobs` are getting scored by `doc_logprobs`:

```
# RAG-sequence marginalization
first_token_scores = seq_logprobs[:, :, :1, :]
second_token_scores = seq_logprobs[:, :, 1:2, :]
remainder = seq_logprobs[:, :, 2:, :]
rag_logprobs = torch.cat([first_token_scores, second_token_scores + doc_logprobs, remainder], dim=2)
```

I wonder if this is intended? I couldn't find this mentioned in the paper.

2. The marginalization for `RagTokenForGeneration` seems more like the Rag-sequence model in the paper, because the doc_scores are the same for all tokens in the same sequence. Is this correct?

```
# RAG-token marginalization
seq_logprobs = torch.nn.functional.log_softmax(seq_logits, dim=-1).view(
            seq_logits.shape[0] // n_docs, n_docs, -1, seq_logits.size(-1)
        )
doc_logprobs = torch.log_softmax(doc_scores, dim=1)
log_prob_sum = seq_logprobs + doc_logprobs.unsqueeze(-1).unsqueeze(-1)
return torch.logsumexp(log_prob_sum, dim=1)
```

Thanks in advance!




Comments:
Hey @XinyuHua,

regarding 1) I'm not really sure either. Maybe the author can give a better answer here (hope it's ok to ping you here @ola13) For 2) Yes it should be correct. You can see from the paper that at every generation step `i` the marginal probability over all tokens `z` is calculated (compare to the equation of RAG-Token Model in 2.1). In the equation, we sum over all `z` which corresponds to summing overall `doc_logprobs` above -> so this looks correct to me. Just the fact that the `marginalize` function (your code above) is executed at every forward pass shows that this has to correspond to `RagToken`.

Thanks for the explanation!

> the marginal probability over all tokens `z` is calculated

`z` is the document instead of token, right? And the paper says "we can draw a different latent document for each
target token and marginalize accordingly", so `p(z|x)` should be a different score for different `y_i`, but the code looks like `p(z|x)` is unchanged for all `y_i`s, whose log form is `doc_logprobs`. So effectively this is how `RagSequence` model is framed. Does that mean the `RagToken` model is actually trained the same way as `RagSequence`, but just the generation is different? 

Another related question, I checked the pre-trained `RagConfig` for `RagToken`, and the `do_marginalize` is actually set to False, so the marginalize method is never called during forward?

`z` is the tensor of containing the logprob of all docs -> this is why it's called `doc_logprobs`. If you check out the dimensions of this tensor you should see that one dimension exactly corresponds to `n_docs`. 

`do_marginalize` is called at every forward pass because it's set to `True` in the function argument here: https://github.com/huggingface/transformers/blob/8e74eca7f2b3235f8d5340d66361ea656c67bac7/src/transformers/models/rag/modeling_rag.py#L1099

Hi @XinyuHua, thanks for the questions! Regarding your first point:

> The marginalization for RagSequenceForGeneration seems a bit strange. From line 998 to line 1001 (link), only the second tokens in seq_logprobs are getting scored by doc_logprobs

Since we're operating in the log-space, multiplications from the formulas in section 2.1. of the paper become additions. So instead of multiplying the `p(z|x) * p(y| x,z)`, we sum: `log p(z|x) + log p(y | z,x)` (or `doc_logprobs` + `seq_logprobs` using our variables names from the code), where x is the input sequence, y is the output sequence and z is the retrieved document. Note that in the logspace, `log p(y | z,x)` decomposes into the sum of logprobs of each token in y.  In the part of the code you linked we perform this summation - we only want to add `doc_logprobs` once per sequence - we don't need to add it to each token of the sequence.

Now the reason we add `doc_logprobs` to the second token is that we want to avoid adding it to the BOS token, in case the target sequence doesn't contain one or in case the `exclude_bos_score` argument is used - otherwise we would effectively do no marginalization at all in these cases.

I hope this helps, but let us know if anything's still unclear!

Hi @ola13 and @patrickvonplaten , thanks for the detailed explanations!  

Hello, 

Thank you for the explanations! However, I'm still having some doubts: in the `get_nll()` function of `RagSequenceForGeneration`, are we basically marginalizing over doc_scores twice? First time is by adding it to `second_token_scores`, second time is by doing `ll = ll.logsumexp(1)  # logsumexp over docs` where dimension 1 is that of docs?

Thank you!