Fine-tuning MLIP foundation models: strategies for accuracy and transferability

Cet article évalue sept stratégies de fine-tuning pour les modèles de fondation de potentiels interatomiques appris par apprentissage automatique (MLIP) à travers divers benchmarks chimiques, révélant que si des prérequis tels que la qualité du modèle de fondation et une initialisation correcte de l'énergie sont primordiaux, le fine-tuning naïf est optimal pour la précision sur un système unique, tandis que le replay multi-têtes préserve de manière unique la robustesse hors distribution pour un déploiement plus large.

L'essentiel

Imaginez que vous avez un chef étoilé qui a passé des années à apprendre à cuisiner des plats parfaits en utilisant uniquement des ingrédients inorganiques comme des roches, des métaux et des sels. Ce chef est un « Modèle de Fondations ». Maintenant, vous voulez que ce chef cuisine un nouveau plat spécifique, comme une soupe organique délicate ou un ragoût biologique, en utilisant une très petite quantité de nouvelles recettes.

La grande question est la suivante : Comment enseigner ce nouveau plat au chef sans qu'il n'oublie ses anciens plats, ou sans gâcher ses compétences existantes ?

Ce document est une expérience culinaire massive testant sept façons différentes de « fine-tuner » (réentraîner) ce chef maître. Les chercheurs ont découvert que la méthode d'enseignement importe moins que trois étapes cruciales de « pré-cuisson » : choisir le bon chef, définir la bonne base et régler la chaleur.

Voici le détail de leurs découvertes en termes simples :

1. Les trois vérifications « Pré-vol » (La partie la plus importante)

Avant même de commencer à enseigner la nouvelle recette, vous devez faire trois choses correctement. Si vous échou-ez sur ces points, aucune méthode d'enseignement ne vous sauvera.

• Choisir le bon chef (Qualité du Modèle de Fondation) :

  • L'analogie : Vous n'engageriez pas un chef qui ne sait que faire bouillir de l'eau pour lui apprendre à cuisiner un soufflé.
  • Le constat : La qualité du modèle original est plus importante que la stratégie de fine-tuning. Un modèle entraîné sur un ensemble de données vaste et diversifié de matériaux inorganiques (comme le modèle « OMat24 ») est bien meilleur pour apprendre de nouvelles chimies étranges qu'un modèle plus ancien et plus petit. Même si vous utilisez la même méthode d'enseignement, un « meilleur » modèle de fondation produira toujours un meilleur plat final.

• Définir le point zéro (Énergie de référence atomique / $E_0$) :

  • L'analogie : Imaginez mesurer la hauteur d'un bâtiment. Si vous commencez la mesure depuis le sous-sol plutôt que depuis le rez-de-chaussée, vos chiffres seront faux, et le bâtiment pourrait sembler flotter ou être enterré. En chimie, vous devez soustraire le « poids » des atomes individuels pour que le modèle n'apprenne que sur la façon dont ils interagissent.
  • Le constat : Les chercheurs ont découvert qu'utiliser une manière intelligente et « consciente du modèle » pour définir ce point zéro est crucial. Si vous utilisez une supposition moyenne et paresseuse, le modèle devient instable. Il peut paraître bon sur le papier (scores d'erreur faibles), mais il s'effondrera lorsque vous essaierez de simuler la physique réelle (comme un bâtiment s'écroulant lors d'un test de soufflerie).

