Data Engineering for Scaling Language Models to 128K Context

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 \textit{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 \textit{quantity} and \textit{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 \textit{domain balance} and \textit{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.

한 줄 요약

LLaMA-2 기반 모델을 128K 컨텍스트 길이로 확장하기 위한 데이터 엔지니어링 전략을 제시하며, 1B~5B 토큰의 적절한 데이터 믹스로 GPT-4 128K 수준 성능을 달성한다.

핵심 기여도

500M~5B 토큰의 데이터로 128K 컨텍스트 내 임의 위치 정보 추출 가능성을 입증
도메인 밸런스와 길이 업샘플링을 강조하며, 기존 방식(예: 책 중심 업샘플링)은 성능 저하를 초래함을 밝힘
1B~5B 토큰의 풀 모델 연속 사전학습으로 GPT-4 128K와 유사한 성능 달성
Needle-in-a-Haystack 테스트에서 기존 오픈소스 모델 대비 10~15% 개선

핵심 아이디어

기존 사전학습 모델은 이미 4K 컨텍스트 내 임의 위치 정보 활용 능력을 갖추고 있으며, 이를 128K로 확장하기 위해 추가적인 연속 사전학습이 필요하다는 가설을 제시한다. 이는 기존 연구가 400B 토큰 이상의 대규모 데이터로 ‘새로운 능력’을 주입하는 방식과 대조된다. 본 연구는 1B~5B 토큰의 적절한 데이터 믹스로도 충분히 성능을 확장할 수 있음을 보인다. 특히, 도메인 밸런스를 유지하면서 각 도메인 내 길이 업샘플링을 적용하는 것이 핵심이며, 이는 기존 방식(예: 책 중심 업샘플링)보다 퍼플렉시티 저하 없이 안정적인 성능 향상 효과를 가져온다.

기술적 접근법

**모델**: LLaMA-2 7B 및 13B 사용, RoPE 기반 조정 외에는 아키텍처 변경 없음
**데이터**: 1B~5B 토큰의 길이별 업샘플링 데이터로 풀 어텐션 연속 사전학습
**데이터 믹스 전략**: 도메인 비율을 사전학습 데이터와 동일하게 유지한 후, 각 도메인 내 길이 업샘플링 적용
**하이퍼파라미터**: 64K~80K 컨텍스트 길이로 훈련, FlashAttention 활용으로 메모리 효율 향상

주요 결과

**Needle-in-a-Haystack 테스트**: 128K 컨텍스트 내 임의 위치 정보 정확히 추출, 기존 오픈소스 모델(YaRN-Mistral 128K, LongLoRA 100K) 대비 10~15% 개선
**데이터 양**: 500M~5B 토큰으로 128K 컨텍스트 내 정보 검색 가능
**도메인 업샘플링 효과**: 책 중심 업샘플링은 퍼플렉시티 상승, 도메인 밸런스 유지 시 안정적 성능 향상
**GPT-4 128K 수준 달성**: GPT-4 128K와 유사한 성능을 오픈소스 모델로 달성

의의 및 한계

본 연구는 대규모 사전학습 모델이 이미 장거리 의존성 모델링 능력을 갖추고 있으며, 이를 확장하기 위해 데이터 엔지니어링이 핵심이라는 통찰을 제시한다. 특히, 도메인 밸런스와 길이 업샘플링을 통해 기존 방식의 한계를 극복한 점이 학술적·실용적 가치를 높인다. 그러나 본 연구는 연속 사전학습만 다루고, 100K 이상 컨텍스트의 지시 학습(Instruction Finetuning) 연구는 제한적이다. 또한, 200K 이상 컨텍스트 확장을 위해서는 시퀀스 병렬화 기술이 필요하다는 한계도 인정한다.

실용적 활용

LLaMA-2와 유사한 아키텍처를 가진 모델에 적용 가능하며, 128K 컨텍스트를 요구하는 다중 문서 QA, 코드 리포지토리 이해, 장기 대화 모델링 등에 활용 가능하다. 특히, 학술 연구 예산 내에서도 GPT-4 128K 수준의 성능을 달성할 수 있어, 장거리 컨텍스트 모델링 연구의 진입 장벽을 낮춘다.