CoLoR-Filter: Filtro de Redução de Perda Condicional para Pré-treinamento Direcionado de Modelos de Linguagem

CoLoR-Filter: Conditional Loss Reduction Filtering for Targeted Language Model Pre-training

June 15, 2024

Autores: David Brandfonbrener, Hanlin Zhang, Andreas Kirsch, Jonathan Richard Schwarz, Sham Kakade

cs.AI

Resumo

A seleção de dados de alta qualidade para pré-treinamento é crucial para moldar o desempenho de modelos de linguagem em tarefas subsequentes. Um grande desafio reside na identificação desse subconjunto ideal, um problema geralmente considerado intratável, o que torna necessário o uso de heurísticas escaláveis e eficazes. Neste trabalho, propomos um método de seleção de dados, o CoLoR-Filter (Filtragem por Redução de Perda Condicional), que utiliza uma abordagem inspirada no método empírico de Bayes para derivar um critério de seleção simples e computacionalmente eficiente com base nos valores de perda relativa de dois modelos auxiliares. Além da fundamentação teórica, avaliamos o CoLoR-Filter empiricamente em duas tarefas de modelagem de linguagem: (1) seleção de dados do C4 para adaptação de domínio em avaliações sobre livros e (2) seleção de dados do C4 para uma série de tarefas subsequentes de resposta a perguntas de múltipla escolha. Demonstramos uma escalabilidade favorável tanto ao selecionar dados de forma mais agressiva quanto ao utilizar modelos auxiliares pequenos para selecionar dados para modelos alvo grandes. Como um dos principais resultados, os dados selecionados pelo CoLoR-Filter usando um par de modelos auxiliares com 150 milhões de parâmetros podem treinar um modelo alvo de 1,2 bilhão de parâmetros para igualar o desempenho de um modelo de 1,2 bilhão de parâmetros treinado com 25 bilhões de tokens selecionados aleatoriamente, utilizando 25 vezes menos dados para livros e 11 vezes menos dados para as tarefas subsequentes. Código: https://github.com/davidbrandfonbrener/color-filter-olmo Dados filtrados: https://huggingface.co/datasets/davidbrandfonbrener/color-filtered-c4

English

Selecting high-quality data for pre-training is crucial in shaping the downstream task performance of language models. A major challenge lies in identifying this optimal subset, a problem generally considered intractable, thus necessitating scalable and effective heuristics. In this work, we propose a data selection method, CoLoR-Filter (Conditional Loss Reduction Filtering), which leverages an empirical Bayes-inspired approach to derive a simple and computationally efficient selection criterion based on the relative loss values of two auxiliary models. In addition to the modeling rationale, we evaluate CoLoR-Filter empirically on two language modeling tasks: (1) selecting data from C4 for domain adaptation to evaluation on Books and (2) selecting data from C4 for a suite of downstream multiple-choice question answering tasks. We demonstrate favorable scaling both as we subselect more aggressively and using small auxiliary models to select data for large target models. As one headline result, CoLoR-Filter data selected using a pair of 150m parameter auxiliary models can train a 1.2b parameter target model to match a 1.2b parameter model trained on 25b randomly selected tokens with 25x less data for Books and 11x less data for the downstream tasks. Code: https://github.com/davidbrandfonbrener/color-filter-olmo Filtered data: https://huggingface.co/datasets/davidbrandfonbrener/color-filtered-c4