ALR$^2$: A Retrieve-then-Reason Framework for Long-context Question Answering

Q1: While the generalization study presented in Section 5.2 is important, it only experiments with datasets that were “modified” to have a long-context, in a method that is similar to those of the training data. I believe the paper could benefit from experimenting with QA datasets that have a long context by design, such as those in the SCROLLS benchmark [1]. Additionally, the paper focuses on reasoning, but experiments with a single multi-hop QA benchmark in each setting (in-domain and out-of-domain). (Also see the first comment in the Questions section)

Following your suggestion, we further evaluated our Retrieve-then-Reason (RR) method on the NarrativeQA dataset, which has the longest context among the QA datasets in the SCROLLS benchmark. The average number of words in the input context for the NarrativeQA dataset is 51,653.

Setup We randomly sampled 100 examples from the validation set of NarrativeQA. We apply both the Retrieve-then-Reason (RR) and Direct-Answer (DA) methods to Command-R model. Following the setting in [1], we input the full story into the model and use corpus-level multi-reference BLEU-1 and BLEU-4 scores as the evaluation metric.

Results As shown in the following table, the RR approach outperforms the DA approach on both BLEU-1 and BLEU-4, demonstrating its potential in enhancing the long-context reasoning in realistic QA datasets.

[1]: The NarrativeQA Reading Comprehension Challenge

Q2: The results in Tab.1 only compare against all baselines for the trained Command-R models, and models prompted to directly generate the answer for other models. I believe the paper could benefit from experimenting with another trained LLM with all variants, or at least adding results with the ‘retrieve-then-reason’ prompt with current strong LLMs (e.g., Llama-405B, GPT4-o, Claude-Sonnet-3.5).

Thanks for your suggestion. We further evaluate the Retrieve-then-Reason (RR) method on GPT-4 and Claude-3. As shown in the following table, the RR method achieves consistent performance gain on GPT-4 and Claude-3, compared with the Direct-Answering (DA) method. Note that the experiments on 128K is still under-running because of the time limit.

(Experiments on HotpotQA dataset)

Q3: Perhaps I am missing something, but as far as I am aware the idea that a LLM is trained to first retrieve information from a long context and then reason over it has been presented previously, for example by training LLMs as re-rankers (e.g., [2]). While there are differences between these two approaches (in  the LLM generates relevant sentences instead of ranking), it seems there are also similarities (the LLM first identifies relevant information from a long context). Hence, the paper might benefit from additional experiments or analysis that highlight what leads to the improvements of  (see questions).

We agree that [2] is a related work, as the retrieved results could be derived from the ranked results. We will cite this work and discuss the differences in our paper. In summary, unlike [2], which focuses on re-ranking at the passage level, our approach trains the LLM to retrieve precise information, such as key sentences from different passages, to support generation. This is because not all information within a passage is necessarily useful. Employing a re-ranking approach becomes burdensome for fine-grained retrieval, such as at the sentence or span level. Additionally, our approach has the potential to extend retrieval to the span level.

[2] RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!

Thank you for the detailed response! It indeed addresses most of my concerns.

My main remaining concern regards the lack of a re-ranking baseline. If I understand correctly, the new RAG baseline is relatively weak (DPR only and performs worse than directly answering the question when context size is smaller than 32K), while stronger baselines exist (e.g., re-ranking). I agree with the authors that the presented method is more fine-grained and could have additional benefits, but I believe the paper could benefit from empirically examining the differences.

Thank you for adding this experiment.

I am a bit surprised that DPR-RAG performs better or similarly than RankZephyr-RAG, as this seems to contradict the original paper (and other works showing that re-ranking with stronger models outperform dense retrievers).

To summarize, I believe the paper could benefit from (a) adding stronger baselines that are commonly used in previous works, (b) discussing these works and how they differ from $ALR^2$, and (c) focusing on comparison to stronger retriever-rerank baselines in addition to direct-answering. Nevertheless, the proposed method has merits and I appreciate the authors' efforts to improve the paper during the discussion period. I updated my score to reflect these changes.

Q1: Details about the experiments are missing. For example, there is insufficient information regarding the hyperparameters used during model training (e.g., learning rate, seed) and the specific training data. It is unclear whether ALR was trained separately on each dataset (e.g., SQuAD, HotpotQA) or trained using both datasets together (according to line 407). If it was the former, which model was used in the experiment described in section 5.2?

Thanks for pointing out. We will add the following training details in the revised version:

Please feel free to let us know if there are any additional hyper-parameters you want to know.

Moreover, our ALR$^2$ is trained on the mixed dataset, i.e., training one model for all the tasks. Therefore, the model used in Sec 5.2 is the same as the one used in the main experiments.

Q2: The experiment setting appears to be incomplete. It seems feasible to apply the RR method (e.g., using multi-turn conversation) to the frontier models utilized in the study (GPT-3.5, GPT-4, GEMINI, CLAUDE-3). Additionally, the benefits of aligning with both objectives are not clearly demonstrated. Why should the LLM be aligned to both objectives? It seems plausible to fine-tune the LLM separately for each objective and then combine the results (e.g., passing the output of the retrieval LLM to the reasoning LLM). An ablation study addressing this approach should be provided.

