ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities

## Abstract

Abstract In this work, we introduce ChatQA 2, an Llama 3.0-based model with a 128K context window, designed to bridge the gap between open-source LLMs and leading proprietary models (e.g., GPT-4-Turbo) in long context understanding and retrieval-augmented generation (RAG) capabilities.
These two capabilities are essential for LLMs to process large volumes of information that cannot fit into a single prompt and are complementary to each other, depending on the downstream tasks and computational budgets.
We present a detailed continued training recipe to extend the context window of Llama3-70B-base from 8K to 128K tokens, along with a three-stage instruction tuning process to enhance the model’s instruction-following, RAG performance, and long-context understanding capabilities.
Our results demonstrate that the 𝙻𝚕𝚊𝚖𝚊𝟹 ⁢ - ⁢ 𝙲𝚑𝚊𝚝𝚀𝙰 ⁢ - ⁢ 𝟸 ⁢ - ⁢ 𝟽𝟶 ⁢ 𝙱 𝙻𝚕𝚊𝚖𝚊𝟹 - 𝙲𝚑𝚊𝚝𝚀𝙰 - 2 - 70 𝙱 \mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B} typewriter_Llama3 - typewriter_ChatQA - typewriter_2 - typewriter_70 typewriter_B model outperforms most existing state-of-the-art models, including GPT-4-Turbo-2024-04-09, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct, on ultra-long tasks beyond 100K tokens, as well as on the RAG benchmark using only a 4K context window, showing the strong long context capability across varying sequence lengths.
We further provide extensive comparisons between direct long-context and RAG solutions using the same state-of-the-art long-context LLMs.
Interestingly, we find that the performance of strong long-context LLMs using RAG improves when retrieving a larger number of chunks. With a large set of top- k chunks, RAG consistently outperforms direct long-context solution using the same state-of-the-art long-context models (e.g., 𝙻𝚕𝚊𝚖𝚊𝟹 ⁢ - ⁢ 𝙲𝚑𝚊𝚝𝚀𝙰 ⁢ - ⁢ 𝟸 ⁢ - ⁢ 𝟽𝟶 ⁢ 𝙱 𝙻𝚕𝚊𝚖𝚊𝟹 - 𝙲𝚑𝚊𝚝𝚀𝙰 - 2 - 70 𝙱 \mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B} typewriter_Llama3 - typewriter_ChatQA - typewriter_2 - typewriter_70 typewriter_B and Qwen2-72B-Instruct) on both 32K benchmarks and real-world 128K tasks.
To advance research in this field, we open-sourced the model weights,
training data, and the evaluation setup for the for the community: https://chatqa2-project.github.io/

## 1 Introduction

The open LLM community has made significant progress in advancing the capabilities of open-access large language models (LLMs), including Llama-3-70B-Instruct (Meta-AI, ), Qwen2-72B-Instruct (Alibaba-QWen, ), Nemotron-4-340B-Instruct (NVIDIA et al., ), and Mixtral-8x22B-Instruct-v0.1 (Mistral, ). However, performance gaps compared to frontier proprietary models, e.g., GPT-4-Turbo (OpenAI, ), still exist in many domains.
Additionally, open-access models focused on key domains have been developed, such as DeepSeek-Coder-V2 (Zhu et al., ) for coding and math, ChatQA 1.5 (Liu et al., ) for conversational QA and retrieval-augmented generation (RAG), and InternVL 1.5 (Chen et al., ) for vision-language tasks, which can be on par with GPT-4-Turbo-2024-04-09 (OpenAI, ) in the certain domains.

In recent developments, the trend of extending the context window length in LLMs has gained remarkable traction within both the industrial and research communities. All leading proprietary LLMs support very large context window, allowing them to accommodate several hundred pages of text in a single prompt. For example, GPT-4 Turbo (OpenAI, ) and Claude 3.5 Sonnet offer a 128K and 200K context window, respectively.
Meanwhile, Gemini 1.5 Pro (Gemini-Team, ) impressively
supports up to a 10M context.
Open-access LLMs have also made significant strides to keep up (01.AI et al., ; Alibaba-QWen, ). For instance, Qwen2-72B-Instruct (Alibaba-QWen, ) and Yi-34B (01.AI et al., ) support 128K and 200K context windows, respectively.
The concurrent work Llama-3.1-70B-Instruct (Meta-AI, ) also supports a 128K context window.
However, the absence of training data and reproduction recipe makes replicating these models challenging.
In addition, these models have mostly been evaluated on synthetic tasks, like Needle in a Haystack (Kamradt, ) test, which does not accurately represent real-world downstream task performance.
For example, previous studies have shown a significant gap between open-access LLMs and leading proprietary models on real-world long-context understanding tasks (Zhang et al., ; Hsieh et al., ). In this work, we aim to bridge the gap between open-source LLMs and proprietary models (e.g., GPT-4-Turbo-2024-04-09) in real-world long-context understanding tasks.

