Rho-1: 모든 토큰이 필요한 것은 아니다

Rho-1: Not All Tokens Are What You Need

April 11, 2024

저자: Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, Weizhu Chen

cs.AI

초록

기존의 언어 모델 사전 학습 방법들은 모든 학습 토큰에 대해 균일하게 다음 토큰 예측 손실을 적용해 왔습니다. 이러한 관행에 도전하며, 우리는 "코퍼스 내 모든 토큰이 언어 모델 학습에 동등하게 중요한 것은 아니다"라는 가설을 제시합니다. 초기 분석에서 우리는 언어 모델의 토큰 수준 학습 동역학을 탐구하며, 서로 다른 토큰에 대한 뚜렷한 손실 패턴을 발견했습니다. 이러한 통찰을 바탕으로, 우리는 Rho-1이라는 새로운 언어 모델을 소개합니다. 전통적인 언어 모델들이 코퍼스의 모든 다음 토큰을 예측하도록 학습하는 것과 달리, Rho-1은 원하는 분포와 일치하는 유용한 토큰을 선택적으로 학습하는 선택적 언어 모델링(Selective Language Modeling, SLM)을 채택합니다. 이 접근법은 참조 모델을 사용해 사전 학습 토큰을 점수화하고, 더 높은 초과 손실을 보이는 토큰에 집중된 손실로 언어 모델을 학습시키는 것을 포함합니다. 15B OpenWebMath 코퍼스에서의 지속적인 사전 학습 시, Rho-1은 9개의 수학 과제에서 최대 30%의 절대적 향상을 보이는 소수 샷 정확도를 달성했습니다. 미세 조정 후, Rho-1-1B와 7B는 각각 MATH 데이터셋에서 40.6%와 51.8%의 최첨단 결과를 기록하며, 사전 학습 토큰의 단 3%만으로 DeepSeekMath와 동등한 성능을 보였습니다. 또한, 80B 일반 토큰에 대한 사전 학습 시 Rho-1은 15개의 다양한 과제에서 평균 6.8%의 향상을 달성하며, 언어 모델 사전 학습의 효율성과 성능을 모두 높였습니다.

English

Previous language model pre-training methods have uniformly applied a next-token prediction loss to all training tokens. Challenging this norm, we posit that "Not all tokens in a corpus are equally important for language model training". Our initial analysis delves into token-level training dynamics of language model, revealing distinct loss patterns for different tokens. Leveraging these insights, we introduce a new language model called Rho-1. Unlike traditional LMs that learn to predict every next token in a corpus, Rho-1 employs Selective Language Modeling (SLM), which selectively trains on useful tokens that aligned with the desired distribution. This approach involves scoring pretraining tokens using a reference model, and then training the language model with a focused loss on tokens with higher excess loss. When continual pretraining on 15B OpenWebMath corpus, Rho-1 yields an absolute improvement in few-shot accuracy of up to 30% in 9 math tasks. After fine-tuning, Rho-1-1B and 7B achieved state-of-the-art results of 40.6% and 51.8% on MATH dataset, respectively - matching DeepSeekMath with only 3% of the pretraining tokens. Furthermore, when pretraining on 80B general tokens, Rho-1 achieves 6.8% average enhancement across 15 diverse tasks, increasing both efficiency and performance of the language model pre-training.