Following your suggestion, we further evaluate the Retrieve-then-Reason (RR) method on GPT-4 and Claude-3. As shown in the following table, the RR method achieves consistent performance gain on GPT-4 and Claude-3, compared with the Direct-Answering (DA) method. Note that the experiments on 128K is still under-running because of the time limit.

(Experiments on HotpotQA dataset)

We agree that it is indeed feasible to train two separate models for retrieval and reasoning, respectively. This approach aligns with the implementation of RAG methods, which typically consist of a dual-encoder-based retriever and a language model. However, enabling a single general-purpose LLM to handle both retrieval and reasoning within a unified framework, namely auto-regressive generation, offers significant advantages in terms of efficiency and practicality in real-world scenarios. This eliminates the need to deploy and manage two LLM services simultaneously. We will clarify this point further in the revised paper.

I understand that time is limited, but as I initially asked, could CMD-R QF be added to Table 2?

Yes, we will add CMD-R QF to Table 2.

Following your suggestion, we conducted the analysis in Table 2 for QF method, and also add the ngram-level precision metric, i.e., the percentage of 4-grams in the retrieved sentences that also appear in the golden-truth supporting facts. Below are the results:

Recall

Precision (4-gram)

We found that RR performs better in terms of recall, while QF excels in precision. This supports our hypothesis that RR is more effective for multi-hop QA, whereas QF is better suited for single-hop QA.

We really appreciate your suggestions to our work, and will integrate those results into the revised version. We hope our results could address your concerns. If you have additional questions, please feel free to let us know.

Q1: I believe an important baseline is missing: How is your retrieve-then-reason framework compared to retrieval-augmented generation for LLM baseline methods including DPR (Dense Passage Retrieval), FLARE (Active Retrieval Augmented Generation), RETA-LLM (RETA-LLM: A Retrieval-Augmented Large Language Model Toolkit). Throughout these methods, retrieval can even be performed with a smaller, faster model on a larger context like a database and are applicable to the long-context understanding scenario. (Also see the first comment in Questions)

Following your helpful suggestion, we add the baseline of RAG method for a more comprehensive comparison. We will also cite and discuss the related works you mentioned in our paper.

RAG Setup We employ the DPR model to encode the question and documents. We use the facebook-dpr-question_encoder-multiset-base and facebook-dpr-ctx_encoder-multiset-base encoders provided by sentence-transformers to encode the passage and question separately. For each question, we build a retrieval datastore independently based on the in-context documents of it. We retrieve the top five documents according to the dot-product score. Afterwards, we assemble the retrieved documents and the original question using DA template.

Results As shown below, the overall performance of the RAG approach surpasses that of the Direct-Answering (DA) method, although the DA approach outperforms RAG in short-context scenarios. The RAG approach demonstrates extremely stable performance across context lengths ranging from 4K to 128K. However, when using the LLM itself for in-context retrieval, i.e., the Retrieve-then-Reason (RR) method, significant performance gains are observed compared to the RAG approach. This indicates that the LLM has a superior understanding of the relationship between context information and the user’s question.

(Experiments Hotpot QA Dataset)

Q2: The contribution is limited. As the author indicated in sec 4.2, the retrieve-then-reason framework is very similar to "generate quotes and citations first" framework. Therefore, the innovation of this work seems just to fine-tune the models based on these data.

We agree that the technical part of this work is simple, but the research problem explored in this paper is important and our proposed approach is very effective.

To summarize, we aim to address several critical questions in this work:

1. What are the common problems with current long-context LLMs? In this preliminary study, we reveal that the long-context capabilities of many commercial LLMs are surprisingly fragile. Specifically, increasing task complexity by requiring additional reasoning significantly degrades the performance of even strong LLMs such as GPT-4, Claude-3, and Command-R. We believe this observation is crucial for the research community.
2. How can the observed problem be solved? We address the issue by decomposing the original problem into retrieval and reasoning sub-problems, employing the RAG formulation, i.e., enabling the LLM itself to perform both retrieval and reasoning.
3. What is the impact of this approach on the community? Our simple training paradigm significantly stabilizes long-context reasoning performance. Furthermore, we highlight practical challenges, such as severe hallucination, when using general-purpose LLMs for retrieval tasks.

Q3: There is lack of a clear explanation for some unexpected parts of the experimental results. For example, the model's performance on squad for some lengths (8k, 16k) in Table 1 is significantly lower than shorter and longer lengths, and I have not seen a clear interpretation for this unusual behavior.

Thanks for the good catch. We hypothesize that this is because the model was trained on a mixed dataset consisting of HotpotQA, SQuAD, and four tasks from NIAH, where SQuAD constitutes only a small portion of the training dataset.

Q4: The term "Retrieval" in this paper is different from previous work like DPR, where the models do not retrieve in-context but retrieve on an external codebase. Considering this may cause misunderstanding for unfamiliar readers, I suggest considering changing names like "in-context retrieval”

Thanks for this good suggestion! This is one thing that our authors are also considering. We will clarify the related parts clearer in the revised paper following your suggestion.