OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

Ce papier présente OmniMol, un potentiel interatomique appris par machine à la pointe de l'art pour les petites molécules, qui exploite une architecture Transformer Point-Edge et un transfert de connaissances depuis la physique des hautes énergies pour obtenir des performances excellentes avec un micro-ajustement minimal et une inférence unique rapide.

L'essentiel

Imaginez deux mondes très différents : l'un est le monde chaotique et ultra-rapide de la physique des particules (où les scientifiques font entrer en collision des atomes pour observer ce qui en émerge), et l'autre est le monde complexe et collant de la chimie moléculaire (où les atomes s'assemblent pour former des médicaments, des matériaux et la vie).

Pendant longtemps, les scientifiques de ces deux domaines ont utilisé des outils complètement différents pour comprendre leurs mondes respectifs. Mais dans cet article, les auteurs présentent OmniMol, un nouvel outil qui tente d'enseigner aux experts de la physique des particules à comprendre la chimie en utilisant un « modèle de fondation » qu'ils avaient déjà construit.

Voici une explication simple de la méthode employée et des résultats obtenus :

1. L'analogie du « Grand Chef »

Imaginez le modèle original, appelé Omnilearned, comme un grand chef qui a passé des années à cuisiner avec des jets de particules.

• Les Ingrédients : En physique des particules, un « jet » est un spray de particules subatomiques (comme des protons et des neutrons) éjectés lors d'une collision.
• La Compétence : Ce chef a appris à reconnaître des motifs dans ces sprays. Il sait comment les particules interagissent, comment elles se regroupent et comment prédire ce qui se produira ensuite. Il a été entraîné sur un milliard de sprays de particules différents.

Maintenant, les auteurs se sont demandé : Ce même chef peut-il cuisiner un repas moléculaire ?

• Les Nouveaux Ingrédients : Au lieu de particules subatomiques, les « ingrédients » sont des atomes (comme le Carbone, l'Oxygène, l'Hydrogène) au sein d'une molécule.
• Le Défi : Les atomes se comportent différemment des particules subatomiques, mais ils partagent une structure similaire : ce sont simplement des points dans l'espace avec des types spécifiques.

2. Le « Traducteur Universel » (L'Architecture)

Pour rendre cela possible, ils n'ont pas construit un nouveau chef à partir de zéro. Ils ont pris le « Grand Chef » existant (Omnilearned) et lui ont fourni un nouvel ensemble d'outils :

• Le Point-Edge Transformer (PET) : Imaginez le chef regardant une assiette de nourriture. Au lieu de regarder un seul ingrédient à la fois, cet outil lui permet de regarder tous les ingrédients à la fois et de voir comment chacun se relie à tous les autres.
• Le « Biais Physique » : C'est la sauce secrète. Le modèle possède un « manuel de règles » intégré qui lui dit : « Hé, ces deux particules/atomes sont proches l'un de l'autre, ils devraient donc prêter plus d'attention l'un à l'autre. » Cela aide le modèle à se concentrer sur les relations les plus importantes sans se laisser troubler par le bruit.

3. L'Expérience : Le Fine-Tuning

Les auteurs ont pris ce modèle entraîné sur des particules et lui ont donné un « cours accéléré » en chimie en utilisant un ensemble de données appelé oMol (une collection de millions de molécules).

• L'Objectif : Ils voulaient que le modèle agisse comme un Potentiel Interatomique Appris par Machine (MLIP). En termes simples, cela signifie que le modèle doit prédire deux choses pour n'importe quel groupe d'atomes :
  1. L'Énergie : Quelle est la quantité de « colle » qui les maintient ensemble ?
  2. La Force : Si vous poussez un atome, avec quelle force va-t-il repousser ?

4. Les Résultats : Rapide et Étonnamment Efficace

L'article a révélé des choses passionnantes :

• Le Super-pouvoir du « Few-Shot » : Habituellement, enseigner la chimie à un ordinateur nécessite d'énormes quantités de données. Mais parce qu'OmniMol a commencé avec la « connaissance » de la physique des particules, il a appris la chimie très rapidement. Même avec une quantité relativement faible de nouvelles données (comme 100 000 molécules), il a performé presque aussi bien que des modèles entraînés sur des millions. C'est comme un grand chef qui peut apprendre une nouvelle cuisine avec seulement quelques recettes parce qu'il comprend déjà les bases des saveurs et de la chaleur.
• Vitesse : OmniMol est incroyablement rapide. Alors que d'autres modèles peuvent prendre beaucoup de temps pour calculer comment une molécule se déplace, OmniMol le fait en un clin d'œil. Les auteurs notent que pour chaque heure de temps de calcul, OmniMol peut simuler trois fois plus de molécules que certains de ses concurrents.
• Le Compromis : Lorsqu'ils disposaient de vastes quantités de données (des millions de molécules), l'avantage de commencer avec la connaissance de la physique des particules s'est un peu estompé. Cela suggère que la « connaissance de la physique des particules » agit comme un départ avantageux, mais si vous avez suffisamment de temps et de données pour entraîner un modèle à partir de zéro, cet avantage initial compte moins.

5. La Vue d'Ensemble

L'article conclut que OmniMol représente la première fois qu'un « modèle de fondation » construit pour une discipline scientifique (la physique des particules) a été transféré avec succès vers une discipline complètement différente (la chimie).

Ils ont prouvé que si vous avez un modèle intelligent qui comprend comment des points dans l'espace interagissent dans un domaine, il peut être adapté pour comprendre comment des points dans l'espace interagissent dans un autre domaine, économisant ainsi du temps et de la puissance de calcul.

En résumé : Les auteurs ont pris une IA super-intelligente entraînée sur des collisions de particules à haute énergie, ajusté son cerveau pour comprendre les atomes plutôt que les particules, et découvert qu'elle est devenue un outil ultra-rapide et extrêmement précis pour prédire le comportement des molécules, en particulier lorsque les données sont rares.

Résumé technique

Résumé Technique : OmniMol

Énoncé du Problème

L'apprentissage automatique (ML) a transformé la représentation et la simulation de systèmes physiques complexes, en particulier en physique des particules et en chimie moléculaire. Bien que ces domaines diffèrent considérablement par leurs échelles d'énergie, ils partagent une structure de données fondamentale : des ensembles de particules (ou d'atomes) de taille variable dans l'espace des phases, formant efficacement des nuages de points structurés.

Le défi principal abordé est le développement de potentiels interatomiques appris par machine (MLIP) efficaces. Les méthodes traditionnelles comme la théorie de la fonctionnelle de la densité (DFT) sont coûteuses en calcul, limitant les simulations de dynamique moléculaire (MD) à grande échelle et sur de longs horizons. Les MLIP visent à approximer les surfaces d'énergie potentielle et les forces à une fraction de ce coût. Cependant, l'entraînement de MLIP robustes nécessite généralement des jeux de données massifs et des ressources computationnelles importantes. L'article émet l'hypothèse qu'un modèle de fondation pré-entraîné sur des nuages de points en physique des particules (spécifiquement des jets de particules) pourrait être transféré à la dynamique moléculaire, accélérant potentiellement l'optimisation et améliorant la précision dans les régimes à faible quantité de données.

Méthodologie

Architecture : Transformer Point-Edge (PET)

OmniMol est construit en adaptant Omnilearned, un modèle de fondation initialement conçu pour classifier et générer des jets de particules en physique des hautes énergies (HEP). L'architecture de base est un Transformer Point-Edge (PET), qui couple une attention locale sur les $k$ plus proches voisins avec des blocs de transformateurs globaux tout-à-tout.

Les composants architecturaux clés incluent :

1. Encodages d'entrée : Les atomes sont encodés dans un espace de tokens combinant des informations de position ($\vec{r}$), des numéros atomiques discrets ($Z$) et des caractéristiques supplémentaires (charge, spin).
2. Bloc d'attention locale : Pour chaque atome, un voisinage local est construit en utilisant les $K$ plus proches voisins ($K=15$ pour les molécules, comparé à $K=10$ pour les jets). Des caractéristiques physiques paires sont calculées, incluant des termes de distance, des puissances inverses de la distance et des fonctions apprises des encodages atomiques. Ceux-ci sont traités par un petit transformateur local pour créer un vecteur d'encodage local.
3. Attention globale avec biais d'interaction : Le mécanisme d'auto-attention global intègre un biais explicite dérivé de caractéristiques physiques paires. Les logits d'attention sont modifiés comme $A^*_{ij} = A_{ij} + B_{ij}$, où $B_{ij}$ est un terme de biais encodé par un MLP. Ce « biais d'attention de matrice d'interaction » injecte directement des priors physiques pairs dans le transformateur, orientant le réseau vers des voisinages physiquement significatifs sans sacrifier l'expressivité.
4. Têtes de sortie : La tête générative d'Omnilearned est réaffectée à deux tâches :
   • Prédiction de forces : Une tête permutation-équivariante prédisant les forces par atome.
   • Prédiction d'énergie : Une tête prédisant des corrections d'énergie par atome, qui sont sommées pour obtenir l'énergie moléculaire totale, préservant ainsi les priors extensifs.

Contraintes d'Invariance et de Conservation

Pour satisfaire les contraintes physiques, les auteurs répondent à deux exigences :

• Conservation de l'énergie : Les forces ne sont pas prédites directement mais sont calculées par rétropropagation de la sortie d'énergie ($\vec{F}_i = \nabla_{\vec{r}_i} E$). Cela assure une conservation exacte de l'énergie mais augmente le coût computationnel pendant l'entraînement (nécessitant une double rétropropagation). Par conséquent, cette contrainte est appliquée uniquement à la variante « petite » du modèle.
• Équivariance rotationnelle : L'architecture standard n'est pas intrinsèquement équivariante car les différences de coordonnées brutes sont alimentées dans des MLP. Pour y remédier, les auteurs introduisent une variante « équivariante et conservatrice ». Cette version supprime les termes de différence de coordonnées directs des caractéristiques paires et incorpore à la place des informations angulaires (cosinus des angles formés par les vecteurs entre atomes voisins) dans le bloc local. Cette modification conserve l'équivariance tout en récupérant significativement les pertes de performance associées à la suppression des termes de coordonnées.

Stratégies d'Entraînement et de Fine-Tuning

Le modèle est affiné sur le jeu de données oMol (spécifiquement les sous-ensembles oMol-25, oMol-4M, oMol-100M et oMol-140M). Deux stratégies sont explorées :

1. LoRA (Adaptation à faible rang) : Les poids du backbone PET pré-entraîné sont figés. Des adaptateurs à faible rang sont introduits uniquement pour les matrices du corps du transformateur ($W_Q, W_K, W_V, W_O, W_{MLP}$), ainsi que l'entraînement des encodeurs d'entrée moléculaire, du MLP de biais et des têtes de tâche. Une couche « d'adaptation d'encodage » est également ajoutée pour modifier les encodages appris.
2. Fine-Tuning complet : Tous les poids du corps et des encodeurs d'entrée sont dégelés et entraînés, tandis que les têtes de tâche sont entraînées à partir de zéro.

L'objectif d'entraînement minimise la somme des erreurs absolues moyennes (MAE) pour les énergies et les forces, les forces étant pondérées plus lourdement ($\lambda_F = 10$).

Résultats Clés

Performance sur oMol

• Entraînement complet : Lorsqu'entraîné sur de grands jeux de données (oMol-4M et oMol-100M/140M), OmniMol atteint des performances compétitives avec les MLIP de l'état de l'art. Par exemple, sur oMol-140M, le modèle OmniMol-large atteint une MAE d'énergie de 1,04 meV/atome et une MAE de force de 13,59 meV/Å.
• Régime à faible quantité de données : Les gains les plus significatifs sont observés lorsque les données d'entraînement sont limitées. Lorsqu'affiné sur seulement 100k molécules ou avec très peu d'époques (2 passages) sur oMol-4M, les variantes pré-entraînées d'OmniMol surpassent nettement les modèles entraînés à partir de zéro.
  • Sur un sous-ensemble de 100k, le pré-entraînement a amélioré la MAE d'énergie de jusqu'à 29,4 % et la MAE de force de 26,9 % pour le modèle moyen.
  • Avec seulement deux époques sur oMol-4M, le modèle moyen a montré une amélioration de 54,6 % de la MAE d'énergie et de 56,9 % de la MAE de force par rapport à son homologue non pré-entraîné.
• Variante équivariante/conservatrice : La variante équivariante et conservatrice du modèle montre des performances nettement améliorées (surtout pour les forces) dans les régimes à faible quantité de données, bien que cet avantage diminue à mesure que la taille du jeu de données augmente.

Mise à l'échelle et Vitesse d'Inférence

• Mise à l'échelle : OmniMol suit une mise à l'échelle en loi de puissance propre avec la taille du modèle, ne montrant aucun signe de saturation jusqu'à 1 milliard de paramètres, ce qui est cohérent avec les découvertes récentes sur les MLIP basés sur des transformateurs.
• Vitesse d'inférence : Malgré de grands nombres de paramètres, OmniMol démontre des vitesses d'inférence uniques grâce à des optimisations matérielles pour les transformateurs. Sur une GPU A100 pour des systèmes d'environ 100 atomes, OmniMol-medium est environ 3 fois plus rapide que des baselines de réseaux de neurones graphiques (GNN) comparables (eSEN-md-d et AllScAIP-md) tout en maintenant une précision compétitive (erreur d'énergie seulement ~0,7 meV/atome plus élevée que AllScAIP-md).

Signification et Revendications

L'article prétend présenter la première démonstration de transfert interdisciplinaire pour des modèles de fondation de nuages de points scientifiques. En adaptant un modèle pré-entraîné sur des jets de particules de physique des hautes énergies à la dynamique moléculaire, les auteurs démontrent que :

1. Le transfert interdomaine est viable : Un modèle de fondation construit pour la physique des particules peut efficacement se transférer à la chimie moléculaire, suggérant que les structures de nuages de points sous-jacentes partagent des caractéristiques apprenables à travers des échelles physiques très différentes.
2. Le biais inductif accélère l'apprentissage : Le pré-entraînement agit comme un fort biais inductif. De la même manière que l'équivariance aide lorsque les données sont rares, la « leçon amère » du pré-entraînement permet une optimisation rapide et une précision améliorée lorsque les données d'entraînement sont limitées.
3. Efficacité : Le transfert architectural permet des vitesses d'inférence uniques, ce qui est critique pour les applications nécessitant une exploration rapide des espaces de conception, telles que la découverte de médicaments à petites molécules.

Les auteurs concluent que, bien que l'étude se concentre sur les MLIP, les leçons concernant les modèles de fondation de nuages de points peuvent avoir une utilité généralisée dans les domaines scientifiques où les systèmes sont décrits comme des ensembles non ordonnés de corps en interaction. Ils ne revendiquent pas une supériorité universelle sur toutes les méthodes existantes dans tous les régimes, mais mettent en évidence les avantages spécifiques dans les scénarios à faible quantité de données et la vitesse d'inférence.