Resumo Técnico: Pré-treinamento de Modelos de Linguagem com Restrição de Dados

Declaração do Problema

As leis de escala clássicas para o pré-treinamento de Grandes Modelos de Linguagem (LLMs), como a Chinchilla, operam sob a premissa de um regime de "dados abundantes", onde os dados de treinamento únicos podem escalar livremente com o computação. Essas leis otimizam a alocação de um orçamento de computação fixo entre parâmetros do modelo e tokens de treinamento, tipicamente assumindo uma única passagem sobre o corpus.

No entanto, o artigo argumenta que o campo está se deslocando para um regime de computação rica e dados restritos. Neste cenário, o suprimento de dados de linguagem natural de alta qualidade é fixo ou limitado, enquanto a computação de treinamento continua a crescer. Consequentemente, as estratégias de treinamento devem se adaptar para utilizar modelos maiores e realizar múltiplas épocas (passagens repetidas) sobre um conjunto de dados finito. Os autores identificam duas lacunas críticas na compreensão atual deste regime:

1. Regularização: Como os modelos podem evitar o overfitting quando a computação aumenta, mas os dados únicos permanecem fixos? Embora o weight decay forte seja conhecido por ajudar, não está claro se outros mecanismos, como o mascaramento aleatório usado em Modelos de Linguagem de Difusão (dLLMs), oferecem benefícios complementares.
2. Leis de Escala: As leis de escala aditivas clássicas (que desacoplam o tamanho do modelo e o tamanho dos dados) descrevem com precisão a superfície de perda neste regime? Os autores hipotetizam que essas leis estão mal especificadas porque falham em capturar a interação entre a capacidade do modelo e os dados únicos limitados.

Metodologia

1. Regularização: Regularização de Entrada Mascarada (MIR)

Para investigar a regularização, os autores introduzem a Regularização de Entrada Mascarada (MIR). Esta é uma modificação mínima ao pré-treinamento autorregressivo (AR) padrão que não altera a arquitetura do modelo ou o procedimento de inferência.

• Mecanismo: Para uma sequência de entrada limpa $x$, uma versão mascarada $\hat{x}$ é gerada substituindo tokens por um token [MASK] com base em uma razão de mascaramento amostrada.
• Objetivo: A perda de treinamento é aumentada para incluir uma perda auxiliar de previsão do próximo token na sequência mascarada:
  $$L = L_{NTP}(x) + \lambda L_{NTP}(\hat{x})$$
• Racional: Os autores hipotetizam que o mascaramento aleatório atua como um regularizador, forçando o modelo a aprender características preditivas generalizáveis em vez de memorizar ruídos específicos do contexto. Isso é testado contra um baseline de modelos AR "fortemente regularizados" (usando weight decay elevado) e comparado a Modelos de Linguagem de Difusão (dLLMs).

2. Escala: A Lei SoftQ

Para abordar a questão da escala, os autores constroem uma grade de experimentos variando o tamanho do modelo ($N \in \{72M, \dots, 1.4B\}$) e o orçamento de dados únicos ($U \in \{100M, \dots, 400M\}$). Eles avaliam várias leis de escala:

• Chinchilla: Uma forma aditiva $L = E + A N^{-\alpha} + B U^{-\beta}$, que assume que a lacuna de perda entre orçamentos de dados é independente do tamanho do modelo.
• Estilo Muennighoff: Utiliza recursos efetivos que saturam com a repetição.
• Quanta: Uma lei acoplada derivada de visões de aprendizado de habilidades.
• SoftQ (Proposta): Uma lei de escala de cinco parâmetros que acopla o tamanho do modelo e o tamanho dos dados através de um "gargalo suave" (soft bottleneck):
  $$L_{SoftQ}(N, U) = E + \left( \frac{A}{N^\rho} + \frac{B}{U^{\rho/(1+\alpha)}} \right)^{\alpha/\rho}$$
  Aqui, $\rho$ controla a nitidez da transição entre os regimes limitado por parâmetros e limitado por dados. À medida que $U \to \infty$, ela recupera a escala de parâmetros; à medida que $N \to \infty$, ela recupera a escala de dados.

