Towards Engineering Scaling Laws with Pretraining Data Composition

Este artigo demonstra que na física de partículas, onde simuladores de alta fidelidade permitem a geração barata de dados sintéticos, o comportamento de escala de modelos de classificação de jatos hadrônicos pode ser projetado para priorizar a diversidade e o alinhamento de dados em vez do tamanho do modelo ao curar estrategicamente conjuntos de dados de pré-treinamento.

O essencial

Imagine que você está tentando ensinar um aluno a reconhecer diferentes tipos de veículos em uma cidade movimentada. Você tem duas maneiras principais de ajudá-lo a aprender: você pode dar a ele um cérebro maior (um modelo maior) ou pode dar a ele mais problemas de prática (mais dados).

Por muito tempo, cientistas que estudam Inteligência Artificial (IA) acreditaram que havia uma "regra de ouro" para isso. Eles pensavam que, se você tivesse uma quantidade fixa de tempo e dinheiro (orçamento de computação), a melhor maneira de obter o aluno mais inteligente seria dividir seus recursos em aproximadamente 50/50 entre construir um cérebro maior e dar a ele mais problemas de prática.

No entanto, este novo artigo sugere que, no mundo da física de partículas, podemos projetar uma regra melhor mudando o que o aluno aprende primeiro.

A Configuração: A Sala de Aula de Física

Os pesquisadores estão trabalhando com "jatos" (jets). Na física de partículas, quando partículas minúsculas colidem, elas espirram fluxos de outras partículas chamadas jatos. É como um fogo de artifício explodindo, mas em vez de faíscas, você tem fluxos de partículas subatômicas.

O objetivo é ensinar uma IA a olhar para esses fluxos e dizer: "Ah, este veio de um tipo específico de explosão!"

O Experimento: Mudando o Livro Didático

Os pesquisadores testaram dois "livros didáticos" diferentes (conjuntos de dados de pré-treinamento) para ver como eles mudavam as regras de aprendizado:

1. O Livro Didático Tedioso (Apenas QCD): Este livro continha apenas exemplos de explosões de partículas "padrão". Era como uma escola de condução que só ensinava você a dirigir um sedan padrão.
2. O Livro Didático Diverso (Melhorado com BSM): Este livro incluía os exemplos padrão mais exemplos complexos, raros e exóticos de explosões que não acontecem em nosso universo normal (simulações de física "Além do Modelo Padrão"). Era como uma escola de condução que ensinava você a dirigir sedans, mas também carros de corrida, caminhões e até veículos voadores.

A Descoberta: Reescrevendo as Regras

Quando treinaram a IA usando o Livro Didático Tedioso, a antiga regra de 50/50 manteve-se verdadeira. Para obter melhores resultados, era necessário equilibrar o aumento do tamanho do cérebro e o fornecimento de mais prática.

Mas quando usaram o Livro Didático Diverso, as regras mudaram completamente. A IA aprendeu que mais problemas de prática eram muito mais valiosos do que um cérebro maior.

• A Analogia: Imagine que a IA treinada com o livro didático diverso é como um aluno que já viu todos os tipos de veículos imagináveis. Quando você aplica um novo teste a ele, ele não precisa de um céreão maior para entender o novo carro; ele só precisa ver mais exemplos dele para atingir a perfeição. Seu "cérebro" não precisa crescer tão rápido porque sua "experiência" é rica.

O Resultado: A Nova Estratégia "Foco nos Dados"

O artigo descobriu que, ao usar os dados diversos e exóticos para o treinamento inicial:

• A estratégia de "cérebro maior" tornou-se menos importante.
• A estratégia de "mais dados" foi a vencedora.

De fato, os pesquisadores descobriram que, para cada unidade de poder de computação que você gasta, deve gastar cerca de 78% dela obtendo mais dados e apenas 22% tornando o modelo maior. Isso é uma mudança enorme em relação à antiga divisão de 50/50.

Por Que Isso Importa para a Física

O artigo destaca uma vantagem única da física: Nós podemos criar nossos próprios dados.

