RAG architecture

RAG architecture from the original paper

Since loss is calculated at the output layer of the generator, how the gradients are back propagated to the retriever model?

Because the input to the Generator is pure text i.e. text of retrieved document + question.

It states that the entire architecture is fine-tuned end-to-end. Notice in the figure "end-to-end backprop through q" where q is the Query Encoder. My question is how does error is back propagated to q? Because while calculating the loss at the output of Generator, the Query Encoder model plays no role.


1 Answer 1


The paper clearly stated that

The retriever (DPR) provides latent documents conditioned on the input, and the seq2seq model (BART [32]) then conditions on these latent documents together with the input to generate the output.

So it was evident that the generator takes text (input ids) as input. But then the question arises that "how gradients are propagated till the query encoder model?"

I had to go through the source code of the RAG project on huggingface to find an answer to this question.

I started with the RagModel class. Going to the source code of the class, I observed that the generator model was indeed taking input_ids as the input in the forward function.

In order to find how gradients are backpropagated till query encoder, we need to find out how loss is being calculated. We can see a variant of nll loss is being used (line 864).


Since query encoder has its skin in the game only because of doc_scores, we will be following it. By reading through the functions (get_nll --> marginalize), we observed that a log-softmax is taken over (a) generator sequence logits, and (b) doc_scores.


Then both these quantities are added and passed through a log sum exponential trick. Why log-sum-exp? I think it is because it will be equivalent to summing up the probabilities ($\sum_{i}{e^{log(p_i)}} = \sum_{i}{p_i}$) and taking the log (for numerical stability). Hence, the gradients are backpropagated till the query encoder model.

Since query encoder model has no ground-truth, the retrieval loss is simply calculated as softmax of doc_scores. It is my intuition that when this retrieval loss is added to the LM loss (seq_logits) then the resulting loss basically represents the following philosophy:

"if the final generated answer is good enough (low LM loss) then it indicates that the fetched documents were also good enough. Hence, good query representation! Not much change is needed in query encoder. Similarly, large LM loss --> bad fetched documents --> big changes (parameter updates) are needed in query encoder."


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.