• Baisser la chaleur (Hyperparamètres) :

  • L'analogie : Lors de l'apprentissage d'une nouvelle compétence, vous ne voulez pas aller trop vite pour ne pas trébucher, mais vous ne voulez pas non plus aller trop lentement pour ne jamais finir.
  • Le constat : Différentes méthodes d'enseignement nécessitent différents « taux d'apprentissage ». Par exemple, une méthode appelée LoRA (qui ne modifie qu'une infime partie du modèle) peut supporter un taux d'apprentissage très rapide, tandis qu'une méthode qui enseigne deux choses à la fois nécessite un rythme beaucoup plus lent et doux.

2. Les sept stratégies d'enseignement

Une fois ces trois vérifications passées, les chercheurs ont testé sept façons d'enseigner la nouvelle recette :

1. Fine-Tuning Naïf : « Continuez simplement à cuisiner. » Vous prenez tout le chef et continuez à l'entraîner sur les nouvelles données.
   • Résultat : Excellent pour apprendre un plat spécifique parfaitement. Mais si vous essayez d'utiliser ce chef pour un autre type de nourriture plus tard, il pourrait avoir oublié ses anciennes compétences (un problème appelé « oubli catastrophique »).
2. Gel des Couches (Layer Freezing) : « Ne touchez pas aux bases. » Vous verrouillez les connaissances du chef sur les techniques de base au couteau et ne le laissez apprendre que la nouvelle sauce.
   • Résultat : Bon, mais parfois trop rigide. Cela limite la capacité du chef à s'adapter aux nouveaux ingrédients.
3. LoRA (Low-Rank Adaptation) : « Ajoutez un aide-mémoire. » Au lieu de réécrire tout le livre de cuisine, vous ajoutez un petit carnet de notes efficace sur le tablier du chef qui ne couvre que les nouvelles règles.
   • Résultat : Très efficace et précis pour des tâches spécifiques, similaire au fine-tuning naïf.
4. Multihead Replay : « Le Chef à deux chapeaux. » Vous donnez deux chapeaux au chef. Un chapeau pour le nouveau plat, et l'autre pour les plats anciens et familiers. Il pratique les deux en même temps.
   • Résultat : C'est le vainqueur pour la sécurité. C'est la seule méthode qui empêche systématiquement le chef d'oublier ses anciennes compétences. Elle permet au chef d'être bon pour le nouveau plat et pour les anciens.
5. Pseudolabel Replay : « Le Chef Synthétique. » Au lieu d'utiliser de vraies anciennes recettes, vous utilisez les propres prédictions du chef sur les anciennes recettes pour s'entraîner.
   • Résultat : Fonctionne bien et est flexible car vous n'avez pas besoin des anciennes données réelles, juste de la mémoire du chef.
6. Replay + LoRA : Combiner l'aide-mémoire avec les deux chapeaux.
   • Résultat : Bon, mais le « Dual Head » seul était souvent suffisant.

3. Les grandes conclusions

• Ne réinventez pas la roue : Si vous avez besoin d'un modèle pour une tâche spécifique et étroite (comme simuler uniquement de l'eau salée), le Fine-Tuning Naïf est le moyen le plus rapide et le plus simple d'obtenir un excellent résultat.
• N'oubliez pas le passé : Si vous avez besoin d'un modèle capable de gérer des situations nouvelles et étranges (comme un nouveau type de batterie ou une molécule biologique complexe) sans oublier son entraînement d'origine, vous devez utiliser le Multihead Replay. C'est la seule stratégie qui a maintenu le modèle robuste et protégé contre « l'oubli ».
• La Qualité prime sur les Astuces : Le document souligne que passer du temps à choisir un modèle de fondation de haute qualité et à définir correctement les références d'énergie est plus important que de choisir l'algorithme de fine-tuning parfait. Si la fondation est faible ou si les mathématiques sont mal configurées, la meilleure stratégie d'enseignement au monde ne servira à rien.

En bref : Pour obtenir le meilleur de l'IA en chimie, commencez par une fondation intelligente, définissez vos règles mathématiques correctement, et si vous voulez que l'IA soit polyvalente et ne soit pas sujette à l'oubli, enseignez-lui en utilisant la méthode des « deux chapeaux » (Multihead Replay).

Résumé technique

Résumé technique : Ajustement fin (fine-tuning) des modèles de fondation MLIP