Principais Contribuições e Resultados

1. Descobertas de Regularização

• Transferência de Weight Decay: Os autores demonstram que o weight decay forte (otimizado para AR) melhora substancialmente o desempenho dos dLLMs, reduzindo a lacuna de perda de validação entre os modelos AR e dLLM. Quando ambos são fortemente regularizados, seus desempenhos tornam-se comparáveis.
• Eficácia da MIR: A MIR fornece regularização complementar sobre o weight decay forte. Em tamanhos de modelo de 72M a 1.4B treinados no dataset DCLM (100M tokens únicos), a MIR consistentemente reduz a perda de validação em comparação com o baseline de apenas weight decay forte.
  • O ganho aumenta com o tamanho do modelo (de uma redução de perda de ~0.006 em 72M para ~0.03 em 1.4B), alinhando-se com a intuição teórica de que modelos superparametrizados se beneficiam mais da regularização baseada em mascaramento.
  • Ganhos em Tarefas Downstream: Em 1.4B de parâmetros, a MIR produz melhorias significativas em tarefas subsequentes, incluindo +10.2 pontos em BoolQ e +2.2 pontos em SciQ.
  • Generalização: Os benefícios transferem-se para dados densos em código (Stack-V2) mesmo quando os hiperparâmetros foram ajustados apenas em dados de linguagem natural.
• Análise de Mecanismo: A análise ao nível de token revela que a MIR ajuda mais em tokens "difíceis", onde o contexto é ruidoso, possui scripts mistos ou contém nomes raros, sugerindo que a MIR melhora a robustez ao contexto de prefixo local, em vez de apenas à predição de frequência.

2. Descobertas de Escala

• Chinchilla está Mal Especificada: Resultados empíricos mostram que a lacuna de perda de validação entre diferentes orçamentos de dados únicos cresce com o tamanho do modelo. A lei aditiva da Chinchilla, que prevê uma lacuna invariante ao tamanho do modelo, falha em capturar este comportamento de "abertura" (fan-out).
• Superioridade da SoftQ: A lei SoftQ proposta, que acopla explicitamente o tamanho do modelo e o tamanho dos dados, ajusta-se significativamente melhor aos experimentos de dados restritos do que as leis Chinchilla, Quanta e de estilo Muennighoff. Ela alcança menor RMSE e MAE tanto em ajustes in-sample quanto em extrapolação de dados não observados (prevendo o desempenho de um orçamento de 400M a partir de orçamentos menores).
• Quantificação de Eficiência de Dados: Usando a SoftQ como baseline, os autores estimam que o ganho de desempenho da MIR é equivalente a aproximadamente 1.3 vezes a quantidade de dados de treinamento únicos (especificamente no intervalo de orçamento de 200M–400M tokens).

Significância e Alegações

O artigo afirma fornecer uma compreensão fundamental do regime de computação rica e dados restritos, que se torna cada vez mais relevante à medida que a computação escala mais rápido do que a disponibilidade de dados.

• Receita Prática: Os autores propõem uma receita prática e de baixo custo (MIR) que melhora a eficiência de dados para o pré-treinamento autorregressivo sem alterações arquiteturais ou overhead de inferência.
• Correção Teórica: Eles demonstram que as leis de escala aditivas clássicas são insuficientes para o regime de dados restritos e propõem a SoftQ como um modelo mais preciso para prever superfícies de perda quando os dados únicos são o recurso limitante.
• Interação de Regularização e Escala: O trabalho destaca que as estratégias de regularização (como weight decay e mascaramento) e as leis de escala estão profundamente interligadas neste regime; a escala eficaz requer modelos que sejam adequadamente regularizados para evitar o overfitting em dados repetidos.

Os autores mantêm a modéstia, observando que seus experimentos são limitados a 1.4B de parâmetros e 400M de tokens únicos, valores pequenos em relação ao pré-treinamento de escala de fronteira. Eles sugerem que trabalhos futuros explorem receitas de transferência de hiperparâmetros e variações de arquiteturas de modelos.