A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attention

Dit paper introduceert AllScAIP, een schaalbaar, op attention gebaseerd machine-learning potentiaalmodel dat door middel van een data-gedreven all-to-all node attention-mechanisme nauwkeurige lange-afstand interacties captureert en daarmee de afhankelijkheid van expliciete fysische termen voor grootschalige systemen zoals biomoleculen en elektrolyten overbodig maakt.

De kern

Stel je voor dat je een superkrachtige keukenchef bent die recepten voor moleculen maakt. In de wereld van de chemie en biologie zijn moleculen net als enorme, ingewikkelde gerechten. Om te voorspellen hoe deze gerechten zich gedragen (bijvoorbeeld of ze koken, bevriezen of ontploffen), hebben wetenschappers "rekenmodellen" nodig. Deze modellen heten MLIPs (Machine Learning Interatomic Potentials).

Vroeger waren deze modellen als een chef die alleen naar de ingrediënten in zijn directe handpalm keek. Hij wist precies hoe een snufje zout en een snufje peper samen proefden, maar hij had geen idee wat er gebeurde als je een heel groot bord met eten had, waar de ingrediënten ver uit elkaar lagen. Als je een eiwit (zoals in je lichaam) of een vloeistof wilde simuleren, miste de chef de "langeafstandscommunicatie" tussen de ingrediënten.

De oplossing: AllScAIP

De onderzoekers in dit paper hebben een nieuwe, slimme chef bedacht: AllScAIP. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude probleem: De "Lokaal Kijkende" Chef

De meeste oude modellen werkten als een chef die alleen met zijn buren praat. Als je in een drukke kamer staat, praat je alleen met de mensen direct naast je. Je hoort niet wat er in de hoek van de kamer gebeurt.

• Het gevolg: Als je een groot molecuul simuleert (zoals een eiwit of een batterijvloeistof), mist de chef de interacties tussen de mensen die ver uit elkaar staan. Dit leidt tot fouten in de voorspellingen.

2. De nieuwe truc: "Iedereen met Iedereen" (All-to-All Attention)

AllScAIP introduceert een revolutionaire regel: Iedereen in de kamer mag met iedereen praten, niet alleen met de buren.

• De Analogie: Stel je voor dat je in een groot concertzaal zit.
  • Oude modellen: Je luistert alleen naar de mensen op de stoel naast jou.
  • AllScAIP: Iedereen in de zaal heeft een super-telefoon. Als iemand in de achterste rij fluistert, kan de persoon in de voorste rij het direct horen. Dit heet "All-to-all node attention".
  • Waarom is dit cool? Het model hoeft niet te raden hoe de krachten werken over lange afstanden; het "hoort" het gewoon direct. Hierdoor wordt het model extreem goed in het simuleren van grote, complexe systemen.

3. De "Recepten" (Inductieve Biases) vs. "Leren van de Wereld"

In de wereld van AI gebruiken wetenschappers vaak "inductieve biases". Dat zijn vooraf opgelegde regels of recepten die de AI moet volgen.

• Voorbeeld: "Je moet altijd weten dat water naar beneden valt" of "Atomen kunnen niet door elkaar heen lopen."

De onderzoekers vroegen zich af: "Moeten we deze regels hard in het model stoppen, of kan de AI ze zelf leren als we genoeg data en rekenkracht geven?"

• De ontdekking:
  • Kleine datasets / Kleine modellen: Hier helpen de vaste regels (recepten) enorm. Het is als een leerling die een receptboek nodig heeft om niet te falen.
  • Grote datasets / Grote modellen: Als je de AI duizenden keren meer data geeft, blijkt dat de AI de regels (zoals hoeken en afstanden) zelf kan leren. Het receptboek wordt zelfs een last! De AI wordt slimmer als je hem de ruimte geeft om zelf te ontdekken hoe de wereld werkt, in plaats van hem te vertellen hoe het moet.

4. De Resultaten: Een Chef die niet faalt

