Measuring Task-Agnostic Training Data Influence Across Language Model Pretraining

Resumo Técnico: Medindo a Influência de Dados de Treinamento Agnóstica à Tarefa em Todo o Pré-treinamento de Modelos de Linguagem

Definição do Problema
Os métodos existentes para medir a influência dos dados de treinamento em modelos de linguagem (LMs) dependem tipicamente de alvos específicos de tarefas, como a perda (loss) em um ponto de teste específico ou o desempenho em um benchmark de aplicação (downstream). Embora úteis para objetivos específicos, essas abordagens enfrentam limitações significativas no contexto do pré-treinamento, que visa desenvolver capacidades amplas e gerais, em vez de otimizar para uma única tarefa. Selecionar tarefas de aplicação ou conjuntos de validação que sejam representativos das capacidades gerais de um modelo ao longo de toda a trajetória de pré-treinamento é difícil. Além disso, as estimativas de influência vinculadas ao desempenho de checkpoints intermediários em tarefas específicas são frequentemente inconsistentes e difíceis de comparar entre diferentes estágios de treinamento, pois os checkpoints iniciais podem possuir capacidades relevantes sem ainda terem resolvido totalmente a tarefa alvo.

Metodologia
Os autores propõem uma medida de influência de dados de treinamento agnóstica à tarefa, definida pela extensão em que um exemplo de treinamento individual move o modelo em direção aos parâmetros finais de uma determinada execução de pré-treinamento.