**Table 1: $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ versus state-of-the-art long context LLMs across tasks with varying context lengths. Notably, *we open-source the training data and reproduction recipe for building 128K long-context LLMs from pretrained base models with 8K context*, resources that are not currently available for open-access (open-weights) 128K context models like Qwen2 and Llama3.1. Our model achieves the highest average score on four real-world ultra-long context (beyond 100K) tasks from InfiniteBench, and nine short context (within 4K) tasks from ChatRAG Bench.**
| Model Type | Long Context LLM (128K) | Ultra-long (>100K) | Long (32K) | Short (4K) |
| --- | --- | --- | --- | --- |
| 4 tasks | 6 tasks | 9 tasks (RAG) |  |  |
|  | GPT-4-Turbo-2024-04-09 | 33.16 | 51.93 | 54.72 |
|  | Qwen2-72B-Instruct | 39.77 | 49.94 | 54.06 |
| Llama3.1-70B-Instruct | 39.81 | 49.92 | 52.12 |  |
|  | $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ | 41.04 | 48.15 | 56.30 |

The long-context capability of LLMs is sometimes considered a rival technique to retrieval-augmented generation (RAG).
In fact, RAG and long-context techniques complement each other from a pragmatic perspective.
An LLM with a long context window can either process large volumes of text as a prompt or utilize retrieval methods to efficiently extract relevant information from the extensive text, depending on the downstream tasks and accuracy vs. efficiency trade-offs.
RAG has efficiency advantages and can easily retrieve relevant contexts for query-based tasks (e.g, QA) from billions of tokens, a feat that long context models cannot achieve so far. Meanwhile, long context models are good at tasks such as summarizing entire documents, where RAG can not perform.
As a result, the state-of-the-art (SOTA) LLM needs to excel at both capabilities, providing options for different downstream tasks and meeting different accuracy-efficiency trade-off requirements.

In our previous study, the open-source ChatQA 1.5 (Liu et al., ) model can surpass GPT-4-Turbo on RAG tasks within 4K context window. For longer context tasks,
Xu et al. () extended the context window of Llama2 (Touvron et al., ) to 16K and 32K tokens, and studied the interplay between RAG and long-context LLMs. It demonstrates that the RAG method can improve the generation accuracy and inference efficiency of a GPT-3.5-turbo-16k level long context model for QA and query-based summarization tasks.

In this work, we present ChatQA 2, which extends the study of long-context and RAG by pushing long-context LLMs to the GPT-4-Turbo level capability together with a 128K context window, and combining it with a state-of-the-art long-context retriever for RAG.
See Table 1 for an overview of the results.
The training data and reproduction recipe for this work are publicly accessible.

Specifically, we make the following contributions:

- 1.
We present a two-step approach to establish the long context capability of Llama3-70B. First, we extend Llama3-70B base’s context window from 8K to 128K by continually pretraining it on a mix of SlimPajama (Soboleva et al., ) with upsampled long sequences (Fu et al., ). Then, we apply a three-stage instruction tuning process on curated datasets to enhance the instruction-following, RAG, and long context understanding capabilities at each respective stage. We find that stage-wise instruction-tuning, by incorporating previous datasets, simplifies experimentation and hyperparameter tuning.
- 2.
We demonstrate that the resulting $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ with 128K context window achieves better accuracy than most existing state-of-the-art models, including GPT-4-Turbo-2024-04-09, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct, on both long-context tasks exceeding 100K tokens and the RAG benchmark for conversational QA tasks within a 4K context window.
- 3.
Previous RAG approachs has two potential issues that can undermine downstream task accuracy: *i)* top-*k* chunk-wise retrieval introduces fragmentation of context, and *ii)* a small top-*k* leads to low recall, while a larger *k* introduces too much irrelevant context to the LLM (e.g., see the analysis in Figure 1 of Yu et al., ).
In contrast, we find that the RAG performance of our $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ is highly robust to variations in chunk size when using a long-context retriever (Wang et al., ; Lee et al., ). Even more promising, accuracy consistently improves as the total number of retrieved tokens for input increases.
- 4.
We conduct extensive comparisons between direct long-context and RAG solutions using the same state-of-the-art long-context LLMs.
We find that RAG can still outperform GPT-4-Turbo-level long-context models (including our $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$) on both 32K benchmarks and real-world 128K tasks, given a sufficient number of top-k chunks.

We organize the rest of the paper as follows.
We discuss related work in § 2.
We introduce the continued pretraining for context window extension in § 3.1 and the three-stage instruction tuning in § 3.2.
We report results in § 5 and conclude the paper in § 7.

## 2 Related Work

### 2.1 Long Context LLM

The trend of extending the context window in LLM starts by Claude with 100K token context (Anthropic, ).
Although the underlying long context techniques behind proprietary models are unclear, the open LLM and research community has developed many methods to extend the context window of LLMs through continued training or fine-tuning (Kaiokendev, ; Nijkamp et al., ; Chen et al., ; Tworkowski et al., ; Chen et al., ; Peng et al., ; Xiong et al., ; Fu et al., ), especially for open-access LLMs (Touvron et al., ; ) based on rotary position embedding (RoPE) (Su et al., ).

There are two popular approaches to adapt RoPE for long-context inputs: position interpolation (Chen et al., ) and increasing the base frequency $\theta$ of RoPE (Xiong et al., ; Liu et al., ).
Recently, Yi-34B (01.AI et al., ) was pretrained with a sequence length of 4K, and its context window was extended to 200K by increasing the RoPE $\theta$ from 10,000 to 5M during continued pretraining. Qwen2-72B-Instruct (Alibaba-QWen, ) was trained on 32K-length contexts and extrapolated to a 128K context length using YaRN (Peng et al., ). Instead of extending the context window of the base model then applying instruction tuning, GradientAI () directly fine-tunes the Llama-3-Instruct, which uses NTK-aware interpolation (Peng et al., ) and the formula from Liu et al. () to scale up $\theta$.

