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

## Abstract

Abstract The context window of large language models (LLMs) has been extended significantly in recent years. However, while the context length that the LLM can process has grown, the capability of the model to accurately reason over that context degrades noticeably.
This occurs because modern LLMs often become overwhelmed by the vast amount of information in the context; when answering questions, the model must identify and reason over relevant evidence sparsely distributed throughout the text.
To alleviate the challenge of long-context reasoning, we develop a retrieve-then-reason framework, enabling LLMs to reason over relevant evidence collected during an intermediate retrieval step.
We find that modern LLMs struggle to accurately retrieve
relevant facts and instead, often hallucinate “retrieved facts” , resulting in flawed reasoning and the production of incorrect answers. To address these issues, we introduce ALR 2 , a method that augments the long-context reasoning capability of LLMs via an explicit two-stage procedure, i.e., aligning LLMs with the objectives of both retrieval and reasoning.
We demonstrate the efficacy of ALR 2 for mitigating performance degradation in long-context reasoning tasks. Through extensive experiments on long-context QA benchmarks, we find our method to outperform competitive baselines by large margins, achieving at least 8.4 and 7.9 EM gains on the long-context versions of HotpotQA and SQuAD datasets, respectively.

## 1 Introduction

The ability for large language models (LLMs) to reason over long contexts is critical in many downstream tasks, e.g., document analysis (Wang et al., ), multi-hop tool use (Yao et al., ), agents with long history (Park et al., ), etc. Significant efforts have been made to extend the effective context length of LLMs, ranging from investigating how different modules in Transformer (Vaswani, ) architecture impact long-context performance (Su et al., ; Chiang & Cholak, ; Xiao et al., ) to proposing new architectures with improved efficiency (Dai et al., ; Beltagy et al., ; Chevalier et al., ).

While these developments are promising, in our preliminary study, we show that the long-context performance of LLMs varied significantly across different tasks. In particular, we design a series of experiments to evaluate the performance of LLMs on long-context retrieval(^1^11Following the definition in prior work (Chen et al., ; Hsieh et al., ), i.e., retrieval is performed by generating the relevant information from a given input text.) and reasoning (Figure 1). We observe that, when tasked to generate answers by directly reasoning over the full context,
performance degrades as the input context grows. In contrast, when tasked with retrieving the set of evidence relevant to the question, the performance of LLMs is only mildly affected by the growth of the input context.

Those findings motivate us to explicitly leverage an intermediate retrieval step to address the challenge of LLMs in long-context reasoning. In our proposed retrieve-then-reason framework,
an LLM functions as both a retriever and a reasoner. When tasked to answer a question, it first collects the relevant evidence sparsely distributed across the context, and then derives the final answer by reasoning over that collected evidence.
We also find that modern LLMs struggle to accurately retrieve “relevant facts”,
i.e., simply
prompting them to retrieve supporting evidence often results in severe hallucinations, which leads to flawed reasoning by the LLMs caused by the reliance on inaccurate information.
To this end, we introduce the ALR^2 approach, which augments the long-context reasoning capability of LLMs via an explicit two-stage procedure, i.e., ALigning LLMs with the objectives of both Retrieval and Reasoning.

Extensive experiments on long-context single-hop and multi-hop question answering tasks demonstrate the effectiveness of our ALR^2 approach.
We compare ALR^2 against various advanced prompting techniques, including direct-answering and other popular prompting methods for long-context reasoning (Weston & Sukhbaatar, ; Lee et al., ). Results show that ALR^2 outperforms these approaches by at least 23.4 and 12.7 EM scores on the long-context versions of HotpotQA (Yang et al., ) and SQuAD (Rajpurkar et al., ) benchmarks. In addition, we also compare ALR^2 with LLMs fine-tuned with direct-answering prompting and observe ALR^2 excelling at 8.4 and 7.9 EM gains on HotpotQA and SQuAD, respectively.
Furthermore, we conduct extensive analysis revealing that
ALR^2 displays favourable generalization capability and significantly mitigates retrieval issues, e.g. hallucinations and low accuracy, that are prevalent in prompting-based approaches (§5).

Our contributions are three-fold:

- •
In our preliminary study, we show that as the input context length grows, LLM reasoning performance degradation is higher than that of retrieval.
- •
We propose ALR^2, a retrieve-then-reason approach to alleviate the challenges occurred when LLMs reasoning over long-contexts by aligning LLMs with both retrieval and reasoning objectives.
ALR^2 substantially mitigates the primary difficulties of long-context reasoning, such as hallucination and low retrieval accuracy.
- •
We conduct extensive experiments to systematically compare our approach with baseline methods, including both prompting-based and fine-tuning-based techniques. Our results show that ALR^2 significantly outperforms these baselines on long-context QA tasks, i.e., achieving at least least 8.4 and 7.9 EM gains on long-context version of the HotpotQA and SQuAD benchmarks, respectively.

Figure: Figure 1: Performance of LLMs on three increasingly challenging long-context tasks. The x-axis represents the number of tokens in the long context, while the y-axis indicates the exact match score.
Refer to caption: extracted/5901090/preliminary_exp.jpg

## 2 Preliminary Study

In this section, we will present the challenges of long-context reasoning and the motivation for our work, drawn from observations on a series of long-context tasks.

#### Setup

We assess the long-context capabilities of LLMs using a series of increasingly challenging tasks, inspired by prior work (Chen et al., ; Hsieh et al., ):

