TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces est un cadre à trois flux agnostique du modèle qui combine l'apprentissage auto-supervisé avec des représentations séparées de la composition et de la structure pour améliorer significativement la transférabilité et l'efficacité des données des réseaux de neurones graphiques atomistiques pour les potentiels interatomiques d'apprentissage automatique.

L'essentiel

Imaginez que vous essayiez d'enseigner à un chef robot comment cuisiner.

Le Problème : Le Chef « Taille Unique »
Actuellement, les scientifiques utilisent des modèles d'intelligence artificielle puissants (appelés MLIP) pour prédire le comportement des atomes, comme l'énergie d'un matériau ou la difficulté à déplacer les atomes. Ces modèles sont entraînés sur d'énormes quantités de données provenant de superordinateurs (DFT).

Cependant, ces modèles présentent un défaut : ils sont comme un chef qui a mémorisé le goût exact d'un plat spécifique, mais a oublié pourquoi il avait ce goût. Si vous leur demandez de cuisiner un plat légèrement différent (un nouveau type de matériau), ils peinent. Ils confondent les ingrédients (composition) avec la forme de la casserole (structure). Si vous changez les ingrédients, ils s'embrouillent sur la forme, et vice versa. Cela les rend mauvais pour apprendre rapidement de nouvelles tâches, surtout lorsque vous n'avez pas beaucoup de données pour les enseigner.

La Solution : TriForces (La Cuisine à Trois Flux)
Les auteurs présentent TriForces, une nouvelle façon de construire ces chefs IA. Au lieu d'un seul cerveau géant essayant de tout mémoriser à la fois, ils divisent le cerveau en trois « flux » ou départements spécialisés :

1. Le Flux Ingrédients (Composition) : Ce département ne regarde que ce qui se trouve dans la casserole (par exemple, « Nous avons 2 Hydrogènes et 1 Oxygène »). Il ignore totalement la forme. Il apprend la chimie.
2. Le Flux Forme (Structure) : Ce département ne regarde que comment les atomes sont arrangés dans l'espace (par exemple, « Ils sont en triangle »). Il ignore ce que sont réellement les atomes. Il apprend la géométrie.
3. Le Flux Interaction : C'est le chef principal qui prend les notes des départements Ingrédients et Forme et les combine pour prédire le résultat final (énergie ou force).

La Sauce Secrète : L'Apprentissage Auto-Supervisé
Avant que le modèle ne soit jamais invité à prédire une propriété spécifique, les auteurs l'entraînent en utilisant un jeu appelé « Apprentissage Auto-Supervisé ». Imaginez cela comme une séance d'entraînement où l'IA doit :

• Dénouer : Regarder une image légèrement brisée ou bruitée d'une molécule et la réparer.
• Masquage : Cacher un ingrédient et deviner ce qu'il était en se basant sur les voisins.
• Appariement : Regarder deux versions légèrement différentes de la même molécule et réaliser qu'elles sont la même chose.

Cet entraînement force l'IA à organiser ses connaissances proprement. Elle apprend que les « ingrédients » appartiennent à un dossier et les « formes » à un autre, plutôt que de les mélanger.

Pourquoi Cela Compte (Les Résultats)
L'article montre que cette nouvelle cuisine « Trois Flux » fonctionne bien mieux que les anciennes cuisines « Un Cerveau » :

• Apprentissage Plus Rapide : Lorsqu'on lui donne une petite quantité de nouvelles données (comme 20 000 exemples au lieu de millions), TriForces apprend beaucoup plus vite et fait moins d'erreurs. C'est comme un chef qui peut apprendre une nouvelle recette après l'avoir goûtée une fois, plutôt que d'avoir besoin de la cuisiner mille fois.
• Meilleure Mémoire : L'IA n'oublie pas ce qu'elle a appris. Elle peut transférer ses connaissances d'un type de matériau à un autre sans s'embrouiller.
• Connaissances Recherchables : Parce que l'IA maintient les « ingrédients » et les « formes » séparés, vous pouvez lui demander de trouver des matériaux qui se ressemblent mais ont des ingrédients différents, ou des matériaux avec les mêmes ingrédients mais des formes différentes. Les anciens modèles ne pouvaient pas faire cela car leurs connaissances étaient trop mélangées.

En Résumé
TriForces est un cadre qui décompose le travail complexe de compréhension des atomes en trois tâches plus simples : connaître les ingrédients, connaître la forme et connaître comment ils fonctionnent ensemble. En entraînant l'IA à maintenir ces tâches séparées et en pratiquant avec des « jeux de devinettes » (apprentissage auto-supervisé), le modèle devient un outil beaucoup plus flexible, efficace et précis pour découvrir de nouveaux matériaux.

