CoLoR-Filter : Filtrage par Réduction Conditionnelle des Pertes pour le Pré-entraînement Ciblé de Modèles de Langage

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

June 15, 2024

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

cs.AI

Résumé

La sélection de données de haute qualité pour le pré-entraînement est cruciale pour influencer les performances des modèles de langage sur les tâches en aval. Un défi majeur réside dans l'identification de ce sous-ensemble optimal, un problème généralement considéré comme insoluble, nécessitant ainsi des heuristiques évolutives et efficaces. Dans ce travail, nous proposons une méthode de sélection de données, CoLoR-Filter (Conditional Loss Reduction Filtering), qui exploite une approche inspirée de l'inférence bayésienne empirique pour dériver un critère de sélection simple et efficace sur le plan computationnel, basé sur les valeurs de perte relative de deux modèles auxiliaires. En plus de la justification théorique, nous évaluons empiriquement CoLoR-Filter sur deux tâches de modélisation du langage : (1) la sélection de données à partir de C4 pour l'adaptation au domaine en vue d'une évaluation sur des livres, et (2) la sélection de données à partir de C4 pour une série de tâches de réponse à des questions à choix multiples en aval. Nous démontrons une mise à l'échelle favorable à la fois lorsque nous sous-sélectionnons de manière plus agressive et en utilisant de petits modèles auxiliaires pour sélectionner des données pour de grands modèles cibles. Comme résultat phare, les données sélectionnées par CoLoR-Filter à l'aide d'une paire de modèles auxiliaires de 150 millions de paramètres peuvent entraîner un modèle cible de 1,2 milliard de paramètres à égaler un modèle de 1,2 milliard de paramètres entraîné sur 25 milliards de tokens sélectionnés aléatoirement, avec 25 fois moins de données pour les livres et 11 fois moins de données pour les tâches en aval. Code : https://github.com/davidbrandfonbrener/color-filter-olmo Données filtrées : 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