Rho-1: Non tutti i token sono ciò di cui hai bisogno

Rho-1: Not All Tokens Are What You Need

April 11, 2024

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

cs.AI

Abstract

I precedenti metodi di pre-addestramento dei modelli linguistici hanno applicato uniformemente una funzione di perdita basata sulla previsione del token successivo a tutti i token di addestramento. Contestando questa norma, ipotizziamo che "Non tutti i token in un corpus siano ugualmente importanti per l'addestramento di un modello linguistico". La nostra analisi iniziale approfondisce le dinamiche di addestramento a livello di token del modello linguistico, rivelando distinti pattern di perdita per diversi token. Sfruttando queste intuizioni, introduciamo un nuovo modello linguistico chiamato Rho-1. A differenza dei modelli linguistici tradizionali che imparano a prevedere ogni token successivo in un corpus, Rho-1 utilizza il Selective Language Modeling (SLM), che addestra selettivamente i token utili allineati con la distribuzione desiderata. Questo approccio prevede la valutazione dei token di pre-addestramento utilizzando un modello di riferimento, seguito dall'addestramento del modello linguistico con una funzione di perdita focalizzata sui token con una perdita eccessiva più elevata. Durante il pre-addestramento continuo sul corpus OpenWebMath da 15B, Rho-1 ottiene un miglioramento assoluto nell'accuratezza few-shot fino al 30% in 9 task matematici. Dopo il fine-tuning, Rho-1-1B e 7B hanno raggiunto risultati all'avanguardia rispettivamente del 40,6% e 51,8% sul dataset MATH, eguagliando DeepSeekMath con solo il 3% dei token di pre-addestramento. Inoltre, durante il pre-addestramento su 80B token generici, Rho-1 ottiene un miglioramento medio del 6,8% su 15 task diversi, aumentando sia l'efficienza che le prestazioni del pre-addestramento del modello linguistico.

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.