Weight Decay Improves Language Model Plasticity

Résumé Technique : Le déclin du poids améliore la plasticité des modèles de langage

Énoncé du problème
Les pratiques actuelles pour l'entraînement des grands modèles de langage (LLM) découpent généralement le pré-entraînement du post-entraînement (ajustement fin, alignement). L'optimisation des hyperparamètres et les lois d'échelle lors du pré-entraînement sont principalement évaluées sur la base de la perte de validation du modèle de base (entropie croisée), selon l'hypothèse qu'une perte de pré-entraînement plus faible est corrélée à un modèle plus capable en aval. Cependant, cette approche néglige la plasticité du modèle — la capacité d'un modèle entraîné à s'adapter avec succès à de nouvelles données et tâches lors d'un entraînement ultérieur sans réinitialisation. L'article examine si l'optimisation des hyperparamètres de pré-entraînement uniquement pour la performance de pré-entraînement compromet par inadvertance l'adaptabilité du modèle lors de l'ajustement fin ultérieur. Plus précisément, les auteurs se concentrent sur le déclin du poids (weight decay), un hyperparamètre de régularisation standard, pour déterminer si sa valeur optimale pour minimiser la perte de pré-entraînement diffère de sa valeur optimale pour maximiser l'adaptabilité en aval.

Méthodologie
Les auteurs mènent une étude empirique systématique sur deux familles de modèles (Llama-2 et OLMo-2) et plusieurs tailles de modèles (jusqu'à 4 milliards de paramètres). Le design expérimental fait varier l'hyperparamètre de déclin du poids ($\lambda$) pendant le pré-entraînement tout en contrôlant les autres facteurs.

1. Régimes d'entraînement : Les modèles sont entraînés dans deux régimes distincts :
   • Optimal en calcul : 20 jetons par paramètre (TPP).
   • Sur-entraîné : 140 TPP.
2. Évaluation par ajustement fin : Les modèles pré-entraînés sont soumis à un ajustement fin sur une suite diversifiée de tâches en aval, incluant :
   • Six jeux de données de raisonnement par chaîne de pensée (CoT) (ex: MetaMathQA, MMLUProCoT).
   • Cinq tests de compréhension du langage et de raisonnement de bon sens (ex: HellaSwag, WinoGrande).
   • Une tâche d'alignement de sécurité.
3. Métriques : La performance en aval est évaluée à l'aide de multiples métriques, notamment Greedy (Pass@1), Pass@16, Majority Vote (Maj@16), Reward Model scoring (RM@16) et Correct Ratio.
4. Analyse mécaniste : Pour comprendre pourquoi le déclin du poids affecte la plasticité, les auteurs analysent :
   • Le sondage linéaire (Linear Probing) : Mesure de la séparabilité linéaire des représentations internes (classification de sentiment et de sujet).
   • Le rang de l'attention (Attention Rank) : Analyse du pseudo-rang des matrices de projection requête-clé ($W_{QK}$) et de valeur ($W_{VP}$).
   • Le surapprentissage (Overfitting) : Calcul de l'écart entre la perte d'entraînement et la perte de validation (Écart Train-Val).

Contributions clés

1. Le déclin du poids comme moteur de plasticité : L'article démontre que le déclin du poids est un facteur critique façonnant la plasticité du modèle. Contrairement aux valeurs par défaut standards, des valeurs de déclin du poids plus élevées pendant le pré-entraînement conduisent systématiquement à de meilleures performances en aval après l'ajustement fin.
2. Découplage des objectifs : L'étude fournit la preuve que la valeur de l'hyperparamètre minimisant la perte de validation de pré-entraînement n'est pas nécessairement la valeur qui maximise la performance en aval. Dans plusieurs régimes, les modèles ayant une perte de pré-entraînement plus élevée (due à un déclin du poids plus important) surpassent les modèles ayant une perte de pré-entraînement plus faible après l'ajustement fin.
3. Aperçus mécanistes : Les auteurs identifient trois mécanismes spécifiques par lesquels le déclin du poids améliore la plasticité :
   • Il favorise des représentations linéairement séparables, rendant les caractéristiques internes plus accessibles pour les tâches en aval.
   • Il régularise les matrices d'attention, réduisant spécifiquement le rang du produit requête-clé ($W_{QK}$), ce qui conduit à des structures plus compressées et robustes.
   • Il réduit le surapprentissage sur les données de pré-entraînement, permettant ainsi d'« oublier » efficacement le bruit spécifique de l'entraînement et de préserver la capacité du modèle à apprendre de nouvelles informations.

Résultats clés

• Décalage de l'optimal du déclin du poids : Pour un entraînement optimal en calcul (20 TPP), le déclin du poids optimal pour la performance en aval est nettement plus élevé que la valeur par défaut standard de 0,1. Par exemple, des valeurs optimales ont été trouvées à 0,5–1,0 pour les modèles Llama-2 et OLMo-2 dans ce régime. Dans le régime sur-entraîné (140 TPP), la valeur optimale pour la performance en aval (0,3) reste supérieure à la valeur qui minimise la perte de pré-entraînement (0,1).
• Compromis contre-intuitifs : L'article souligne un compromis où de « moins bons » modèles de base (perte de pré-entraînement plus élevée) peuvent devenir de « meilleurs » modèles finaux après l'ajustement fin. Par exemple, un modèle OLMo-2-1B-140x pré-entraîné avec $\lambda=0,3$ ou $1,0$ présentait une perte de pré-entraînement plus élevée que la référence $\lambda=0,1$, mais a obtenu une précision en aval nettement meilleure.
• Robustesse : Ces conclusions se vérifient à travers différentes familles de modèles, tailles, et même lorsque d'autres hyperparamètres (taux d'apprentissage, taille de lot) sont variés pendant le pré-entraînement et l'ajustement fin.
• Corrélations mécanistes :
  • Un déclin du poids plus élevé est corrélé à une précision de sondage linéaire plus élevée, indiquant des représentations plus structurées.
  • Un déclin du poids plus élevé est corrélé à un pseudo-rang plus faible dans les matrices d'attention (particulièrement $W_{QK}$), suggérant une transition vers des hypothèses de bas rang et robustes.
  • Un déclin du poids plus élevé est corrélé à un écart Train-Val plus faible, indiquant un surapprentissage réduit.