Em campos como medicina ou linguagem, obter novos dados é difícil, caro ou impossível (você não pode simplesmente "simular" um novo paciente humano). Mas na física de partículas, os cientistas usam computadores poderosos para simular colisões de partículas. Eles podem gerar quantidades infinitas de dados de alta qualidade e diversos gratuitamente (uma vez que a simulação esteja rodando).

A Conclusão:
Se você está construindo uma IA superinteligente para a física, não tente apenas construir o maior cérebro possível. Em vez disso, gaste seu tempo e dinheiro projetando um currículo melhor e mais diverso para a IA aprender primeiro. Uma vez que a IA tenha visto uma ampla variedade de exemplos "exóticos", ela aprenderá mais rápida e melhor a tarefa específica que você lhe der, e você obterá melhores resultados ao alimentá-la com mais dados em vez de tornar o modelo maior.

Em resumo: Uma dieta bem escolhida e diversa de dados de treinamento é mais poderosa do que um cérebro maior.

Resumo técnico

Resumo Técnico: Rumo à Engenharia de Leis de Escalonamento com Composição de Dados de Pré-treinamento

Enunciado do Problema
As leis de escalonamento neural descrevem como o desempenho do modelo melhora como uma lei de potência em relação ao computação, tamanho do modelo e tamanho do conjunto de dados. Embora bem estabelecidas para grandes modelos de linguagem (LLMs), essas relações estão emergindo na física de partículas. Uma distinção fundamental na física é a capacidade de gerar dados sintéticos de alta fidelidade via simuladores a um custo relativamente baixo em comparação com o gasto computacional de treinar modelos maiores. Isso cria uma oportunidade única de projetar o próprio conjunto de dados de pré-treinamento para influenciar o comportamento de escalonamento. A questão central abordada é se a composição dos dados de pré-treinamento — especificamente sua diversidade e alinhamento com tarefas de downstream — pode ser projetada para deslocar o regime de escalonamento computacionalmente ótimo de favorecer modelos maiores para favorecer conjuntos de dados maiores.

Metodologia
O estudo foca na tarefa de classificar jatos hadrônicos produzidos em colisões de alta energia. Os autores utilizam uma arquitetura transformer genérica que processa dados de jatos como uma nuvem de pontos, variando os tamanhos dos modelos de aproximadamente 3.000 a 10,5 milhões de parâmetros (abrangendo três ordens de magnitude) enquanto mantém a profundidade e as dimensões das cabeças de atenção fixas.

O design experimental envolve um protocolo de treinamento de dois estágios:

1. Pré-treinamento: Os modelos são pré-treinados em subconjuntos do conjunto de dados JetClass-II, que contém 188 classes de jatos simulados. Os autores definem quatro subconjuntos distintos de pré-treinamento para manipular diversidade e alinhamento:
   • QCD: Apenas jatos iniciados por quarks leves ou glúons (17 classes).
   • QCD + res2p: Jatos QCD mais jatos provenientes de decaimentos de dois corpos de ressonâncias Além do Modelo Padrão (BSM).
   • QCD + res34p: Jatos QCD mais jatos provenientes de decaimentos de três ou quatro corpos de ressonâncias BSM.
   • QCD + res2p + res34p: O conjunto de dados completo, incluindo todos os decaimentos de ressonância BSM.
   • Nota: Os subconjuntos BSM introduzem maior diversidade (mais classes de processos, cobertura de espaço de fase mais ampla) e melhor alinhamento com a tarefa de downstream (topologias multi-prong) comparado aos dados apenas de QCD.
2. Ajuste Fino (Fine-tuning): Os modelos pré-treinados são ajustados finamente no conjunto de dados original JetClass para uma tarefa de classificação de jatos de 10 classes (identificando quarks leves/glúons, quarks top, bósons W/Z e partículas Higgs). Esta tarefa requer a identificação de multiplicidade de prongs e escalas de massa, que são bem representadas nos dados aumentados por BSM, mas mal representadas nos dados apenas de QCD.

