PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Peng Xu, Wei Ping, Xianchao Wu, Chejian Xu, Zihan Liu, Mohammad Shoeybi, Bryan Catanzaro

arXiv:2407.14482v3cs.CLcs.AIcs.IRcs.LG

TL;DR

Open LLMs still face gaps against proprietary models in long-context understanding, while RAG and long context serve complementary tasks. ChatQA 2 extends Llama3 to 128K and applies staged instruction tuning; it reports stronger results than most SOTA models on ultra-long tasks and short-context RAG.

  • Problem

    Open models still show performance gaps from frontier proprietary systems, and long-context evaluation has often relied on synthetic tests rather than real-world downstream tasks.

    Source: §1 · p. 1

  • Method

    ChatQA 2 extends Llama3-70B-base from 8K to 128K by continued pretraining, applies three instruction-tuning stages, and pairs the model with a long-context retriever for RAG.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    On InfiniteBench, ChatQA-2-70B scores 41.04 versus GPT4-Turbo-2024-04-09 at 33.16, and outperforms most SOTA models on tasks beyond 100K tokens and 4K conversational-QA RAG.

    Source: §1 · p. 2, §5.2 · p. 7

  • Takeaways & Limitations

    Long context and RAG offer task-dependent options: long context suits whole-document summarization, while RAG efficiently retrieves relevant context for query-based tasks.

    Source: §1 · p. 2

    The authors report relatively low summarization performance and attribute it to a lack of summarization data in the SFT recipe.

    Source: §5.2 · p. 7

Abstract

from arXiv · show

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-2024-04-09) in long context understanding and retrieval-augmented generation (RAG) capabilities. These two capabilities are complementary to each other and essential for LLMs to process large volumes of information that cannot fit into a single prompt. 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 Llama3-ChatQA-2-70B 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., Llama3-ChatQA-2-70B and Qwen2-72B-Instruct) on both 32K and 128K benchmarks. We open-source the model weights, training data, and the evaluation setup for the for the community: https://chatqa2-project.github.io/

1 INTRODUCTION

Open-access LLMs still trail frontier proprietary models in many domains, while reproducible recipes and realistic long-context evaluation remain limited. ChatQA 2 combines long-context modeling and RAG, and reports strong results across both settings.

Motivation

  • The authors identify a gap in open long-context work: training data and reproduction recipes are often unavailable, and synthetic needle tests do not represent real-world downstream performance well.They aim to provide an open recipe targeting proprietary-model-level capability on real-world long-context tasks.Source: §1 · p. 2
  • The paper frames long context and RAG as complementary: long context suits whole-document summarization, while RAG efficiently retrieves relevant material for query-based tasks.The choice depends on downstream tasks and accuracy-efficiency trade-offs.Source: §1 · p. 2

Contributions

  • The paper makes its training data and reproduction recipe publicly accessible.Table 1 also presents the model comparison across tasks with varying context lengths.Source: §1 · p. 2, Table 1 · p. 2
  • ChatQA 2 extends Llama3-70B-base from 8K to 128K through continued pretraining, then uses three instruction-tuning stages for instruction following, RAG, and long-context understanding.The authors say incorporating previous datasets in stage-wise tuning simplifies experimentation and hyperparameter tuning.Source: §1 · p. 2
  • Llama3-ChatQA-2-70B outperforms most SOTA models on tasks beyond 100K tokens and on conversational-QA RAG within 4K tokens.The comparisons include GPT-4-Turbo-2024-04-09, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct.Source: §1 · p. 2
  • With sufficient top-k chunks, RAG outperforms direct long-context solutions on 32K benchmarks and real-world 128K tasks; RAG accuracy also improves as retrieved tokens increase.The reported RAG performance is robust to chunk-size variation when using a long-context retriever.Source: §1 · p. 3

2 RELATED WORK

Related work describes continued training and fine-tuning approaches for extending LLM context, alongside retrieval systems that address inputs too long for model context windows.

2.1 LONG CONTEXT LLM

  • Long-context methods extend LLM windows through continued training or fine-tuning, including RoPE adaptations such as position interpolation and increasing the base frequency θ.Examples include Yi-34B extending from 4K to 200K by raising θ, and Qwen2 using YaRN to extrapolate from 32K to 128K.Source: §2.1 · p. 3
  • Llama-3.1 models (Dubey et al., 2024), with a 128K context window, were introduced after this work’s release; the authors added a comprehensive experimental comparison with ChatQA-2.Source: §2.1 · p. 3

2.2 RETRIEVAL-AUGMENTED GENERATION (RAG)

  • Standalone retrieval is an established way to handle long texts that do not fit in an LLM context window.Prior work has also proposed retrieval-augmented language models.Source: §2.2 · p. 3
  • Dense retrievers map text chunks to embedding vectors for nearest-neighbor retrieval; short chunk limits can fragment context, motivating retrievers supporting longer chunks.Examples cited support chunks of 8K tokens for Jina Embeddings 2 and Nomic Embed, and 32K for E5-mistral-7B and NV-Embed.Source: §2.2 · p. 3

