RAVEN: In-Context Learning with Retrieval-Augmented Encoder-Decoder Language Models

Thank you for highlighting the writing and extensiveness of our experiments.

novelty

The research problem of “in-context learning with retrieval-augmented encoder-decoder LMs” is an important one that has been overlooked in existing studies, making the problem itself novel. Regarding techniques, it is crucial to emphasize that our methodologies stem from an in-depth understanding of the limitations of current models and a careful consideration of how to overcome these issues, as discussed in Sections 3. The design is also non-trivial; it was based on the exploration of different strategies, as detailed in, e.g., Section 3.1 and Appendix C1-3.

In summary, our work offers contributions from both analytical and technological perspectives:

• From an analytical standpoint, we provide a thorough analysis of the in-context learning capability of state-of-the-art retrieval-augmented encoder-decoder language models. Our insights also highlight potential avenues for improvement. This analysis serves as the foundation for our proposed RAVEN training strategy and Fusion-in-Context prompting approach.
• From a technological perspective, we introduce a confluence of retrieval-augmented masked and prefix language modeling, coupled with our Fusion-in-Context Prompting and In-Context Example Retrieval strategies. These techniques not only enhance the base model's capabilities (by mitigating the mismatch between pretraining and testing) but also maximize the potential of in-context learning for these models (either by scaling the number of in-context examples with FiCL or by retrieving better in-context examples).

interpretation of Table 3

In this table, we aim to demonstrate how in-context example retrieval can further improve the performance of RAVEN. Yes, the baseline is RAVEN, and the baseline score for RAVEN 3B 1-shot on NQ is 31.7%. With in-context example retrieval, we achieve a further improvement of 9.1%, resulting in an accuracy of 40.8%.

Thank you for your positive and constructive feedback. We are happy that the reviewer likes our writing, methods, and experiments.

W1: evaluation on tasks that may not require retrieval augmentation

Thank you for your suggestion! Our primary focus in this paper is on “in-context learning with retrieval-augmented encoder-decoder models”, hence our emphasis on tasks benefiting from retrieval augmentation. We recognize the value of evaluating RAVEN on tasks that do not heavily rely on retrieval augmentation. In the MMLU benchmark, which includes a variety of tasks, some do not require retrieval augmentation. We observe RAVEN demonstrates performance improvements over ATLAS. We also observe performance gains when retrieval is turned off (i.e., retrieval of 0 passages) during evaluation.

W2: analysis on how ICL/FiCL behaviors (wrt number of demonstrations)

Actually, we have already provided some analysis on this in our paper. For example, in Figure 2 and Figure 5, as well as in the results in the tables, we report the performance of RAVEN/ATLAS with different numbers of in-context examples. For decoder-only LMs, we may refer to the results reported in previous papers. The performance varies significantly; for example, on TriviaQA, 1-shot is optimal for PaLM, while GPT-3 uses a 64-shot setting.

It’s true that the improvement of FiCL may also have a "ceiling", but we do not observe a performance drop like standard ICL when continuing to increase the number of examples (as shown in Figure 5).

W3: comparison to retrieval-augmented decoder-only LMs

We have compared RAVEN with Replug, Retro, Retro++ in Table 2, which are retrieval-augmented decoder-only LMs. From our experiments, when the capability of the base model is close, RAVEN can significantly outperform retrieval-augmented decoder-only LMs.

