On The Finetuning of MLIPs Through the Lens of Iterated Maps With BPTT

Résumé technique : Ajustement fin des MLIP par le prisme des applications itérées avec BPTT

Énoncé du problème
La relaxation structurelle précise — le processus consistant à trouver les configurations atomiques correspondant aux minima locaux sur la surface d'énergie potentielle (PES) — constitue un goulot d'étranglement en science des matériaux computationnelle. Les méthodes traditionnelles reposent sur la théorie de la fonctionnelle de la densité (DFT) pour calculer les forces interatomiques, ce qui est coûteux en calcul et présente une montée en charge abrupte avec la taille du système. Les potentiels interatomiques fondés sur l'apprentissage automatique (MLIP) sont apparus comme des substituts efficaces pour approximer les forces de la DFT, étant typiquement utilisés dans des boucles d'optimisation itératives pour émuler la relaxation. Cependant, un défi fondamental dans le développement des MLIP est la rareté des données ; générer de nouveaux exemples d'entraînement nécessite des calculs de premier principe coûteux. Par conséquent, le simple agrandissement des jeux de données est souvent impraticable. De plus, l'entraînement conventionnel des MLIP optimise l'exactitude de la force par étape de manière indépendante, ignorant comment les erreurs s'accumulent au cours de la trajectoire de relaxation, ce qui conduit souvent à des écarts significatifs dans les structures finales prédites.

Méthodologie
Les auteurs proposent un cadre d'ajustement fin (fine-tuning) qui traite la relaxation structurelle comme une boucle de simulation entièrement différentiable et de bout en bout. Au lieu d'entraîner les MLIP uniquement sur des paires structure-force statiques, la méthode déroule des trajectoires de relaxation complètes et applique la rétropropagation à travers le temps (BPTT - Backpropagation Through Time).

Les composantes clés de la méthodologie incluent :

• Entraînement au niveau de la trajectoire : Le processus de relaxation est modélisé comme une séquence de « cadres » (frames), où chaque cadre consiste en une prédiction de force par le MLIP suivie d'une étape de mise à jour structurelle. L'ensemble de la trajectoire est déroulé et les gradients sont suivis à travers la séquence pour mettre à jour les paramètres du modèle en fonction de la qualité de la structure relaxée finale, plutôt que des erreurs de force intermédiaires.
• Fonction de perte : L'objectif d'optimisation est le « Delta Q » ($D_q$), une métrique de déplacement pondérée par la masse entre la structure finale prédite et la structure relaxée de référence (ground-truth). Cette métrique est préférée à l'erreur quadratique moyenne (MSE) dans les cas de défauts afin d'éviter de suraccentuer les erreurs de réseau de la structure globale (bulk).
• Applications itérées et fonctions de substitution : Les auteurs interprètent l'étape de relaxation comme une application itérée. La procédure BPTT affine le MLIP pour qu'il agisse comme une fonction de substitution (proxy function) capable d'approximer la dynamique de contraction de la PES, apprenant à préserver les emplacements des points fixes (structures stables) et leurs bassins d'attraction, même si l'exactitude locale de la force est légèrement compromise.
• Contrôle de la taille du pas : L'étude examine si la taille du pas ($\eta$) dans la descente de gradient doit être fixe, apprise en tant que scalaire ou prédite par un réseau de neurones. Les expériences indiquent qu'une taille de pas fixe ou apprise sous forme de scalaire est suffisante, et que les principaux gains de performance proviennent de la modification des poids du MLIP lui-même pour l'aligner sur la procédure de descente.

Contributions clés

1. Cadre d'ajustement fin basé sur BPTT : Introduction d'une méthode d'ajustement fin par trajectoire complète pour les MLIP pré-entraînés qui optimise directement le résultat du processus de relaxation.
2. Ablation et analyse : Analyse complète des composants d'optimisation au niveau de la PES, démontrant que la méthode est robuste aux variations des hyperparamètres et des modifications procédurales (ex: initialisation de la taille du pas, longueur de la trajectoire).
3. Connexion théorique : Liaison de l'entraînement par BPTT à la théorie des applications itérées et des fonctions de substitution, suggérant que la méthode apprend une contraction simplifiée des dynamiques réelles pilotées par la DFT, adaptée à des variétés structurelles spécifiques.
4. Validation de la généralisabilité : Validation à travers de multiples domaines structurels (défauts du silicium, cristaux purs, catalyseurs) et architectures (ADAPT, ResMLP), montrant des améliorations de performance constantes.

Résultats
La méthode proposée améliore systématiquement la précision des structures relaxées pour tous les modèles pré-entraînés évalués :

• Gains de performance : L'approche produit une réduction moyenne d'environ 32 % de l'erreur de prédiction ($D_q$) sur les jeux de données. Dans certains cas spécifiques, comme les défauts du silicium, la réduction de l'erreur atteint environ 50 % par rapport aux bases non ajustées.
• Précision paradoxale : Un résultat notable est que l'ajustement fin par BPTT dégrade souvent l'exactitude brute de la prédiction de force (les erreurs de force L2 augmentent) tout en améliant simultanément l'exactitude structurelle finale. Cela suggère que le modèle apprend un biais structurel qui donne la priorité à l'état final correct plutôt qu'à la fidélité de la force locale.
• Robustesse : La méthode produit des résultats très similaires malgré diverses configurations d'hyperparamètres et est robuste aux initialisations non optimales de la taille du pas.
• Indépendance de l'architecture : Des améliorations ont été observées tant dans l'architecture ADAPT (basée sur les Transformers, sans graphe) que dans ResMLP, indiquant que la stratégie n'est pas limitée à un type de modèle spécifique.

Signification et affirmations
L'article affirme que cette approche offre une solution pragmatique au goulot d'étranglement de la rareté des données dans le développement des MLIP. En extrayant plus de valeur des données existantes grâce à une supervision au niveau de la trajectoire, elle permet de créer des MLIP hautement efficaces et spécifiques à un domaine sans nécessiter de nouvelles données coûteuses de premier principe.

Les auteurs positionnent le BPTT non pas comme une méthode pour « résoudre la physique » ou récupérer des dynamiques physiques universelles, mais comme une étape finale d'un pipeline d'entraînement par étapes. Elle affine un MLIP largement applicable et pré-entraîné pour qu'il fonctionne de manière fiable sur des classes structurelles spécifiques en apprenant une application de contraction qui dirige les trajectoires vers les états métastables corrects. Cela est particulièrement précieux pour les flux de travail à haut débit où l'amélioration de la fidélité de la relaxation réduit le besoin d'évaluations DFT coûteuses. Ce travail établit un parallèle avec l'apprentissage par renforcement à partir de la rétroaction humaine (RLHF), où des objectifs au niveau de la séquence améliorent le comportement en aval sans nécessairement minimiser la perte d'entraînement au niveau du jeton (token).