Rho-1: Niet alle tokens zijn wat je nodig hebt

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

Samenvatting

Vorige methoden voor het vooraf trainen van taalmodellen pasten uniform een volgende-token-voorspellingsverlies toe op alle trainings-tokens. Deze norm uitdagend, stellen wij dat "Niet alle tokens in een corpus even belangrijk zijn voor het trainen van taalmodellen". Onze initiële analyse verdiept zich in de trainingsdynamiek op tokenniveau van taalmodellen, waarbij verschillende verliespatronen voor verschillende tokens worden onthuld. Gebruikmakend van deze inzichten introduceren we een nieuw taalmodel genaamd Rho-1. In tegenstelling tot traditionele taalmodellen die leren om elk volgend token in een corpus te voorspellen, gebruikt Rho-1 Selective Language Modeling (SLM), waarbij selectief getraind wordt op nuttige tokens die overeenkomen met de gewenste distributie. Deze aanpak omvat het scoren van vooraf getrainde tokens met behulp van een referentiemodel, en vervolgens het trainen van het taalmodel met een gericht verlies op tokens met een hoger excessief verlies. Bij voortgezet vooraf trainen op het 15B OpenWebMath-corpus behaalt Rho-1 een absolute verbetering in few-shot nauwkeurigheid van tot 30% in 9 wiskundige taken. Na fine-tuning behaalden Rho-1-1B en 7B state-of-the-art resultaten van respectievelijk 40,6% en 51,8% op de MATH dataset - wat overeenkomt met DeepSeekMath met slechts 3% van de vooraf getrainde tokens. Bovendien behaalt Rho-1 bij het vooraf trainen op 80B algemene tokens een gemiddelde verbetering van 6,8% over 15 diverse taken, waardoor zowel de efficiëntie als de prestaties van het vooraf trainen van het taalmodel worden verhoogd.

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.