Towards Engineering Scaling Laws with Pretraining Data Composition

Résumé Technique : Vers l'ingénierie des lois d'échelle de mise à l'échelle par la composition des données de préentraînement

Énoncé du problème
Les lois d'échelle neuronales décrivent comment la performance d'un modèle s'améliore selon une loi de puissance par rapport au calcul, à la taille du modèle et à la taille du jeu de données. Bien que ces relations soient bien établies pour les grands modèles de langage (LLM), elles émergent dans le domaine de la physique des particules. Une distinction clé en physique fondamentale est la capacité de générer des données synthétiques de haute fidélité via des simulateurs à un coût relativement faible par rapport à la dépense de calcul nécessaire pour entraîner des modèles plus grands. Cela crée une opportunité unique d'ingénierer le jeu de données de préentraînement lui-même pour influencer le comportement de mise à l'échelle. La question centrale abordée est de savoir si la composition des données de préentraînement — spécifiquement sa diversité et son alignement avec les tâches en aval — peut être ingéniérée pour déplacer le régime de mise à l'échelle optimal en termes de calcul en faveur de jeux de données plus larges plutôt que de modèles plus grands.

Méthodologie
L'étude se concentre sur la tâche de classification des jets hadroniques produits lors de collisions de particules à haute énergie. Les auteurs utilisent une architecture transformer générique qui traite les données de jets comme un nuage de points, faisant varier la taille des modèles d'environ 3 000 à 10,5 millions de paramètres (couvrant trois ordres de grandeur) tout en maintenant fixes la profondeur et les dimensions des têtes d'attention.

Le design expérimental repose sur un protocole d'entraînement en deux étapes :

1. Préentraînement : Les modèles sont préentraînés sur des sous-ensembles du jeu de données JetClass-II, qui contient 188 classes de jets simulés. Les auteurs définissent quatre sous-ensembles de préentraînement distincts pour manipuler la diversité et l'alignement :
   • QCD : Uniquement des jets initiés par des quarks légers ou des gluons (17 classes).
   • QCD + res2p : Jets QCD plus des jets provenant de désintégrations de résonances de type "au-delà du Modèle Standard" (BSM) à deux corps.
   • QCD + res34p : Jets QCD plus des jets provenant de désintégrations de résonances BSM à trois ou quatre corps.
   • QCD + res2p + res34p : Le jeu de données complet incluant toutes les désintégrations de résonances BSM.
   • Note : Les sous-ensembles BSM introduisent une plus grande diversité (plus de classes de processus, couverture plus large de l'espace des phases) et un meilleur alignement avec la tâche en aval (topologies multi-branches) par rapport aux données uniquement QCD.
2. Ajustement fin (Fine-tuning) : Les modèles préentraînés sont ajustés sur le jeu de données JetClass original pour une tâche de classification de jets à 10 classes (identification des quarks légers/gluons, des quarks top, des bosons W/Z et des particules de Higgs). Cette tâche nécessite l'identification de la multiplicité des branches et des échelles de masse, qui sont bien représentées dans les données de préentraînement augmentées par le BSM mais mal représentées dans les données uniquement QCD.

Les exposants de mise à l'échelle sont extraits en ajustant des lois de puissance à la taille de modèle optimale pour le calcul ($N^*$) et à la taille du jeu de données ($D^*$) en fonction du calcul total ($C$). L'étude compare ces exposants entre l'entraînement "à partir de zéro" (sans préentraînement) et les diverses configurations de préentraînement.

Résultats clés
L'étude démontre que la composition des données de préentraînement modifie de manière significative les exposants de mise à l'échelle optimaux pour le calcul :

• Entraînement à partir de zéro : L'entraînement à partir de zéro produit des exposants de $a \approx 0,52$ (taille du modèle) et $b \approx 0,48$ (taille du jeu de données), indiquant une allocation relativement équilibrée des ressources de calcul entre la taille du modèle et les données, cohérente avec les observations dans les LLM.
• Préentraînement uniquement QCD : Le préentraînement uniquement sur des jets QCD entraîne un décalage marginal ($a \approx 0,53, b \approx 0,47$), suggérant que le préentraînement seul, sans alignement ou diversité spécifique, ne change pas fondamentalement le régime de mise à l'échelle.
• Préentraînement augmenté par le BSM : L'inclusion de désintégrations de résonances BSM dans le corpus de préentraînement provoque un décalage spectaculaire. Avec le jeu de données complet augmenté par le BSM, les exposants passent à $a \approx 0,22$ et $b \approx 0,78$.
  • Cela indique un régime où la stratégie de calcul optimale favorise lourdement l'augmentation de la taille du jeu de données par rapport à l'augmentation de la taille du modèle.
  • Ce décalage représente une réduction d'un facteur de 2,3 de l'exposant de mise à l'échelle pour la taille du modèle par rapport à la référence "à partir de zéro".
  • Les courbes de perte de l'ajustement fin confirment que le préentraînement enrichi par le BSM réduit systématiquement la perte pour toutes les tailles de modèles, avec des bénéfices croissants pour les modèles plus grands.

Principales contributions

1. Ingénierie des lois d'échelle : L'article fournit la première étude systématique montrant que la composition des données de préentraînement peut être ingéniérée pour déplacer les exposants de mise à l'échelle en physique fondamentale. Il démontre que la diversité et l'alignement en aval dans le corpus de préentraînement peuvent déplacer le régime de mise à l'échelle optimal vers des stratégies privilégiant les données.
2. Décalage quantitatif : Le travail quantifie le passage d'un régime de mise à l'échelle équilibré ($a \approx b \approx 0,5$) à un régime fortement axé sur les données ($a \approx 0,22, b \approx 0,78$) en incorporant la physique BSM dans le préentraînement.
3. Implications pour les modèles de fondation : Les résultats suggent que les modèles de fondation préentraînés sur des données synthétiques diverses et alignées peuvent atteindre des performances optimales avec des nombres de paramètres plus faibles, permettant de rediriger les budgets de calcul économisés vers la génération de données d'entraînement supplémentaires.

Signification et affirmations
Les auteurs affirment que ce travail identifie un nouvel espace de conception pour l'apprentissage automatique scientifique : les entrées physiques pour l'entraînement des modèles de fondation. Contra à les domaines du langage naturel ou de l'image où la curation des données est limitée par la disponibilité, la physique fondamentale peut exploiter des simulateurs peu coûteux et de haute fidélité pour construire des corpus de préentraînement qui façonnent activement les lois d'échelle.

L'article conclut modestement que, bien que le préentraînement sur des corpus bien composés permette un régime où le calcul en aval est mieux dépensé en plus de données plutôt qu'en modèles plus grands, des travaux supplémentaires sont nécessaires pour vérifier si ces résultats se généralisent à différentes tâches d'ajustement fin, à des échelles de modèles plus grandes et à des tailles de jeux de données différentes. L'étude ne prétend pas avoir résolu tous les défis de mise à l'échelle, mais souligne l'ingénierie de la composition du préentraînement comme un levier sous-exploré pour maximiser le potentiel de découverte des modèles de fondation scientifiques.