Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

1. Problématique

Le pré-entraînement continu (Continual Pre-training) est une méthode efficace pour adapter les grands modèles de langage (LLM) à de nouveaux domaines spécifiques (comme le raisonnement mathématique ou la génération de code) en utilisant de petites quantités de données cibles. Cependant, cette approche présente un risque majeur : l'oubli catastrophique des capacités originales du modèle acquises lors du pré-entraînement initial.

La solution courante consiste à mélanger les données de la source (domaine original) et de la cible (nouveau domaine) pour maintenir un équilibre de performance. Les stratégies existantes reposent sur des heuristiques manuelles (définies par l'intuition humaine ou des résultats empiriques) pour déterminer les poids de mélange des données. Ces méthodes sont souvent rigides, coûteuses à ajuster pour chaque nouveau scénario et ne parviennent pas à exploiter pleinement les relations complexes entre les distributions de données et les performances du modèle.

2. Méthodologie : Data Mixing Agent

Les auteurs proposent le Data Mixing Agent, le premier cadre basé sur un modèle, fonctionnant de bout en bout, qui apprend à ré-pondérer les domaines de manière autonome.

A. Formulation comme Processus de Décision Markovien (MDP)

Le ré-pondérage des domaines est formalisé comme un MDP :

• État ($S$) : La distribution des données des étapes de ré-pondérage précédentes.
• Action ($A$) : La distribution de probabilité des données à mélanger pour l'étape courante (sur un espace de domaines, par exemple 52 dimensions).
• Politique ($f$) : L'agent lui-même, qui prédit la prochaine distribution.
• Récompense ($r$) : Un signal de feedback provenant d'un environnement d'évaluation léger.

B. Pipeline d'Apprentissage

1. Échantillonnage de Trajectoires : Des milliers de trajectoires de mélange de données sont générées aléatoirement. Pour chaque trajectoire, de petits modèles "proxy" sont entraînés de zéro.
2. Collecte de Feedback : À chaque étape de ré-pondérage, les checkpoints des modèles proxy sont évalués dans un environnement léger (utilisant des sous-ensembles de benchmarks comme MMLU pour les capacités générales et MATH pour le raisonnement mathématique). Cela génère des paires (distribution de données, feedback de performance).
3. Modélisation de l'Agent :
   • L'agent est un modèle léger basé sur l'architecture Transformer (2 couches, ~2,1M de paramètres).
   • Il prend en entrée l'historique des distributions et les feedbacks environnementaux pour prédire la prochaine distribution.
4. Optimisation par Apprentissage par Renforcement (RL) :
   • L'agent est d'abord pré-entraîné par Fine-Tuning Supervisé (SFT) sur les trajectoires de haute qualité.
   • Il est ensuite optimisé via Conservative Q-Learning (CQL), un algorithme de RL hors politique (off-policy). Le CQL est crucial car il permet d'apprendre non seulement des actions "bonnes" mais aussi d'éviter les actions "mauvaises" en pénalisant la sur-estimation des valeurs Q pour les actions hors distribution.

C. Généralisation

Une fois entraîné, l'agent peut prédire les mélanges de données pour de nouveaux modèles cibles, de nouvelles sources de données ou de nouveaux domaines sans nécessiter de ré-entraînement, en exploitant des heuristiques généralisables.

3. Contributions Clés

• Premier Agent de Mélange de Données : Introduction d'une méthode basée sur un modèle pour apprendre les heuristiques de ré-pondérage de domaine de manière end-to-end.
• Apprentissage par Renforcement Efficace : Utilisation du CQL pour apprendre des stratégies complexes à partir d'un large éventail de trajectoires, dépassant les simples biais inductifs du SFT.
• Généralisation Sans Ré-entraînement : Démonstration que l'agent apprend des heuristiques agnostiques aux données et aux modèles, fonctionnant bien sur des domaines non vus (ex: passer du raisonnement mathématique à la génération de code).
• Efficacité des Données : Capacité à atteindre des performances supérieures en utilisant moins de données du domaine source, réduisant ainsi les coûts de calcul.

4. Résultats Expérimentaux

Les expériences ont été menées sur des tâches de raisonnement mathématique et de génération de code, en utilisant des modèles cibles de 3B paramètres (architectures LLaMA) et des benchmarks standards (MMLU, GSM8K, MATH, HumanEval, etc.).

• Performance Équilibrée : Le Data Mixing Agent (DataAgentRL) surpasse significativement les méthodes de référence (RegMix, DBL, entraînement naïf).
  • Sur le raisonnement mathématique, il améliore la performance moyenne de 3,02 % par rapport à RegMix et de 8,88 % par rapport au modèle de base, tout en préservant les capacités générales.
  • Il réduit la variance des performances, indiquant une stabilité supérieure entre les différents domaines.
• Généralisation Transversale :
  • L'agent entraîné sur des données mathématiques et un domaine de définition spécifique (52 dimensions) fonctionne efficacement sur d'autres modèles cibles (LLaMA-FWE, LLaMA-Nemo) et d'autres espaces de domaines (2 dimensions) sans ré-entraînement.
  • Application directe réussie sur le domaine de la génération de code (un domaine non vu pendant l'entraînement), surpassant les méthodes optimisées spécifiquement pour le code.
• Alignement avec l'Intuition Humaine : L'analyse des trajectoires montre que l'agent apprend des stratégies cohérentes avec les conclusions empiriques (ex: augmenter les données des domaines "Science" et "Santé" pour MMLU, réduire les domaines "Mode" et "Beauté").
• Efficacité : L'agent permet d'atteindre des performances supérieures avec moins de tokens du domaine source (économie de ~2,14B de tokens dans les expériences mathématiques) et nécessite moins d'heures GPU totales que les méthodes basées sur la régression (RegMix).

5. Signification et Impact

Ce travail marque un changement de paradigme dans la gestion des données pour le pré-entraînement continu. Au lieu de dépendre de règles fixes ou de l'ajustement manuel coûteux, il introduit un agent autonome capable d'apprendre des stratégies de mélange de données optimales et généralisables.

• Réduction des Coûts : En apprenant à utiliser moins de données source tout en évitant l'oubli catastrophique, la méthode rend l'adaptation des LLM plus économique.
• Adaptabilité : La capacité de transférer les connaissances d'un domaine à un autre (ex: Math -> Code) sans ré-entraînement ouvre la voie à des systèmes d'adaptation de modèles plus flexibles et évolutifs.
• Fondation pour l'IA Autonome : La démonstration qu'un petit modèle peut apprendre des heuristiques complexes de gestion de données pour guider l'entraînement de grands modèles valide le concept d'agents IA pour l'optimisation de processus d'apprentissage profonds.

En résumé, le Data Mixing Agent offre une solution robuste, efficace et généralisable au problème de l'oubli catastrophique lors du pré-entraînement continu, surpassant les méthodes de l'état de l'art grâce à l'apprentissage par renforcement et à la découverte automatique d'heuristiques de mélange de données.