[Sessões do Paper-Club] Rho-1: nem todos os tokens são o que você precisa

Tradicionalmente, os grandes modelos de linguagem (LLMs) são treinados com uma abordagem única para todos, prevendo cada próximo token em um corpus usando a mesma função de perda. Mas e se nem todos os tokens contribuírem igualmente para o aprendizado? No artigo RHO-1, RHO-1: Nem todos os tokens são o que você precisa pesquisadores da Universidade de Xiamen, da Universidade de Tsinghua, do Laboratório de IA de Xangai e da Microsoft desafiam essa norma ao introduzir um novo paradigma de treinamento chamado Modelagem Seletiva de Linguagem (SLM), que concentra o aprendizado do modelo nos tokens que realmente importam, conforme ilustrado na Figura 1. Notavelmente, este trabalho foi reconhecido como o melhor vice-campeão em papel no NeuRIPS 2024, ressaltando seu impacto significativo no campo.

Figura 1 (retirada do artigo RHO-1): Comparação entre Modelagem de Linguagem Causal (CLM) e Modelagem Seletiva de Linguagem (SLM). Mesmo com uma filtragem rigorosa em nível de documento, conjuntos de dados de alta qualidade ainda contêm tokens ruidosos que podem atrapalhar o treinamento. Em contraste com a abordagem CLM padrão, que calcula a perda uniformemente para cada token, o SLM processa toda a sequência e descarta seletivamente a perda de tokens indesejados, levando a um aprendizado mais eficiente.

Os métodos convencionais de pré-treinamento aplicam a perda de previsão do próximo token de maneira uniforme em todos os tokens, mas a equipe do RHO-1 demonstra que essa abordagem é ineficiente. Ao analisar a dinâmica do treinamento em nível de token, os autores descobriram que os tokens em um corpus exibem padrões de perda marcadamente diferentes — alguns tokens são “fáceis” e aprendidos rapidamente, enquanto outros permanecem persistentemente desafiadores ou até flutuam durante o treinamento, conforme mostrado na Figura 2. Essa percepção sugere que tratar cada token de forma igual desperdiça computação valiosa em partes dos dados que oferecem poucas informações novas.

‍

Figura 2 (adaptada do artigo RHO-1): Dinâmica de perda em nível de token durante o pré-treinamento. Essa figura avalia a perda de tokens individuais a cada 1 bilhão de tokens de treinamento usando um conjunto de validação de 320 mil tokens. Ele categoriza os tokens em quatro grupos: 26% apresentam redução significativa de perda (H → L), 51% permanecem consistentemente baixos (L → L), 11% são persistentemente desafiadores (H → H) e 12% mostram um aumento inesperado de perda (L → H). Esses resultados destacam que as perdas individuais de tokens não diminuem uniformemente, revelando dinâmicas de treinamento complexas que motivam o foco seletivo do SLM.

No cerne do RHO-1 está a ideia de que apenas um subconjunto de tokens — aqueles que apresentam maior “perda excessiva” em relação a um modelo de referência de alta qualidade — deve conduzir o processo de aprendizado. A abordagem começa com o treinamento de um modelo de referência em um corpus cuidadosamente selecionado para capturar a distribuição de tokens desejada. Cada token no conjunto de dados maior de pré-treinamento é então pontuado comparando a perda no modelo atual com a do modelo de referência, gerando uma “perda excessiva” que indica o potencial do token de melhorar o modelo se focado. Em vez de computar a perda em todo o corpus, o SLM aplica seletivamente a perda somente a tokens com alta perda excessiva, filtrando efetivamente o ruído e concentrando-se nas partes mais informativas dos dados.

Essa abordagem seletiva não apenas melhora a eficiência dos dados, mas também aprimora o desempenho geral do modelo. Experimentos no corpus do OpenWebMath mostraram que o RHO-1 atinge um aumento absoluto de precisão de até 30% em nove tarefas matemáticas. Além disso, após o ajuste fino, os modelos de parâmetros 1B e 7B alcançaram resultados de última geração no conjunto de dados MATH, igualando ou até mesmo excedendo o desempenho de modelos treinados em quantidades muito maiores de dados. Quando estendido para um corpus diversificado de 80B de tokens do domínio geral, o RHO-1 produziu uma melhoria média de 6,8% em 15 benchmarks, com ganhos especialmente notáveis em tarefas envolvendo código e matemática. Além disso, ao se concentrar em tokens de alto impacto, o SLM acelera o aprendizado, atingindo níveis de precisão comparáveis de 5 a 10 vezes mais rápido do que os métodos tradicionais.

Observações

O RHO-1 representa uma nova direção promissora para o pré-treinamento do LLM, desafiando a suposição de que “todos os tokens são criados da mesma forma”. Em vez de desperdiçar a computação em tokens já aprendidos ou irrelevantes, o SLM se concentra seletivamente naqueles que realmente aumentam as capacidades do modelo. Embora os resultados nos domínios matemático e geral sejam impressionantes, esse trabalho também abre vários caminhos interessantes:

• Aplicações mais amplas: O treinamento seletivo poderia ser estendido a tarefas multimodais ou até mesmo a cenários de aprendizado por reforço?

• Opções de modelo de referência: Qual é o impacto do uso de diferentes modelos de referência? As APIs proprietárias poderiam ser aproveitadas para uma pontuação de token mais robusta?

• Adaptação dinâmica: Como a taxa de seleção de tokens pode ser adaptada dinamicamente durante o treinamento para otimizar ainda mais o desempenho?

Embora o RHO-1 demonstre melhorias significativas em tarefas matemáticas e de domínio geral com modelos e conjuntos de dados menores, sua escalabilidade para modelos muito grandes e corpora extensos ainda precisa ser demonstrada. Além disso, confiar em um modelo de referência para pontuação de tokens introduz uma dependência adicional que pode limitar sua aplicabilidade mais ampla em determinadas configurações. O RHO-1 não apenas alcança resultados de última geração com menos tokens, mas também ressalta a importância de repensar nossa abordagem aos dados no pré-treinamento do LLM. Ao nos concentrarmos no que realmente importa, poderemos criar modelos mais eficientes e capazes no futuro.

‍

Referências

[1] Lin, Zhenghao, et al. “Rho-1: Nem todos os tokens são o que você precisa.” pré-impressão arXiv: 2404.07965 (2024).

‍