Resumo Técnico: Desajuste de Repetição em Experimentos de Mistura de Dados

Definição do Problema

No pré-treinamento de modelos de linguagem (LMs), determinar a composição ideal de misturas de dados de múltiplas fontes é crítico para o desempenho downstream. Uma estratégia comum envolve a execução de experimentos em pequena escala e a extrapolação dos resultados para o orçamento de treinamento alvo. No entanto, profissionais frequentemente observam que misturas otimizadas em pequenas escalas falham em se transferir efetivamente para regimes maiores.

Este artigo identifica o desajuste de repetição (repetition mismatch) como a causa primária dessa falha de extrapolação. Quando dados de alta qualidade são escassos em relação ao orçamento de treinamento, eles devem ser repetidos múltiplas vezes durante o treinamento. Crucialmente, o número de repetições necessárias para uma determinada base de dados muda conforme o orçamento total de treinamento cresce. Os métodos padrão de seleção de mistura baseados em escalonamento assumem que a relação entre proporções de mistura e desempenho é invariante à escala. Os autores argumentam que essa suposição é falha porque experimentos de controle em pequena escala impõem um regime de repetição fundamentalmente diferente para os dados de alta qualidade em comparação ao alvo. Esse descompasso distorce o panorama de perda (loss landscape), deslocando a mistura aparentemente ótima de maneiras que experimentos de pequena escala não conseguem antecipar.

Metodologia

Os autores propõem um procedimento de subamostragem consciente da repetição para isolar a repetição como uma variável na previsão da mistura. Esta abordagem, baseada no trabalho de Li et al. (2025), envolve o subamostragem de todas as fontes de dados de modo que os dados de alta qualidade passem pelo exato mesmo número de repetições no experimento de pequena escala que passariam na execução alvo de escala total, enquanto utilizam apenas uma fração do total de tokens.

Configuração Experimental:

• Datasets: O estudo combina conjuntos de dados de alta qualidade limitados (WikiText ou PubMed) com um corpus de web crawl de larga escala (FineWeb).
• Modelos: Os experimentos abrangem quatro tamanhos de modelo (30M, 124M, 345M e 757M de parâmetros) usando uma arquitetura NanoGPT modificada com otimizador Muon e Rotary Positional Embeddings.
• Comparação: O método proposto com controle de repetição é comparado contra um baseline de lei de escala (scaling-law) padrão, que extrapola proporções de mistura ideais a partir de execuções mais curtas sem igualar as taxas de repetição.
• Métricas: O desempenho é medido pela distância da mistura prevista em relação à mistura ótima real (medida por erro de razão de mistura) e pela perda de validação em domínios de alta qualidade retidos.

Principais Contribuições e Resultados

1. O Desajuste de Repetição é um Confundidor Dominante

O artigo demonstra que igualar a taxa de repetição da execução alvo é suficiente para recuperar previsões de mistura precisas a partir de experimentos de pequena escala.

• Resultados de Duas Fontes: Em um cenário com uma fonte de alta qualidade e FineWeb, um único experimento controlado por repetição usando apenas 1/16 dos tokens alvo recupera uma mistura dentro de 0,05 do ótimo para um modelo de 757M de parâmetros. Em contraste, a abordagem padrão de lei de escala sem controle de repetição apresenta um erro de 0,75.
• Eficiência de Custo: Alcançar precisão comparável sem o controle de repetição requer três a quatro horizontes de treinamento, consumindo de 44% a 94% do orçamento de tokens alvo. O método controlado por repetição alcança isso com um único horizonte.
• Interação com a Capacidade do Modelo: O benefício do controle de repetição escala com o tamanho do modelo. A redução do erro é mais pronunciada para modelos maiores (757M), enquanto modelos menores (30M) mostram menos benefício ou até superam o método controlado por repetição com leis de escala padrão, sugerindo um limite inferior de escala de modelo onde este método é mais eficaz.

2. Eficácia em Configurações de Múltiplas Fontes

Os autores estendem os achados para uma configuração de três fontes (WikiText, PubMed e FineWeb).

• Embora o espaço de mistura maior exija mais experimentos para ser restringido, o controle de repetição permanece altamente eficaz.
• Na escala de 757M, dois horizontes controlados por repetição são suficientes para recuperar a mistura ótima alvo, superando baselines que exigem experimentos completos de duas fontes para construção.
• Na escala de 124M, previsões de múltiplos horizontes controladas por repetição (até quatro horizontes) convergem para o ótimo, igualando efetivamente a perda real (2,91950 vs. 2,91820) enquanto reduzem significativamente o custo computacional em comparação aos baselines não controlados.

3. Repetição como uma Variável de Primeira Classe

Os resultados indicam que a dinâmica de repetição de dados, e não apenas a escala, determina se os experimentos de mistura em pequena escala generalizam. O artigo argumenta que a repetição deve ser tratada como uma variável explícita na otimização de mistura, em vez de um efeito colateral incidental da escassez de dados.

Significância e Alegações

O artigo afirma que controlar o desajuste de repetição resolve o problema de extrapolação em regimes de dados limitados. Ao preservar a dinâmica de repetição da execução alvo através da subamostragem, pesquisadores podem prever misturas de dados ideais com alta precisão usando uma fração do orçamento computacional.

Os autores posicionam isso como uma intervenção simples e ortogonal que não requer modelagem paramétrica ou execuções de treinamento de proxy, tornando-a aplicável a qualquer método existente de previsão de mistura. Eles afirmam que, à medida que os profissionais dependem cada vez mais de execuções de ablação menores para informar decisões de mistura para modelos de bilhões de parâmetros, considerar o desajuste de repetição se tornará crítico para evitar erros de previsão sistemáticos que crescem com a capacidade do modelo e o hiato entre as escalas de proxy e alvo.

O estudo é modesto em seu escopo, observando que os experimentos foram conduzidos até 757M de parâmetros e ~3,8B de tokens, menores do que o pré-treinamento de LLMs modernos. No entanto, a tendência consistente de que o controle de repetição se torna mais eficaz com o tamanho do modelo sugere que esses achados provavelmente se aplicarão a escalas maiores. Os autores também observam que sua avaliação foca na perda de validação sobre domínios de alta qualidade, e confirmar a transferência para benchmarks downstream permanece como uma direção para trabalhos futuros.