Résumé technique : Mise à l'échelle de la répétition des données de domaine dans le pré-entraînement des LLM

Énoncé du problème

À mesure que les modèles de langage de grande taille (LLM) passent à l'échelle, le maintien d'un ratio approprié de jetons par paramètre (TPP - tokens-per-parameter) nécessite d'augmenter le budget total de jetons d'entraînement. Si les données web générales peuvent être augmentées relativement facilement, les données spécifiques à un domaine de haute qualité (ex: code, mathématiques, textes médicaux) sont rares et difficiles à faire croître au même rythme. Par conséquent, lorsque la taille du modèle et le budget total de jetons augmentent, le ratio de mélange des données de haute qualité par domaine a tendance à diminuer, ce qui peut entraîner un échec de l'acquisition des connaissances cibles si le ratio tombe en dessous d'un seuil critique.

Pour contrer cette dilution, les praticiens répètent souvent les données de haute qualité disponibles. Cependant, cela introduit un compromis : la répétition aide à maintenir la fraction nécessaire de données pour l'acquisition des connaissances, mais augmente le risque de surapprentissage (overfitting). La littérature existante suggère que les modèles plus grands sont plus susceptibles de subir un surapprentissage dû aux données répétées, impliquant que la répétition devrait être minimisée. Toutefois, ces conclusions sont souvent tirées dans des configurations expérimentales où la taille totale des données d'entraînement est fixe à travers les échelles de modèles, plutôt que de préserver le TPP. Cet article étudie le nombre de répétitions optimal sous un régime de mise à l'échelle plus pratique où le budget de jetons d'entraînement croît proportionnellement à la taille du modèle (TPP fixe).

Méthodologie

Les auteurs ont mené une étude empirique systématique sur quatre domaines de haute qualité : le Code, les Mathématiques, Wikipédia et les données Médicales.

• Régime de mise à l'échelle : L'étude respecte un TPP fixe ($D/N$), où $D$ est le nombre total de jetons d'entraînement et $N$ est la taille du modèle. Le budget de jetons $D_N$ évolue linéairement avec la taille du modèle ($D_N = \text{TPP} \cdot N$).
• Configuration expérimentale : Pour chaque taille de modèle et chaque domaine, les auteurs ont fait varier deux paramètres :
  1. Fraction de données uniques ($\alpha$) : La proportion de jetons uniques de haute qualité dans le budget total (testée à 1/40, 1/20, 1/10).
  2. Nombre de répétitions ($e$) : Le nombre de fois qu'un sous-ensemble fixe de jetons uniques de haute qualité est revisité (testé de 1 à 7).
• Évaluation : Les modèles ont été évalués via la perte de validation en distribution (sur des données de réserve issues du même domaine) et la perte de validation hors distribution (OOD) du pré-entraînement (sur des données web générales comme ArXiv et News).
• Analyse théorique : Pour expliquer les résultats empiriques, les auteurs ont employé un modèle de régression linéaire "one-hot" dans $\mathbb{R}^\infty$ avec des unités de connaissance distribuées selon une loi de puissance. Ils ont analysé la décomposition du risque en "erreur d'acquisition de connaissances" et "erreur d'ajustement du bruit" pour dériver les temps d'arrêt (comptes de répétition optimaux) sous différentes conditions de mise à l'échelle.

Principaux résultats

1. Corrélation avec la perte de validation

Le nombre de répétitions optimal est fortement corrélé négativement avec la perte de validation finale d'un domaine (corrélation de Pearson de -0,944).

• Les domaines qui atteignent des pertes de validation plus faibles (indiquant que le modèle apprend bien le domaine) peuvent tolérer et bénéficier de plus de répétitions.
• Les domaines avec des pertes de validation plus élevées ont tendance à surapprendre plus tôt, nécessitant moins de répétitions.
• Cela suggère que la perte de validation sert de proxy pratique pour déterminer les limites de répétition sécurisées.

2. Impact de la taille du modèle sous TPP fixe

Contrairement aux conclusions des contextes à budget de données fixe (où la répétition optimale diminue à mesure que les modèles croissent), cette étude démontre que sous un TPP fixe, le nombre de répétitions optimal augmente légèrement avec la taille du modèle.

