Rho-1: Not All Tokens Are What You Need

TL;DR

Rho-1 introduces Selective Language Modeling, training on useful tokens identified by scoring, which improves efficiency and accuracy in math and diverse tasks compared to traditional methods.

Contribution

The paper proposes a novel selective training approach that focuses on useful tokens, significantly enhancing language model performance and efficiency.

Findings

Up to 30% improvement in few-shot accuracy on math tasks

State-of-the-art results on MATH dataset with fewer tokens

6.8% average performance boost across diverse tasks

Abstract

Previous language model pre-training methods have uniformly applied a next-token prediction loss to all training tokens. Challenging this norm, we posit that "9l training". Our initial analysis examines 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 scores. 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 continual 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.