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

Authors: jie  huang
wei  ping
peng  xu
mohammad  shoeybi
kevin  chen-chuan
  chang
bryan  catanzaro
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-23 01:31:39

Abstract: In this paper, we investigate the in-context learning ability of retrieval-augmented encoder-decoder language models. We first conduct a comprehensive analysis of existing models and identify their limitations in in-context learning, primarily due to a mismatch between pretraining and inference, as well as a restricted context length. To address these issues, we propose RAVEN, a model that combines retrieval-augmented masked language modeling and prefix language modeling. We further introduce Fusion-in-Context Learning to enhance the few-shot performance by enabling the model to leverage more in-context examples without requiring additional training. Through extensive experiments, we demonstrate that our simple yet effective design significantly improves performance, achieving results comparable to the most advanced language models in certain scenarios, despite having substantially fewer parameters. Our work underscores the potential of retrieval-augmented encoder-decoder language models for in-context learning and encourages further research in this direction.

Main paper: 
 introduction: 
 Recent advancements in natural language processing have been predominantly driven by the development of large language models (LLMs) (Brown et al., 2020; OpenAI, 2022; 2023; Chowdhery et al., 2023; Smith et al., 2022) . These models have demonstrated remarkable performance across a wide range of tasks (Qin et al., 2023; Bubeck et al., 2023; Huang & Chang, 2023) . One of the key features that enables these models to excel is their ability to perform in-context learning (Dong et al., 2022) . By conditioning on given context, LLMs can adapt to new tasks and domains without the need for task-specific fine-tuning. This enables LLMs to perform well on zero-shot or few-shot learning tasks, where only a limited number of examples are available.While in-context learning has been extensively studied for decoder-only language models like GPT-3 (Brown et al., 2020) and PaLM (Chowdhery et al., 2023) , research on encoderdecoder language models, which have shown to learn stronger representations (Devlin et al., 2019; Raffel et al., 2020) , remains limited. Notably, Patel et al. (2023) tap into the potential of mT5 (Xue et al., 2021) , a multilingual encoder-decoder LM, by iteratively prompting the model to produce long generations with in-context examples. Chung et al. (2022) ; Longpre et al. (2023) finetune T5 (Raffel et al., 2020) with a large mixture of tasks using instruction tuning (Mishra et al., 2022; Wei et al., 2022; Sanh et al., 2022) to improve model performance and generalization to unseen tasks in both zero-shot and few-shot settings.On the other hand, LLMs still face challenges such as hallucination and limitations in representing the long-tail and most recent knowledge (Mallen et al., 2022; Huang et al., 2022; Luu et al., 2022; Jang et al., 2022; Zheng et al., 2023) . Retrieval-augmented language models (Izacard et al., 2023; Borgeaud et al., 2022; Wang et al., 2023; Shi et al., 2023) have emerged as a powerful approach to address these issues by retrieving relevant knowledge from an external corpus. Among these, the encoder-decoder models, such as ATLAS (Izacard et al., 2023) , stand out. They benefit from the strong representation ability of a bidirectional encoder, coupled with of the efficacy of a Fusion-in-Decoder architecture (Izacard & Grave, 2021) , enabling the effective integration of multiple retrieved passages. Despite these advancements, in-context learning with these models remains underexplored.In this regard, we first conduct a comprehensive analysis of the state-of-the-art retrievalaugmented encoder-decoder language models by designing and experimenting with different prompting strategies. We find that these models exhibit a certain in-context learning ability; however, due to a mismatch between pretraining and inference and a limited context length-issues that are common to existing encoder-decoder LMs trained with masked language modeling-its few-shot performance is not stable and providing more than, e.g., 8-shot, examples does not lead to further improvement.Based on the analysis, we develop RAVEN1 by first mitigating the mismatch between pretraining and inference through a combination of retrieval-augmented masked language modeling and prefix language modeling. Moreover, to enable the model to learn from more in-context examples, we propose Fusion-in-Context Learning, a novel approach that allows the model to utilize more in-context examples without modifying the model configuration or requiring additional training. Furthermore, we suggest using the retriever of the model to obtain relevant in-context examples to further enhance few-shot performance. Our empirical results demonstrate that RAVEN significantly outperforms previous retrieval-augmented encoder-decoder LMs in both zero-shot and few-shot settings, even achieving comparable results to decoder-only LLMs in some settings despite having 180 times fewer parameters.The main contributions of this paper are twofold:• From an analytical standpoint, we provide a thorough analysis of the in-context learning ability of retrieval-augmented encoder-decoder language models. We demonstrate the possibilities and offer insights for future development. • From a technological perspective, we introduce RAVEN, coupled with our Fusionin-Context Learning and In-Context Example Retrieval strategies, building upon the analytical groundwork. These techniques, though simple, are highly effective. They not only enhance the base model's capabilities but also highlight the potential of in-context learning with retrieval-augmented encoder-decoder LMs. 
background and related work: 
 Retrieval-augmented language models are a class of language models designed to enhance their performance by incorporating external knowledge. These models typically employ an information retrieval mechanism to access relevant information from a large corpus, which is then integrated into the model's prediction process. Retrieval-augmented LMs can be based on both encoder-decoder (Izacard et al., 2023; Lewis et al., 2020) and decoderonly (Khandelwal et al., 2020; Borgeaud et al., 2022; Shi et al., 2022) architectures. 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 concept is also investigated by Ye et al. (2023) for more efficient in-context learning.While there has been some research on in-context learning with retrieval-augmented decoder-only LMs, which can be straightforwardly implemented by concatenating retrieved passages with the query as the input of the LM (Mallen et al., 2022; Shi et al., 2023; Khattab et al., 2022) , in-context learning with retrieval-augmented encoder-decoder LMs remains unexplored to the best of our knowledge. This is despite the fact that encoder-decoder LMs can be more efficient at incorporating multiple (e.g., 40) retrieved passages.  