3 METHOD

The method extends Llama3’s context window from 8K to 128K, tunes it in three stages with long-context data, and uses a long-context retriever to support RAG.

3.1 EXTENDING CONTEXT WINDOW TO 128K

  • The continued-pretraining corpus uses upsampled SlimPajama documents to produce 10 billion tokens of 128K sequences; training runs for 2,000 steps on 8 billion tokens.The method raises the RoPE base frequency to 150M, uses batch size 32 and learning rate 3e-5, and separates documents with special characters such as <s>.Source: §3.1 · p. 4

3.2 INSTRUCTION-TUNING WITH LONG CONTEXT DATA

  • Three-stage instruction tuning starts with instruction-following data, then conversational QA with context, before long-context SFT supports sequences up to 128K.Long SFT data below 32K uses existing datasets; 32K–128K samples are synthesized from NarrativeQA by inserting summaries into source web pages, and NarrativeQA is excluded from evaluation to avoid potential contamination.Source: §3.2 · p. 4

3.3 LONG CONTEXT RETRIEVER MEETS LONG CONTEXT LLM

  • To address fragmentation and recall trade-offs in chunk-wise RAG, the method uses a long-context retriever that supports thousands of tokens per chunk.E5-mistral embeds document chunks and questions; indexed chunk embeddings are retrieved by nearest-neighbor search using the question embedding as query.Source: §3.3 · pp. 4–5

4 BASELINES AND EVALUATION BENCHMARKS

The evaluation compares long-context models and RAG across ultra-long, long, and short context tasks, with benchmarks and metrics tailored to each setting.

4 BASELINES AND EVALUATION BENCHMARKS

  • The baselines include GPT-4-Turbo, Qwen2-72B-Instruct, Llama-3-70B-Instruct-Gradient-262k, and Llama-3.1 8B/70B; default RAG uses E5-mistral with 1,200-token chunks and top 5.The authors describe this RAG configuration as a cost-performance trade-off supported by an ablation study.Source: §4 · p. 5

4.1 ULTRA-LONG CONTEXT BENCHMARKS BEYOND 100K TOKENS

  • InfiniteBench evaluates ultra-long context beyond 100K tokens; the study uses four English tasks: novel summarization, long-book QA, multiple choice, and dialogue.The tasks use ROUGE-L-Sum, F1, or exact match according to task type; QA questions require long-range dependencies and reasoning.Source: §4.1 · p. 5

4.2 LONG CONTEXT BENCHMARKS WITHIN 32K TOKENS

  • The within-32K benchmark comprises six datasets from SCROLLS and LongBench, with ROUGE, exact match, or F1 metrics following the official evaluations.QMSum uses geometric mean ROUGE; QuALITY uses exact match; Qasper, MuSiQue, HotpotQA, and MultiFieldQA-en use F1.Source: §4.2 · p. 5

4.3 SHORT CONTEXT WITHIN 4K TOKENS

  • The short-context ChatRAG evaluation uses nine datasets after excluding HDial; long-document tasks retrieve top 5 or top 20 chunks, and average F1 is reported.The Needle In A Haystack test compares Llama3.1-Instruct 8B/70B and ChatQA-2 8B/70B up to 128K using the same needle.Source: §4.3 · pp. 5–6, Figure 1 · p. 6

5 RESULTS

The results compare synthetic needle retrieval with real-world long-context and RAG benchmarks, followed by ablations of retrieval and training choices.

5.1 NEEDLE IN A HAYSTACK

  • 100% accuracy through 128K tokens on both needle and passkey tests; Llama3.1-Instruct 8B and 70B fail the harder needle despite passing passkey.Source: §5.1 · p. 6

5.2 ULTRA-LONG CONTEXT EVALUATION BEYOND 100K TOKENS

  • 41.04 on InfiniteBench: ChatQA-2-70B outscored GPT-4-Turbo-2024-04-09 (33.16), Llama3.1-70B-Instruct (39.81), and Qwen2-72B-Instruct (39.77) beyond 100K tokens.The evaluation covers real-world long-context understanding tasks beyond a 100K context window.Source: Table 2 · p. 7, §5.2 · p. 7

5.3 LONG CONTEXT EVALUATION WITHIN 32K TOKENS

  • 51.93 was GPT-4-Turbo-2024-04-09’s top score within 32K; ChatQA-2 scored 48.15, below Qwen2 (49.94) and Llama3.1 (49.92).The default RAG setting—top-5 chunks of 1,200 tokens—performed worse than top long-context solutions.Source: §5.3 · p. 7, Table 3 · p. 8

