Learning Inter-Atomic Potentials without Explicit Equivariance

Titel: Leren van Inter-Atomaire Potentiaals zonder Expliciete Equivariantie

Auteurs: Ahmed A. Elhag, Arun Raja, et al. (Universiteit van Oxford, LBNL, AstraZeneca, AITHYRA)

---

1. Het Probleem

In de chemie en materiaalkunde zijn Machine-Learned Inter-Atomic Potentials (MLIPs) essentieel voor het simuleren van moleculaire systemen, van drugontwikkeling tot nieuw materiaalontwerp. Traditionele methoden zoals Dichtheidsfunctionaaltheorie (DFT) zijn nauwkeurig maar computationally zeer duur (schalen kubisch met systeemgrootte), waardoor grootschalige simulaties onmogelijk zijn.

Bestaande state-of-the-art MLIP-modellen lossen dit op door equivariante neurale netwerken te gebruiken. Deze architecturen zijn "hard-wired" om rotatie- en translatiesymmetrieën (SE(3)-equivariantie) te respecteren, vaak via complexe mechanismen zoals sferische harmonischen.

• Nadeel: Deze specifieke architecturen kunnen minder flexibel zijn, computationally inefficiënter zijn en minder goed schalen dan generieke modellen.
• Vraag: Is het mogelijk om een generiek, niet-equivariant model (zoals een Transformer) te trainen dat symmetrieën leert uit de data, zonder de architectuur expliciet te beperken?

2. Methodologie: TransIP

De auteurs introduceren TransIP (Transformer-based Inter-Atomic Potentials), een nieuw trainingsparadigma dat symmetriecompliance bereikt zonder expliciete architecturale constraints.

Kerncomponenten:

1. Ongeconstrueerde Transformer-Backbone:

   • Het model behandelt elk atoom als een "token" in een variabele lengte sequentie.
   • Het gebruikt een standaard Transformer-encoder met self-attention (geen afstands-cutoffs of specifieke graf-structuur nodig).
   • Atomaire eigenschappen (coördinaten, atoomnummer, lading, spin) worden geëmbed en door de lagen verwerkt.
   • De output is een energievoorspelling (invariant) en krachten (equivariant), berekend via de gradiënt van de energie.

2. Gestudeerde Latente Equivariantie (Learned Latent Equivariance):

   • In plaats van de architectuur te forceren om equivariant te zijn, wordt een transformatienetwerk $T_\tau$ geleerd.
   • Dit netwerk leert hoe rotaties ($g \in SO(3)$) in de inputruimte corresponderen met transformaties in de latente embeddingruimte ($\mathbb{R}^d$).
   • Het doel is om te leren dat: $f(\phi(g)(m)) \approx T_\tau(\phi(g), f(m))$.

3. Contrastieve Doelfunctie (Loss Function):

   • Het trainingsdoel combineert drie termen:
     1. Energieverlies ($L_E$): Voorspelling van totale energie.
     2. Krachtverlies ($L_F$): Voorspelling van atomaire krachten.
     3. Latente Equivariantie-verlies ($L_{leq}$): Een contrastieve loss die de afstand minimaliseert tussen de embedding van een geroteerd molecuul en de transformatie van de originele embedding via het geleerde netwerk $T_\tau$.
   • Formule: $L_{total} = \lambda_E L_E + \lambda_F L_F + \lambda_{leq} L_{leq}$.

3. Belangrijkste Bijdragen

• Single-Stage Training: Een pipeline die SO(3)-equivariantie bereikt door training, zonder dure pre-training/fine-tuning frameworks of hard-wired equivariante lagen.
• Architectuur-onafhankelijke Loss: Een nieuwe contrastieve loss die symmetrie bevordert in de embeddingruimte van een ongeconstrueerd model.
• Superieure Schaalbaarheid: Het bewijst dat een generieke Transformer, gecombineerd met deze leerstrategie, beter schaalt dan traditionele data-augmentatie methoden en vergelijkbaar presteert met state-of-the-art equivariante modellen.

4. Resultaten

Het model is getraind en geëvalueerd op de Open Molecules 2025 (OMol25) dataset, een grote en diverse collectie (organische moleculen, biomoleculen, elektrolyten, metaalcomplexen).

• Vergelijking met Data Augmentatie (TransAug):

  • TransIP presteert 40% tot 60% beter dan een ongeconstrueerde Transformer die gebruikmaakt van SO(3)-data-augmentatie (TransAug).
  • Het voordeel is het grootst bij beperkte datasets (1M moleculen), waar TransIP de krachten (Force MAE) aanzienlijk nauwkeuriger voorspelt (255 vs 600 meV/Å).
  • Zelfs bij grotere datasets (4M) behoudt TransIP een voorsprong in krachtvoorspelling.

• Vergelijking met State-of-the-Art (SOTA):

  • TransIP (grootste variant, 302M parameters) bereikt vergelijkbare prestaties als de huidige SOTA equivariante modellen (zoals eSCN en GemNet-OC) op de OMol25 benchmark.
  • Het presteert vooral sterk in de voorspelling van totale energie en krachten over diverse chemische categorieën.

• Efficiëntie:

  • TransIP is sneller in inferentie dan equivariante baselines (eSEN). Bijvoorbeeld, TransIP-S verwerkt ~160.000 atomen/sec, terwijl eSEN slechts ~15.000 atomen/sec haalt.
  • De prestaties schalen soepel met modelgrootte, terwijl data-augmentatie modellen (TransAug) vaak degradatie vertonen bij grotere modellen.

• Analyse van de Leerstrategie:

  • Er is een sterke correlatie gevonden tussen de fout in de latente equivariantie en de algehele voorspellingsnauwkeurigheid.
  • Analyse toont aan dat het model leert dat rotaties in de inputruimte corresponderen met een gedeelde orthogonale transformatie in de latente ruimte, zelfs zonder dat dit expliciet in de architectuur is ingebouwd.

5. Betekenis en Conclusie

Dit werk demonstreert dat geleerde equivariantie een krachtig en efficiënt alternatief is voor expliciet equivariante architecturen of data-augmentatie.

• Flexibiliteit: Het maakt het mogelijk om de schaalbaarheid en hardware-efficiëntie van moderne Transformers (zoals die gebruikt in NLP en visie) te benutten voor moleculaire dynamica.
• Toekomstperspectief: Met de toenemende beschikbaarheid van grote datasets (zoals OMol25) en de behoefte aan grotere modellen, biedt TransIP een pad naar schaalbare, nauwkeurige MLIP's zonder de complexiteit van sferische harmonischen of tensor-producten.
• De code is open-source beschikbaar gesteld, wat de gemeenschap in staat stelt deze aanpak verder te onderzoeken en te verbeteren.

Kortom, TransIP toont aan dat symmetrie niet per se in de "hard-wiring" van het netwerk moet zitten, maar effectief kan worden geleerd als een inductieve bias via een goed ontworpen trainingsdoel.