Rho-1: No todos los tokens son lo que necesitas

Rho-1: Not All Tokens Are What You Need

April 11, 2024

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

cs.AI

Resumen

Los métodos anteriores de preentrenamiento de modelos de lenguaje han aplicado uniformemente una pérdida de predicción del siguiente token a todos los tokens de entrenamiento. Desafiando esta norma, proponemos que "No todos los tokens en un corpus son igualmente importantes para el entrenamiento de modelos de lenguaje". Nuestro análisis inicial profundiza en la dinámica de entrenamiento a nivel de token de los modelos de lenguaje, revelando patrones de pérdida distintos para diferentes tokens. Aprovechando estas ideas, presentamos un nuevo modelo de lenguaje llamado Rho-1. A diferencia de los modelos de lenguaje tradicionales que aprenden a predecir cada siguiente token en un corpus, Rho-1 emplea Modelado de Lenguaje Selectivo (SLM), que entrena selectivamente en tokens útiles alineados con la distribución deseada. Este enfoque implica puntuar los tokens de preentrenamiento utilizando un modelo de referencia y luego entrenar el modelo de lenguaje con una pérdida enfocada en tokens con mayor pérdida excesiva. Al realizar un preentrenamiento continuo en el corpus OpenWebMath de 15B, Rho-1 logra una mejora absoluta en la precisión de pocos ejemplos de hasta un 30% en 9 tareas matemáticas. Después del ajuste fino, Rho-1-1B y 7B alcanzaron resultados de vanguardia del 40.6% y 51.8% en el conjunto de datos MATH, respectivamente, igualando a DeepSeekMath con solo el 3% de los tokens de preentrenamiento. Además, al preentrenar con 80B tokens generales, Rho-1 logra una mejora promedio del 6.8% en 15 tareas diversas, aumentando tanto la eficiencia como el rendimiento del preentrenamiento de modelos de lenguaje.

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.