AutoMixer: Checkpoint Artifacts as Automatic Data Mixers

Résumé Technique : AutoMixer

1. Problématique (The Problem)

L'entraînement de modèles de langage (LLM) performants repose sur la composition optimale de leurs données d'entraînement. Cependant, deux défis majeurs subsistent :

• L'absence de correspondance directe : Il est difficile de mapper précisément des ensembles de données brutes (souvent hétérogènes et massifs comme Common Crawl) à des compétences spécifiques (raisonnement, mathématiques, etc.), car les frontières entre les domaines de connaissances sont floues et se chevauchent.
• La dynamique d'acquisition des compétences : Les capacités d'un modèle n'émergent pas de manière linéaire ; elles suivent des trajectoires non monotones au cours de l'entraînement. Identifier les données "idéales" pour une compétence est un problème de type "l'œuf ou la poule" : il faut que le modèle possède déjà la compétence pour savoir quelles données l'ont favorisée.

2. Méthodologie (Methodology)

Les auteurs proposent AutoMixer, un cadre qui utilise les points de contrôle (checkpoints) intermédiaires du modèle comme des "signaux" pour regrouper et pondérer les données. La méthode se décompose en deux étapes principales :

A. Regroupement des données (Data Regrouping) via l'influence des échantillons :

• Simulation par modèles proxys : Au lieu d'utiliser le modèle final (trop coûteux), les auteurs entraînent des modèles plus petits (proxys de 75M ou 350M de paramètres) pour simuler la trajectoire d'apprentissage.
• Sélection de checkpoints par tâche : Pour chaque tâche cible (ex: ARC, BoolQ), ils identifient le checkpoint qui présente la meilleure performance. Ces checkpoints sont considérés comme des "artefacts" capturant l'état optimal pour une compétence donnée.
• Estimation de l'influence : Ils utilisent une approximation efficace de la fonction d'influence (via DataInf) pour calculer l'impact de chaque échantillon de donnée sur la performance de ces checkpoints spécifiques.
• Calcul de l'influence jointe ($I_{joint}$) : Pour chaque échantillon, une influence agrégée est calculée en combinant les scores de plusieurs checkpoints, pondérés par un facteur de mélange ($\alpha_j$) basé sur la vitesse d'acquisition de la tâche. Les échantillons sont ensuite regroupés en "groupes de données" basés sur leur utilité.

B. Repondération du mélange de données (Data-mix Reweighting) :

• Une fois les données regroupées, AutoMixer détermine la probabilité de tirage (poids $w_g$) de chaque groupe lors de l'entraînement final.
• Le poids est proportionnel à la densité d'influence du groupe ($\rho_g$), ce qui maximise l'influence globale sur toutes les tâches cibles tout en respectant le budget total de tokens.

3. Contributions Clés (Key Contributions)

1. Cadre AutoMixer : Une méthode innovante qui transforme les checkpoints de modèles intermédiaires en outils de sélection de données automatiques.
2. Utilisation de modèles proxys : Démontre que des modèles de petite taille peuvent servir de guides efficaces pour l'entraînement de modèles beaucoup plus grands.
3. Approche multi-checkpoint : Contrairement aux méthodes classiques qui ne regardent que l'état final du modèle, AutoMixer capture l'évolution temporelle des compétences.

4. Résultats (Results)

Les expérimentations ont été menées sur le jeu de données FineWeb-Edu avec des modèles de 350M, 1.5B et 3B de paramètres :

• Amélioration de la performance : AutoMixer-350M surpasse systématiquement l'échantillonnage uniforme et les méthodes basées sur la perplexité (PPL) ou les n-grammes.
• Gain de précision : Sur l'échelle 350M, AutoMixer-350M affiche une amélioration moyenne de 1,93 % sur huit benchmarks de raisonnement.
• Efficacité du proxy : Les résultats montrent que la taille du modèle proxy est cruciale : un modèle de 350M est bien plus performant qu'un modèle de 75M pour identifier les échantillons à haute influence, bien que le modèle de 75M reste utile pour le filtrage de base.
• Trajectoire d'apprentissage : Les courbes de performance montrent qu'AutoMixer maintient un avantage constant sur l'échantillonnage uniforme tout au long des étapes d'entraînement.

5. Signification (Significance)

Ce travail démontre que les "déchets" de l'entraînement (les checkpoints intermédiaires souvent ignorés) sont en réalité des sources d'informations précieuses pour optimiser la qualité des données. AutoMixer offre une voie vers un entraînement de modèles plus efficace et plus ciblé, permettant de maximiser l'acquisition de compétences spécifiques sans nécessiter de curations manuelles de données complexes ou de méthodes de force brute coûteuses.