Rho-1 : Tous les tokens ne sont pas ce dont vous avez besoin

Rho-1: Not All Tokens Are What You Need

April 11, 2024

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

cs.AI

Résumé

Les méthodes précédentes de pré-entraînement des modèles de langage appliquaient uniformément une perte de prédiction du token suivant à tous les tokens d'entraînement. En remettant en question cette norme, nous postulons que "Tous les tokens d'un corpus ne sont pas également importants pour l'entraînement des modèles de langage". Notre analyse initiale explore la dynamique d'entraînement au niveau des tokens des modèles de langage, révélant des motifs de perte distincts pour différents tokens. En tirant parti de ces insights, nous introduisons un nouveau modèle de langage appelé Rho-1. Contrairement aux modèles de langage traditionnels qui apprennent à prédire chaque token suivant dans un corpus, Rho-1 utilise le Modélisation Sélective du Langage (Selective Language Modeling, SLM), qui entraîne sélectivement les tokens utiles alignés avec la distribution souhaitée. Cette approche consiste à évaluer les tokens de pré-entraînement à l'aide d'un modèle de référence, puis à entraîner le modèle de langage avec une perte focalisée sur les tokens présentant une perte excédentaire plus élevée. Lors d'un pré-entraînement continu sur le corpus OpenWebMath de 15 milliards de tokens, Rho-1 obtient une amélioration absolue de la précision en few-shot allant jusqu'à 30% sur 9 tâches mathématiques. Après un fine-tuning, Rho-1-1B et 7B atteignent des résultats de pointe de 40,6% et 51,8% sur le jeu de données MATH, respectivement - égalant DeepSeekMath avec seulement 3% des tokens de pré-entraînement. De plus, lors d'un pré-entraînement sur 80 milliards de tokens généraux, Rho-1 réalise une amélioration moyenne de 6,8% sur 15 tâches diverses, augmentant à la fois l'efficacité et la performance du pré-entraînement des modèles de langage.

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.