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

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.

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

* The paper is well-written and organized. * Experiments are thorough, compared with a variety of baseline models.

Reasons to reject

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 to authors

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?

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.

Rating

7

Confidence

5

Ethics flag

1

Reasons to accept

1. The paper is well written, and well motivated, and the experiments are very clear and clean. 2. The analysis on Atlas' in-context learning performance is insightful and well motivated the proposed techniques. The proposed techniques are interesting, effective, and novel. 3. The improvement on the selected tasks is significant. The ablation study is comprehensive.

Reasons to reject

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 to authors

Please see "reasons to reject".

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.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

- The proposed method is technical sound. The resulted LM improves upon the previous encoder-decoder LM. - The effectiveness of the proposed method is empirically verified on NQ and TQA.

Reasons to reject

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.

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.

Rating

5

Confidence

4

Ethics flag

1

Reasons to accept

1. The paper’s main strength is demonstrating that with a couple modifications to the train/inference procedure, that it outperforms Atlas by a respectable margin. 2. The paper proposes a method that does not (at least on the surface) seem overly compute intensive, making it easier to apply. 3. The paper looks solid in methodology, the methods are clear enough to understand and the experiments seem executed well.

Reasons to reject

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 to authors

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.

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).

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

* The proposed method, RAVEN, appears to work well across multiple tasks and addresses the concern that ATLAS needs additional finetuning. * The addition method, FiCL, is a novel method to more efficiently combine the notion of in-context learning to encoder-decoder models. The fact that this method works so well is a cool result. * The experiments are extensive and cover lots of ablations like model size, prompting strategy, multiple shots, etc. Overall, they are convincing.

Reasons to reject

* 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 to authors

* 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?

Follow-up questions

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.)

Ack

Thanks for answering my question! I will keep the score unchanged.