After the release of the first version of this work, the Llama-3.1 models with a 128K context window were introduced. We add a comprehensive comparison between ChatQA-2 and Llama-3.1 in the experiments.

### 2.2 Retrieval-augmented Generation (RAG)

Retrieval with a standalone retriever (e.g., Karpukhin et al., ; Wang et al., ; Lin et al., ; Lee et al., ) is a long-standing solution for handling long texts that cannot fit into the context window of language models.
In previous work, various retrieval-augmented language models have been proposed (Nakano et al., ; Borgeaud et al., ; Wang et al., ; ; Guu et al., ; Izacard & Grave, ; Izacard et al., ; Lewis et al., ; Huang et al., ; Khandelwal et al., ; Liu et al., ).

Previous dense-embedding-based retrievers only supported limited context windows (e.g., 512 tokens) (e.g., Karpukhin et al., ; Wang et al., ; Lin et al., ).
In top-*k* chunk-wise retrieval, the short chunk size increases context fragmentation. As a result, extending the context window of retrievers has become popular. For example, Jina Embeddings 2(Günther et al., ) and Nomic Embed (Nussbaum et al., ) support 8K tokens, while E5-mistral-7B (Wang et al., ) and NV-Embed (Lee et al., ) support 32K tokens.

## 3 Method

### 3.1 Extending Context Window to 128K

In this section, we present the method to extend the context window from 8K to 128K for Llama3.
Note that, the Llama-3.1 models with 128K context window are released after the first version of this work. We compare our $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ and $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}8B}$ with Llama-3.1-70B-Instruct and Llama-3.1-8B-Instruct in experiments.

We prepare our long context pretraining corpus from the Slimpajama (Soboleva et al., ) following Fu et al. (). We upsample long-context documents with the hyperparameter set as 0.1 to produce 10 billion tokens with sequence length of 128k.
Since Llama3 is pretrained with a much higher RoPE base frequency of 500,000 compared to Llama2, we increased the RoPE base frequency to 150M accordingly. We set the batch size to 32 to fit 4 million tokens in a batch and use a learning rate of 3e-5 to train 2000 steps (8B tokens in total).

Interestingly, we found it more effective to separate different documents using special characters, such as "
", rather than the reserved beginning and ending tokens <BOS> and <EOS>. We hypothesize that the <BOS> and <EOS> tokens in Llama3 signal the model to ignore previous chunks of text after pretraining, which is not helpful for the LLMs to adapt for longer context inputs.

### 3.2 Instruction-Tuning with Long Context Data

In this section, we present the instruction tuning method designed to enhance both long context understanding capability and RAG performance.

Specifically, we implement three stages of instruction-tuning. For the first two stages, we follow ChatQA 1.5 (Liu et al., ), where the model is initially trained on 128k high-quality instruction-following datasets, and then trained on a blend of conversational QA data with provided context.
However, these two stages involve relatively short contexts, with a maximum sequence length of only 4K tokens.
To enhance our model’s capability to handle very long context sequences up to 128k tokens, we collect a long SFT dataset.

This dataset is collected through two categories:
1) For SFT data sequences less than 32k: We leverage existing long-context datasets from LongAlpaca12k, GPT-4 samples from Open Orca (^1^11[https://huggingface.co/datasets/Open-Orca/OpenOrca](https://huggingface.co/datasets/Open-Orca/OpenOrca)), and Long Data Collections (^2^22[https://huggingface.co/datasets/togethercomputer/Long-Data-Collections](https://huggingface.co/datasets/togethercomputer/Long-Data-Collections)).
2) For sequence lengths between 32k and 128k: Since it is challenging to collect such SFT samples, we rely on synthetic datasets. We utilize NarrativeQA, which contains both the ground truth summary and semantically related paragraphs. We assemble all the related paragraphs and randomly insert the ground truth summary to simulate a real long document for its QA pairs. Both the full long SFT dataset and the short SFT dataset from the first two stages are then blended for training. We set the learning rate at 3e-5 and the batch size at 32.

### 3.3 Long Context Retriever meets Long Context LLM

As we mentioned in previous section, the current RAG pipeline for LLM has the following issues:
*i)* The top-*k* chunk-wise retrieval introduces non-negligible fragmentation of context for generating accurate answers. For example, previous state-of-the-art dense-embedding based retrievers (e.g., Li et al., ; Lin et al., ) only support 512 tokens.
*ii)* Small top-*k* (e.g., 5 or 10) usually leads to relatively low recall, while much larger *k* (e.g., 100) can lead to worse generation (see Table 5 in Xu et al. ()) as the previous LLMs could not utilize too many chunked context very well (Liu et al., ).
To address the issue, we propose to use the most recent long-context retriever (Wang et al., ; Lee et al., ), which can support thousands of tokens.
In our setting, we use the E5-mistral embedding model (Wang et al., ) as the retriever.

## 4 Baselines and Evaluation Benchmarks

#### Long context models

We compare our models against SOTA long context models: 1) GPT-4-Turbo-2024-04-09 (128K context window) (OpenAI, ), 2) Qwen2-72B-Instruct (128K context window) (Alibaba-QWen, ), 3) Llama-3-70B-Instruct-Gradient-262k (256K context window) (GradientAI, ), 4) Llama-3.1-8B-Instruct and Llama-3.1-70B-Instruct (128K context window).
Note that, the Llama-3.1 models are released after the first version of this work.

