Retrieval Collapse when fine-tuning RAG
Author: JamesDeAntonis
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: latest production version
- Platform:
- Python version: 3.8
- PyTorch version (GPU?): 1.7
- Tensorflow version (GPU?):
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 RAG: @patrickvonplaten, @lhoestq
 -->

## Information

Model I am using (Bert, XLNet ...):

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Fine-tune RAG on FEVER dataset
2. Notice that the same documents are retrieved every time

We are trying to fine-tune RAG on the FEVER dataset.  We follow the steps in the paper (we think) to a 'T' (RAG plus BART classifier head).  However, when we fine-tune, retrieval "collapses" (a term used in the paper) so that all queries retrieve the same irrelevant documents.  As a sanity check, we fine-tuned with a frozen retriever, and achieved similar results (72%) to what the paper achieves with frozen retriever.  Thus, it appears that perhaps there is a bug in HF's implementation of the retriever (and its gradients) that is causing this.  Alternatively, perhaps there is an obvious mistake in our config of the retriever.  Do you have any insights into this?  Thanks!


Comments:
Maybe @patrickvonplaten or @lhoestq can chime in here.

It may very well be that this feature is not implemented. @ola13, @lhoestq  do you have more insight here maybe?

The gradient does propagate to the question encoder. What configuration did you use ?

We used the default configuration and 'compressed' DPR.  We are indeed seeing that the gradient is propagating to the question encoder; before training starts, a retriever retrieves related articles, but after training, the retriever universally retrieves the same few documents regardless of the query.  That shows that the gradients are propagating, but not well.

Following up on this.  Any thoughts?  We are curious whether there is an issue with the implementation of DPR training in HF

Hi @JamesDeAntonis! Retrieval collapse is a problem we encountered in some setups and not necessarily caused by a bug in the retriever - basically what it means is that the passages retrieved at the beginning of the training are not useful enough so the models learns to ignore them.

We experienced collapse when training a RAG-Sequence model for FEVER, but we were successful with RAG-Token and RAG-Classifier. An option to move forward here could be:
- try training a RAG-Token generative model (it'd be generating the labels)
- share the classification code, maybe there's some issue there? Are you performing marginalization on top of the BART classification head logprobs?

Thanks for the response!  Good to hear that you were able to train successfully.  When you trained, did you use the two-label or three-label dataset?  (we are currently using the three-label)  I'm curious whether the inconclusive samples are contributing to the collapse.

* We are using the final hidden state of RAG-token as input into a classification head, and the model properly trains with the generator and classifier heads unfrozen (just the retriever is frozen in this case).  This gets to 72% accuracy, same as the paper.  I think this implies that the generator and classifier head are configured properly.
* We are indeed marginalizing on top of the BART classification head logprobs

Here is our classification head, mostly taken from HF:

```{python}
class BartClassificationHead(Module):
    """Head for sentence-level classification tasks."""

    def __init__(
        self,
        input_dim: int,
        inner_dim: int,
        num_classes: int,
        pooler_dropout: float,
        **config_kwargs
    ):
        super().__init__(**config_kwargs)
        self.dense = Linear(input_dim, inner_dim)
        self.dropout = Dropout(p=pooler_dropout)
        self.out_proj = Linear(inner_dim, num_classes)

    def forward(self, hidden_states: torch.Tensor):
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.dense(hidden_states)
        hidden_states = torch.tanh(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.out_proj(hidden_states)
        return hidden_states
```

..and here is the high-level model code:

```{python}
        # input_ids shape = (batch_size, 512)
        outputs = super().forward(input_ids=input_ids, attention_mask=attention_mask, **rag_kwargs)

        ### the following code is inspired by BartForSequenceClassification forward method
        # best practice for bart classification is to use the last hidden state
        # hidden.shape=(batch_size * n_documents, 300, 1024)
        hidden = outputs.generator_dec_hidden_states[-1]  # last hidden state;
        #print (hidden)

        # eos_mask.shape = (batch_size * n_documents, 300)
        eos_mask = outputs.context_input_ids.eq(self.rag.generator.config.eos_token_id)

        if len(torch.unique(eos_mask.sum(1))) > 1:
            raise ValueError("All examples must have the same number of <eos> tokens.")

        # pass along the hidden state at the eos token
        # (batch_size * n_documents, 1024)
        sentence_representation = hidden[eos_mask, :].view(hidden.size(0), -1, hidden.size(-1))[:, -1, :]

        # (batch_size * n_documents, 1, 3)
        document_level_logits = self.classification_head(sentence_representation)

        # finally, marginalize across all the retrieved documents
        # (batch_size, 1, 3)
        logits = self.marginalize(document_level_logits, outputs.doc_scores)

        # (batch_size, 3)
        logits = logits.squeeze(1)
```

We were able to train RAG-Token and RAG-Classifier successfully both on 2-way and the 3-way variant of FEVER. One important thing to note though is that those were on our internal `fairseq` implementation.

> try training a `RagToken` generative model (it'd be generating the labels)

What I meant when suggesting to use `RagToken` would be to use it as-is, without a classification head - it might seem counterintuitive but the generative model is actually able to learn to generate the labels.

As for the classification implementation - what you're proposing is quite different from our implementation in `fairseq`. What happens currently in your implementation is that you marginalize twice - once inside the forward pass on `RagToken`, and then again after applying your classification head. What we do instead is the following: 
1) take the generator hidden states (not marginalized)
2) apply BART-like classification head on top of that
3) marginalize