Énoncé du problème
Les modèles de fondation de potentiels interatomiques appris par apprentissage automatique (MLIP) ont démontré leur capacité à se transférer à travers divers systèmes chimiques, offrant un flux de travail qui évite le processus gourmand en ressources consistant à entraîner des potentiels spécifiques à une tâche à partir de zéro. Cependant, la communauté manque de directives systématiques sur comment et quand ajuster ces modèles. Les premiers rapports suggéraient qu'un ajustement fin naïf conduit souvent à un « oubli catastrophique », incitant l'adoption de techniques contraintes (par exemple, le gel des couches, l'adaptation de bas rang ou LoRA) initialement développées pour les grands modèles de langage. Cet article examine si ces contraintes sont nécessaires ou si les échecs précoces étaient dus à d'autres facteurs, tels que des modèles de fondation plus faibles, une initialisation inappropriée de l'énergie de référence atomique ($E_0$) ou des procédures d'entraînement instables. L'étude vise à caractériser les principaux facteurs façonnant les résultats de l'ajustement fin, spécifiquement la précision sur la tâche cible et la robustesse hors distribution (OOD).

Méthodologie
Les auteurs évaluent sept stratégies d'ajustement fin distinctes à travers cinq benchmarks chimiquement divers, trois générations de modèles de fondation et des ensembles d'entraînement couvrant cinq ordres de grandeur en taille.

• Stratégies d'ajustement fin évaluées :

  1. Naïve : Mises à jour complètes des paramètres via une descente de gradient continue.
  2. Gel des couches (Variantes) : Gel des couches d'embedding/de passage de messages tout en entraînant les lectures (readouts) ; ou gel de l'embedding et de la première couche de passage de messages.
  3. Adaptation de bas rang (LoRA) : Injection de décompositions de bas rang entraînables dans les couches linéaires scalaires et équivariantes tout en gelant les poids préentraînés.
  4. Rejeu multi-têtes (Multihead Replay) : Optimisation simultanée sur les données cibles et un ensemble de données de rejeu (provenant du préentraînement ou étiqueté par pseudo-étiquetage) en utilisant des têtes de lecture distinctes.
  5. Rejeu par pseudo-étiquetage (Pseudolabel Replay) : Une variante du rejeu multi-têtes où les étiquettes de rejeu sont générées par le modèle de fondation lui-même, découplant la source de rejeu du corpus de préentraînement original.
  6. Rejeu + LoRA : Combinaison du rejeu multi-têtes avec LoRA.

• Benchmarks : L'étude couvre des systèmes avec un écart croissant par rapport au domaine de préentraînement OMat24 (solide inorganique périodique) :

  • Électrolytes d'argyrodite de lithium (solide périodique inorganique).
  • NaCl aqueux (solution ionique).
  • Polymorphes de glace (solide moléculaire).
  • Réactions SN2 (chimie réactive en phase gazeuse).
  • Biomolécules SPICE (conformères organiques/biomoléculaires).

• Implémentations techniques : Les auteurs ont implémenté trois nouvelles capacités dans le code de MACE :

  • LoRA adapté aux architectures de passage de messages équivariantes (couvrant à la fois les couches linéaires scalaires et équivariantes).
  • Pseudo-étiquetage pour le rejeu afin de découpler les sources de données de rejeu.
  • Réestimation du modèle d'énergie de référence atomique ($E_0$) consciente du modèle pour aligner les bases préentraînées avec les données cibles.

• Métriques d'évaluation : Au-delà des erreurs standard de point sur l'énergie et les forces, l'étude sonde les comportements dynamiques et extrapolatifs, incluant les fonctions de distribution radiale (RDF) issues de la dynamique moléculaire (MD), les profils de réaction NEB (Nudged Elastic Band), les tests de stabilité MD et la recherche de structures aléatoires (RSS) pour détecter les échecs de répulsion à courte portée.

Résultats clés

1. Les prérequis dominent le choix de la stratégie : L'étude conclut que la qualité du modèle de fondation, une initialisation correcte de $E_0$ et des hyperparamètres bien choisis sont des prérequis dont l'impact dépasse systématiquement celui de la stratégie d'ajustement fin spécifique.

   • Qualité de la fondation : Les nouveaux modèles de fondation (ex: basés sur OMat24) surpassent systématiquement les anciens (basés sur MPTraj) en termes de transfert OOD, même avec des recettes d'ajustement fin fixes.
   • Initialisation de $E_0$ : L'utilisation de valeurs $E_0$ « moyennées » entraîne des erreurs nettement plus élevées et une instabilité de la MD (ex: les modèles de glace échouant en moins de 50 ps). Les $E_0$ « réestimés » (alignant le zéro de référence du modèle préentraîné sur les données cibles) sont critiques pour la stabilité et la transférabilité, offrant souvent de meilleurs résultats que le choix de l'algorithme d'ajustement fin lui-même.
   • Hyperparamètres : L'ajustement fin naïf nécessite des taux d'apprentissage réduits et une décroissance EMA accrue. LoRA tolère des taux d'apprentissage plus élevés. Le rejeu multi-têtes nécessite des taux d'apprentissage substantiellement plus bas pour éviter des signaux de mise à jour concurrents. Le décrochage de poids (weight decay) doit être réglé à zéro pour éviter d'éloigner les paramètres de la solution préentraînée.

2. Performance par objectif :

   • Spécialisation en distribution (Système unique) : Pour les tâches étroites (ex: barrières SN2, solvatation NaCl aqueux), la plupart des stratégies (Naïve, LoRA, Multihead) atteignent une grande précision, surpassant systématiquement les modèles entraînés à partir de zéro. L'ajustement fin naïf offre la meilleure convergence pour les applications à système unique.
   • Robustesse hors distribution (OOD) : Lors de l'évaluation du transfert vers des compositions connexes mais non vues (ex: électrolytes non-argyrodite) ou des chimies différentes (ex: biomolécules), le Rejeu multi-têtes (avec des données originales ou pseudo-étiquetées) est la seule approche qui préserve systématiquement la robustesse OOD. Il maintient la précision sur la distribution de préentraînement tout en apprenant la tâche cible, empêchant efficacement l'oubli catastrophique.
   • Gel et LoRA : Bien qu'efficaces pour l'efficacité des paramètres, le gel des couches et LoRA ont montré des limites pour s'adapter aux caractéristiques de solvatation ou pour maintenir une large robustesse chimique par rapport au rejeu multi-têtes dans les scénarios testés.

Signification et affirmations
L'article affirme que la fragilité perçue de l'ajustement fin naïf dans les MLIP est largement le résultat d'une configuration sous-optimale plutôt que d'une limitation intrinsèque de la méthode. Les auteurs soutiennent que :

• L'ajustement fin naïf est une option viable et souvent supérieure pour les applications à système unique, à condition que le modèle de fondation soit de haute qualité et que les $E_0$ soient correctement réestimés.
• Le rejeu multi-têtes est la stratégie nécessaire pour un déploiement plus large où la préservation du comportement du modèle de fondation en dehors de la distribution d'ajustement fin est requise.
• Le rejeu par pseudo-étiquetage offre un avantage pratique en permettant l'utilisation de n'importe quel ensemble de données structurellement diversifié pour le rejeu, supprimant la dépendance vis-à-vis de l'accès au corpus de préentraînement original.

Ce travail établit que pour les praticiens, investir dans le modèle de fondation le plus performant et assurer l'alignement correct de l'énergie de référence atomique sont des choix de conception plus critiques que la sélection d'un algorithme d'ajustement fin contraint spécifique. L'étude fournit un cadre systématique pour déployer les modèles de fondation MLIP, faisant de l'ajustement fin non plus une option de niche, mais un point de départ par défaut pour le développement spécifique à un système.