[논문 리뷰] Data Engineering for Scaling Language Models to 128K Context

Yao Fu, Rameswar Panda외 5명|arXiv (Cornell University)|2024. 02. 15.

Network Packet Processing and Optimization인용 수 4

한 줄 요약

이 논문은 LLaMA-2 언어 모델을 128K 컨텍스트 길이로 확장하기 위한 데이터 엔지니어링 레시피를 제안한다. 이는 각 소스에서 길이를 업샘플링한 데이터를 10억~50억 토큰 분량의 연속 미세조정을 통해 실현된다. 이를 통해 도메인 비율을 균형 있게 유지하면서 각 도메인 내에서 장문의 시퀀스를 강조함으로써, 나침반 속 바늘 검색 태스크에서 우수한 성능을 달성하며, GPT-4 128K에 가까운 성능을 달성하면서도 짧은 컨텍스트 성능을 유지함을 입증한다.

ABSTRACT

We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular extit{the ability to utilize information at arbitrary input locations}, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the extit{quantity} and extit{quality} of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize extit{domain balance} and extit{length upsampling}. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.

연구 동기 및 목표

• 모델 아키텍처 변경 없이도 언어 모델을 128K 컨텍스트 길이로 효과적으로 확장하기 위한 데이터 엔지니어링 전략을 탐구하는 것.
• 특히 모델의 원래 사전학습 길이를 초월한 장문 컨텍스트 내에서 임의의 위치에 있는 정확한 정보 검색을 가능하게 하는 과제를 해결하는 것.
• 대규모 사전학습 과정에서 이미 부분적으로 확보된 장문 컨텍스트 능력을, 정제된 데이터에 대한 최소한의 연속 미세조정을 통해 해제할 수 있는지 확인하는 것.
• 장문 컨텍스트 적응 과정에서 데이터 양, 도메인 균형, 시퀀스 길이 분포 간의 상호 교환 관계를 평가하는 것.
• 향후 지시 테이닝을 가능하게 하기 위해, 오픈소스 모델을 128K 컨텍스트로 확장하는 실용적이고 저비용이며 효과적인 레시피를 제공하는 것.

제안 방법

• LLaMA-2 7B 및 13B 모델을 10억~50억 토큰 분량의 정제된 데이터에 대해 연속 미세조정하며, 64K~80K 컨텍스트 창을 사용한 전체 어텐션 메커니즘을 적용한다.
• 모델 아키텍처 변경 없이도 더 긴 컨텍스트 길이를 지원하기 위해 RoPE 위치 임베딩 기반 조정 기법을 적용한다.
• 각 도메인 내에서 장문 시퀀스의 비율을 높이되, 원래의 도메인 혼합 비율을 유지하는 방식의 각 소스 길이 업샘플링 전략을 설계한다.
• 학계 수준의 하드웨어에서도 장문 컨텍스트 훈련을 가능하게 하기 위해 FlashAttention를 활용하여 메모리 사용량을 감소시킨다.
• 4K에서 128K에 이르는 다양한 컨텍스트 길이에서 정밀한 검색 성능을 측정하기 위해 나침반 속 바늘 검색 벤치마크를 평가에 활용한다.
• 데이터 혼합 설계의 영향을 분리하기 위해 YaRN-Mistral 128K 및 LongLoRA 100K와 같은 강력한 베이스라인과의 비교를 수행한다.

Figure 1: Needle-in-a-Haystack test (Kamradt, 2023 ) performance comparison. The x-axis denotes the length of the document (the “haystack”); the y-axis indicates the position that the “needle” (a short sentence) is located within the document, from 1K to 128K tokens. For example, 50% indicates that

실험 결과

연구 질문

• RQ1 4K 컨텍스트로 사전학습된 모델이, 정제된 데이터에 대해 단지 10억~50억 토큰 분량의 연속 미세조정만으로도 128K 컨텍스트 능력을 해제할 수 있는가?
• RQ2 예를 들어 책과 같은 장문 도메인에서 장문 시퀀스를 단순히 업샘플링하는 것은 코드나 웹 콘텐츠와 같은 다른 도메인에서 성능 저하를 초래하는가?
• RQ3 각 도메인 내에서 장문 시퀀스를 업샘플링하면서도 원래의 도메인 혼합 비율을 유지하는 것이 도메인 기반 업샘플링보다 더 효과적인가?
• RQ4 데이터 혼합 설계가 장문 컨텍스트 성능와 짧은 컨텍스트 성능 간의 상호 교환 관계에 어떻게 영향을 미치는가?
• RQ5 데이터 중심 접근 방식이 장문 컨텍스트 모델링의 효율성과 성능 측면에서 대규모 사전학습 레시피를 능가할 수 있는가?

주요 결과

• 각 소스의 길이를 업샘플링한 데이터에 대해 단지 10억~50억 토큰 분량의 연속 미세조정만으로도 LLaMA-2 7B 및 13B 모델이 128K 컨텍스트 길이까지 나침반 속 바늘 검색 테스트에서 뛰어난 성능을 달성한다.
• 각 소스 기반 길이 업샘플링 전략—원래의 도메인 비율을 유지하면서 각 도메인 내에서 장문 시퀀스의 비율을 높이는 방식—은 도메인 간 균형 잡힌 성능 향상을 가져오며, 단순 도메인 기반 업샘플링보다 뛰어난 성능을 보인다.
• 단순 도메인 기반 업샘플링(예: 책을 과도하게 강조)은 코드나 웹 콘텐츠 등의 다른 도메인에서 심각한 성능 저하를 초래하며, 이는 비효율적인 데이터 혼합 설계를 시사한다.
• 제안된 레시피로 훈련된 모델은 장문 컨텍스트 검색에서 GPT-4 128K와 유사한 성능을 달성하며, YaRN-Mistral 128K 및 LongLoRA 100K와 같은 강력한 오픈소스 모델을 크게 능가한다.
• 검증 손실만으로는 성능 격차를 감지하기 부족하며, 나침반 속 바늘 테스트는 손실이 반영하지 못하는 장문 컨텍스트 검색 능력의 핵심 격차를 드러낸다.
• 본 연구는 데이터 엔지니어링이 최첨단 장문 컨텍스트 성능를 달성하기 위해 모델 아키텍처나 알고리즘 혁신만큼이나 핵심적인 역할을 한다는 점을 입증한다.

Figure 2: Length and domain distributions of the various data mixture strategies. We use SlimPajama (Soboleva et al., 2023 ) , an open-source reproduction of LLaMA (Touvron et al., 2023a ) pretraining data mixture, as the source dataset. The original data mixture has about 30% of documents that are