Resumo Técnico: Escalonamento da Repetição de Dados de Domínio em Pré-treinamento de LLMs

Declaração do Problema

À medida que os Grandes Modelos de Linguagem (LLMs) escalam, manter uma proporção apropriada de tokens por parâmetro (TPP - tokens-per-parameter) exige o aumento do orçamento total de tokens de treinamento. Embora dados da web de uso geral possam ser escalados relativamente fácil, dados de domínios específicos de alta qualidade (ex: código, matemática, textos médicos) são escassos e difíceis de escalar na mesma taxa. Consequentemente, conforme o tamanho do modelo e o orçamento total de tokens crescem, a razão de mistura de dados de domínio de alta qualidade tende a diminuir, o que pode levar a uma falha na aquisição do conhecimento pretendido se a razão cair abaixo de um limiar crítico.

Para combater essa diluição, os profissionais frequentemente repetem os dados de alta qualidade disponíveis. No entanto, isso introduz um compromisso (trade-off): a repetição ajuda a manter a fração necessária de dados para a aquisição de conhecimento, mas aumenta o risco de sobreajuste (overfitting). A literatura existente sugere que modelos maiores são mais suscetíveis ao overfitting decorrente de dados repetidos, implicando que a repetição deve ser minimizada. No entanto, essas conclusões são frequentemente extraídas sob configurações experimentais onde o tamanho total dos dados de treinamento é fixo entre as escalas de modelos, em vez de preservar o TPP. Este artigo investiga a contagem de repetição ideal sob um regime de escalonamento mais prático, onde o orçamento de tokens de treinamento cresce proporcionalmente com o tamanho do modelo (TPP fixo).

Metodologia

Os autores conduziram um estudo empírico sistemático através de quatro domínios de alta qualidade: Código, Matemática, Wikipedia e Dados Médicos.

• Regime de Escalonamento: O estudo adere a um TPP ($D/N$) fixo, onde $D$ é o total de tokens de treinamento e $N$ é o tamanho do modelo. O orçamento de tokens $D_N$ escala linearmente com o tamanho do modelo ($D_N = \text{TPP} \cdot N$).
• Configuração Experimental: Para cada tamanho de modelo e domínio, os autores variaram dois parâmetros:
  1. Fração de Dados Únicos ($\alpha$): A proporção de tokens únicos de alta qualidade no orçamento total (testada em 1/40, 1/20, 1/10).
  2. Contagem de Repetição ($e$): O número de vezes que um subconjunto fixo de tokens únicos de alta qualidade é revisitado (testado de 1 a 7).
• Avaliação: Os modelos foram avaliados usando a perda de validação intra-domínio (em dados retidos do mesmo domínio) e a perda de validação fora de distribuição (OOD) de pré-treinamento (em dados gerais da web como ArXiv e Notícias).
• Análise Teórica: Para explicar os achados empíricos, os autores empregaram um modelo de regressão linear one-hot em $\mathbb{R}^\infty$ com unidades de conhecimento distribuídas por lei de potência. Eles analisaram a decomposição do risco em "erro de aquisição de conhecimento" e "erro de ajuste de ruído" para derivar tempos de parada (contagens de repetição ideais) sob diferentes condições de escalonamento.

Principais Achados

1. Correlação com a Perda de Validação

A contagem de repetição ideal possui uma forte correlação negativa com a perda de validação final de um domínio (correlação de Pearson de -0,944).

• Domínios que alcançam perdas de validação menores (indicando que o modelo aprende bem o domínio) podem tolerar e beneficiar-se de mais repetições.
• Domínios com perdas de validação mais altas tendem a sofrer overfitting mais cedo, exigindo menos repetições.
• Isso sugere que a perda de validação serve como um proxy prático para determinar limites seguros de repetição.

2. Impacto do Tamanho do Modelo sob TPP Fixo

Ao contrário das descobertas em configurações de orçamento de dados fixo (onde a repetição ideal diminui conforme os modelos crescem), este estudo descobre que, sob um TPP fixo, a contagem de repetição ideal aumenta levemente com o tamanho do modelo.