#### Retrieval-augmented generation (RAG)

We apply RAG to our ChatQA 2 models and baseline long-context LLMs when it is applicable to the downstream tasks, such as those with informative prompts or queries. By default, we use the E5-mistral retriever (Wang et al., ), as it supports both short and long contexts. In our experiments, we use a chunk size of 1200 and select the top 5 chunks, as this configuration provides a good trade-off between cost and performance, as demonstrated in the ablation study in Section 6.

To give a comprehensive study of different context lengths, our evaluation benchmarks covers three categories, 1) ultra-long context benchmarks beyond 100K tokens, 2) long context benchmarks within 32K tokens, and 3) short context benchmarks within 4K tokens.

### 4.1 Ultra-long Context Benchmarks Beyond 100K Tokens

InfiniteBench (Zhang et al., ) is proposed to evaluate the ultra-long context capability of LLMs over 100K sequence length. As we focus on real-world english tasks, we only take the four related tasks from the InfiniteBench, i.e. longbook summarization (En.Sum), longbook qa (En.QA), longbook multiple choice (En.MC), and longbook dialogue (En.Dia). En.Sum is a task that requires models to generate a concise summary of the given novel and is evaluated using the ROUGE-L-Sum metric (Lin, ). En.QA is annotated by a pipeline that ensures the questions’ necessitating of long-range dependencies and reasoning, beyond simple short passage retrieval. Aggregation reasoning and filtering reasoning are the two primary reasoning categories. F1 score is used to evaluate the quality of the answer. En.MC is annotated with the same pipeline of En.QA except that four answer choices are provided and exact matching scores are reported. En.Dia leverages movie and drama scripts from a designated online database (^3^33[https://imsdb.com](https://imsdb.com)) with long, multi-role dialogues. Exact matching score is used again to evaluate the prediction accuracy.

### 4.2 Long Context Benchmarks within 32K Tokens

We use the long context datasets (except NarrativeQA as it is included in our training) from Xu et al. () as our benchmark for long datasets within 32K. There are six datasets in total, where QMSum (QM), Qasper (QASP), QuALITY (QLTY) are token from SCROLLS and HotpotQA (HQA)
MuSiQue (MSQ), MultiFieldQA-en (MFQA) are token from LongBench. Following the official metrics, we report the geometric mean of ROUGE scores (i.e., ROUGE1/2/L) (Lin, ) for QM, the exact matching (EM) score for QLTY, and F1 scores for the remaining four datasets QASP, MSQ, HQA and MFQA.

### 4.3 Short Context within 4K Tokens

We use ChatRAG Bench (Liu et al., ) as our benchmark for short context within 4k. ChatRAG bench consists of 10 datasets and we exclude HDial as it is vasincluded in our training.
Following the setup of Liu et al. (), for Doc2Dial (D2D), QuAC, and QReCC task with long documents, each document is divided into segments of roughly 300 words. The top 5 relevant chunks are then retrieved as context for each user question.
For TopiOCQA and INSCIT, top-20 chunks were retrieved to obtain similar context length to the first three datasets. The other four datasets are CoQA, DoQA, ConvFinQA (CFQA), and SQA, which cover a wide range of domains like finance, children’s stories, literature, mid/high school exams, news, Wikipedia and etc. We use F1 score as the metric to evaluate the generations and report the average score without HDial as it is a fair zero-shot comparisons over different models.

## 5 Results

In this section, we present the results and comparisons from extensive benchmark evaluations.
We begin with the synthetic Needle in a Haystack test, then focus on real-world long context understanding and RAG tasks.

Figure: (a) Llama3.1-Instruct 8B
Refer to caption: extracted/5841941/figures/llama3_1/llama3_1_instruct_8b_tp8_mcore.png

### 5.1 Needle In A Haystack

We evaluate our $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ and $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}8B}$ model on the Needle In A Haystack (NIAH) test (Kamradt, ). This synthetic task is popular for testing the long-context capability of LLMs, and can be considered as a threshold level evaluation.
We use the needle test: *“The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.”*
We find this is generally a harder pressure test than the “passkey” in a haystack test, which is commonly used by other long-context LLM works.
Figure 3 demonstrates the performance of our models up to 128K tokens. It shows that our model achieves 100% accuracy under both test cases. This test confirms our model’s excellent long-context retrieval capability.

In addition, we also include the NIAH test for Llama3.1-Instruct 8B and Llama3.1-Instruct-70B. From Figures 1(a) and 1(b), we find that neither model could pass the NIAH test with the needle: *“The best thing to do in San Francisco…”*, although they could pass the easier *Passkey* retrieval test, as shown in Figures 3(a) and 3(b) in Appendix.

