기억할 것인가, 검색할 것인가: RAG(Retrieval-Augmented Generation)에 적합한 사전 훈련을 위한 확장 법칙

To Memorize or to Retrieve: Scaling the Interaction Between Pretraining and Retrieval

검색 기반 증강 생성(RAG)은 테스트 시 관련 문맥을 제공하여 지식 집약적인 상황에서 언어 모델(LM)의 성능을 향상시킵니다. 그러나 사전 훈련 중에 획득된 매개변수 지식과 검색을 통해 접근되는 비매개변수 지식 간의 관계는 특히 고정된 데이터 예산 하에서 잘 이해되지 않습니다. 본 연구에서는 다양한 모델 및 데이터 크기에서 사전 훈련 코퍼스 크기와 검색 저장소 크기 간의 균형을 체계적으로 연구합니다. 우리는 30M에서 3B개의 매개변수를 가진 OLMo-2 기반 LM을 최대 100B개의 DCLM 데이터로 훈련하고, 사전 훈련 데이터 크기(매개변수 수의 1~150배)와 검색 저장소 크기(1~20배)를 모두 변경하며, 추론, 과학 QA 및 개방형 도메인 QA를 포괄하는 다양한 벤치마크에서 성능을 평가합니다. 연구 결과, 검색은 모델 크기 전반에 걸쳐 매개변수만 사용하는 기준 모델보다 성능을 꾸준히 향상시키며, 모델 크기, 사전 훈련 토큰 수 및 검색 코퍼스 크기를 함수로 모델링하는 3차원 확장 프레임워크를 제시합니다. 이 확장 공간은 고정된 데이터 예산을 사전 훈련과 검색 간에 최적으로 할당하는 데 도움이 되며, 검색의 한계 효용은 모델 크기, 작업 유형 및 사전 훈련 포화 정도에 따라 크게 달라짐을 보여줍니다. 본 연구 결과는 검색이 언제, 어떻게 사전 훈련을 보완해야 하는지에 대한 정량적인 기반을 제공하며, 확장 가능한 언어 모델링 시스템 설계 시 데이터 리소스를 할당하는 데 실질적인 지침을 제공합니다.

Retrieval-augmented generation (RAG) improves language model (LM) performance by providing relevant context at test time for knowledge-intensive situations. In this work, we systematically study the trade-off between pretraining and retrieval by training OLMo-2-based LMs ranging from 30M to 3B parameters on up to 100B DCLM tokens, while varying pretraining data scale, retrieval store size, and retrieval store source (pretraining vs. new data) across reasoning, scientific QA, and open-domain QA benchmarks. We find that retrieval gains depend on model capacity and pretraining exposure and are strongly front-loaded, with a median 91% of the largest observed improvement realized by one retrieval token per model parameter. However, the interaction is objective-dependent: smaller models gain more in gold-answer perplexity, whereas larger, more-pretrained models gain more in accuracy. Retrieval from previously seen data also preserves most of the held-out retrieval gain. Retrieval is therefore a task-, regime-, and metric-dependent complement to parametric learning whose value also depends on datastore size and information novelty. Overall, this motivates the explicit partitioning of data between internalization and external access for LM design.