• Modelos maiores treinados com orçamentos de tokens proporcionalmente maiores beneficiam-se de ligeiramente mais repetições antes que o overfitting ocorra.
• Isso apoia uma estratégia de transferência conservadora: uma contagem de repetição que não causa overfitting em um modelo auxiliar menor (treinado com o mesmo TPP) dificilmente causará overfitting em um modelo alvo maior.

3. Insensibilidade à Fração de Dados Únicos

A contagem de repetição ideal é amplamente insensível à fração de dados únicos de alta qualidade ($\alpha$) dentro do intervalo testado.

• Embora o aumento de $\alpha$ reduza consistentemente a perda de validação absoluta, a localização da contagem de repetição ideal (o $e$ que minimiza a perda) permanece quase inalterada.
• Isso implica que, ao ajustar a repetição em um modelo auxiliar, a fração específica de dados únicos utilizada não precisa corresponder exatamente ao alvo de implantação; qualquer fração representativa dentro do intervalo testado é suficiente.

4. Dependência de Domínio e Efeitos OOD

• Especificidade de Domínio: A eficácia da repetição é altamente dependente do domínio. Por exemplo, dados de Matemática são robustos à repetição (até 4-5 vezes), enquanto dados da Wikipedia degradam bruscamente após 2 repetições.
• Desempenho OOD: Quando a fração total de tokens de alta qualidade é fixa, substituir tokens únicos por tokens repetidos tem um efeito limitado no desempenho fora de distribuição (OOD). A degradação é concentrada principalmente na perda de validação intra-domínio.

5. Influência do Cronograma de Taxa de Aprendizado (Learning Rate Schedule)

A contagem de repetição ideal depende do cronograma da taxa de aprendizado. Cronogramas com decaimento precoce (ex: iniciando o decaimento em 0% do treinamento) levam à degradação após menos repetições. Por outro lado, atrasar o decaimento ou usar uma taxa de aprendizado constante permite que o modelo tolere mais repetições, provavelmente devido ao ruído de otimização mais forte atuando como regularização implícita contra a memorização.

Explicação Teórica

Os autores fornecem um arcabouço teórico explicando essas observações através do compromisso entre aquisição de conhecimento e ajuste de ruído:

• Ajuste de Ruído (Noise Fitting): A otimização repetida reduz o erro de aquisição de conhecimento, mas aumenta o erro de ajuste de ruído. O ponto de parada ideal é onde o benefício marginal do sinal de aprendizado é superado pelo custo de ajustar o ruído.
• Perda de Validação: Uma perda de validação menor implica um nível de ruído ($\sigma^2$) menor em relação ao sinal. Teoricamente, um ruído menor atrasa o início do overfitting, permitindo mais repetições (Teorema 4.2).
• Escalonamento: Sob TPP fixo, o orçamento de dados cresce com o tamanho do modelo. Isso permite que a redução no erro de aquisição de conhecimento domine por um período mais longo à medida que o modelo escala, justificando o aumento nas contagens de repetição ideais (Teorema 4.4).

Significância e Implicações Práticas

O artigo conclui que contagens de repetição ajustadas em modelos auxiliares menores, desde que compartilhem o mesmo TPP do modelo alvo, podem servir como uma estimativa prática e conservadora para modelos maiores.

O fluxo de trabalho proposto para configurar dados de domínio de alta qualidade em pré-treinamento de larga escala é:

1. Treinar um modelo auxiliar menor com o mesmo TPP.
2. Realizar um teste de variação (sweep) das contagens de repetição para o domínio alvo usando qualquer fração representativa de dados únicos.
3. Identificar a contagem de repetição que minimiza a perda de validação (ou que para logo antes da degradação).
4. Aplicar esta contagem ao modelo alvo maior, pois a tendência sugere que modelos maiores não sofrerão overfitting mais cedo do que o modelo auxiliar sob o mesmo regime de escalonamento.

Esta abordagem oferece uma estratégia eficiente de dados para manter a representação de domínio de alta qualidade no pré-treinamento de LLMs sem exigir conjuntos de dados únicos massivos, que muitas vezes não estão disponíveis.