Quest: Query-centric Data Synthesis Approach for Long-context Scaling of Large Language Model

2024-05-30

arXiv.org

AI Summary

To address performance bottlenecks in training long-context large language models—stemming from data domain shift and semantic fragmentation—this paper proposes a query-centric data synthesis paradigm. Methodologically, it introduces the first query-prediction-driven document clustering mechanism, integrating generative query prediction, keyword-query joint embedding matching, semantic similarity-based grouping, and dynamic synthesis to jointly preserve semantic coherence and enhance contextual diversity. The approach supports modeling sequences up to one million tokens and scales across model sizes. Empirically, it substantially outperforms standard fine-tuning, KNN-based retrieval, and in-context learning (ICLM) baselines on multiple long-context benchmarks. Crucially, it provides the first empirical validation of both effectiveness and scalability for million-token-context modeling, demonstrating robust generalization across diverse long-context reasoning and retrieval tasks.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Recent advancements in large language models (LLMs) have highlighted the importance of extending context lengths for handling complex tasks. While traditional methods for training on long contexts often use filtered long documents, these approaches lead to domain imbalances, limiting model performance. To address this, techniques like random document concatenation (Standard) and similarity-based methods (KNN, ICLM) have been developed. However, they either sacrifice semantic coherence or diversity. To balance both aspects, we introduce Quest, a query-centric data synthesis method aggregating semantically relevant yet diverse documents. Quest uses a generative model to predict potential queries for each document, grouping documents with similar queries and keywords. Extensive experiments demonstrate Quest's superior performance on long-context tasks, achieving remarkable results with context lengths of up to 1M tokens and confirming its scalability across various model sizes.

Problem

• Extend context lengths for LLMs
• Balance semantic coherence and diversity
• Improve long-context task performance

Innovation

• Query-centric data synthesis
• Semantic relevance and diversity
• Scalability up to 1M tokens