- •
Task 1: Passkey Retrieval: A sentence, such as ”The passkey of Alice is KEY,” is inserted into the long context, and the LLM is asked to retrieve Alice’s passkey. The ”KEY” is a randomly generated string of 10 characters, sampled from the alphabet (A–Z).
- •
Task 2: Multi-hop Passkey Retrieval: In this task, the 10-character passkey is split into two halves, with each half randomly assigned to “Alice” and “Bob.” The LLM is then prompted to retrieve either Alice’s or Bob’s passkey. This task requires an LLM to identify the most relevant one from multiple passkeys, making it more complex than the first.
- •
Task 3: Multi-hop Passkey Retrieval + Reasoning: This task involves both retrieval and reasoning, and Task 1 & 2 are the pre-requisite tasks, i.e., the passkey retrieval is critical for the final reasoning. Given the passkeys of both Alice and Bob, the LLM is asked to concatenate them, e.g., if “Alice has ABC” and “Bob has DEF”, then the expected answer is “ABCDEF.” The order of concatenation is critical, and the final concatenated answer is controlled to match the passkeys used in Task 1. Notably, the reasoning is designed to be simple purposefully, and a model is expected to achieve good performance on this task if it performs well on Task 1 & 2.

For all tasks, the LLMs are prompted to generate the correct passkey from the long context. We generate 100 random passkeys and use the Wikitext-103 dataset (Grave et al., ) to augment the input context. The length of the Wikitext-103 context is varied between 0 and 128K tokens, with the passkeys randomly inserted at 30% and 60% positions within the context. Exact match (EM) is used to evaluate the LLMs’ performance. We test three modern LLMs in this experiment: Command-R-Plus(^2^22[https://cohere.com/blog/command-r-plus-microsoft-azure](https://cohere.com/blog/command-r-plus-microsoft-azure)), GPT-4 (Achiam et al., ), and Claude-3-Haiku ().

#### Results

From Figure 1, we make several key observations. First, most LLMs perform well on the first two passkey retrieval tasks (Tasks 1 & 2), achieving high EM scores across context lengths ranging from 0K to 128K. Second, while the simple reasoning required by Task 3 is not an issue for most LLMs within short context, their performance drops significantly as the context length increases. This decline is more pronounced in Task 3 compared to the retrieval-only tasks in Task 1 and 2. Results of Task 3 reveal the challenge facing by modern LLMs in directly reasoning over the long-context input(^3^33Similar trends are observed in more complex, realistic question-answering tasks, such as HotpotQA and SQuAD (see §4 for more details).). It further drives our motivation in leveraging an intermediate retrieval step as a stepping stone for the subsequent reasoning step to address the challenge.

## 3 Methods

### 3.1 Overview

To address the aforementioned challenges, we introduce a retrieve-then-reason approach which augments the long-context reasoning capability of LLMs via a two-stage breakdown, i.e., (i) first explicitly retrieving relevant facts pertaining to the task; and (ii) then reasoning over the collected facts to derive the final answer. Moreover, as manifested in Figure 2,
modern LLMs struggled to retrieve with high accuracy in realistic scenarios and frequently hallucinate “retrieved facts”, leading to faulty reasoning and the generation of incorrect answers. To this end, we introduce a simple and effective method, namely ALR^2, to augment the long-context reasoning capability of LLMs by an explicit two-stage procedure, i.e., ALigning LLMs with the objectives of both Retrieval and Reasoning. Our approach is inspired by the formulation of Retrieval-Augmented Generation (RAG) approach (Lewis et al., ), which was originally designed for open-domain QA.

### 3.2 Formulation of RAG

The standard open-domain QA problem can be formalized as:

$$ ${\bm{y}}^{*}=\operatorname*{arg\,max}_{{\bm{y}}}p_{\theta}({\bm{y}}|{\bm{q}}),$ $$

where ${\bm{q}}$ is the user question, ${\bm{y}}^{*}$ is the answer with the highest probability, and $\theta$ denotes parameters of the whole system. RAG introduces a latent variable ${\bm{z}}$ and reformulates the conditional probability $p_{\theta}({\bm{y}}|{\bm{q}})$ as follows:

$$ $p_{\theta}({\bm{y}}|{\bm{q}})=\sum_{{\bm{z}}}p_{\phi}({\bm{z}}|{\bm{q}})p_{\mu }({\bm{y}}|{\bm{z}},{\bm{q}}),$ (1) $$

where ${\bm{z}}$ represents the retrieved passage, while $p_{\phi}({\bm{z}}|{\bm{q}})$ and $p_{\mu}({\bm{y}}|{\bm{z}},{\bm{q}})$ are the probability distributions of the retriever and generator, respectively. The retriever, parameterized by $\phi$, builds an external retrieval index and retrieves the top-$K$ most relevant passages for each question from the index (Karpukhin et al., ; Izacard et al., ; Ni et al., ). The generator, typically a language model (Vaswani, ) parameterized by $\mu$, then takes the retrieved passages and the question ${\bm{q}}$ as input to generate the final answer.

### 3.3 Adapting RAG Formulation to Long-Context Reasoning

Given that LLMs perform well in explicit long-context retrieval, a natural solution is to adapt the RAG formulation (Eq. 2) to re-organize the sparsely distributed relevant facts in intermediate steps, and then reason over those retrieved facts. We define our approach for long-context reasoning as follows:

$$ $p_{\theta}({\bm{y}}|{\bm{q}},{\bm{c}})=\sum_{{\bm{z}}}p_{\theta}({\bm{z}}|{\bm {q}},{\bm{c}})p_{\theta}({\bm{y}}|{\bm{z}},{\bm{q}},{\bm{c}}),$ (2) $$

where ${\bm{c}}$ represents the long context, and both the retriever and the generator are parameterized by $\theta$. Since the key information ${\bm{z}}$ is essential for answering the question, the term $p_{\theta}({\bm{z}}|{\bm{q}},{\bm{c}})$ models the process of explicitly retrieving ${\bm{z}}$ from the long context ${\bm{c}}$. Once the key information is retrieved, the generator $p_{\theta}({\bm{y}}|{\bm{z}},{\bm{q}},{\bm{c}})$ can utilize it to answer the question.
By leveraging an intermediate retrieval step, our approach simplifies the long-context reasoning to its short-context counterpart, i.e, the reasoning procedure operates on the condensed evidence collected from the retrieval step.(^4^44In Figure 3, we provide a concrete prompt to implement our approach. As shown in Appendix A, it can already significantly alleviate the performance drop on the reasoning task introduced in preliminary study (§2).)

Figure: Figure 2: Error case of the prompting-based retrieve-then-reason approach. The Command-R model and prompt in Figure 3 are used, and more details are in §4.1. The text with under-wave is the information matched between the golden facts and retrieved facts. We use the text with underline to represent information that is in the long context but not in golden facts. The information hallucinated by LLM are marked by red. The italic text in supporting facts shows the appearance of the answer.

### 3.4 Aligning Generation to Retrieval & Reasoning

Though Differentiable Search Index (DSI) (Tay et al., ; Bevilacqua et al., ) demonstrates that replacing the retrieval process with generation is feasible, pre-trained LLMs are generally not well-aligned with the retrieval objective. As illustrated in Figure 2, simply prompting the LLM to retrieve key information often results in low retrieval accuracy and severe hallucination(^5^55See §5.1 for more quantitative analysis for the retrieval.), which in turn leads to flawed reasoning and the production of incorrect answers.

To address these challenges, we
align the LLM with both retrieval and reasoning objectives as defined in Eq. 2. To build the training data, we collect datasets that provide the necessary supporting facts (a list of statements crucial for answering the question) and use these facts as the target ${\bm{z}}^{*}$. We then jointly optimize the LLM for both long-context retrieval and answering with the following objective:

$$ $\mathcal{L}_{align}=\frac{1}{N}\sum_{i=1}^{N}\big{(}\log p_{\theta}({\bm{z}}^{ *}_{i}|{\bm{q}}_{i},{\bm{c}}_{i})+\log p_{\theta}({\bm{y}}^{*}|{\bm{z}}^{*}_{i },{\bm{q}}_{i},{\bm{c}}_{i})\big{)},$ (3) $$

where $N$ is the number of training examples, and ${\bm{z}}^{*}_{i}$ and ${\bm{y}}^{*}_{i}$ are the golden supporting facts and answers for the $i$-th example, respectively. Notably, for questions with multiple supporting facts, such as in multi-hop QA datasets (Yang et al., ), the LLM is trained to retrieve all supporting facts by formatting ${\bm{z}}^{*}$ as a concatenated string of bulleted facts. We adopt the prompting template in Figure 3 for our model training.

### 3.5 Discussion

Our retrieval approach differs from others in two key aspects. First, most dense retrieval methods (Karpukhin et al., ), which use collections of dense vectors as the index, or Differentiable Search Index (DSI), which leverages model parameters (Tay et al., ) for indexing. In contrast, we treat the long input context itself as the retrieval index which supports the update of index on-the-fly. Second, most prior methods do not account for the relationships between retrieved statements and the correct number of statements to retrieve, often relying on large top-$K$ values to ensure all necessary information are retrieved. Differently, our retriever is explicitly trained to retrieve a coherent set of facts that are relevant for answering the question, which is particularly valuable in tasks like multi-hop QA. A better retrieval quality is not only important for the subsequent reasoning, but also provides a more trustworthy rationale, helping users better understand the decision-making process of LLMs.

## 4 Experiments

### 4.1 Setup

We follow the Ruler benchmark (Hsieh et al., ) to evaluate the long-context performance of our approach and the baselines. Our focus is on question-answering (QA) tasks, but we also assess performance on retrieval tasks, such as variants of the Needle-in-a-Haystack (NIAH) task:

- •
QA: Ruler utilizes two QA datasets, i.e., SQuAD (Rajpurkar et al., ) and HotpotQA (Yang et al., ), to build long-context evaluation datasets, which test the single-hop and multi-hop QA capabilities of LLMs in the long-context scenario.
- •
NIAH: The NIAH task is a synthetic retrieval task commonly used for long-context evaluation. The simplest version involves extracting a “needle”, such as an 8-digit number, from a long context. Ruler includes four variants of the NIAH task: Single NIAH, Multi-keys NIAH, Multi-values NIAH, and Multi-queries NIAH.

More details about those datasets are shown in Appendix B.1.

We prepared the training data for each task accordingly. For the QA tasks, we used the supporting facts provided by HotpotQA and the sentences containing the answers in SQuAD as the golden retrieval results. Notably, each question in the SQuAD dataset only has one retrieved sentence, while HotpotQA questions typically have 2 to 6 supporting facts. These supporting facts are organized into a string of bulleted items. For the NIAH tasks, we treated the needle sentences as the retrieved results. The number of training examples from HotpotQA, SQuAD, and the four NIAH variants is 5,000, 25,000, and 1,600, respectively. The context lengths (in tokens) of the training examples are 4K, 8K, 16K, and 32K, with the data ratios for each being 1/10, 2/10, 3/10, and 4/10, respectively. Inference for longer contexts is done by extrapolation. More training details are in Appendix B.2. We used Exact Match (EM) as our evaluation metric. Following the setup of Ruler, we report the average EM score on 500 test cases for each dataset or sub-task. Importantly, to avoid data contamination, we ensured that both the context and the question-answer pairs for testing were unseen during training.

**Table 1: Experiments on question-answering (QA) and needle-in-a-haystack (NIAH) tasks. The best results are highlighted by bold text. The evaluation metric is Extract Match (EM).**
| Model | Prompt | 4K | 8K | 16K | 32K | 64K | 128K | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| HotpotQA |  |  |  |  |  |  |  |  |
| GPT-3.5 | DA | 52.4 | 49.4 | 46.5 | – | – | – | – |
| GPT-4 | DA | 59.8 | 58.1 | 57.8 | 53.8 | 51.0 | 48.4 | 54.8 |
| Gemini | DA | 59.4 | 57.3 | 58.72 | 54.5 | 56.7 | 52.6 | 56.5 |
| Claude-3-Haiku | DA | 49.2 | 48.0 | 42.4 | 43.0 | 38.8 | 49.4 | 45.1 |
| \hdashlineCMD-R | DA | 51.6 | 51.2 | 47.4 | 47.4 | 43.6 | 34.29 | 45.9 |
| RR | 58.8 | 55.6 | 54.8 | 54.6 | 51.4 | 44.8 | 53.3 |  |
| QF | 53.4 | 54.4 | 49.4 | 48.6 | 44.8 | 32.8 | 46.0 |  |
| S2A | 43.0 | 44.6 | 42.8 | 40.8 | 34.0 | 30.5 | 39.2 |  |
| \hdashlineCMD-R-FT | DA | 72.8 | 71.39 | 72.0 | 69.1 | 67.4 | 57.4 | 68.3 |
| ALR^2 | RR | 77.2 | 76.8 | 77.2 | 76.6 | 77.5 | 75.2 | 76.7 |
| SQuAD |  |  |  |  |  |  |  |  |
| GPT-3.5 | DA | 63.8 | 60.0 | 53.1 | — | — | — | — |
| GPT-4 | DA | 82.3 | 75.6 | 72.3 | 70.1 | 63.0 | 59.4 | 70.4 |
| Gemini | DA | 65.9 | 67.2 | 66.5 | 65.6 | 64.6 | 62.5 | 65.4 |
| Claude-3-Haiku | DA | 67.0 | 63.6 | 60.6 | 53.2 | 53.8 | 55.0 | 58.8 |
| \hdashlineCMD-R | DA | 65.6 | 61.4 | 63.8 | 62.4 | 54.0 | 49.6 | 59.4 |
| RR | 58.8 | 55.1 | 52.4 | 51.8 | 51.3 | 49.2 | 53.1 |  |
| QF | 63.6 | 61.6 | 60.6 | 57.9 | 46.6 | 43.2 | 55.5 |  |
| S2A | 42.0 | 42.4 | 49.4 | 46.6 | 41.1 | 38.3 | 43.3 |  |
| \hdashlineCMD-R-FT | DA | 80.8 | 54.2 | 68.6 | 78.2 | 50.6 | 53.2 | 64.2 |
| ALR^2 | RR | 89.8 | 63.0 | 64.0 | 78.2 | 66.6 | 71.2 | 72.1 |
| NIAH |  |  |  |  |  |  |  |  |
| GPT-3.5 | DA | 99.85 | 98.55 | 93.05 | - | - | - | - |
| GPT-4 | DA | 100.0 | 99.9 | 99.65 | 98.5 | 96.1 | 91.4 | 97.60 |
| CMD-R | DA | 94.6 | 91.9 | 89.29 | 94.37 | 93.84 | 87.99 | 92.00 |
| \hdashlineCMD-R-FT | DA | 100.0 | 100.0 | 99.95 | 99.95 | 99.85 | 98.79 | 99.75 |
| ALR^2 | RR | 100.0 | 100.0 | 100.0 | 99.95 | 99.85 | 98.99 | 99.79 |

### 4.2 Baselines

We employ several prompting techniques for the long-context tasks:

- •
DA: This prompting approach is used for Direct Answering, where the model directly generates an answer without explicit retrieval. The template is illustrated in Figure 4.
- •
RR: This prompting method, shown in Figure 3, is used for the Retrieve-then-Reason framework in our work. It first prompts the model to retrieve the key information from the long context and then complete the reasoning to answer the question.
- •
QF: Similar to the RR, Anthropic(^6^66[https://docs.anthropic.com/claude/docs/advanced-text-analysis](https://docs.anthropic.com/claude/docs/advanced-text-analysis)) and Lee et al. () empirically find that guiding the LLM to generate Quotes and citations First, before providing answers, improves long-context reasoning. We evaluate a variant of this approach, as shown in Figure 5.
- •
S2A: Instead of retrieving information directly from the long context, the System-2 Attention approach (Weston & Sukhbaatar, ) focuses on answering questions based on summarized and reorganized information extracted from the long context. Our prompt template for this approach is presented in Figure 6.

For our experiments, we use the Command-R (CMD-R) model(^7^77[https://cohere.com/blog/command-r](https://cohere.com/blog/command-r)) with 35 billion parameters as our backbone. We evaluate its performance using the four prompting techniques: DA, RR, QF, and S2A. These prompting-based baselines are referred to as CMD-R + X, where X is the name of the prompt used. Additionally, to ensure a fair assessment on the performance gains otained from model fine-tuning, we include another CMD-R-FT + DA baseline. This model is built by fine-tuning CMD-R on the same training data, described in §4.1, to directly produce the answer, i.e., without a retrieval step, given the input context and the question.
For a broader set of evaluations, we also test the DA prompt on other frontier models, including GPT-3.5(^8^88Since thye maximum context length of GPT-3.5 is 16K, we only show the performance of GPT-3.5 within this context length.) (Ouyang et al., ), GPT-4 (Achiam et al., ), Gemini (Team et al., ), and Claude-3-Haiku ().

### 4.3 Results

As shown in Table 1, ALR^2 achieves significantly better performance on the two question-answering (QA) benchmarks. Notably, on the multi-hop QA benchmark, HotpotQA, ALR^2 maintains consistent performance from 4K to 128K tokens of context. Additionally, the retrieve-then-reason prompting baseline, i.e., CMD-R+RR, shows a smoother performance curve on HotpotQA compared to the direct-answer baseline, i.e., CMD-R+DA. However, on the single-hop SQuAD benchmark, CMD-R+RR performs worse than CMD-R+DA. We hypothesize that this is due to over-retrieval by CMD-R+RR, as it tends to retrieve excessive information that becomes redundant for single-hop questions. In contrast, after aligning the LLM to handle both single-hop and multi-hop questions, our ALR^2 approach achieves superior overall performance.

We also compare different two-stage prompting techniques, namely RR, QF, and S2A. As shown in Table 1, CMD-R+RR outperforms CMD-R+QF on HotpotQA but falls short on SQuAD. We find that this discrepancy is closely related to the number of sentences retrieved by each prompt. Specifically, CMD-R+RR retrieves an average of 7.5 sentences across both QA datasets, while CMD-R+QF retrieves only 2.2 sentences. Since HotpotQA is a multi-hop QA dataset that requires more supporting facts, CMD-R+RR performs better on this dataset. Both CMD-R+QF and CMD-R+RR outperform the system-2 attention baseline, i.e., CMD-R+S2A, in our experiments.

In addition, on the NIAH tasks, all evaluated approaches obtain comparable performances.
This is unsurprising given the fact that NIAH tasks only require straightforward information retrieval, which is relatively easy for LLMs to handle. When evaluating long-context performance, we caution researchers against averaging the performance of long-context retrieval and reasoning tasks, such as NIAH and QA, due to the significant performance gap between them.

## 5 Analyses

**Table 2: Evaluation of retrieval quality on HotpotQA dataset. The hallucination rate is the percentage of retrieved sentences that are not presented in the context. The recall score is the percentage of golden facts that are displayed in retrieved sentences. The best results are highlighted by bold text.**
| Model | Prompt | 4K | 8K | 16K | 32K | 64K | 128K | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Hallucination Rate $\downarrow$ |  |  |  |  |  |  |  |  |
| CMD-R | RR | 66.57 | 61.04 | 54.54 | 57.01 | 60.79 | 66.74 | 61.1 |
| ALR^2 | RR | 0.26 | 0.17 | 0.17 | 0.17 | 0.36 | 0.61 | 0.29 |
| Recall $\uparrow$ |  |  |  |  |  |  |  |  |
| CMD-R | RR | 27.37 | 32.50 | 39.53 | 41.35 | 41.34 | 22.28 | 34.06 |
| ALR^2 | RR | 69.47 | 69.39 | 69.64 | 69.47 | 67.78 | 66.99 | 68.79 |

### 5.1 Retrieval Quality by LLM

#### Setup.

We evaluate the hallucination rate of our retriever, $p_{\theta}({\bm{z}}|{\bm{q}},{\bm{c}})$, on the HotpotQA dataset. The dataset construction for the long-context setting follows the same procedure as described in the main experiments (§4). The hallucination rate is defined as the percentage of retrieved sentences that do not match any sentences in the context, calculated as:

$$ $\mathrm{Hallucination}(\mathcal{Z},\mathcal{C})=1-\frac{1}{N}\sum_{i=1}^{N} \mathds{1}_{\{z_{i}\in{\bm{c}}_{i}\}},$ $$

where $\mathcal{Z}=\{z_{i}\}^{N}_{i=1}$ is the collection of all retrieved sentences, $z_{i}$ is a single retrieved sentence, and $\mathcal{C}=\{{\bm{c}}_{i}\}^{N}_{i=1}$ is the corresponding collection of input context. Additionally, since HotpotQA provides golden supporting facts for each question, we evaluate retrieval quality using the recall score, which measures the fraction of golden facts successfully retrieved by the LLM.

#### Results.

As shown in Table 2, by aligning generation with retrieval objective, the hallucination rate of our approach is significantly better than the baseline.
Additionally, ALR^2 also achieves notably higher retrieval recall scores compared to the prompting-based method, highlighting the importance of generation-retrieval alignment. The reduced hallucination rate and improved retrieval recall are not only crucial for enhancing the generator’s accuracy, but they also provide a more trustworthy rationale, helping users better understand the decision-making process of LLMs.

### 5.2 Generalization to Unseen Data

#### Setup.

We are also interested in whether ALR^2, trained on the SQuAD (Rajpurkar et al., ) and HotpotQA (Yang et al., ) datasets, can generalize to unseen datasets. To explore this, we follow the strategy of Ruler (Hsieh et al., ) to pre-process two additional QA datasets: StrategyQA (Geva et al., ) and TriviaQA (Joshi et al., ). We compare the performance of the LLM fine-tuned on the direct-answer template (CMD-R-FT + DA) with our approach (ALR^2).

#### Results.

As shown in Table 3, ALR^2 consistently outperforms CMD-R-FT + DA across context lengths ranging from 4K to 128K. These results indicate that ALR^2 provides a more robust and stable framework for solving long-context QA problems, even when applied to unseen datasets.

**Table 3: EM scores on StrategtQA and TriviaQA tasks. The best results are highlighted by bold text.**
| Model | Prompt | 4K | 8K | 16K | 32K | 64K | 128K | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| StrategyQA |  |  |  |  |  |  |  |  |
| CMD-R-FT | DA | 61.57 | 60.26 | 60.69 | 60.69 | 58.95 | 56.33 | 59.74 |
| ALR^2 | RR | 72.48 | 71.61 | 70.74 | 70.16 | 71.61 | 72.48 | 71.51 |
| TriviaQA |  |  |  |  |  |  |  |  |
| CMD-R-FT | DA | 74.6 | 74.2 | 73.2 | 70.0 | 69.0 | 69.19 | 71.69 |
| ALR^2 | RR | 77.0 | 76.4 | 74.4 | 72.60 | 73.94 | 71.98 | 74.38 |

## 6 Related Work

### 6.1 Long-context Large Language Model

Long-context ability is critical for employing a large language model (LLM) to solve real-world problems, e.g., document analysis (Wang et al., ), tool use (Yao et al., ; Schick et al., ), etc. Many components of an LLM may affect the long-context performance. For instance, researchers find that the design of the position embedding is important for the extrapolation performance of an LLM (Su et al., ; Press et al., ; Chen et al., ; Ruoss et al., ; Men et al., ), i.e., testing the LLM on context length that is larger than the training window size. In addition, the attention distribution and scale of the Transformer model (Chiang & Cholak, ; Xiao et al., ) may also affect the long-context ability.

To extend the context length of LLM, some works try to compress the long context into shorter text (Jiang et al., ; ; Xu et al., ), limited hidden representations (Chevalier et al., ; Ge et al., ), or model parameters (Wang et al., ), saving the number of input tokens in the context. Moreover, splitting long context into fixed segments and caching them for re-use is also widely used for increasing the effective context length of LLMs (Dai et al., ). Different from Transformer-XL (Dai et al., ) that only leverages local attention, Munkhdalai et al. () introduce an additional global momentum to store the knowledge in all the previous context. Yen et al. () leverages a small bi-directional model (Liu, ) to encode the segments independently, and allows the LLM to attend to those segments during generation.

In our work, the target is not to extend the context length. In contrast, we find that the performance of QA will degenerate more significantly with longer context for most popular LLMs, similar to the findings in Liu et al. (). Therefore, we propose a new formulation inspired by retrieval-augmented generation to alleviate this common issue of LLMs.

Recently, there are also many benchmarks are proposed to evaluate the long-context performance of LLMs (Hsieh et al., ; Wang et al., ; Lee et al., ; Zhang et al., ; Krishna et al., ; Karpinska et al., ; Kuratov et al., ). Considering our focus is the general QA problems and the construction of long-context QA data is similar in many benchmarks, we use Ruler (Hsieh et al., ) as our main benchmark.

### 6.2 Retrieval-Augmented Generation

Retrieval-Augmented generation (RAG) has been extensively used to solve NLP problems, e.g., open-domain question-answering (Lewis et al., ; Asai et al., ), language modeling (Guu et al., ; Lan et al., ; Shi et al., ), machine translation (Gu et al., ; Cai et al., ), etc. For most RAG works, the dual-encoder-based retriever (Karpukhin et al., ; Izacard et al., ; Ni et al., ), which measures the similarity between inputs and retrieval documents by their encoded embeddings, has became the de facto approach for retrieval. The retrieved top-$K$ documents will be integrated into the generator, e.g., a language model (Radford et al., ), to enhance the generation. The most related work in this research line is Xu et al. (), which aims to save the number of input tokens to an LLM by compressing the retrieved $K$ documents. However, our work focuses on enhancing the long-context reasoning of LLMs.

### 6.3 Citation Generation

There exist many works that aim to attribute the prediction of an LLM to the information source that the LLM is based on, i.e., citation generation. Some train the model to learn this ability on a standardized citation format (Menick et al., ; Taylor et al., ; Khalifa et al., ; Ye et al., ). While others use the instruction-following ability and post-processing techniques to enable this feature at inference (Gao et al., ; Kamalloo et al., ; Gao et al., ; Liu et al., ).

Our work aims to precisely retrieve the bag of facts that can support the question-answering in long-context scenario, which is also beneficial for the explainability of the prediction. However, the focus of our work is more about enhancing the long-context QA ability of LLMs, which is orthogonal to works in citation generation.

## 7 Conclusion

In this work, we addressed the challenge of long-context reasoning in large language models (LLMs) by introducing ALR^2, a retrieve-then-reason approach that aligns LLMs with both retrieval and reasoning objectives. Through extensive experiments on both single-hop and multi-hop QA tasks, we demonstrated that ALR^2 significantly outperforms existing baselines, especially in long-context scenarios. Our analysis also highlighted the importance of aligning generation with retrieval to reduce hallucinations and improve retrieval accuracy. Moreover, we showed that ALR^2 generalizes well to unseen datasets, offering a robust solution for long-context QA problems.

#### Limitations.

We also note some of the limitations of our work. For example, as most works based on the RAG formulation, our method is hard to enhance the summarization tasks, in which all the information in the long context is important for the final prediction. In addition, only considering one granularity, i.e., sentence, in the retrieval may be limited in some scenarios. A better way is to allow the users to choose the retrieval granularity, e.g., phrase, sentence, or passage, in the instruction. We leave the addressing of those limitations in future works.

## Appendix A Additional Results for Preliminary Study

We also evaluate the effectiveness of the retrieve-then-reason framework on the challenging task in the section of preliminary study (§2). As shown in Table 4, simply using the RR prompt (Figure 3) for the Command-R-Plus model can improve the overall EM score from 0.21 to 0.56, significantly alleviating the performance drop on the task that requires long-context reasoning.

**Table 4: EM scores on Task 3, i.e., Multi-hop Passkey Retrieval + Reasoning, in our preliminary study (§2). The best results are highlighted by bold text.**
| Model | Prompt | 16K | 32K | 64K | 128K | Overall |
| --- | --- | --- | --- | --- | --- | --- |
| CMD-R-Plus | DA | 0.70 | 0.16 | 0.0 | 0.0 | 0.21 |
| CMD-R-Plus | RR | 0.84 | 0.86 | 0.56 | 0.0 | 0.56 |

## Appendix B Details of Experiment Setup

### B.1 Ruler Tasks

#### QA Tasks

We use the validation sets of SQuAD (Rajpurkar et al., ) and HotpotQA (Yang et al., ) datasets to build the long-context benchmarks, where each question in those datasets is associated with some relevant passages. Following Hsieh et al. (), We use concatenated and shuffled passages as the context, which include both relevant passages and randomly sampled ones.

#### NIAH Tasks

Following Hsieh et al. (), we consider four kinds of Needle-in-a-Haystack (NIAH) tasks in our work and report the averaged performance of them:

- •
Single NIAH (S-NIAH): The “haystack” text is from Paul Graham essays (Kamradt, ) is from “needle” sentence of this task is “One of the special magic numbers for {KEY} is: {VALUE}”, where the “{KEY}” and “{VALUE}” are placeholders for English words and 7-digit numbers, respectively. For S-NIAH, we only insert one needle sentence into the long-context.
- •
Multi-keys NIAH (MK-NIAH): In this setting, we insert four needle sentences into the context and ask an LLM to retrieve only one of them. Each of the four needle sentences has unique “{KEY}” and “{VALUE}”.
- •
Multi-values NIAH (MV-NIAH): In this setting, part of the needle sentences may share the same key. We will ask an LLM to retrieve all the values associated with the same key.
- •
Multi-queries NIAH (MQ-NIAH): Compared with MK-NIAH, we need the LLM to retrieve multiple values in needle sentences associated with specific keys.

Figure: Figure 5: The quotes-and-citation-first (QF) prompt for long-context QA. The {CONTEXT} is the placeholder for long context and {QUERY} is for user question. The format of {QUOTES_AND_ANSWER} is similar to the example embraced by <example> and </example>.

Figure: Figure 6: The system-2 attention (S2A) prompt for long-context QA. The {CONTEXT} is the placeholder for long context and {QUERY} is for user question. The {COMPRESSED_TEXT} contains the copied question and the compressed long context based on the question. In stage-2, we only input the {COMPRESSED_TEXT} to LLM, discarding the previous chat history.

### B.2 Training Details

During training, we set the base value of rotary position embedding to $8e6$ (Men et al., ; Su et al., ). The training batch size is 64 on V4-512 TPU. For every 3 updates on the prepared training dataset, we will also update our parameters on pre-training data with 64K tokens for one step. For both of the two-stage and one-stage fine-tuning methods, we will train the command-r model with 35 billion parameters for 3 epochs.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. As mentioned in the first weakness, can you evaluate both efficiency and effectiveness comparison between these methods, and show how in-context retrieve-then-reason can bring extra benefits compared to RAG?

2. 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"
soundness: 3
strengths: 1. The research shows that breaking down long-context reasoning into explicit retrieval and reasoning steps can enhance LLM performance on long document understanding tasks and also provide more transparent reasoning processes. These combined strengths position the paper with higher value.

2. The model can help with hallucination mitigation, where the system dramatically reduces the tendency to generate fictional content by explicitly aligning the LLM with retrieval objectives.

3. The paper is easy to follow and is well-written.
summary: This paper focus on addressing the challenge of LLMs' degrading performance when reasoning over long contexts. This paper proposes a two-stage retrieve-then-reason approach: first retrieving relevant information from the long context, then reasoning over the retrieved information to produce answers. The model is explicitly trained to align with both retrieval and reasoning objectives, leading to improvements over baselines.
weaknesses: 1. 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. 

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

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

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Lines 357-359 point out that QF only retrieves 2.2 sentences on average (compared to 7.5 for RR). However, in HotpotQA, it is common for only a small number of supporting facts (1-2) to be necessary [1] (even if 2-6 supporting facts are provided). To assess whether the retrieval was truly effective, could the performance of CMD-R QF also be included in the experiment results of Section 5.1? It is particularly important to include not only the recall score but also precision. Relying solely on the recall score could bias the evaluation toward retrieval methods that select a larger number of facts.

[1] [MuSiQue: Multihop Questions via Single-hop Question Composition] (Trivedi et al., TACL 2022)
soundness: 2
strengths: - The paper is generally well-written, easy to understand, and clearly presented.
- The preliminary study provides interesting insights into where issues arise when LLMs handle long contexts, showing that while retrieval itself is not problematic, the issues occur during reasoning.
summary: This paper proposes a retrieve-then-reason method to address the issue that LLMs (Large Language Models) struggle to handle long contexts properly, especially in reasoning. This proposed approach improves the handling of long contexts by explicitly retrieving supporting facts from the given long context as well as by aligning LLMs with retrieval and reasoning objectives at the same time.
weaknesses: In my opinion, there seem to be some areas for improvement in the design of the experiments aimed at proving the claims of this paper. If these below points can be clarified through additional responses from the authors, I would be glad to adjust my evaluation accordingly.

1. 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?

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

3. There may not have been a fair comparison with the baselines. It appears that there is a difference in the answer template between the DA prompt (figure 4) and the RR prompt (figure 3), as the "Answer:" is missing in figure 3. Given that LLMs can be sensitive to prompts, this could lead to performance differences. To provide evidence that the comparisons in the experiment were fair, results using the same answer template should be included. Additionally, based on figure 3 and equation (2), it seems that the given long context was used alongside the retrieved facts. A comparison showing the performance when the given long context (i.e., prior chat history) is removed would be necessary for a fair comparison with S2A, which also omits prior records.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Regarding weakness 1 - do you think there are QA datasets that have long-context by design and might be relevant for evaluation?

Regarding weakness 3 - do you think there are other experiments/ablations that are similar to current methods that can be added? For example, the model can be trained to first re-rank the long context (variants that generate relevant indexes for passages like in RankZephyr or relevant snippets as in the paper can be explored). Additionally, an ablation where the final answer is generated by a short-context QA model can help understand if the full context is necessary for reasoning after the retrieval step, or if only the information in the retrieval step is sufficient to answer the questions.

Additionally, perhaps I missed something, but will the model and data be publicly available? 

Also, I believe the paper could benefit from a discussion regarding reproducibility and an ethics statement.
soundness: 2
strengths: - The paper deals with an important subject, long-range QA.
- The improvements over a variant trained to directly generate the answer are significant, and improvements generalize to unseen datasets.
- The paper is well-written and easy to follow.
summary: The paper presents $ALR^2$, a retrieve-and-reason framework for QA over long contexts. $ALR^2$ is trained to first retrieve relevant information from a long context, and then reason over it to arrive at an answer. With Command-R, $ALR^2$ significantly outperforms prompted variants and a model directly trained to generate an answer on HotpotQA and SQuAD (in a variant where gold facts are augmented to create a long context), and generalizes to unseen StrategyQA and TriviaQA datasets.
weaknesses: - 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).

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

- 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 $ALR^2$ 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 $ALR^2$ (see questions).

- While the paper is overall well-written, there are a few parts that I think could be further improved. Mainly, the preliminary study provides motivation for the work, but lacks a qualitative analysis, so it was not clear to me why models struggle on Task 3. Additionally, I felt like the discussion in Section 3.5 was a bit misleading, as dense retrievers are often much more efficient than the presented approach.


[1] SCROLLS: Standardized CompaRison Over Long Language Sequences

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

Meta-review: This submission introduces ALR2 designed to enhance the long-context reasoning ability of LLMs. While LLMs have larger context windows, the authors argue that their reasoning capability over long contexts often degrades due to difficulty in identifying and reasoning over relevant information. The idea of ALR2 is to address this by designing a retrieve-then-reason framework. First, it retrieves relevant information; Second, it starts to "reason" with related background information. The proposed method is shown to improve retrieval accuracy and reduce hallucination. 

The reviewers identified the strengths of this work as:
- this paper is easy to follow and clearly presented (Reviewers 4gL5, Q28w, WCoe).
- the retrieve-then-reason framework can enhances LLM performance on long-context understanding tasks, helping reduce hallucinations (Reviewers 4gL5, Q28w).
- the proposed r-and-r method shows improvements over some baseline models, especially on long-context QA tasks (Reviewers 4gL5, WCoe).

The points that this work could improve upon include the following: 
- Several reviewers in the original reviews and discussions mentioned that the experimental design could be improved, particularly on lacks comparisons to important baselines, e.g., DPR, RETA-LLM etc.,  and lacks certain important datasets, and 2) it does not provide fair comparisons due to differences in answer templates and context usage (Reviewers 4gL5, Q28w).
- the experiments lack details, e.g., hyperparameters, training data, ablation studies (Reviewers Q28w, WCoe).
- The idea of using retrieval and reasoning has been presented before in similar methods, and the authors did not offer evidence / arguments to highlight its unique contributions (Reviewers WCoe, Q28w).

During the rebuttal, the authors successfully convinced two reviewers to update their ratings. Relatively speaking, the engagement between reviewers and the authors are enthusiastic, the reviewers were also active during discussion. 
- Reviewer WCoe updated the score from 5 to 6. The reviewer also acknowledged that many of the questions were answered by the authors' rebuttal. However, the reviewer remained concerned about its empirical setup and the lack of importance baselines (as mentioned above). 
- Reviewer Q28w raised the rating from 3 to 5 as some issues were addressed by the authors during the rebuttal. This reviewer had extensive back-and-forth discussions with the authors regarding the datasets and evaluation metrics used. Ultimately, the reviewer decided  not to increase the score further. 
- Reviewer 4gL5 (an initial & final rating of 6) did not respond to the authors' rebuttal. 


Overall, the final ratings for this work are 6, 6 (up from 5), 5 (up from 3). The reviewers' main concerns on technical novelty, empirical demonstration (missing baselines and important datasets / metrics) remained unaddressed after the rebuttal. From the AC's standpoint, the idea of retrieval-and-reasoning presented in this work does not outweigh these issues. Based on these factors, this work in its current form is not recommended for acceptance, since it would significantly benefit from another round of revision.

Decision: Reject