5.4 CHATRAG BENCH: SHORT CONTEXT EVALUATION WITHIN 4K TOKENS

  • 56.30 average on ChatRAG Bench: ChatQA-2-70B outperformed most 128K-context models but remained slightly behind ChatQA-1.5-70B (57.14).The benchmark includes nine datasets, excludes HDial, and caps context at 4K tokens.Source: §5.4 · p. 7, Table 4 · p. 8

5.5 ABLATION STUDY

  • Within 32K, accuracy improved monotonically as the total retrieved tokens increased across the tested top-k and chunk-size settings.Source: Figure 2 · p. 9
  • Retriever changes produced relatively small performance variations, with no clear overall winner; NV-Embed-v2 performed better on ultra-long tasks.Source: Table 6 · p. 8, §5.5 · p. 9
  • Three-stage tuning outperformed all-in-one tuning on RAG by a significant margin, with marginally better or worse results on ultra-long and long tasks.Source: Table 7 · p. 9, §5.5 · p. 9
  • With 1.6M additional SFT samples, ChatQA-2-8B was comparable to Llama3.1-8B on 32K tasks and better on 4K and ultra-long tasks, but weaker on MMLU and HumanEval.The paper attributes its lower MT-Bench score to the absence of DPO/PPO alignment and proposes RLHF as future work.Source: §5.5 · p. 9

6 RAG VERSUS LONG CONTEXT

The comparison finds that RAG can outperform direct long-context evaluation at both up to 32K and beyond 100K tokens, with larger retrieved contexts helping at 32K and lower inference cost noted beyond 100K.

6 RAG VERSUS LONG CONTEXT

  • For the beyond-100K comparison, the reported averages cover En.QA and En.MC because RAG does not directly apply to En.Sum and En.Dia.Table 5 compares RAG with long-context evaluation on applicable tasks.Source: §6 · p. 9
  • At least 12,000 retrieved tokens lets RAG marginally outperform direct long-context processing for downstream tasks up to 32K; more tokens from 3,000 to 24,000 consistently improve results.The number of retrieved chunks can be selected to trade off efficiency and accuracy.Source: §6 · p. 10
  • Beyond 100K tokens, RAG with top-30 ChatQA-2 or top-20 Qwen2 chunks outperforms full-context evaluation and is recommended when applicable for accuracy and lower inference cost.The comparison uses average En.QA and En.MC scores, since RAG does not directly apply to En.Sum and En.Dia.Source: §6 · pp. 9–10

7 CONCLUSION

The conclusion presents Llama3-ChatQA-2-70B as a reproducible long-context model and reports GPT-4-Turbo-level accuracy across short-context RAG, long-context, and ultra-long-context evaluations.

7 CONCLUSION

  • Llama3-ChatQA-2-70B achieves GPT-4-Turbo-2024-04-09-level accuracy across benchmarks spanning different sequence lengths.Evaluations include short-context ChatRAG within 4K tokens, SCROLLS and LongBench within 32K, and InfiniteBench beyond 100K.Source: §7 · p. 10
  • The model is presented as supporting both understanding contexts up to 128K and generation from retrieved contexts.Source: §7 · p. 10
  • The conclusion provides a detailed, reproducible technical recipe covering methods, training data, and evaluation benchmarks.Source: §7 · p. 10

A APPENDIX

The appendix reports strong InfiniteBench retrieval-task results for ChatQA-2 and describes a separate passkey-retrieval probe of Llama 3.1-Instruct models up to 128K context.

A APPENDIX

  • Figure 3 probes passkey retrieval for Llama3.1-Instruct 8B and 70B models across context windows up to 128K.The needle repeats pass key 385243, and the question asks for the pass key.Source: Figure 3 · p. 14
  • All evaluated models score 100% on InfiniteBench number_string and passkey, while Llama3-ChatQA-2-70B scores better than GPT-4-Turbo on kv_retrieval.Table 9 covers these three retrieval-based tasks.Source: A APPENDIX · p. 14, Table 9 · p. 14

B APPENDIX

The appendix reports that NIAH results favor the <s> document breaker over <EOS> and <BOS> in a 2B-token separator comparison.

B APPENDIX

C APPENDIX

The appendix notes that three-stage SFT blends datasets with potentially different epoch counts and identifies a comparison with all-in-one-stage training in Table 10.

  • Dataset epoch counts can differ according to blending ratio and sample count in the three-stage SFT process.Source: C APPENDIX · p. 15
  • Stage 2 dataset statistics are provided as an example of the analysis complexity introduced by the blended training process.Source: C APPENDIX · p. 15
  • Table 10 compares three-stage training with all-in-one-stage training.Source: Table 10 · p. 15

Related papers

Loading 2407.14482v3…