Signification et affirmations
L'article soutient que le domaine doit reconsidérer l'reliance exclusive sur la perte de validation du pré-entraînement comme unique métrique d'optimisation des hyperparamètres. Les auteurs affirment que la plasticité du modèle est une propriété distincte et cruciale qui est souvent sacrifiée lorsque l'on optimise purement pour la perte de pré-entraînement.

La portée de ce travail réside dans :

1. Réévaluation des standards : Il suggère que les choix d'hyperparamètres standards (ex: $\lambda=0,1$) peuvent être sous-optimaux pour le cycle de vie complet du développement des LLM, particulièrement quand l'adaptabilité en aval est l'objectif ultime.
2. Limites de l'entropie croisée : Il souligne que la perte d'entropie croisée est un proxy imparfait pour la capacité en aval, car des modèles ayant des pertes similaires peuvent présenter des plasticités très différentes.
3. Rôle multifacette des hyperparamètres : Il souligne qu'un hyperparamètre unique tel que le déclin du poids joue un rôle complexe et multifacette dans le façonnage du comportement du modèle, influençant non seulement la généralisation mais aussi les propriétés structurelles (linéarité, rang) qui facilitent l'apprentissage futur.

Les auteurs restent modestes, notant que bien que leurs résultats soient robustes à travers les échelles testées, les valeurs optimales spécifiques peuvent varier dans des scénarios de très grande échelle ou fortement sur-entraînés au-delà de leur portée expérimentale actuelle. Ils appellent à des travaux futurs pour approfondir l'étude du compromis stabilité-plasticité et l'applicabilité de ces découvertes à d'autres modalités de modèles de fondation.