So basically - you don't want to just add a `BartClassificationHead` on top of `RagToken` hidden states. You want to implement something similar to [`BartForSequenceClassification`](https://github.com/huggingface/transformers/blob/master/src/transformers/models/bart/modeling_bart.py#L1254-L1346) - a `RagForSequenceClassification` of sorts, doing what I outlined above - if you're interested in implementing that I think it'd be a great contribution to the repo, cc @patrickvonplaten :)

Let me know if this makes sense!



Thanks for the advice!  I have a few questions

(1) what makes you think we marginalize twice?  `do_marginalize=False` in `RagToken` by default.
(2) What's the difference between a `BartClassificationHead` on top of `Bart`, and `BartClassificationHead` on top of `RagToken`?  Isn't the generator of `RagToken` simply `Bart` already, so the final hidden state of `RagToken` == final hidden state of `Bart`?
(3) Did you use 'adam' as your optimizer

And yes, I would be delighted to contribute it to the repo :)

Hi there! I'm a colleague of @JamesDeAntonis and I just wanted to chime in and clarify that we are using RagTokenForGeneration only as a neat wrapper to have a `self.rag` variable if we ever needed it. During training we take `outputs.generator_dec_hidden_states[-1]` as posted in the code above. Then we proceed with the 3 steps you listed above, essentially ending up with a SequenceClassification head from the rag generator hidden state outputs. We don't interact with the generation aspect at all, as you correctly identified.

> (1) what makes you think we marginalize twice? do_marginalize=False in RagToken by default.

Hey James and @suriyakode, I didn't realize that was the default configuration, in such case indeed your implementation should be equivalent to what I was suggesting.

And yes, we did use Adam as optimizer.

In such case I don't see anything obvious unfortunately. What accuracy do you get with the collapse?

Bummer :/

With frozen retriever, we achieved 72% accuracy, then unfroze the retriever and the accuracy fell to 68% after collapse

To be clear, you used `RagToken`, aka the current `RagTokenForGeneration` object that was implemented in HF?  My original thought was that something could be wrong with the gradients or something in this specific implementation

@JamesDeAntonis all of my FEVER experiments were done on `fairseq`, but I have been able to replicate RAG paper results training HF `RagToken` models on Natural Questions, which gives me some level of confidence in the implementation.

Ok, thanks.  What was your learning rate?

It was 1e-05 for training the classifier.

I also noticed an issue with the finetuning script. I ended up printing `doc_scores` while using the RAG finetuning scripts and saw that there was no gradient. Is there no gradient passed to the question encoder from the generator?

The `doc_scores` is supposed to have gradients. And the gradients are propagated to the weights of the RAG question encoder.

I cloned the transformers repo and pip installed transformers using the repo cloned to verify. I printed `doc_scores` in `RagModel` and got the following for what `doc_scores` was:
![Screen Shot 2021-01-29 at 3 20 34 PM](https://user-images.githubusercontent.com/18504534/106323530-69eaa300-622c-11eb-831c-19a2c5ba7f87.png)
I don't see gradients in the tensor.


Is there a fix to the `doc_scores` gradient?

It looks like there's no grad_fn on your `doc_scores`. Are the weights of the question encoder updated during finetuning ?

Nope the weights aren't updated.

Though should there even be a `grad_fn` when it's running "Validation sanity check" (as pictured above)?

Good catch @dblakely indeed during validation it makes sense to not have any gradient

Thanks for the clarification! I continued printing after and got a `grad_fn` in the tensor.

@ola13 is there any code that we can see regarding the internal `fairseq` implementation of RAG and the training you did with it?  I don't think there's any RAG in the public `fairseq` repo, but would be useful for me to be able to compare the two implementations

@ola13 I have a couple of questions on this
(1) how many docs did you retrieve?
(2) regarding the `RagForFeverClassification` idea, what is meant by "[we] first regenerate the claim" in section C of the RAG paper's appendix?  Does that mean that we should only provide tokens from the claim as `decoder_input_ids` in the generator?  Does it mean we should run multiple passes of the generator?  Curious what the correct interpretation is

Hi @JamesDeAntonis - sorry just noticed your previous commend, the `fairseq` implementation is not available publicly. 

Regarding your latest questions:
1) we used 5 docs at training time and tried evaluation with up to 50 docs
2) This idea is adapted from BART - https://arxiv.org/pdf/1910.13461.pdf - section 3.3 or Figure 3.a - in case of RAG, we don't copy all contextualized input, just the claim tokens. Since BART model will copy all input to `decoder_input_ids` if you just leave it at `None`  ([here](https://github.com/huggingface/transformers/blob/master/src/transformers/models/bart/modeling_bart.py#L1135-L1140)) you can try adding extra logic to only pass claim tokens to the BART decoder with RAG (that may mean explicitly setting `decoder_input_ids` like you mention).  This does not require multiple runs of the generator.

@JamesDeAntonis if you remember, could you please let me know what was the value of starting value of training loss, when using RAG -token or sequence models.

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.

@JamesDeAntonis Hello, do you have address this problem, or any thoughts? I think I have encounter a similar problem, so I would like to know how you deal with it afterwards?