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

2024-07-19

arXiv.org

AI Summary

Open-source large language models (LLMs) significantly underperform closed-source counterparts like GPT-4-Turbo in ultra-long-context understanding and retrieval-augmented generation (RAG). Method: We introduce a Llama 3.0-based LLM with native 128K-context support, proposing a novel three-stage joint training paradigm: (i) context window expansion (from 8K to 128K tokens), (ii) RAG-aware supervised fine-tuning, and (iii) multi-granularity long-text evaluation. Contribution/Results: Empirical results demonstrate that robust long-context capability enables RAG performance to consistently surpass pure long-context inference as the number of retrieved chunks increases. Our model achieves state-of-the-art results on benchmarks involving >100K-token documents and 4K-RAG tasks, outperforming GPT-4-Turbo, Qwen2-72B, and Llama-3.1-70B. To foster reproducibility and community advancement, we fully open-source the model weights, training code, and evaluation framework.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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-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/

Problem

• Enhance long-context understanding in open-source LLMs.
• Improve retrieval-augmented generation performance.
• Bridge gap to proprietary models like GPT-4-Turbo.

Innovation

• Extended context window to 128K tokens
• Three-stage instruction tuning process
• Open-sourced model weights and data