**Table 2: Evaluation results on InfiniteBench includes real-world long-context understanding tasks beyond a 100K context window. For RAG, we use top-5 retrieved chunks, each with 1200 tokens from E5-mistral retriever (Wang et al., ).**
| Model | Avg. | En.Sum | En.QA | En.MC | En.Dia |
| --- | --- | --- | --- | --- | --- |
| GPT-4-Turbo-2024-04-09 | 33.16 | 17.62 | 19.29 | 77.73 | 18.00 |
| Claude 2 | 33.96 | 14.50 | 11.97 | 62.88 | 46.50 |
| Yi-34B-200K | < 15.15 | < 5 | 12.17 | 38.43 | < 5 |
| w/ RAG | N/A | N/A | 17.69 | 77.29 | N/A |
| Qwen2-72B-Instruct | 39.77 | 31.66 | 21.46 | 82.97 | 23.00 |
| w/ RAG | N/A | N/A | 16.48 | 76.86 | N/A |
| Llama-3-70B-Instruct-Gradient-262k | 32.57 | 14.27 | 29.52 | 69.00 | 17.50 |
| \hdashlineLlama3.1-8B-Instruct | 33.17 | 29.20 | 31.51 | 58.95 | 13.00 |
| w/ RAG | N/A | N/A | 15.41 | 64.63 | N/A |
| Llama3.1-70B-Instruct | 39.81 | 30.94 | 38.49 | 75.55 | 14.29 |
| w/ RAG | N/A | N/A | 23.44 | 79.04 | N/A |
| $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}8B}$ | 35.59 | 17.14 | 43.51 | 64.19 | 17.50 |
| w/ RAG | N/A | N/A | 41.48 | 63.76 | N/A |
| $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ | 41.04 | 16.08 | 48.22 | 80.35 | 19.50 |
| w/ RAG | N/A | N/A | 40.66 | 73.80 | N/A |

### 5.2 Ultra-long Context Evaluation Beyond 100K Tokens

In this subsection, we evaluate the ultra-long context capability beyond 100K tokens on the real-world tasks from InfiniteBench (Zhang et al., ).
Table 2 shows that our model (41.04) outperforms many existing state-of-the-art models, such as GPT4-Turbo-2024-04-09 (33.16), GPT4-1106 preview (28.23), Llama3.1-70B-Instruct (39.81), Qwen2-72B-Instruct (39.77), Llama-3-70B-Instruct-Gradient-262k (32.57) and Claude 2 (33.96), confirming the effectiveness of our recipe and state-of-the-art long-context capability of our model. We find that our model performs extremely well at QA tasks while performs relatively low for summarization. This can be attributed to the lack of summarization data in our SFT recipe. Additionally, we compare $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}8B}$ to Llama3.1-8B-Instruct and our model again shows better results.

### 5.3 Long Context Evaluation within 32K Tokens

In this subsection, we evaluate the long context capability within 32K tokens. Table 3 shows that GPT-4-Turbo-2024-04-09 achieves the highest score of 51.93 among all models. Our model scores 48.15, which is higher than Llama-3-70B-Instruct-Gradient-262k but is a bit lower than Qwen2-72B-Instruct (49.94) and Llama3.1-70B-Instruct (49.92). This difference can be attributed to the extensive 32K pretraining and large SFT data implemented by Qwen2-72B-Instruct and Llama3.1-70B-Instruct, while we used a much smaller continued pretraining corpus and SFT datasets.

Additionally, we found that all the RAG solutions (top-5 with chunk size as 1200) perform worse than the long context solution, which suggest all these SOTA long context LLMs can handle 32K tokens well within their context window.

**Table 3: Evaluation results on the long context benchmarks within 32K tokens. For RAG, we use top-5 retrieved chunks, each with 1200 tokens from E5-mistral retriever (Wang et al., ). Note that here, we do not intend to optimize the RAG scores by increasing the number of total retrieved tokens and we leave the ablation in Section 6.**
| Model | Avg. | QM | QASP | QLTY | MSQ | HQA | MFQA |
| --- | --- | --- | --- | --- | --- | --- | --- |
| GPT-4-Turbo-2024-04-09 | 51.93 | 16.37 | 38.96 | 88.45 | 44.88 | 70.65 | 52.26 |
| w/ RAG | 49.84 | 16.07 | 36.18 | 85.85 | 42.17 | 67.85 | 50.94 |
| Qwen2-72B-Instruct | 49.94 | 17.06 | 34.84 | 84.85 | 46.80 | 65.98 | 50.12 |
| w/ RAG | 48.08 | 17.63 | 35.19 | 83.05 | 39.92 | 64.58 | 48.10 |
| Llama-3-70B-Instruct-Gradient-262k | 40.51 | 20.72 | 30.64 | 74.35 | 20.20 | 45.82 | 51.33 |
| w/ RAG | 40.57 | 20.04 | 30.68 | 72.35 | 22.19 | 46.85 | 51.31 |
| \hdashlineLlama3.1-8B-Instruct | 42.42 | 15.16 | 36.59 | 68.20 | 25.89 | 52.42 | 56.25 |
| w/ RAG | 41.35 | 14.08 | 35.05 | 65.35 | 28.30 | 50.76 | 54.54 |
| Llama3.1-70B-Instruct | 49.92 | 14.85 | 35.75 | 87.40 | 45.65 | 59.64 | 56.23 |
| w/ RAG | 47.63 | 13.59 | 32.11 | 85.65 | 39.53 | 58.30 | 56.58 |
| $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}8B}$ | 39.41 | 11.64 | 28.85 | 63.35 | 27.81 | 53.81 | 51.02 |
| w/ RAG | 40.31 | 13.20 | 28.85 | 61.10 | 29.77 | 57.81 | 51.15 |
| $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ | 48.15 | 15.99 | 32.84 | 82.20 | 39.26 | 62.24 | 56.39 |
| w/ RAG | 47.59 | 16.16 | 33.27 | 81.70 | 38.16 | 61.13 | 55.14 |