For decoder-only LMs like GPT-4 (stacking all the demonstrations and passages), there is still a gap, but the comparison may not be meaningful since the capability of the base model differs significantly (T5 vs. GPT-4). We notice there is a family of high-performing encoder-decoder LMs recently developed by Reka (https://arxiv.org/abs/2404.12387). It would be very interesting to see the performance of RAVEN built on Reka (though it’s not open source so far). We feel this is a very promising area for future exploration.

We thank the reviewer for highlighting the soundness and effectiveness of our methods.

“It's unclear whether the proposed methods remain effective when they are applied to today's popular LLMs, despite that most of them are based on decoder-only architectures.”

While decoder-only LMs are currently more popular, our research aims to explore and highlight the potential of encoder-decoder architectures. From a research standpoint, demonstrating the efficacy of alternative approaches is crucial for driving innovation and may encourage the development of more powerful encoder-decoder LMs. Notably, recent Reka models (https://arxiv.org/abs/2404.12387) are based on encoder-decoder architectures and achieve performance comparable to GPT-4.

Additionally, it is also interesting to explore the combination of the Fusion-in-Decoder architecture with existing decoder-only LMs. In this way, we can use a bidirectional architecture to effectively encode retrieved passages for powerful decoder-only LLMs.

“There is no evidence whether in-context learning is still necessary when fine-tuned LLMs are applied for RAG.”

This comment seems confusing. In-context learning applies to both pretrained and fine-tuned models with or without RAG. It remains important, as shown in leaderboards and technical reports of recent models like Gemini, where many results (of both pretrained or finetuned models) are still reported using in-context learning.

“Besides, efficiency is another concern of the proposed method, considering that in-context examples in the RAG settings will take a great deal of token consumption.”

The additional cost of RAG applies to all existing LMs. And encoder-decoder LMs are actually more efficient than decoder-only LMs in RAG settings. As explained in Section 2: “For decoder-only LMs, the computational cost typically increases quadratically with the input length, as well as with the number of retrieval passages. In contrast, for encoder-decoder LMs with a Fusion-in-Decoder architecture, the computation cost grows linearly with the number of retrieved passages, as they only perform self-attention over one passage at a time (Izacard & Grave, 2021).”

This efficiency also applies to FiCL. If we encode different in-context examples separately with FiCL, the computation cost also grows linearly with the number of in-context examples. Additionally, we don’t retrieve passages for in-context examples, which further improves efficiency and reduces token consumption.

Thank you for highlighting the solidity of our methods and experiments.

W1: novel insight

Our methodologies stem from an in-depth understanding of existing limitations and how to overcome them, as discussed in Section 3. The proposed methods are designed carefully based on novel insights and comprehensive analysis (see our response to Q1 below as a simple example).

1. There might be a misunderstanding regarding Section 3.2. Our design differs from previous work by combining retrieval-augmented masked and prefix language modeling to enhance in-context learning. In Section 3.3, we also propose a novel prompting strategy named FiCL that further boosts performance.

2. We don’t understand this comment. Strategy 1 puts all the input into the encoder. Should "decoder" in your comment be "encoder"?

W2: error analysis

Great suggestion! Briefly, for short-form QA, RAVEN achieves higher answer accuracy due to improved language modeling and better utilization of retrieved passages. For long-form QA, RAVEN is able to generate more reasonable and informative responses (examples in Appendix C4). We will provide more error analysis in the final version.

W3: MMLU

The gain is more significant for factual-seeking questions, where retrieval is crucial, but the improvement is relatively smaller for tasks such as answering math questions.

Q1: variant of strategy 1

We designed strategy 1 to align closely with the pretraining objective. The alternative prompting strategies suggested do not match this objective.

In retrieval-augmented prefix language modeling, the order is (prefix, <extra_id_0>, passages). strategy 1 aligns with this order, i.e., (examples + query, <extra_id_0>, passages).

• If the passage precedes q0 (a typo by the reviewer?), the order is (examples, passages, query, <extra_id_0>).

• If the passage is before all the QA pairs, the order is (passages, examples + query, <extra_id_0>).

Both alternatives mismatch pretraining, where passages are not after <extra_id_0>, thus are unlikely to perform well.

Q2: comparison

As shown in Tables 2 and 4, we also compare our model to other retrieval-augmented LMs such as RePlug, Retro, and Retro++, as well as several popular LMs. Performance varies significantly depending on the base model's capability. But in general, when the capability of the base model is close, RAVEN can significantly outperform the baselines.

Thank you for your suggestions. We will take into account all of them in the final version.

We thank the reviewer for highlighting the innovation and soundness of our methods, as well as the extensiveness of our experiments.

Q1: Section 3.1

Thank you for your feedback! We believe Section 3.1 is important as it is the first analysis of the in-context learning ability of Retrieval-Augmented Encoder-Decoder LMs. This analysis inspired the design of RAVEN and FiCL.

Yes, Appendix C1 (first/last/random) can also be treated as different prompting strategies. Regarding the design of “putting the retrieved docs into the prompt either”, one reason we don’t apply it is exactly due to the limited context length. Another reason is that this prompting strategy (i.e., q1 + doc1 + q2 + doc2 + … + q0 + doc0 => masked tokens) does not align well with the objective of pretraining (i.e., corrupted text + doc => masked tokens). Designing a pretraining objective for this prompting approach (as well as extend the context length) is a promising future direction.

Q2: masked spans

In the first training stage, masked spans are determined similarly to Atlas, with 15% of tokens randomly masked. In the second training stage, we mask 10% of the tokens on average at the end of the sequence, based on preliminary experiments. The p_reader is related to the Likelihood Distillation described in previous literature. We will add more explanations.

Q3: dynamic retrieval

Following ATLAS, we use Contriever (Izacard et al., 2022), where a single encoder embeds both queries and documents (different from DPR). Since we use a single encoder, embeddings for similar/related queries should be close. Therefore, we can simply use the same retriever.

Q4: prompting vs finetuning

Yes, we report the few-shot prompting results, not fine-tuning. Actually, fine-tuning for RAVEN can also yield better performance, but it’s not the focus of this paper since we want to disregard it, considering it’s not the common few-shot usage nowadays.

Q5: FiCL

• For standard ICL, we perform hyperparameter search in {2,3,4,5,8,16}; for FiCL, we search in {[40,1], [64,5], [64,8]}. We follow previous literature (e.g., GPT-3, PaLM) by evaluating on the dev/test sets directly.
• According to our results, FiCL [k,m] is (almost) always better than m-shot without FiCL. Settings like k=16 or m=1 also show FiCL outperforms standard ICL.

Q6: reference

This is explained in the previous paragraph: “...the computation cost grows linearly with the number of retrieved passages…(Izacard & Grave, 2021).”