Les auteurs ont publié leur code et leurs modèles pré-entraînés afin que d'autres scientifiques puissent utiliser cette « cuisine à trois flux » pour construire une meilleure IA pour la science des matériaux.

Résumé technique

Résumé Technique : TriForces

Énoncé du Problème

Les Potentiels Interatomiques basés sur l'Apprentissage Automatique (MLIPs) fondés sur des Réseaux de Neurones à Graphes Géométriques (GNN) ont atteint une haute précision dans la prédiction des propriétés atomistiques lorsqu'ils sont entraînés sur des ensembles de données à grande échelle issus de la Théorie de la Fonctionnelle de la Densité (DFT). Cependant, leur utilité pratique est entravée par deux limitations principales :

1. Transférabilité Incohérente : Les MLIPs échouent souvent à se transférer efficacement vers de nouvelles chimies ou tâches, particulièrement lors du fine-tuning sur de petits ensembles de données coûteux et spécifiques à la tâche. Les modèles pré-entraînés éprouvent fréquemment des difficultés avec des tâches diagnostiques simples (par exemple, identifier les systèmes cristallins ou les éléments majoritaires) et présentent un « oubli catastrophique » ou une performance de transfert instable à travers les domaines.
2. Mauvaise Réutilisabilité des Représentations : Les MLIPs actuels optimisent les représentations pour des cibles de prédiction spécifiques (énergie et forces) plutôt que pour une réutilisation générale. Par conséquent, ces représentations entremêlent souvent la composition (chimie) et la structure (géométrie), les rendant inadaptées à l'analyse exploratoire, à la recherche de voisins les plus proches, ou à la recherche de similarité décomposée. Les objectifs supervisés standards encouragent des représentations suffisantes pour la régression mais non organisées pour préserver les informations de composition et de structure accessibles.

Méthodologie : Le Cadre TriForces

Les auteurs proposent TriForces, un cadre agnostique au modèle qui enrichit les GNN géométriques existants avec une architecture à trois flux et un apprentissage auto-supervisé (SSL) multi-objectifs. L'innovation centrale est la factorisation explicite des représentations atomistiques en trois composantes distinctes :

1. Architecture à Trois Flux

Au lieu d'un seul vecteur latent, TriForces décompose la représentation au niveau des nœuds $h_i$ en trois flux concaténés :

• Flux de Composition ($h^{comp}$) : Encode l'information chimique sans coordonnées. Il traite l'ensemble des éléments atomiques uniques et leurs comptes stœchiométriques en utilisant un Transformer avec une attention pondérée par le compte. Ce flux préserve les comptes absolus d'éléments pour encoder la taille du système et l'échelle d'énergie, indépendamment de la géométrie.
• Flux de Structure ($h^{struct}$) : Encode la géométrie sans identité d'élément (agnostique au type). Il construit des descripteurs locaux invariants par rotation inspirés de SOAP (Smooth Overlap of Atomic Positions), utilisant des fonctions de base radiales, des harmoniques sphériques et des coupures multi-échelles. Ce flux capture des motifs géométriques réutilisables et des motifs topologiques, renforcés par un passage de messages invariant.
• Flux d'Interaction ($h^{int}$) : Un GNN géométrique de base standard (par exemple, MACE, eSEN, Orb-v3) qui capture le couplage entre la composition et la géométrie, préservant l'expressivité de l'architecture originale.

2. Pré-entraînement Auto-Supervisé

Pour organiser l'espace latent et améliorer la transférabilité, TriForces emploie une stratégie de pré-entraînement SSL multi-objectifs utilisant des augmentations stochastiques (bruit de position, masquage d'éléments, variation de graphe et rotations). Le cadre combine trois objectifs complémentaires :

• Non-Reconstruction (LeJEPA) : Aligne les embeddings de deux vues augmentées d'une même structure aux niveaux nœud et graphe. Cela impose l'invariance aux augmentations et organise l'espace latent global sans nécessiter de gradients arrêtés ou d'encodeurs à momentum.
• Débruitage : Entraîne le modèle à récupérer les positions atomiques propres à partir d'entrées bruitées. Cela stabilise les représentations géométriques et fournit implicitement une augmentation par rotation.
• Masquage : Prédit les types atomiques masqués basés sur la géométrie et la composition environnantes. Cela encourage le modèle à apprendre des motifs compositionnels et le contexte.