Dit nieuwe model (AllScAIP) is getest op enorme datasets met miljoenen moleculen.

• Precisie: Het is de beste in zijn klas voor het voorspellen van energie en krachten in moleculen.
• Stabiliteit: Het kan simulaties draaien die dagenlang duren (zoals het simuleren van hoe een vloeistof zich gedraagt in een batterij) zonder dat de simulatie "ontspoort" of onrealistische resultaten geeft.
• Realiteit: Het model voorspelde de dichtheid en verdampingswarmte van vloeistoffen zo nauwkeurig, dat het bijna perfect overeenkwam met echte experimenten in het lab.

Samenvattend in één zin:

AllScAIP is een slimme AI die stopt met het volgen van starre regels en in plaats daarvan leert door met iedereen in het systeem tegelijk te communiceren, waardoor het de perfecte voorspeller wordt voor de grootste en meest complexe moleculaire gerechten die we kunnen bedenken.

Het bewijst dat als je genoeg data en rekenkracht hebt, je de AI meer vrijheid kunt geven om de natuurwetten zelf te ontdekken, in plaats van ze voor te schrijven.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attention" in het Nederlands.

Probleemstelling

Machine Learning Interatomische Potentiaalmodellen (MLIPs) hebben grote vooruitgang geboekt in het voorspellen van energie en krachten op atomaire schaal, vaak met een nauwkeurigheid die dicht bij ab initio (DFT) berekeningen ligt. Echter, bestaande topmodellen vertrouwen sterk op sterke, fysiek gebaseerde inductieve biases (zoals lokale afkortingen en symmetrie-eisen).

Het grootste probleem ontstaat bij het schalen naar grotere systemen, zoals biomoleculen en elektrolyten. In deze regimes zijn langeafstandsinteracties (long-range interactions, LR) cruciaal voor het correct beschrijven van subtiele effecten zoals elektrostatische krachten, polarisatie en dispersie. Bestaande schaalbare modellen, die vaak gebaseerd zijn op lokale graafstructuren met straalafkortingen (cutoffs), worstelen om deze LR-interacties nauwkeurig te vangen. De huidige oplossing is vaak het handmatig toevoegen van expliciete fysieke termen (zoals Coulomb-termen of polarisatie-oplossers), wat de modellen complex maakt en niet altijd goed schaalbaar is voor zeer grote, heterogene datasets.

Methodologie: AllScAIP

De auteurs stellen AllScAIP (All-to-all Scalable Attention Interatomic Potential) voor. Dit is een MLIP-model dat een "prior-light" aanpak hanteert: het minimaliseert handmatige fysieke biases en vertrouwt op schaalbaarheid en data-gedreven leren, met behoud van energiebehoud.

Het modelarchitectuur bestaat uit twee hoofdcomponenten die beide gebruikmaken van geoptimaliseerde Multi-Head Self-Attention (MHSA) CUDA-kernels:

1. Lokale Zelf-Attention (Neighborhood Attention):

   • Werkt op vaste lokale stencils (burenlijsten) rondom elk atoom.
   • Lost fijne, anisotrope geometrische interacties op.
   • Complexiteit: $O(N \cdot k)$, waarbij $N$ het aantal atomen is en $k$ het maximale aantal buren.
   • Optioneel gebruik van Legendre Angular Encoding (LAE) om hoekinformatie (richting) in te voegen.

2. Globale Zelf-Attention (All-to-all Node Attention):

   • Een tweede stap waarbij alle atomen in het systeem met elkaar communiceren.
   • Hierdoor kunnen signalen over het hele graf bewegen, wat langeafstandsinteracties mogelijk maakt zonder expliciete fysieke formules.
   • Complexiteit: $O(N^2)$.
   • Optioneel gebruik van Euclidean Rotary Position Encoding (ERoPE) om afstandsgegevens (radiale informatie) in te voegen.

Geometrische Encodings:
De auteurs testen of specifieke geometrische biases (LAE voor hoeken, ERoPE voor afstanden) nodig zijn of dat het model deze kan leren.