• Les modèles plus grands entraînés avec des budgets de jetons proportionnellement plus importants bénéficient de quelques répétitions supplémentaires avant que le surapprentissage ne survient.
• Cela soutient une stratégie de transfert conservatrice : un nombre de répétitions qui ne provoque pas de surapprentissage dans un modèle proxy plus petit (entraîné au même TPP) est peu susceptible de provoquer un surapprentissage dans un modèle cible plus grand.

3. Insensibilité à la fraction de données uniques

Le nombre de répétitions optimal est largement insensible à la fraction de données de haute qualité uniques ($\alpha$) dans la plage testée.

• Bien que l'augmentation de $\alpha$ réduise systématiquement la perte de validation absolue, l'emplacement du nombre de répétitions optimal (le $e$ qui minimise la perte) reste presque inchangé.
• Cela implique que lors du réglage de la répétition sur un modèle proxy, la fraction spécifique de données uniques utilisées n'a pas besoin de correspondre exactement à la cible de déploiement ; toute fraction représentative dans la plage testée suffit.

4. Dépendance au domaine et effets OOD

• Spécificité du domaine : L'efficacité de la répétition dépend fortement du domaine. Par exemple, les données de Mathématiques sont robustes à la répétition (jusqu'à 4-5 fois), tandis que les données Wikipédia se dégradent brutalement après 2 répétitions.
• Performance OOD : Lorsque la fraction totale de jetons de haute qualité est fixe, remplacer les jetons uniques par des jetons répétés a un effet limité sur la performance hors distribution (OOD). La dégradation est principalement confinée à la perte de validation intra-domaine.

5. Influence du programme de taux d'apprentissage (Learning Rate)

Le nombre de répétitions optimal dépend du programme de taux d'apprentissage. Les programmes avec une décroissance précoce (ex: commençant la décroissance à 0 % de l'entraînement) entraînent une dégradation après moins de répétitions. Inversement, retarder la décroissance ou utiliser un taux d'apprentissage constant permet au modèle de tolérer plus de répétitions, probablement grâce à un bruit d'optimisation plus fort agissant comme une régularisation implicite contre la mémorisation.

Explication théorique

Les auteurs fournissent un cadre théorique expliquant ces observations à travers le compromis entre l'acquisition de connaissances et l'ajustement du bruit :

• Ajustement du bruit (Noise Fitting) : L'optimisation répétée réduit l'erreur d'acquisition de connaissances mais augmente l'erreur d'ajustement du bruit. Le point d'arrêt optimal est celui où le bénéfice marginal du signal d'apprentissage est dépassé par le coût de l'ajustement du bruit.
• Perte de validation : Une perte de validation plus faible implique un niveau de bruit ($\sigma^2$) plus bas par rapport au signal. Théoriquement, un bruit plus faible retarde l'apparition du surapprentissage, permettant plus de répétitions (Théorème 4.2).
• Mise à l'échelle : Sous un TPP fixe, le budget de données croît avec la taille du modèle. Cela permet à la réduction de l'erreur d'acquisition de connaissances de dominer pendant une durée plus longue à mesure que le modèle change d'échelle, justifiant l'augmentation des comptes de répétition optimaux (Théorème 4.4).

Signification et implications pratiques

L'article conclut que les comptes de répétition réglés sur des modèles proxy plus petits, à condition qu'ils partagent le même TPP que le modèle cible, peuvent servir d'estimation pratique et conservatrice pour les modèles plus grands.

Le flux de travail proposé pour configurer les données de domaine de haute qualité dans le pré-entraînement à grande échelle est le suivant :

1. Entraîner un modèle proxy plus petit avec le même TPP.
2. Effectuer une recherche (sweep) des comptes de répétition pour le domaine cible en utilisant n'importe quelle fraction représentative de données uniques.
3. Identifier le nombre de répétitions qui minimise la perte de validation (ou qui s'arrête juste avant la dégradation).
4. Appliquer ce compte au modèle cible plus grand, car la tendance suggère que les modèles plus grands ne surapprendront pas plus tôt que le proxy sous le même régime de mise à l'échelle.

Cette approche offre une stratégie efficace en termes de données pour maintenir la représentation des domaines de haute qualité dans le pré-entraînement des LLM sans nécessiter de jeux de données uniques massifs, qui sont souvent indisponibles.