• Definição de Influência: A influência de um mini-lote $B_t$ é quantificada como a redução na distância $L_2$ ao quadrado entre os parâmetros atuais $\theta_t$ e os parâmetros pré-treinados finais $\theta^*$ induzida pela atualização no passo $t$. Matematicamente, $Cont(B_t) = \|\theta^* - \theta_t\|_2^2 - \|\theta^* - \theta_{t+1}\|_2^2$. Isso se expande para $2\Delta_t^\top(\theta^* - \theta_t) - \|\Delta_t\|_2^2$, onde o primeiro termo mede o alinhamento com a direção em direção aos parâmetros finais, e o segundo termo penaliza a norma ao quadrado da atualização.
• Decomposição ao Nível de Exemplo: Assumindo o Gradiente Descendente Estocástico (SGD) padrão, a influência do mini-lote é decomposta em contribuições individuais de exemplos. A influência de um exemplo $x_k^t$ é definida como $Cont(x_k^t) = 2\Delta_{t,k}^\top(\theta^* - \theta_t) - \Delta_{t,k}^\top \Delta_t$. Esta formulação leva em conta a interação entre exemplos dentro de um mini-lote.
• Aproximação Baseada em Checkpoints: Para tornar o método tratável para pré-treinamentos de larga escala sem a necessidade de retreinamento, os autores derivam uma aproximação usando checkpoints salvos periodicamente. Em vez de computar atualizações em cada passo, eles aproximam o estado $\theta_t$ e a atualização $\Delta_t$ usando os parâmetros no início ($\theta_c$) e no fim ($\theta_{c'}$) de um intervalo de checkpoint. Isso permite a estimativa post-hoc da influência a partir de checkpoints publicamente disponíveis.

Principais Contribuições

1. Reformulação Agnóstica à Tarefa: O artigo reformula a influência dos dados de treinamento ao longo do pré-treinamento em termos de progresso em direção aos parâmetros finais, removendo a dependência de tarefas ou conjuntos de validação específicos.
2. Estimativa Baseada em Checkpoints: Os autores derivam um método de aproximação prático que permite a estimação da influência ao nível de exemplo a partir de checkpoints intermediários sem precisar retreinar o modelo do zero.
3. Caracterização Empírica: O estudo caracteriza a evolução da contribuição de dados ao longo do pré-treinamento, revelando variações sistemáticas e dependentes de estágio entre diferentes tipos de dados.
4. Análise de Confiabilidade: O trabalho estabelece a confiabilidade da análise proposta através de aproximações de checkpoint, pontos de referência finais e várias configurações de modelos (tamanhos, inicializações e ordenações de dados).

Resultados Experimentais
O método foi aplicado a 18 configurações das suítes Pythia e PolyPythia, cobrindo tamanhos de modelo de 70M a 12B de parâmetros e utilizando 154 checkpoints por execução.

• Dinâmica de Contribuição: A contribuição média aumenta significamente após o estágio inicial de treinamento, atinge o pico por volta de 40k passos e depois diminui gradualmente. O desvio padrão segue um padrão semelhante. Notavelmente, a parcela de exemplos "oponentes" (aqueles com contribuição negativa, movendo o modelo para longe dos parâmetros finais) permanece baixa (<1%) durante os estágios iniciais e intermediários, mas sobe para 8–10% nos estágios finais.
• Dificuldade do Texto: Textos com maior perplexidade (PPL) (mais difíceis para o modelo) representam uma parcela maior da contribuição normalizada durante o estágio intermediário do pré-treinamento, enquanto a contribuição é distribuída de forma mais uniforme entre os níveis de dificuldade nos estágios iniciais e finais.
• Mudança de Domínio (O Cruzamento): Uma mudança temporal sistemática nos domínios de dados influentes foi observada:
  • Treinamento Inicial: Dados relacionados à literatura (ex: Livros, Hobbies) estão mais fortemente alinhados com a trajetória em direção aos parâmetros finais.
  • Treinamento Tardio: Dados relacionados a STEM (Ciência, Tecnologia, Engenharia e Matemática, como Ciências, Computadores e Eletrônicos) tornam-se cada vez mais alinhados com a trajetória.
  • Este cruzamento qualitativo é amplamente consistente em várias configurações de modelos, embora o tempo das transições possa variar com a escala do modelo (modelos maiores apresentam as transições mais cedo).
• Validação: A aproximação baseada em checkpoint mostra uma forte correlação (Pearson $r > 0.94$ nos estágios tardios) com as contribuições exatas ao nível de exemplo computadas via retreinamento em modelos menores. Os resultados são robustos a variações na inicialização de pesos e na ordenação de dados, embora a ordenação de dados introduza um pouco mais de variação do que a inicialização.
• Comparação com Métodos Específicos de Tarefa: Ao contrário das análises do tipo TracIn focadas em tarefas específicas, que geram padrões de atribuição fortemente dependentes do domínio de validação escolhido, a medida agnóstica à tarefa proposta revela consistentemente o cruzamento Literatura-para-STEM, independentemente do alvo específico de aplicação.

Significância e Alegações
O artigo afirma fornecer uma visão tratável e de nível de trajetória de como os dados influentes mudam ao longo do pré-treinamento. Ao desacoplar a análise de influência de tarefas específicas, os autores oferecem uma perspectiva que complementa as análises existentes focadas em tarefas.

Os resultados sugerem que os exemplos de treinamento que mais movem o modelo em direção aos seus parâmetros finais não são estáticos ao longo do pré-treinamento. Essa descoberta motiva abordagens conscientes de estágio para a composição de dados de pré-treinamento e design de currículo. Especificamente, a mudança observada em direção aos dados de STEM nos estágios finais fornece suporte empírico para as recentes receitas de pré-treinamento de LLMs que aumentam a alocação de dados de matemática e código nas fases finais de treinamento. Os autores posicionam seu trabalho como uma ferramenta para guiar estratégias de pré-treinamento mais sutis que adaptam a composição de dados ao longo dos estágios de treinamento, em vez de tratar uma única mistura de dados como igualmente adequada durante todo o processo.

Os autores mantêm modéstia quanto a alegações causais, observando que sua análise caracteriza como a contribuição varia ao longo do treinamento, em vez de estabelecer diretamente o efeito causal da mudança de misturas de dados no desempenho de tarefas específicas. Eles também reconhecem limitações, como a dependência da distância ao quadrado no espaço de parâmetros (que é sensível à parametrização) e a suposição de SGD para a decomposição, apesar do uso de otimizadores adaptativos na prática.