Table: Table 4: Evaluation results on ChatRAG Bench with 9 datasets. Following Liu et al. (), we exclude HDial as it is included in the instruction tuning datasets. The maximum context lengths are 4K tokens.

### 5.4 ChatRAG Bench : Short Context Evaluation within 4K Tokens

In this subsection, we evaluate the models on the short context tasks within 4K tokens from ChatRAG Bench (Liu et al., ).
The results can be found in Table 4.
Our model achieves an average score of 56.30, outperforming most existing SOTA models with 128K context window, including GPT-4-Turbo-2024-04-09, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct. It is also comparable, though slightly behind, Llama3-ChatQA-1.5-70B (57.14), which only supports 4K context length.
This suggests that extending short-context models to long-context has trade-offs. Exploring how to effectively extend the context window to even larger scales, e.g., millions of tokens as in Gemini 1.5 Pro (Gemini-Team, ), without degrading performance on regular short-context tasks is an exciting research direction.

Figure: Figure 2: Ablation of $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ with RAG given different top-*k* = {5, 10, 20, 40} retrieval, and chunk-size = {300, 600, 1200} on long context benchmarks within 32K tokens (see Section 4.2 for more details). The accuracy can be monotonically improved with more retrieved tokens (i.e., *k*$\times$ chunk_size) in the context window.
Refer to caption: x2.png

**Table 5: We compare RAG vs. long context evaluation using our $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ on tasks beyond 100k. Here we use average accuracy of En.QA and En.MC tasks that can apply RAG. The RAG-based result is still better than long context evaluation.**
|  | RAG (top-*k*) | Long context |
| --- | --- | --- |
| $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ | 64.55 (k=30) | 64.29 |
| Qwen2-72B-Instruct | 52.95 (k=20) | 52.22 |

## 6 RAG versus Long Context

In Figure 2 and Table 5,
we compare direct long-context and RAG solutions
on 32K and beyond 100K benchmarks.
For sequences beyond 100K tokens, we report only the average scores for En.QA and En.MC, as the RAG setting is not directly applicable to En.Sum and En.Dia.

Figure 2 compares different chunk sizes for top-*k* retrieval and the total number of tokens fed into our long context LLM. Comparing total tokens from 3000 to 24000, we found that more tokens consistently yield better results, confirming the strong long-context capability of our model.
Also, for downstream tasks with sequence lengths up to 32K, RAG can still outperform the long-context solution if a sufficient number of top-k chunks are used, i.e., at least 12,000 tokens as shown in Figure 2. Therefore, the hyperparameter k can be selected based on the tradeoff between efficiency and accuracy in real-world applications.

On the other hand, we found that for context lengths beyond 100K, RAG (using top-30 chunks for our $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$, and top-20 for Qwen2-72B-Instruct) outperforms the full long-context solution. This indicates that even SOTA long-context LLMs may struggle to effectively understand and reason over 100K tokens. In such scenarios, RAG is recommended for better accuracy and much lower inference cost, provided it is applicable to the downstream tasks.

## 7 Conclusion

We introduce $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$, a long context LLM that possess GPT-4 Turbo-level capabilities for both understanding up to 128K long contexts and utilizing retrieved contexts for generation.
This provides the flexible options for different downstream tasks with specific accuracy and efficiency requirements.
We present a detailed and reproducible technical recipe for building and evaluating the model, including the methods, training data, and evaluation benchmarks.
In particular, we evaluate ChatQA 2 on short-context RAG benchmark (ChatRAG) (within 4K tokens), long context tasks from SCROLLS and LongBench (within 32K tokens), and ultra-long context tasks from InfiniteBench (beyond 100K tokens).
Our results demonstrate that the $\mathtt{Llama3\text{-}ChatQA\text{-}2\text{-}70B}$ can achieve GPT-4-Turbo-2024-04-09 level accuracy on these benchmarks across different sequence lengths.

## Appendix A Appendix

Figure: (a) Llama3.1-Instruct 8B (passkey retrieval)
Refer to caption: extracted/5841941/figures/llama3_1/llama3_1_instruct_8b_tp8_mcore_pass_sent.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: - More details in the method. The 3.1 and 3.2 looks quite common. Since the authors emphasizes RAG, what specific changes have the authors made? And how RAG is used in the training?
- I am a bit confused by the highlight of this paper. I suppose the authors want to present a training pipeline for long context. But I don't quite understand the ablation of why certain choices are made.
soundness: 1
strengths: The performance looks good.
summary: This paper presents a way to train a short context base model to handle long context. In involves positional encoding (RoPE) to allow long context handling, pretraining on long context data and RAG. The trained performance on llama 70b shows comparable results to commercial LLMs in long context.
weaknesses: - The method is not clearly presented. I don't understand what fundamental changes of RAG or finetuning that makes the long context performance better
- The ablation study of the training process / design is missing. I don't understand the takeaway message from this paper

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Will the collected long SFT dataset utilized in the three-stage instruction tuning be open source?

2. Will the ChatQA 2 model be open source?
soundness: 2
strengths: 1. This paper introduces a significant enhancement to Llama3-70B’s long-context processing abilities, extending its context window from 8K to 128K tokens. Compared to Llama3.1with 128K context window, the proposed ChatQA2 achieves better performance. 