Os expoentes de escalonamento são extraídos ajustando leis de potência ao tamanho computacionalmente ótimo do modelo ($N^*$) e ao tamanho do conjunto de dados ($D^*$) como uma função da computação total ($C$). O estudo compara esses expoentes entre o treinamento "do zero" (sem pré-treinamento) e as várias configurações de pré-treinamento.

Resultos Principais
O estudo demonstra que a composição dos dados de pré-treinamento altera significativamente os expoentes de escalonamento computacionalmente ótimos:

• Treinamento do Zero: O treinamento a partir do zero produz expoentes de $a \approx 0,52$ (tamanho do modelo) e $b \approx 0,48$ (tamanho do conjunto de dados), indicando uma alocação de recursos computacionais aproximadamente equilibrada entre o tamanho do modelo e os dados, consistente com achados em LLMs.
• Pré-treinamento Apenas com QCD: O pré-treinamento exclusivamente em jatos QCD resulta em um deslocamento marginal ($a \approx 0,53, b \approx 0,47$), sugerindo que o pré-treinamento sozinho, sem alinhamento ou diversidade específicos, não altera fundamentalmente o regime de escalonamento.
• Pré-treinamento Aumentado por BSM: Incluir decaimentos de ressonância BSM no corpus de pré-treinamento causa um deslocamento dramático. Com o conjunto de dados completo aumentado por BSM, os expoentes mudam para $a \approx 0,22$ e $b \approx 0,78$.
  • Isso indica um regime onde a estratégia computacionalmente ótima favorece fortemente o aumento do tamanho do conjunto de dados em detrimento do aumento do tamanho do modelo.
  • O deslocamento representa uma redução de um fator de 2,3 no expoente de escalonamento para o tamanho do modelo em comparação com a linha de base do treinamento do zero.
  • As curvas de perda de ajuste fino confirmam que o pré-treinamento aumentado por BSM reduz consistentemente a perda em todos os tamanhos de modelo, com benefícios aumentando para modelos maiores.

Principais Contribuições

1. Engenharia de Leis de Escalonamento: O artigo fornece o primeiro estudo sistemático mostrando que a composição dos dados de pré-treinamento pode ser projetada para deslocar os expoentes de escalonamento em física fundamental. Demonstra que a diversidade e o alinhamento de downstream no corpus de pré-treinamento podem mover o regime de escalonamento ótimo para estratégias que favorecem os dados.
2. Deslocamento Quantitativo: O trabalho quantifica o deslocamento de um regime de escalonamento equilibrado ($a \approx b \approx 0,5$) para um regime fortemente favorável aos dados ($a \approx 0,22, b \approx 0,78$) ao incorporar física BSM no pré-treinamento.
3. Implicações para Modelos de Fundação: Os resultados sugerem que modelos de fundação pré-treinados em dados sintéticos diversos e alinhados podem alcançar desempenho ideal com contagens menores de parâmetros, permitindo que orçamentos de computação economizados sejam redirecionados para a geração de dados de treinamento adicionais.

Significância e Alegações
Os autores alegam que este trabalho identifica um novo espaço de design para o aprendizado de máquina científico: os inputs de física para o treinamento de modelos de fundação. Diferente dos domínios de linguagem natural ou imagem, onde a curadoria de dados é limitada pela disponibilidade, a física fundamental pode alavancar simuladores de alta fidelidade e baixo custo para construir corpora de pré-treinamento que moldam ativamente as leis de escalonamento.

O artigo conclui modestamente que, embora o pré-treinamento em corpora bem compostos permita um regime onde a computação de downstream é melhor gasta em mais dados do que em modelos maiores, trabalhos adicionais são necessários para verificar se esses resultados se generalizam para diferentes tarefas de ajuste fino, escalas de modelos maiores e diferentes tamanhos de conjuntos de dados. O estudo não afirma ter resolvido todos os desafios de escalonamento, mas destaca a engenharia de composição de pré-treinamento como uma alavanca subexplorada para maximizar o potencial de descoberta de modelos de fundação científicos.