• LAE: Gebruikt reële sferische harmonischen om richtingsinformatie aan de attention-logits toe te voegen.
• ERoPE: Voegt een isotrope radiale bias toe op basis van afstanden tussen atomen.

Energiebehoud:
Het model garandeert energiebehoud door krachten te berekenen als de negatieve gradiënt van de voorspelde energie ($F = -\nabla E$), gebruikmakend van een differentieerbare kNN-grafiek.

Belangrijkste Bijdragen en Ablatiestudies

De kernbijdrage is de empirische validatie van de hypothese dat veel inductieve biases leerbaar zijn onder schaal (data- en modelgrootte).

• Low-Data / Small-Model Regime: In situaties met weinig data of kleine modellen zijn de inductieve biases (LAE, ERoPE) zeer waardevol. Ze verbeteren de sample-efficiency aanzienlijk.
• High-Data / Large-Model Regime: Naarmate de dataset (tot 102 miljoen samples) en het modelgrootte (tot 1 miljard parameters) toenemen, nemen de voordelen van de vaste geometrische encodings af, en soms zelfs om. Het model leert de richtings- en afstandsafhankelijkheden end-to-end uit de data.
• De Rol van All-to-all Attention: In tegenstelling tot de geometrische encodings, blijft de all-to-all node attention cruciaal en voordelig op alle schalen. Dit is de belangrijkste architectonische component voor het vangen van langeafstandsinteracties.

Resultaten

AllScAIP presteert state-of-the-art op meerdere benchmarks:

1. Nauwkeurigheid (OMol25):

   • Op het enorme OMol25-dataset (102M samples) bereikt AllScAIP de beste energie- en krachtnauwkeurigheid, vooral op biomoleculen waar langeafstandsinteracties kritiek zijn.
   • Het model presteert ook competitief op materialen (OMat24) en katalysatoren (OC20).

2. Langeafstandsvalidatie:

   • Distance Scaling Test: Wanneer moleculen worden uitgerekt, behoudt AllScAIP een lage foutmarge, terwijl concurrenten (zoals eSEN en UMA) sterk degradëren. Dit bewijst het vermogen om langeafstandsinteracties correct te modelleren.
   • Moleculaire Dynamica (MD): AllScAIP kan stabiele, lange tijdschaal MD-simulaties uitvoeren. De resultaten tonen een uitstekende overeenkomst met experimentele waarden voor dichtheid en verdampingswarmte, zonder dat er expliciete langeafstands-termen in het model zijn verwerkt.

3. Efficiëntie:

   • Het model schaalbaar van $O(Nk)$ naar $O(N^2)$. Hoewel de $O(N^2)$-fase (globale attention) rekentijd kost, blijft het efficiënt genoeg voor systemen met $10^3$tot$10^5$ atomen, wat het bereik van veel interessante wetenschappelijke toepassingen dekt.

Significantie en Conclusie

Dit paper biedt een nieuw paradigma voor het ontwikkelen van MLIPs:

• Prior-Light Recipe: In plaats van complexe fysieke biases in te bouwen, kan men vertrouwen op schaalbare architecturale componenten (zoals all-to-all attention) en grote datasets.
• Lerend vs. Enkoderend: Veel fysieke eigenschappen (rotatie-equivalentie, langeafstandsinteracties) hoeven niet handmatig te worden gecodeerd; ze kunnen door het model worden geleerd als er voldoende capaciteit en data is.
• Toekomstvisie: De auteurs suggereren dat de toekomst van MLIPs ligt in het prioriteren van schaalbare componenten en het laten "zorgen" van schaal voor de rest, in plaats van het beperken van modellen door starre inductieve biases.

Kortom, AllScAIP bewijst dat een data-gedreven aanpak met all-to-all attention superieur is in het modelleren van complexe, langeafstandsinteracties in grote moleculaire systemen, terwijl het tegelijkertijd energiebehoud en fysieke consistentie garandeert.