Rho-1: Nem Todos os Tokens São o que Você Precisa

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

Resumo

Métodos anteriores de pré-treinamento de modelos de linguagem aplicaram uniformemente uma perda de previsão do próximo token a todos os tokens de treinamento. Desafiando essa norma, propomos que "Nem todos os tokens em um corpus são igualmente importantes para o treinamento de modelos de linguagem". Nossa análise inicial explora a dinâmica de treinamento em nível de token de modelos de linguagem, revelando padrões distintos de perda para diferentes tokens. Aproveitando esses insights, introduzimos um novo modelo de linguagem chamado Rho-1. Diferente dos modelos de linguagem tradicionais que aprendem a prever cada próximo token em um corpus, o Rho-1 emprega a Modelagem de Linguagem Seletiva (SLM), que treina seletivamente em tokens úteis alinhados com a distribuição desejada. Essa abordagem envolve a pontuação de tokens de pré-treinamento usando um modelo de referência e, em seguida, o treinamento do modelo de linguagem com uma perda focada em tokens com maior perda excessiva. Quando pré-treinado continuamente no corpus OpenWebMath de 15 bilhões de tokens, o Rho-1 proporciona uma melhoria absoluta na precisão de poucos exemplos de até 30% em 9 tarefas matemáticas. Após o ajuste fino, o Rho-1-1B e o 7B alcançaram resultados de ponta de 40,6% e 51,8% no conjunto de dados MATH, respectivamente - igualando o desempenho do DeepSeekMath com apenas 3% dos tokens de pré-treinamento. Além disso, quando pré-treinado em 80 bilhões de tokens gerais, o Rho-1 alcança um aprimoramento médio de 6,8% em 15 tarefas diversas, aumentando tanto a eficiência quanto o desempenho do pré-treinamento de modelos de linguagem.

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.