2. This paper proposes a three-stage instruction tuning, and collect a long SFT dataset to enhance the model's capability to handle very long context sequences up to 128k tokens. 

3. The authors propose utilizing the latest long-context retriever to address challenges in the current RAG pipeline for LLMs, specifically the limitations of top-k chunk-wise retrieval.
summary: This paper introduces ChatQA 2, an open-source long-context language model with a 128K token context window, aiming to achieve performance comparable to GPT-4-Turbo. By leveraging a continue pretraining and a three-stage instruction tuning process, ChatQA 2 demonstrates improved accuracy and robustness in handling large volumes of context and consistently outperforms direct long-context models when using RAG with sufficient top-k chunks.
weaknesses: 1. Compared to Llama3.1, ChatQA 2 shows improved performance on InfiniteBench, which includes real-world long-context understanding tasks beyond a 100K context window. However, on long-context benchmarks within 32K tokens, ChatQA 2 does not outperform Llama3.1, possibly due to its instruction tuning being focused on long SFT dataset. The author should also pay attention to the performance of 32K tokens.

2. The long-context open-source model ChatQA 2 presented in this paper achieves strong performance, though it lacks some innovation. The synthetic dataset used in the long SFT dataset could be described in more detail.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. At several places, e.g. line 144, 168, 222, in this manuscript, the authors mentioned "the first version of this work.", what does that mean? Please clarify.

2. What are the key new insights that this work offer, considering that the Llama3.1-128k model was published late July? I am not complaining about anything, but mainly want to make sure the technical contributions of this work. I will appreciate the answer from the authors.

3. Around line #188, the authors mentioned "where the model is initially trained on 128k high-quality instruction following datasets", while also mentioning " these two stages involve relatively short contexts", what does that mean, why 128k high-quality datasets are short contexts? Is there a typo?

4. In the long sequence SFT data, the authors used NarrativeQA for synthetic data generation. Are there potential data contamination in the benchmarks?


5. Minor questions: line # 263: why are the documents segmented into 300-word chunks? Does the size matter for the benchmark results?
soundness: 3
strengths: - The developed ChatQA2 model shows promising performances than existing open-source baselines and the April version of OpenAI GPT model. 

- The evaluation from short to long contexts are comprehensive and convincing.

- This work presents a detailed data and training recipe that enables a good-performing long-context model, bridging the gap between the commercial sector and open-source community for long context modeling.
summary: This work develops ChatQA 2, a long context LLM based on llama3-8k model that matches the performance of the April version of OpenAI GPT4 models and outperform SOTA open-source models. This work introduces a two-stage adaptation pipeline including long context continued training and multi-stage instruction finetuning. The developed model is then evaluated on short-context RAG, medium context up to 32k and long context above 100k benchmarks to show the superior performances than the existing open-source models.
weaknesses: - The authors should clearly articulate the relationship of this work and ChatQA 1.5, right now it was described as both a baseline (as shown in Table 4) and confusingly as "the first version of this work", i.e. the same work. If ChatQA1.5 is indeed a baseline, then it should be listed in all the tables and result analysis.

- Some details of the computational experiment are missing in section 3: how many epochs and tokens were trained in each stage? How much does this impact the results? For example, there is a discussion to the potential reason of low longbench scores due to under-training. The authors should present the trend of higher scores with more training epochs (in supplementary information).  Please also provide the data, and training for reproducibility purposes.

- The authors should improve the comprehensiveness of their ablation studies and results analysis: (1) The authors should also evaluate the retrieval-based tasks in infinitebench to further confirm the NIAH superior performances as shown at the beginning. (2) The section 5.2 is a bit lacking: the RAG numbers seem not being discussed. (3) The authors should add an ablation study to support their claim that started at line #177, i.e. using an alternative document breaker helps the training.

- Minor comments: (1) the authors should add citations of scrolls and longbench when they are mentioned in line # 252 and line # 253. (2) There seems to be a typo in line #261: "as it is vasincluded" seems a bit confusing. (3) The authors should keep the bold texts consistent in all tables, right now some Llama-3.1-128k models are made bold but some are not. For example, I am not sure why 48.15 is made bold in Table 3, but not 49.92 (for 3.1-70b-instruct).

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: See weakness section for questions.
soundness: 3
strengths: ## Originality
The technique of extending the context lengths of LLMs is not new. The most originality lies in the creation of the synthetic, long-context instruction dataset highlighted in Section 3.2. The authors seem to suggest that utilizing long-context retrievers for RAG is a contribution in Section 3.3.

## Quality
The paper is well written. Experiments cover a wide range of tasks and reflect actual usage of long-context LLMs in real life. Baselines covered are extensive. Nothing to pick here.

## Clarity
The recipe is simple and presented clearly in the paper. Particularly, technical details are well-documented in the Method Section (Sec 3).

## Significance
The topic of extending LLMs to support longer context than it was original pretrained is of itself significant, which in turn makes potential solutions significant.
summary: This paper introduces a recipe to extend Llama-3.0 models to long context windows (128K). The recipe involves increasing the RoPE frequency, continual pretraining on longer length corpus, then supervised fine-tuning on a mix of short and long context length instructions. They benchmark the model on tasks (synthetic and real-world) with various context lengths and found consistent gains over the original Llama-3 models, with the 70B model exceeding the performance of GPT-4-Turbo on some tasks. Ablation studies demonstrate tradeoffs involved with extending LLMs to handle long-context inputs.
weaknesses: ## Discussion on the key factors influencing performance tradeoff amongst tasks with different context lengths in long-context models