La perte finale de pré-entraînement est une somme pondérée de ces trois objectifs.

Contributions Clés

1. Décomposition Architecturale : Une conception à trois flux qui sépare explicitement la composition, la structure et l'interaction, garantissant que les deux facteurs sont préservés par conception plutôt qu'entremêlés.
2. Stratégie de Pré-entraînement Hybride : Une approche auto-supervisée combinant des objectifs basés sur la reconstruction (débruitage, masquage) avec un apprentissage de prédiction latente (LeJEPA) pour structurer l'espace d'embedding afin d'une meilleure transfert en aval.
3. Recherche Interprétable : La capacité d'effectuer une recherche de similarité ciblée dans les espaces d'embedding compositionnels, structurels ou conjoints, permettant la comparaison de matériaux basée sur des critères spécifiques (par exemple, chimie uniquement ou structure uniquement).
4. Validation Empirique : Des expériences extensives à travers plusieurs architectures (Orb-v3, eSEN, MACE) et benchmarks (OMat24, MatBench, QM9) démontrant une efficacité des données améliorée, une performance de transfert et une qualité de représentation.

Résultats

• Performance de Transfert (OMat24) : Dans des régimes à données limitées, TriForces surpasse significativement les modèles de référence. À 20K échantillons, il réduit l'Erreur Absolue Moyenne (MAE) de l'énergie de 57 % par rapport au modèle de base. Il améliore la MAE des forces pour toutes les tailles d'échantillons et réduit les erreurs de contrainte.
• Efficacité des Données : TriForces atteint des erreurs plus faibles à chaque taille d'ensemble de données (de 20K à 2M d'échantillons), avec les gains les plus significatifs observés dans les configurations à faible quantité de données.
• Performance sur les Benchmarks :
  • MatBench : Les variantes de TriForces obtiennent les meilleurs résultats globaux sur 6 tâches sur 8, surpassant à la fois les modèles de pré-entraînement auto-supervisés et ceux étiquetés par DFT. Par exemple, la MAE des phonons s'est améliorée de 57,8 à 19,5 cm$^{-1}$.
  • Découverte MatBench : TriForces eSEN-sm atteint une MAE d'énergie comparable à un modèle eSEN-30M-OAM beaucoup plus grand tout en utilisant 60 % de paramètres en moins et en s'entraînant jusqu'à 5 fois plus vite.
  • QM9 : Le pré-entraînement sur des entrées chimiques diversifiées (solides + molécules) réduit systématiquement la MAE par rapport aux modèles de référence basés uniquement sur les solides ou sans SSL.
• Qualité de Représentation : Le sondage linéaire sur des embeddings figés montre que TriForces préserve des informations fondamentales (système cristallin, élément majoritaire, nombre de coordination) que les MLIPs standards perdent. TriForces atteint une précision de 96 à 100 % pour la classification du système cristallin et de l'élément majoritaire, tandis que les modèles de référence peinent (55–73 %).
• Recherche : Le cadre permet une recherche efficace par k-NN où le flux de composition excelle dans le rappel des ensembles d'éléments et le flux de structure excelle dans le rappel des groupes d'espace, une capacité absente dans les modèles à flux unique.

Importance et Revendications

L'article positionne TriForces non pas simplement comme une méthode auto-supervisée, mais comme un cadre architectural dont les représentations sont davantage améliorées par le SSL.

• Dépendance au Régime : Les auteurs affirment que la factorisation de flux fournit les gains dominants dans les contextes supervisés à grande échelle, tandis que le SSL est le plus précieux pour le transfert à faible quantité de données, l'organisation des représentations et les tâches de recherche.
• Découplage : En séparant la composition et la structure, TriForces adresse les problèmes de « transfert fragile » et de « difficulté de réutilisation » des modèles de fondation actuels. Il permet aux modèles d'apprendre des représentations organisées pour l'analyse (par exemple, sondage, recherche) en plus de la prédiction.
• Praticité : Le cadre est agnostique au modèle et plug-in, le rendant immédiatement applicable aux architectures atomistiques existantes ou nouvelles. Les auteurs publient des points de contrôle pré-entraînés et du code pour faciliter la réutilisation dans la modélisation atomistique en aval.

Le travail suggère que les futurs modèles de fondation atomistiques devraient dépasser les architectures de prédiction à flux unique pour se diriger vers des représentations factorisées qui préservent explicitement les facteurs physiques distincts des systèmes chimiques.