methodology: 
 In this section, we first explore in-context learning with retrieval-augmented encoderdecoder language models in the literature. Building upon the analysis, we develop models with enhanced zero-shot performance and improved in-context learning abilities.To investigate the in-context learning ability of retrieval-augmented encoder-decoder language models, we first aim to gain insights from the state-of-the-art designs in the literature. Among them, the design of ATLAS (Izacard et al., 2023) stands out; it combines a general-purpose dense retriever with a sequence-to-sequence reader (i.e., T5 (Raffel et al., 2020) ) using the Fusion-in-Decoder architecture (Izacard & Grave, 2021) . The retriever, encoder and decoder are jointly trained during the pretraining process. In this process, the dense retriever, based on the Contriever model (Izacard et al., 2022) , is responsible for selecting relevant passages from an external knowledge source, e.g., Wikipedia, based on the given corrupted context. The retrieved passages are then processed along with the context by the encoder, which generates the corresponding output, i.e., the masked spans, at the decoder (Figure 1 , left). ATLAS demonstrates exceptional few-shot performance on knowledge-intensive language tasks (Petroni et al., 2021) , despite having a lower parameter count compared to other recent LLMs.However, in Izacard et al. (2023) , the few-shot performance is achieved by finetuning the model with few-shot examples, which requires additional training and may limit its applications, such as dealing with dynamic and diverse real-time user queries like GPT-3/4 (Brown et al., 2020; OpenAI, 2023) , where in-context learning plays a vital role. Therefore, we take the initiative to explore the in-context learning ability of this type of models, using open-domain question answering (Chen et al., 2017) as a representative task for some preliminary experiments.Prompting Strategies. To facilitate in-context learning, an effective prompting strategy is paramount. In contrast to decoder-only LMs, where the input can only be fed to the decoder, encoder-decoder LMs can take input in either the encoder or the decoder. In alignment with the pretraining objective, we identify two prompting strategies for in-context learning:Strategy 1. The first strategy involves feeding all example question-answer pairs and the target question to the encoder, without any input to the decoder. The prompt is designed as:2 TriviaQA 0-shot 1-shot 5-shot 8-shot 0-shot 1-shot 5-shot 8-shot Enc: Question: q 1 Answer: a 1 . . . Question: q k Answer: a k Question: q 0 Answer:<extra id 0> dwhere (q 1 , a 1 ), . . . , (q k , a k ) represent example QA pairs, q 0 denotes the target question, <extra id 0> is a sentinel token (Raffel et al., 2020) , and d is the relevant passage retrieved with q 0 . An example in a 2-shot setting is illusated in Figure 1 (middle).Strategy 2. As the decoder of the encoder-decoder model can also accept input, we can feed the answers of in-context examples to the decoder and only feed the questions to the encoder, using multiple sentinel tokens:Enc: Question: q 1 Answer:<extra id 0> . . . Question: q k Answer:<extra id (k -1)> Question: q 0 Answer:<extra id k> d We select two widely-used datasets in the domain of open-domain question answering for the preliminary study: Natural Questions (NQ) (Kwiatkowski et al., 2019) and TriviaQA (TQA) (Joshi et al., 2017) 3 . Table 1 summarizes the results. We find that the model struggles to learn from in-context examples using strategy 2, as the few-shot performance is worse than the zero-shot performance. We hypothesize that this is because the model has difficulty learning the pattern of S2 with masked language modeling during its pretraining, since it is unlikely to obtain several consecutive question-answer pairs (or something similar) in the form of strategy 2 by randomly masking several spans in a sequence.Dec: <extra id 0> a 1 . . . <extra id (k -1)> a kOn the other hand, we observe that with strategy 1, the model does exhibit some in-context learning ability, where the 5-shot and 8-shot performance is significantly better than the zero-shot performance on both NQ and TriviaQA. Therefore, we choose to focus on strategy 1 for further study and disregard strategy 2 for the remainder of the paper.As the encoder of encoder-decoder language models is bidirectional, it can also examine in-context examples that follow the target question to fill in the masked token. This means that we may position the target question at the beginning or middle of a sequence, for example:Question: q 0 Answer:<extra id 0> Question: q 1 Answer: a 1 . . . Question: q k Answer: a k dQuestion: q 1 Answer: a 1 . . . Question: q 0 Answer:<extra id 0>. . . Question: q k Answer: a k dTable 2 summarizes the results. We denote the target question's position as "first" for the beginning of the sequence, "random" for a random position, and "last" for the original setting (S1). Interestingly, placing the target question anywhere other than the last position results in a significant performance drop. Upon examining the generated answers, we observe that when the target question is placed at the beginning or in the middle, the model tends to repeat the answer or generate additional text. For example, for the prompt "Question: What number in Bingo is sometimes referred to as Heinz varieties? Answer:<extra id 0> Question: . . . ". The generated text is "57 'Heinz varieties' is a term used in Bingo to describe". This indicates that the model does not fully understand and follow the style of in-context examples. Therefore, by default, we position the target question after all the in-context examples.The number of in-context examples is a crucial hyperparameter for in-context learning. Generally, we expect better performance from a model with more in-context examples, but there is an upper limit due to 1) the maximum context length setup, e.g., 512 tokens, during the pretraining process, and 2) the point at which the model has received sufficient examples and cannot gain additional information from more examples. The optimal number of in-context examples also varies between models. For instance, on TriviaQA, PaLM (Chowdhery et al., 2023) exhibits better 1-shot performance than settings with more examples, while this is not the case for GPT-3 (Brown et al., 2020) .Figure 2 illustrates the impact of varying the number of in-context examples across different model sizes. Interestingly, the 11B model demonstrates poor performance in low-shot settings, e.g., 1-shot, but improves significantly after 4-shot and 5-shot. Upon examining the generated responses, we find that the model tends to produce answers with more tokens in low-shot settings, while the ground truth typically consists of shorter answers with fewer than 5 tokens. By relaxing the criteria for a correct prediction to include instances where the ground-truth answer is a substring of the model output, we find that the 1-shot performance surpasses that of the 0-shot setting (38.3 vs 32.1 on NQ).All models perform well in the 5-shot and 8-shot settings, but their performance does not continue to improve with more in-context examples (e.g., 16-shot). We believe this plateau may be attributed to two factors: 1) the sequence length constraints during pretraining, where the maximum input length to the encoder is set to 384 tokens, and the average input sequence length (excluding passages) is around 130 tokens; 2) the model's ability to learn adequately with 5 or 8 examples, making additional examples less beneficial.Effect of Number of Retrieved Passages. Figure 3 illustrates the impact of the number of retrieved passages on model performance. We observe that for both 0-shot and 5-shot settings, the performance of the models increases significantly with the number of retrieved passages. This highlights the effectiveness of the Fusion-in-Decoder architecture (Izacard & Grave, 2021) for knowledge-intensive tasks like open-domain question answering, and underscores the importance of pretraining language models with retrieval augmentation. Additionally, the 5-shot performance consistently outperforms the 0-shot setting. This observation further emphasizes the value of providing in-context examples to improve the performance of retrieval-augmented encoder-decoder language models.Input to Encoder: Machine learning algorithms build a model based on sample data, known as training data, in order to make predictions or decisions without being explicitly programmed to do so. Machine learning algorithms are used in a wide variety of applications, such as in medicine, email filtering, speech recognition, agriculture, and computer vision, where it is difficult or<extra_id_0> Passage: … machine learning models require a high quantity of reliable data in order for the models …Output: <extra_id_0> unfeasible to develop conventional algorithms to perform the needed tasks. In §3.1, we observe that retrieval-augmented encoder-decoder LMs exhibit a certain ability for incontext learning, which has been overlooked in previous studies. However, there are also some limitations such as unstable performance in low-shot settings, and the fact that providing more in-context examples does not consistently improve performance.To learn a better retriever and enhance the bidirectional understanding ability of the reader, as demonstrated in Izacard et al. (2023) , a practical choice is to pretrain the model with the masked language modeling objective, where the input is a corrupted text with several masked spans placed randomly within the sequence (refer to Figure 1 (left) for an example). However, in testing, based on our analysis in §3.1, it is most effective to place the target question after all the in-context examples, with a masked token (i.e., <extra id 0>) following the question (Figure 1 , middle)). Thus, there exists a mismatch between pretraining and inference.To solve this issue, we propose combining retrieval-augmented masked and prefix language modeling. Specifically, in the first stage, following Izacard et al. (2023) , the retriever and reader are trained jointly with retrieval-augmented masked language modeling. The training objective for the retriever is to minimize the KL divergence KL(p READER ∥ p RETRIEVER ) between the passage posterior distribution according to the reader and the passage distribution from the retriever over the top-K retrieved passages, i.e., p READER (d) =exp(log p LM (a | d,q)) ∑ K i=1 exp(log p LM (a | d i ,q)) , p RETRIEVER (d) = exp(s(d,q)/T) ∑ K i=1 exp(s(d i ,q)/T), where s(•) calculates the dot product between the query q and passage d vectors, and T is a hyperparameter. The training objective for the reader is to maximize the likelihood of the masked spans with n retrieved passages: ∑ i log p(a i | q, {d k } 1,...,n , a 1:i-1 ).In the second stage, for each sequence, we mask 10% of the tokens on average at the end of the sequence with the <extra id 0> token. Then, we use the retriever obtained from the first stage to retrieve relevant passages using the prefix and train the reader to recover the suffix of this sequence with the prefix and the passages as input. An example of input and output for retrieval-augmented prefix language modeling is shown in Figure 4 . We can observe that the pretraining objective aligns more closely with prompting strategy 1 in Figure 1 . We refer to the model trained with this combined objective as RAVEN.RAVEN benefits from both the retrieval-augmented masked language modeling, which contributes to a better reader and retriever, and retrieval-augmented prefix language modeling, which mitigates the gap between pretraining and inference. This design is non-trivial. In Appendix C.1, we verify the effectiveness of it by exploring different training strategies.In §3.1, we observe that the performance does not further improve with more in-context examples after 8-shot. One major reason for this is the limited sequence length during the pretraining process, which makes it difficult for the model to handle long sequences during inference. Pretraining models with longer contexts would be a potential solution, but it would significantly increase computation cost. Additionally, the maximum input length is also constrained by the maximum sequence length of the retriever, i.e., Contriever, which is based on BERT (Devlin et al., 2019) and has a maximum length of 512 tokens.As an alternative, we propose an approach to enable models to learn from more in-context examples without requiring additional training. As described in §3.1, the reader is based on the Fusion-in-Decoder architecture (Izacard & Grave, 2021) , where multiple passages are retrieved, and each passage, concatenated with the in-context examples and target question, is fed to the encoder separately (Figure 5  
experimental setup: 
 Datasets. Following the setup in §3.1, we first evaluate on two widely-used open-domain question answering datasets: Natural Questions (Kwiatkowski et al., 2019) and Trivi-aQA (Joshi et al., 2017) . Additionally, we conduct a case study on long-form question answering using the ELI5 dataset (Fan et al., 2019) . Furthermore, we assess the models' language understanding ability using the Massively Multitask Language Understanding (MMLU) benchmark (Hendrycks et al., 2021) . Detailed information regarding the MMLU evaluation is in Appendix B.5. Other evaluation settings are the same as §B.1.Baselines. Since both RAVEN and ATLAS (Izacard et al., 2023) are trained starting from T5, we choose ATLAS as a primary baseline for comparison. We also compare our model with decoder-only LLMs such as GPT-3 (Brown et al., 2020) , PaLM (Chowdhery et al., 2023) , and LLaMA (Touvron et al., 2023) (in a closed-book setting). Additionally, for open-domain QA, we evaluate our approach against REPLUG (Shi et al., 2023) and RETRO (Borgeaud et al., 2022) , as well as its improved version RETRO++ (Wang et al., 2023) . These models are decoder-only language models augmented with retrieval. REPLUG is based on Codex (Chen et al., 2021) and Contriever (Izacard et al., 2022) , where the passages are retrieved by Contriever (using ensemble and additional adaptation) and fed directly to Codex. RETRO is a GPT model (Radford et al., 2019) augmented with a transformer encoder to encode the retrieved passages. RETRO++ is a variant of RETRO that feeds the most relevant retrieved passage into the GPT decoder while providing other passages to its encoder. For MMLU, we also include FLAN-T5 (Chung et al., 2022) , an enhanced version of T5 that has been trained on a large mixture of tasks with instruction finetuning.4 We choose open-domain QA as our primary evaluation task, as it effectively represents knowledge-intensive challenges and is widely employed in real-world applications. 3 present the exact match (EM) scores for ATLAS and RAVEN on the NQ and TriviaQA datasets. Both the 3B and 11B RAVEN models significantly outperform ATLAS. For instance, on TriviaQA, RAVEN 11B achieves an improvement of 8.8%, 30.7%, and 2.8% in the 0-shot, 1-shot, and few-shot settings respectively, compared to ATLAS 11B. Furthermore, the performance of RAVEN increases steadily with the number of in-context examples, while the performance of ATLAS experiences a substantial decline in low-shot settings, demonstrating the effectiveness of RAVEN across various shot settings.Learning. We also report the results of models with Fusion-in-Context Learning (FiCL) in Table 3 . For both ATLAS and RAVEN, FiCL contributes to approximately a 1% improvement, which is not attainable by standard in-context learning, where performance does not further improve (or even decreases) with more than 8 in-context examples. This demonstrates the superiority of FiCL for enabling models to learn from more examples.Comparison to Other Models. In Table 3 , we further compare RAVEN to other baselines. On NQ, RAVEN's zero-shot and one-shot performance surpasses all the baselines, including PaLM, even though RAVEN 3B has 180 times fewer parameters than PaLM 540B. The zero-shot performance of RAVEN on TriviaQA is also on par with PaLM 62B. Furthermore, RAVEN's zero-shot performance significantly exceeds that of both RETRO and RETRO++, which are retrieval-augmented language models of a similar scale.In the few-shot setting, with FiCL, RAVEN achieves performance comparable to GPT-3 175B and PaLM 62B. However, there remains a gap between RAVEN and the larger PaLM 540B * For TriviaQA, PaLM's 1-shot performance surpasses other settings. We follow the original paper to report the 1-shot result.For other models, we select the best k-shot (k ∈ {2, 3, 4, 5, 8, 16}) performance or report the number in the original paper.Table 3 : Results on NQ and TriviaQA. Since the performance varies significantly depending on the capability of the base model, the results from models other than ATLAS should only be used for reference to gauge the position. And we assume RAVEN can achieve significant performance improvement when based on a stronger base model. and Codex 175B models. Nevertheless, given the considerably smaller scale of RAVEN in comparison to PaLM and Codex, its performance can be considered impressive. The performance of RAVEN may be further improved if it is built upon a larger model, in which case its few-shot performance is likely to surpass that of PaLM and Codex.Effect of Number of Retrieved Passages. Figure 7 illustrates the effect of the number of retrieved passages. As the number of retrieved passages increases, we observe a significant performance improvement of RAVEN 11B in both the 0-shot and 5-shot settings.In-Context Example Retrieval. §3.4 suggests using RAVEN's retriever for in-context example retrieval. Results in Table 4 show that this approach improves RAVEN's few-shot results, especially on NQ where a ∼10% improvement is observed. This indicates the positive impact of incorporating more relevant in-context examples. Table 5 summarizes the results (accuracy) on Massive Multitask Language Understanding (MMLU). We find that the zeroshot performance of RAVEN is impressive, surpassing the few-shot performance of GPT-3 175B and being slightly worse than PaLM 62B, despite having a significantly smaller number of parameters. Furthermore, with the same number of parameters, the performance of RAVEN is far superior to T5. Additionally, even without instruction finetuning, RAVEN achieves performance comparable to FLAN-T5, a model finetuned on a large collection of tasks. We expect further improvement of RAVEN by applying instruction tuning as well and leave it for future study.Interestingly, with standard in-context learning, the few-shot performance of RAVEN is worse than zero-shot, possibly due to the longer questions and answer options in MMLU causing context length issues in the 5-shot setting. Also, in the one-shot setting, since MMLU is a multiple-choice QA task, providing only one example might introduce bias in the model's prediction, favoring a specific option. However, with Fusion-in-Context Learning, the performance improves significantly, leading to better few-shot performance for the 11B model compared to its zero-shot performance, further demonstrating the effectiveness of FiCL. 
conclusion: 
 In this study, we have delved into the in-context learning ability of retrieval-augmented encoder-decoder language models. We commenced with a comprehensive analysis of the models in the literature and subsequently developed our model based on the analysis. Our extensive experimental results demonstrated that our model significantly outperforms previous models and achieves results on par with some of the most advanced language models, even with substantially fewer parameters. These findings highlight the potential of retrieval-augmented encoder-decoder language models in the realm of in-context learning. 
 
 Appendix: While the performance of RAVEN is impressive considering its scale and training budget, there are also some limitations. One limitation arises from the constrained context length inherent to the base models (i.e., T5 or ATLAS) we employed. This restriction poses challenges to the scalability of in-context learning, especially as the number of in-context examples increases. While our Fusion-in-Context Learning (FiCL) strategy does offer a mitigative approach to this constraint, an alternative and possibly more optimal solution might involve extending the context length. This would be particularly beneficial for tasks requiring extensive inputs.Furthermore, when compared to some of the prevailing decoder-only language models, particularly those exceeding 100B parameters, the models deployed in our research might appear relatively diminutive in scale (in terms of both the number of parameters and the amount of training data). Our endeavor partially seeks to catalyze further investigations into more powerful encoder-decoder models.Nonetheless, the insights and methods proposed are transferable and have the potential to enhance other models, including those that are domain-specialized or more powerful, such as mT5 (Xue et al., 2021) and UL2 (Tay et al., 2023) . Future work focusing on scaling up the model, applying these methods, and further studying its in-context learning ability is encouraged. Drawing on the benefits of scaling up and combining this with our proposed approaches, we believe that there is potential to develop even more powerful retrievalaugmented language models in the future. Another promising future direction is exploring how to combine the Fusion-in-Decoder architecture with existing decoder-only language models. By doing so, we can harness the advantages of both architectures-employing a bidirectional architecture to effectively encode retrieved passages for the most powerful decoder-only LLMs.B.1 Experimental Setup for §3.1We select two widely-used datasets in the domain of open-domain question answering: Natural Questions (NQ) (Kwiatkowski et al., 2019) and TriviaQA (TQA) (Joshi et al., 2017) .To assess the performance, we follow the previous work (Izacard et al., 2023) to employ the standard exact match (EM) metric. For the few-shot settings, we follow Brown et al. (2020) We train two versions of RAVEN: 3B and 11B. To isolate the effect of training variance with masked language modeling, we initialize both the retriever and the reader of the models with the weights of ATLAS (3B and 11B) and continue to pretrain the model with prefix language modeling. To isolate the effect of retrieval, we do not update the retriever during the training process for prefix language modeling. We pretrain the reader using the December 2021 Wikipedia corpora preprocessed by Izacard et al. (2023) , where the index is also constructed using the same corpora. In accordance with Izacard et al. (2023) , we retrieve 20 passages for each masked sequence (excluding passages identical to the original sequence). Both the 3B and 11B models are trained for 5,000 steps, using AdamW optimizer (Loshchilov & Hutter, 2019) with a batch size of 64. We employ a learning rate of 4 × 10 -5 for the 3B model and 1 × 10 -5 for the 11B model, with linear decay and 100 warmup steps. All the models are trained on NVIDIA A100 GPUs (80 GB). For the 3B model, we utilize 8 GPUs, whereas for the 11B model, we employ 32 GPUs. The prompt used for prefix language modeling is detailed in Appendix B.3. During testing, we default to retrieving 40 documents for all tasks. The prompts used can be found in Appendix B.4 and Appendix B.5.In alignment with the pretraining of ATLAS, we design the prompt for prefix language modeling as {prefix}<extra id 0> title: {title} context: {text} where {prefix} represents the prefix of an input sequence. The {title} and {text} elements are retrieved by the model's retriever using the prefix as a query. Here, {text} signifies the retrieved passage, while {title} denotes the corresponding article and section title of the passage. The model is trained to generate <extra id 0>{suffix}where {suffix} is the suffix (masked by <extra id 0>) of the input sequence.In accordance with pretraining, we use the following prompt for open-domain question answering:Question: {question} Answer:<extra id 0> title: {title} context: {text} For example, Question: In which country was the first permanent bungee jumping site situated? Answer:<extra id 0> title: Bungee jumping: Modern sport context: first permanent commercial bungee site, the Kawarau Bridge Bungy at the Kawarau Gorge Suspension Bridge near Queenstown in the South Island of New Zealand. Hackett remains one of the largest commercial operators, with concerns in several countries. Several million successful jumps have taken place since 1980. This safety record is attributable to bungee operators rigorously conforming to standards and guidelines governing jumps, such as double checking calculations and fittings for every jump.As with any sport, injuries can still occur (see below), and there have been fatalities. A relatively common mistake in fatality cases is to use a cord thatMMLU comprises 57 multiple-choice question answering datasets that span various domains, including elementary mathematics, US history, computer science, and more. For the evaluation on MMLU, we report the accuracy and use an index composed of December 2021 Wikipedia corpora. We follow Izacard et al. (2023) to apply the "de-biased" inference. Specifically, during inference, we execute four forward passes, each corresponding to a cyclic permutation of the answer letter-option assignment within the question. For instance, the answer option designated to letter 'A' is shifted to 'B', 'B' to 'C', 'C' to 'D', and 'D' to 'A'. The final prediction is obtained by summing up the probabilities from these four forward passes.We design the prompt in the following format: We conduct an ablation study by training ATLAS and RAVEN with different pretraining strategies. First, to isolate the effect of more training steps of RAVEN, we also train ATLAS for 5,000 more steps using the masked language modeling objective. Results in Table 6 (row 2) show that the performance does not improve, indicating that the performance improvement of RAVEN compared to ATLAS is not simply due to training for more steps.Second, to verify the effectiveness of RAVEN's training strategy (i.e., first masked language modeling, and then prefix language modeling), we train two variants of RAVEN, starting from the T5-lm-adapt checkpoint6 , which is the checkpoint that ATLAS starts from. For the first variant, we use the same prefix language modeling objective of RAVEN. For the second variant, we train the model with a mixture of masked and prefix language modeling. Specifically, we construct corrupted texts by both masking 15% spans in the sequence (same as ATLAS) and replacing the suffix with a special mask token <extra id 99> (used in testing). We train the model for 10,000 steps and update the retriever and refresh the index during training with the optimal strategy described in Izacard et al. (2023) . Table 6 (RAVEN -in row 3 and 4) summarizes the results. We find that the performance of these two variants is superior to ATLAS, but inferior to RAVEN when trained using the strategy described in §3.2. An explanation for this is that, by training with masked language modeling first, the model can achieve better language understanding ability and is equipped with a more effective retriever (as emperically verified in Izacard et al. (2023) ). Subsequently, by training with prefix language modeling, the mismatch between pretraining and inference is mitigated, resulting in improved zero-shot and few-shot performance.Table 7 presents some example outputs of ATLAS and RAVEN 11B on long-form question answering. The questions are sampled from the ELI5 dataset (Fan et al., 2019) . An examination of these results reveals that ATLAS typically generates concise answers, while the output from RAVEN generally encompasses more information. This is a predictable outcome given that ATLAS is pretrained solely with masked language modeling, where each masked span usually contains only a handful of tokens. Besides, while RAVEN's answers are not always entirely accurate, they generally exhibit higher quality compared to ATLAS. Furthermore, the use of Fusion-in-Context Learning in RAVEN appears to contribute to a more coherent and informative generation.RAVEN, a bird known for its intelligence and adaptability, has the letters "RA" in its name, which represents "Retrieval-Augmented" in our context.Here we present a format designed for better demonstration. The actual prompt, which follows the template used in pretraining, can be found in Appendix B.4.Experimental setup is detailed in the Appendix B.1.Implementation details are described in Appendix B.2.https://github.com/facebookresearch/atlashttps://huggingface.co/google/t5-xl-lm-adapt 


Decision: This paper investigates the combination of using an encoder-decoder model, with in-context examples, and retrieval augmentation, looking at ATLAS as a main comparison point (without considering finetuning of the model on downstream tasks)

The authors add a pretraining stage that makes train and test mismatch smaller, and also use different examples at test time, and improve prompting and the retrievers. The reviewers mostly agree the paper is written well, experiments are convincing and empirical gains are sizable so this is a good addition to the literature on retrieval-augmented language modeling.

Reviews: review: paper_summary: paper_summary: This paper studies the problem of retrieval and in-context learning for encoder-decoder languages models. Specifically, they find limitations of prior methods that arise from being not truly few-shot in that they require additional fine-tuning. This work then proposes a new model (RAVEN) that obviates the need for fine-tuning by modifying the prompt/encoding format. The retriever part is trained by a KL objective while the reader is trained using masked language modeling for prefix-based LM. Additionally, this work introduces “Fusion-in-context learning” (FiCL), which distributes a high number of examples (e.g. 64) across multiple embeddings that get encoded. Both RAVEN and FiCL are shown to be effective on Natural Questions, TriviaQA, and MMLU compared to the prior state of the art model. 

The main strengths are that the method appears to work well and also are well-suited solutions to the disadvantages of prior work (fine-tuning); in some sense, this work modernizes past work with more attention to prompts and ICL. The main weakness, which could be a result of my own misunderstanding, is the comparison against baseline, ATLAS, does not seem to correspond with the original paper (more explained below).
strength_weakness: strength_weakness: * The paper presentation is a little bit confusing; the first contribution of “thorough analysis” does not seem too conclusive.
* Some technical details were confusing, especially as there is a strong reliance on the reader understanding all the technical details and hyper parameters of ATLAS first. 
   * In particular, the ATLAS baselines reported here are lower than the numbers in the ATLAS paper.
questions: questions: * Section 3.1 feels out of place relative to the rest of the paper. The conclusions of 3.1 are that prompt 1 is better than prompt 2 and that going up to 8-shot is sufficient given the small context sizes (512). Both of these findings could belong in an appendix in favor of clearer explanation and analysis of the main method RAVEN and secondary method contribution FiCL, both of which are, in my opinion, more convincing and stronger results. 
    * To nitpick a little more, strategy 1 and 2 are not the only strategies: C.1. actually outlines 3 (first/last/random) more, and none of these consider putting the retrieved (maybe limited to top-1 title only due to seq. size) docs into the prompt either -- which would be more akin to what RAG-like approaches would do to fully mimic the data point for ICL. 
    * I stared at Figure 1 for a long time at the beginning trying to understand what it had to do with encoder/decoder architecture, expecting it to be a compare/contrast of ATLAS and RAVEN, but later on realized it was unrelated to 3.2.

* This might be a naive question: how are masked spans determined? (i.e. what constitutes a span in this model?) If these decisions follow those of ATLAS, perhaps it should be made even clearer in Sec 3.2 that a lot of the modeling objectives/decisions come from ATLAS. e.g. How does p_reader differ from ATLAS's p_attn?

* The dynamic retrieval confused me a lot. The retriever is retrieving text documents, and the embeddings are embeddings of the documents. The few-show examples are (q, a) pairs, which is a different "type" than documents. How, or why, should we use the same retriever? And does this retriever get trained with an auxiliary objective somehow to retrieve good (q, a) pairs?

* ATLAS's numbers in their paper are because they fine-tuned their model further -- is that why the scores reported in this paper are lower than the ones in their paper? Since one of the objectives of this work was to fix some of the shortcomings of ATLAS, it should probably be also compared to the models like ATLAS + finetuning. If it comes close but doesn't beat ATLAS + finetuning yet, then that would be useful to know as practitioners can be informed about which model to use based on their compute budget, etc. If it surpasses ATLAS+finetuning, that would be a stronger result.
 
* I would have liked to see a little more discussion around FiCL. The tables are picking the best numbers based on the [k, m] hyperparameters
    * Are [k, m] picked based on a validation set?
    * What are the tradeoffs between higher [k, m] and no FiCL? Is FiCL always strictly better than few-shot, or if k (or m) is small, would it still be preferable to do few-shot without FiCL?

* "Encoder-decoder LMs are more efficient at multiple (e.g. 40) retrieved passages" -- this seems to be missing a citation?
paper_summary: paper_summary: The paper investigates to combine retrieval augmentation, in context learning, and encoder-decoder models. After establishing the primary grounding literature (e.g., Atlas), the paper begins by presenting two strategies for in-context learning + retrieval augmentation + encoder-decoder LMs. Provided sufficient in context examples (4 or 5 more depending on QA set), one strategy shows improvement over a zero shot Atlas baseline.
strength_weakness: strength_weakness: 1. The paper overall seems too incremental and this isn’t helped by lack of analysis beyond headline performance numbers. While there is a nuanced difference between simple yet powerful methods based on novel insight and incrementally building on existing work, I feel the paper falls into the latter category.
    1. In 3.2, the main difference between prior work seems to be possibly expanding the number of concatenated examples and instead of repeating the same, shuffling them (compared to Izzacard 2023, it looks like the KL divergence minimization objective is similar/same modulo attention difference).
    2. Strategy 1 in 3.1 seems like the simplest possible way (ie would be a good baseline) to do retrieval augmented QA with an encoder-decoder model (put QA pairs for context learning in decoder, Q + retrieved passage + blank answer in decoder).
2. While there are a couple demonstrative examples, the paper does not include much analysis beyond headline performance numbers (or performance as number if IC examples varies). It would have been nice to see something like an error analysis that convincingly argues that method fixes what it aims to fix.
3. Along similar lines, although there was deeper discussion of two QA benchmarks, the discussion of MMLU was quite short. For example, I would be interested in knowing if the gains are uniform or from specific tasks in MMLU.
questions: questions: 1. Did you try a variant of strategy 1 in Figure 1, where the passage either precede the A0, or the passage is before all the QA pairs?
2. The tables primarily compare to Atlas, are there other recent models that are worth comparing to?

Suggestions that didn’t influence my scores
1. I realize that LLMs are relatively new, but aside from datasets, the literature cited is almost exclusively 2020 or after, surely there is relevant literature to draw from before that.
2. I’m not certain if COLM style is different, but at least for many other venus figures (ie Figure 3, Table 4 should be at the top of the page (unless they are preceded by other tables/figures).
3. Figure 2: The y-axis scales should show where the red lines for Atlas 11B bottoms out, it looks truncated right now. Same with Figure 5
4. In 3.1, since you discard strategy 2 past 3.1, to make the paper more focused, I would suggest moving most/all the experiments from the main paper into the appendix and include a textual mention/footnote that notes another strategy was attempted, but it did not work as well as S1, with details in appendix.
5. In Figure 2, you should indicate the results are for S1 in the caption.
6. In Table 2, there is a mix between open weights and API only models. I would suggest distinguishing between the two somehow, as I don’t feel confident that it is a fair comparison to open weights models whose training data is known to a larger degree. A similar argument would apply to making it easier to compare small model results to large model results.
7. The paper spends quite a bit of time discussing details in Atlas, it would be stronger by discussing only the minimal necessary and focusing on the new work being introduced.
paper_summary: paper_summary: This paper tackles the problem of in-context learning with retrieval-augmented encoder-decoder LMs. There are two main techniques introduced in this paper. One is the pre-training method called RAVEN, which combines retrieval-augmented masked language modeling and prefix language modeling for the training of retriever and generator. The other one the fusion-in-context learning, which enables the model to leverage more in-context examples. The proposed method shows improved performances over Altas and other decoder-only LLMs like GPT-3 and PaLM on NQ and TQA.
strength_weakness: strength_weakness: There are two serious concerns regarding the motivation of the paper. 
- 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. 
- There is no evidence whether in-context learning is still necessary when fine-tuned LLMs are applied for RAG. 

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.
paper_summary: paper_summary: This paper aims to improve the in-context learning ability of retrieval-augmented encoder-decoder models---to be more specific, Atlas. The authors first conducted some analysis on the in-context learning behavior of Atlas with different prompting strategies, which shows that Atlas can do in-context learning but is unstable and does not benefit from more in-context examples. Then the authors proposed three ways to improve it: 

1. (Training): add a second stage training for Atlas which uses prefix language modeling, to handle the training/inference mismatch (in inference it's almost always a continual generation task). 

2. (Inference): instead of using the same group of ICL examples in each fusion-in-decoder context, the authors propose to use different ones so that the model can use more ICL examples without the limit of encoder context window. This is called "Fusion-in-Context Learning (FiCL)".

3. Select better ICL examples by using the Atlas retriever following previous work.

Those techniques combined form the new model, Raven. The authors mainly demonstrate its efficacy on open-domain QA tasks (NaturalQuestion, TriviaQA), as well as MMLU. The experiments show that Raven can improve significantly over Atlas. Thorough ablations also demonstrate the efficacy of each of the above techniques.
strength_weakness: strength_weakness: 1. This is more of a suggestion than critique: I wish there were more tasks analyzed and evaluated. The authors only explored tasks that significantly benefit from retrieval-augmentation, such as open-domain QA tasks and MMLU. To demonstrate the general ability of the language model, I believe people would also be interested to see how this suite of techniques work on tasks that may not require retrieval augmentation. This would show that FiD/encoder-decoder style LMs not only work for RAG, but also general ICL.

2. It would be great to have some analysis on how ICL behaviors (wrt numbers of demonstrations) differ in FiD/decoder-only models. Even with FiCL, I imagine the improvement has a "ceiling" given there is no attention across different encoder context windows. 

3. How about comparing this to using retrieval-augmentation in decoder-only LMs too? You can just stack all the demonstrations as well as all the passages in the decoder context as well. Also, what about [Replug](https://arxiv.org/abs/2301.12652)?
questions: questions: Please see "reasons to reject".
paper_summary: paper_summary: This paper investigates the in-context learning ability of encoder-decoder models and further proposes an improved version of the studied model. RAVEN, the proposed model, is built on top of ATLAS by further pre-training with the suggested objectives, with an aim to shorten the gap between the masked language modeling during pre-training and in-context learning during inference time. In addition, the few-shots examples are key to the in-context learning performance. The authors further present to use different examples for each retrieved document (Fusion-in-Context Learning) and use the retrieval component to obtain the demonstration examples (In-Context Example Retrieval). Results indicate that RAVEN outperforms ATLAS and even surpasses (or performs on par with) much larger models.
strength_weakness: strength_weakness: Novelty of the proposed methods is somewhat limited. Essentially, RAVEN is trained by further pre-training ASLAS with objectives that were discussed in the original paper of ASLAS (the prefix LM used in the second stage can be seen as a variant of the standard MLM and standard prefix LM).
questions: questions: What is the baseline for Table 3 and how to interpret them? If the baseline is RAVEN, does it mean the baseline score, for example, 31.7% for RAVEN 3B 1-shot on NQ can be improved by 9.1%, as reported on Table 3?

Comments: title: Thank you!
Thank you!
Thank you!
Ack
Follow-up questions
comment: Thank you for your acknowledgment! We are happy that our response addressed your concerns.
Thanks for providing the detailed explanation. This answers my questions and concerns. I increased my score from 6 to 7.
Thank you for your acknowledgment! We are happy that our response answered your questions :)
Thank you for highlighting FiCL and offering your thoughts! This is a very interesting point for further analysis. We really appreciate the engagement and responsibility of the reviewer—this is incredible.

We are also happy that our responses addressed your confusion.
Thank you for answering the questions.

I think there is already enough substance in this paper, however I still find FiCL an interesting approach and this paper only scratches the surface of what is possible. It would be an even stronger paper with a more thorough study on FiCL specifically to learn why it helpful: this study explores the three parameters [k, m, n], but to me, it feels like there are 4: [k, m, num_docs, num_encodings], where currently it is assumed each doc gets a single encoding. But if the argument is that fusion specifically helps with aggregating over more "shots"/examples, then something like [64, 5, 1, 40] should strongly improve over [1, 1, 1, 40] or more realistically, [1,1,1,1]. That is, 64-shot, 5 sampled randomly 40 times for a *single* document, then fused should outperform one-shot with no fusion (still 1 doc) -- this hypothetical but be a little degenerate but something like [64, 5, 10, 40] vs. [5, 5, 10, 40] would be a more reasonable setting to demonstrate and ablate between whether it's the diversity of examples or the fusion mechanism itself which leads to performance gains.

Nonetheless, I am increasing my score slightly upon the authors' response which answered my confusion around the ATLAS eval, and their committal for clarifying the relation between this work and ATLAS, and on expanding Section 3.1.
Thanks for answering my question! I will keep the score unchanged.
Dear Reviewer k2PV, thank you again for reviewing our paper! We hope our response addresses your questions. Please let us know your thoughts, and we are more than happy to answer any further questions.
Dear Reviewer SfUy, since the deadline of the discussion period is approaching, we would like to follow up to see if our response addresses your comments and to ensure we have enough time to respond if you have any follow-up questions.

Thank you again for your review!
Thank you for your feedback! From the developer's perspective, we do believe there are many use cases for ICL in RAG, and we have observed significant performance improvements in our experiments. The cost of ICL here is also not high (the FiCL design even makes it cheaper than traditional ICL).

Thank you for agreeing with the quality of our study! We think one interesting future exploration is to apply our methods to more advanced encoder-decoder LMs, such as Reka (though it’s not open source so far).
Thank you for your feedback and suggestions! Appendices C.1 and C.2 were initially part of the main text, but we moved them to the appendix to keep the paper within 9 pages. We plan to move them back to Section 3.1 in the final version. We will also take into account your suggestions for adjusting the space of other content.

For FiCL, [k, m] refers to [k-shot, m-fusion]. Using Figure 4 as an illustration, it's [j-shot, 2-fusion] (there are j in-context examples in total, and the encoder encodes 2 examples each time). We will make it more clear (e.g., replacing j with k to be consistent and adding more explanations in the caption).

For the number of retrieved docs (i.e., n), as detailed in Appendix B.1, it was set to 40 by default.

Thank you for carefully reviewing our paper and considering our responses. We greatly appreciate it and are happy with the feedback. We are also more than happy to answer any additional questions you may have.
Thank you for the response. 

I'm still not convinced by the findings/strength of the current section 3.1 -- if it is intended to be the first analysis for ICL for these models, then it should definitely be expanded to additionally include C.1 and C.2, and there would be 3 (equally important) subquestions, which together inform the design of RAVEN. With extra space, this may be possible, but may also require condensing the existing 3.1 a little (perhaps the discussion of "effect of number of in-context examples" can be reduced, since the findings there aren't too unexpected: that there's some threshold at which ICL plateaus.

For FiCL, I'm re-reading this section trying to recall what [k,m] referred to again, and I'm a little confused by Table 2. In table 2, how many docs are retrieved (i.e. what is n)? As a comment, it would be helpful to illustrate or re-use the variables k, m in Figure 4 too.

(I'll leave this comment here for now so authors get a chance to respond, but will look more closely at this section later).
Thanks for the authors' response! I'm increasing my overall assessment to 6 after having calibration with other papers in my batch. I think the paper basically presents a quality study in this field.  

However, I'm still uncertain about the necessity of using ICL for RAG. The improvement can be very marginal with the latest LLMs whereas the additional cost from using ICL will be substantial. The authors may provide a comprehensive investigation in their future research.
Dear Reviewer SfUy, we hope our response addresses your concerns. Please let us know your thoughts, and we are more than happy to answer any further questions.

BTW, as suggested by Reviewer A7yJ, we plan to highlight more about the distinction between our work and ATLAS in the final version. One significant difference between them is the problem itself: finetuning (ATLAS) vs in-context learning (ours). Besides that, there are also several technical innovations (combined retrieval-augmented masked and prefix language modeling, Fusion-in-Context Learning, In-Context Example Retrieval) in our paper that improved the in-context learning performance of retrieval-augmented encoder-decoder LMs.
Thank you for your feedback! We would like to further clarify and strengthen our response to your concerns.

> Effectiveness of ICL in RAG:

The motivation for using ICL with retrieval is aligned with its application in models without retrieval. In most cases, retrieval just adds external knowledge to the model's context but does not fundamentally change the task being performed. If ICL can work on tasks for a model without retrieval, there is no reason why it should not be applied to a model with retrieval. Both ICL and RAG are important and common usages of LLMs; they are orthogonal and don't conflict with each other. Therefore, there must be many use cases of ICL with retrieval, such as answering a factual-seeking question with the help of examples, solving knowledge-intensive tasks with the help of examples, writing a passage (which may require retrieval to collect the knowledge) based on some examples, etc. And empirical evidence in our paper exactly demonstrates that ICL enhances performance in models with retrieval on popular tasks such as Open-domain QA and MMLU.

> Running Cost of ICL:

While it's true that ICL introduces additional costs, we designed our approach to minimize these costs effectively. Here are the key points regarding cost management:

- **Efficiency in Prompt Design:** In our implementation, we do not retrieve passages for in-context examples (i.e., only retrieve passages for the target query). Therefore, the additional tokens introduced by ICL in our setting are the same as traditional ICL.

- **Linear Cost Growth with FiCL:** The computation cost of FiCL grows linearly with the number of in-context examples, as it performs self-attention over one (or $m$) example(s) at a time. This is more efficient than decoder-only LMs, where the time complexity is quadratic. This efficiency allows us to leverage the benefits of ICL without incurring prohibitive costs.

In summary, our paper provides strong empirical evidence and methodological innovations that support the use of ICL in RAG settings, demonstrating meaningful performance improvements without incurring unsustainable costs. We believe that these contributions are significant and address both the effectiveness and cost concerns you raised. 

We hope our response addresses your concerns. Please let us know your thoughts, and we are more than happy to answer any further questions.
Thanks for the authors' response! I'm sorry for being confused. Let me try to clarify my questions. From my perspective, ICL is not a very common practice for RAG due to the following concerns: 
- **Effectiveness:** I haven't seen many cases where ICL contributes substantially to RAG, though it's useful in many other situations. If there are examples where it does, please remind me. 
- **Running cost:** Applying RAG is already costly due to the encoding of retrieved information. Introducing ICL will add considerable extra costs.

As a result, I still cannot agree with the motivation of the paper, though it made meaningful progress in improving upon Altas.
Thank you for your follow-up questions!

Yes, your understanding is correct. All the few-shot results reported in the ATLAS paper are from finetuning (FT) performance. In our paper, all results are from in-context learning on the base models, as we want to disregard FT since FT is not the common few-shot usage nowadays. This is also a significant distinction between these two papers, i.e., finetuning vs in-context learning.

Regarding the 0-shot results of ATLAS, we ran the evaluation ourselves. There is a slight difference in the prompt. In their code, they add a space between "Answer:" and "<extra_id_0>", but we found there is no such space in their template for pretraining, so we removed it (i.e., "Answer:<extra_id_0>", see our prompt in Appendix B5). Therefore, our reported 0-shot results of ATLAS are slightly higher than the number in the original paper.

Thank you for your suggestion. The rebuttal is limited to 2500 characters, so we need to keep the above response precise to cover all of your questions. Feel free to let us know if you would like any additional explanation. We did a lot of work to shorten this paper to 9 pages, lol. We will definitely use the extra page to incorporate your suggestions in the final version.
Hi! Thanks for the response. I still have a couple questions that I'm confused about around how ATLAS was used in this work, and would like to see if my understanding is right.

ATLAS is a modification of the Contriever/T5 architecture that was pretrained on various "pretext" tasks on general datasets. They released several models and this work uses the 3B and 11B models base models (not the ones fine-tuned on NQ). RAVEN then does additional (pre)training of the encoder/decoder with the modified objective, and no finetuning at all on the downstream datasets. Meanwhile both ATLAS and ATLAS + FiCL scores reported in your paper use those 3B/11B base models directly.

So the reason why the numbers reported in this work do not match any of the numbers in the original ATLAS paper is because almost all of the numbers in the ATLAS paper are FT. The only exception would be the 0-shot results (e.g. table 7). For that, the numbers aren't exactly the same either? (47.1 vs. 47.4), and that's because I assume you ran eval again and got slightly different numbers? Or is the dataset slightly different? Or prompt/post-processing slightly different?

(as a general comment: while you stated you will add more explanations, I think the main thing missing right now is a clearer distinction between this work and ATLAS, and may reduce the impression of it being "too incremental" as reviewer SfUy wrote. The extra page you'll get would be helpful towards this, and can help lift some of the appendices and analysis into the main paper while keeping Sec 3.1 in if you really insist on it.)