A major part of the paper is dedicated to compare the performance of long-context models on various (shorter) context lengths, specifically 4K, 32K, and 100K+. As the authors stated in P8L429 "This suggests that extending short-context models to long-context has trade-offs.", it is important to pinpoint what exactly controls the trade-off. The authors claimed in P7L372 that "This difference (in 32K performance) can be attributed to the extensive 32K pretraining and large SFT data implemented by Qwen2-72B-Instruct and Llama3.1-70B-Instruct.", which seems to hint at the data distribution being mainly responsible for the tradeoff. However, to the extent of my knowledge, Llama3.1 didn't released their training recipe nor data distribution. Can the tradeoff be fully attributed to the context-length distribution in the continual pretraining and SFT dataset? More discussion would greatly benefit the work.

## Relevance of stage-wise instruction-tuning

The authors highlighted in P2L98 and P4L187 that the instruction-tuning is split into 3 stages: instruction-following on short context, QA on short context, then a mix of long context instruction-following and synthetic QA. Stated in P4L199 "Both the full long SFT dataset and the short SFT dataset from the first two stages are then blended for training." What is the relevance of breaking down the instruction-tuning into stages when the final stage includes all data from the first two stages? An ablation study on training the model all-at-once with all the available data would help justify the multi-stage approach. Perhaps there are benefits in gradually increasing the context-length when instruction-finetuning and is actually key to the success of this recipe.

## Evaluating tradeoffs with other LLM capabilities with standard, general benchmarks

Section 5.2 highlights the difference in performance of the proposed model on QA and summarization. Quote "We find that our model performs extremely well at QA tasks while performs relatively low for summarization. This can be attributed to the lack of summarization data in our SFT recipe." Indeed the SFT data is heavy on QA and thus accounting for the performance bias. It is concerning that the good performance on this model on the benchmarked tasks might merely be a tradeoff with other LLM capabilities (e.g. reasoning, ICL, commonsense, creative writing, coding). Providing results on more general benchmarks (e.g. MT-Bench, MMLU, HumanEval) would help users evaluate the entire spectrum of tradeoffs of the proposed recipe.

---

confidence: 2
contribution: 4
initial_rating: 6
presentation: 1
questions: 1. Under what circumstances does RAG improve LLM performance? For instance, why does RAG cause a performance drop for Llama3-ChatQA-2 on En.MC in Table 2, while improving the performance of Llama3.1-Instruct? The SOTA results in Tables 2 and 3 were achieved without RAG. Does this suggest that the effectiveness of RAG depends on the specific benchmark rather than the context length?
2. Could you test the results with retrievers other than E5-mistral?
soundness: 4
strengths: 1. The paper provides open access to the trained model and training recipe, which facilitates further research in the field of LLM + RAG.
2. The paper focuses on increasing the context length of LLMs and has conducted extensive testing across benchmarks of varying lengths. It compares the performance with both open-source and proprietary models, achieving promising results.
summary: This paper presents Llama3-ChatQA-2, a model based on LLama-3-base that is trained with an extended context window. The authors first collect a long SFT dataset for instruction-tuning and train the model with modified PE. On top of this, they integrate a long-context retriever to implement RAG. Extensive testing are conducted across a variety of tasks involving short, long, and ultra-long token lengths, achieving state-of-the-art (SOTA) results.
weaknesses: Some sections of the paper are too brief, presuming the reader’s familiarity with certain concepts, particularly those related to RAG. For example, the meaning of “chunk size” and a simple introduction to the inputs and outputs of the E5-mistral retriever are not sufficiently explained. Since the paper emphasizes *open-source training data and reproduction recipe* compared to the open-access (Open Weights) Llama3.1-70B-Instruct, it would benefit from providing more detailed explanations, like [1].

> [1] The llama 3 herd of models. *arXiv preprint arXiv:2407.21783*.

Meta-review: (a) Scientific Claims and Findings
This paper introduces ChatQA 2, an open-source long-context language model with a 128K token context window, demonstrating improved accuracy and robustness in handling large volumes of context and consistently outperforming direct long-context models when using RAG with sufficient top-k chunks.    

(b) Strengths of the Paper
The paper introduces a significant enhancement to Llama3-70B's long-context processing abilities, extending its context window from 8K to 128K tokens.    
The paper proposes a three-stage instruction tuning and collects a long SFT dataset to enhance the model's capability to handle very long context sequences up to 128k tokens.    
The authors propose utilizing the latest long-context retriever to address challenges in the current RAG pipeline for LLMs, specifically the limitations of top-k chunk-wise retrieval.    

(c) Weaknesses of the Paper
Compared to Llama3.1, ChatQA 2 shows improved performance on InfiniteBench, which includes real-world long-context understanding tasks beyond a 100K context window.    
The long-context open-source model ChatQA 2 presented in this paper achieves strong performance, though it lacks some innovation.    

(d) Reasons for Accepting the Paper
The paper presents a detailed data and training recipe that enables a good-performing long-context model, bridging the gap between the commercial sector and the open-source community for long context modeling.    
The developed ChatQA2 model shows promising performances than existing open-source baselines and the April version of the OpenAI GPT model.    
The evaluation from short to long contexts is comprehensive and convincing.

Decision: Accept (Poster)