TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces is een model-agnostisch, drie-stromenkader dat zelftoezichthoudend leren combineert met gescheiden samenstellings- en structuurrepresentaties om de overdraagbaarheid en data-efficiëntie van atomaire grafische neurale netwerken voor machine learning-krachtpotentialen tussen atomen aanzienlijk te verbeteren.

De kern

Stel je voor dat je probeert een robotchef te leren koken.

Het Probleem: De "Eén-Oplossing-Voor-Allen"-Chef
Momenteel gebruiken wetenschappers krachtige AI-modellen (zogenaamde MLIP's) om te voorspellen hoe atomen zich gedragen, zoals hoeveel energie een materiaal heeft of hoe moeilijk het is om atomen te verplaatsen. Deze modellen zijn getraind op enorme hoeveelheden data van supercomputers (DFT).

Echter, deze modellen hebben een gebrek: ze zijn als een chef die de exacte smaak van een specifiek gerecht heeft onthouden, maar vergeten is waarom het zo smaakte. Als je ze vraagt om een iets ander gerecht te koken (een nieuw type materiaal), hebben ze moeite. Ze verwarren de ingrediënten (samenstelling) met de vorm van de pan (structuur). Als je de ingrediënten verandert, raken ze in de war over de vorm, en andersom. Dit maakt ze slecht in het snel leren van nieuwe taken, vooral wanneer je niet veel data hebt om ze te leren.

De Oplossing: TriForces (De Drie-Stroom Keuken)
De auteurs introduceren TriForces, een nieuwe manier om deze AI-chefs te bouwen. In plaats van één groot brein dat probeert alles tegelijk te onthouden, splitsen ze het brein op in drie gespecialiseerde "stromen" of afdelingen:

1. De Ingrediëntenstroom (Samenstelling): Deze afdeling kijkt alleen naar wat er in de pan zit (bijvoorbeeld: "We hebben 2 Waterstofatomen en 1 Zuurstof"). Ze negeert de vorm volledig. Ze leert de chemie.
2. De Vormstroom (Structuur): Deze afdeling kijkt alleen naar hoe de atomen in de ruimte zijn gerangschikt (bijvoorbeeld: "Ze vormen een driehoek"). Ze negeert wat de atomen eigenlijk zijn. Ze leert de geometrie.
3. De Interactiestroom: Dit is de hoofdschef die de notities van de Ingrediënten- en Vormafdelingen neemt en ze combineert om het eindresultaat te voorspellen (energie of kracht).

De Geheime Saus: Zelftoezichtend Leren
Voordat het model ooit wordt gevraagd om een specifieke eigenschap te voorspellen, trainen de auteurs het met een spel genaamd "Zelftoezichtend Leren". Denk hierbij aan een oefensessie waarbij de AI moet:

• Ruis verwijderen: Kijken naar een licht beschadigde of ruizige afbeelding van een molecuul en deze herstellen.
• Maskeren: Een ingrediënt bedekken en raden wat het was op basis van de buren.
• Aanpassen: Kijken naar twee licht verschillende versies van hetzelfde molecuul en beseffen dat het hetzelfde ding is.

Deze training dwingt de AI om haar kennis netjes te organiseren. Ze leert dat "ingrediënten" in één map horen en "vormen" in een andere, in plaats van ze door elkaar te gooien.

Waarom Dit Belangrijk Is (De Resultaten)
Het artikel toont aan dat deze nieuwe "Drie-Stroom" keuken veel beter werkt dan de oude "Eén-Brein" keukens:

• Sneller Leren: Wanneer ze een kleine hoeveelheid nieuwe data krijgen (zoals 20.000 voorbeelden in plaats van miljoenen), leert TriForces veel sneller en maakt minder fouten. Het is als een chef die een nieuw recept kan leren na het één keer te hebben geproefd, in plaats van het duizend keer te moeten koken.
• Beter Geheugen: De AI vergeet niet wat ze heeft geleerd. Ze kan haar kennis overdragen van het ene type materiaal naar het andere zonder in de war te raken.
• Zoekbare Kennis: Omdat de AI "ingrediënten" en "vormen" gescheiden houdt, kun je haar vragen om materialen te vinden die er hetzelfde uitzien maar verschillende ingrediënten hebben, of materialen met dezelfde ingrediënten maar verschillende vormen. De oude modellen konden dit niet doen omdat hun kennis te veel door elkaar zat.

Samenvattend
TriForces is een raamwerk dat de complexe taak van het begrijpen van atomen opdeelt in drie eenvoudigere taken: de ingrediënten kennen, de vorm kennen en weten hoe ze samenwerken. Door de AI te trainen om deze taken gescheiden te houden en te oefenen met "raadselspellen" (zelftoezichtend leren), wordt het model een veel flexibeler, efficiënter en accurater hulpmiddel voor het ontdekken van nieuwe materialen.

De auteurs hebben hun code en voorgetrainde modellen vrijgegeven zodat andere wetenschappers deze "drie-stroom keuken" kunnen gebruiken om betere AI voor materiaalkunde te bouwen.

Technische samenvatting

Technische Samenvatting: TriForces

Probleemstelling

Machine Learning Interatomische Potentiaals (MLIPs) gebaseerd op geometrische Graph Neural Networks (GNNs) hebben hoge nauwkeurigheid bereikt bij het voorspellen van atomaire eigenschappen wanneer ze getraind zijn op grote datasets van Dichtheidsfunctionaaltheorie (DFT). Hun praktische bruikbaarheid wordt echter gehinderd door twee primaire beperkingen:

1. Inconsistente Transferbaarheid: MLIPs slagen er vaak niet in om effectief over te gaan naar nieuwe chemieën of taken, vooral bij fijnafstemming op kleine, dure, taakspecifieke datasets. Voorgetrainde modellen worstelen vaak met eenvoudige diagnostische taken (bijvoorbeeld het identificeren van kristalstelsels of meerderheidselementen) en vertonen "catastrophic forgetting" of instabiele transferprestaties over domeinen heen.
2. Slechte Herbruikbaarheid van Representaties: Huidige MLIPs optimaliseren representaties voor specifieke voorspellingsdoelen (energie en krachten) in plaats van voor algemeen hergebruik. Bijgevolg verstrengelen deze representaties vaak samenstelling (chemie) en structuur (geometrie), waardoor ze ongeschikt zijn voor exploratieve analyse, nearest-neighbor retrieval of gedecomponeerde gelijkeniszoekopdrachten. Standaard toezichtobjectieven stimuleren representaties die voldoende zijn voor regressie, maar niet georganiseerd om toegankelijke samenstellings- en structurele informatie te behouden.

Methodologie: Het TriForces-kader

De auteurs stellen TriForces voor, een model-onafhankelijk kader dat bestaande geometrische GNNs aanvult met een drie-stroom architectuur en multi-objectief zelftoezichtleren (SSL). De kerninnovatie is de expliciete factorisatie van atomaire representaties in drie distincte componenten:

1. Drie-Stroom Architectuur

In plaats van één latente vector decomposeert TriForces de node-niveau representatie $h_i$ in drie samengevoegde stromen:

• Samenstellingsstroom ($h^{comp}$): Codeert chemische informatie zonder coördinaten. Deze verwerkt de set unieke atomaire elementen en hun stoichiometrische aantallen met behulp van een Transformer met op aantallen gewogen aandacht. Deze stroom behoudt absolute elementaantallen om systeemgrootte en energieschaal te coderen, onafhankelijk van geometrie.
• Structuurstroom ($h^{struct}$): Codeert geometrie zonder elementidentiteit (type-agnostisch). Deze bouwt rotatie-invariante lokale beschrijvers op, geïnspireerd door SOAP (Smooth Overlap of Atomic Positions), met gebruik van radiale basisfuncties, sferische harmonischen en multi-schaal afsnijdingen. Deze stroom vangt herbruikbare geometrische motieven en topologische patronen, versterkt door invariante berichtdoorvoer.
• Interactiestroom ($h^{int}$): Een standaard basis geometrische GNN (bijvoorbeeld MACE, eSEN, Orb-v3) die de koppeling tussen samenstelling en geometrie vastlegt, waarbij de expressiviteit van de originele architectuur behouden blijft.

2. Zelftoezicht Pretraining

Om de latente ruimte te organiseren en de transferbaarheid te verbeteren, hanteert TriForces een multi-objectief SSL pretrainingstrategie met stochastische augmentaties (positieruis, elementmaskering, grafvariatie en rotaties). Het kader combineert drie complementaire objectieven:

• Niet-herstelling (LeJEPA): Aligneert embeddings van twee versterkte weergaven van dezelfde structuur op zowel node- als graf-niveau. Dit dwingt invariantie ten opzichte van augmentaties af en organiseert de globale latente ruimte zonder stop-gradients of momentum-encoders te vereisen.
• Denoising: Traint het model om schone atomaire posities te herstellen uit ruisige invoer. Dit stabiliseert geometrische representaties en biedt impliciet rotatie-augmentatie.
• Maskering: Voorspelt gemaskeerde atomaire types op basis van omringende geometrie en samenstelling. Dit stimuleert het model om samenstellingspatronen en context te leren.

De uiteindelijke pretraining-los is een gewogen som van deze drie objectieven.

Belangrijkste Bijdragen

1. Architecturale Decompositie: Een drie-stroom ontwerp dat samenstelling, structuur en interactie expliciet scheidt, zodat beide factoren per ontwerp behouden blijven in plaats van verstrengeld.
2. Hybride Pretrainingstrategie: Een zelftoezichtbenadering die reconstructie-gebaseerde objectieven (denoising, masking) combineert met latente-voorspellingsleren (LeJEPA) om de embeddingruimte te structureren voor betere downstream transfer.
3. Interpreteerbare Retrieval: Het vermogen om gerichte gelijkeniszoekopdrachten uit te voeren in de samenstellings-, structurele of gezamenlijke embeddingruimten, waardoor materiaalsovereenkomsten mogelijk zijn op basis van specifieke criteria (bijvoorbeeld alleen chemie of alleen structuur).
4. Empirische Validatie: Uitgebreide experimenten over meerdere architecturen (Orb-v3, eSEN, MACE) en benchmarks (OMat24, MatBench, QM9) die verbeterde data-efficiëntie, transferprestaties en representatiekwaliteit aantonen.

Resultaten

• Transferprestaties (OMat24): In regimes met beperkte data presteert TriForces aanzienlijk beter dan baselines. Bij 20K samples verlaagt het de Mean Absolute Error (MAE) voor energie met 57% ten opzichte van het basismodel. Het verbetert de force MAE over alle steekproefgroottes en verlaagt stressfouten.
• Data-efficiëntie: TriForces bereikt lagere fouten bij elke datasetgrootte (20K tot 2M samples), met de meest significante winsten waargenomen in settings met weinig data.
• Benchmarkprestaties:
  • MatBench: TriForces-varianten behalen de beste algemene resultaten op 6 van de 8 taken, en presteren beter dan zowel zelftoezicht- als DFT-gelabelde pretraining-baselines. Bijvoorbeeld, Phonons MAE verbeterde van 57,8 naar 19,5 cm$^{-1}$.
  • MatBench Discovery: TriForces eSEN-sm bereikt een vergelijkbare energie MAE als een veel groter eSEN-30M-OAM-model, terwijl het 60% minder parameters gebruikt en tot 5$\times$ sneller traint.
  • QM9: Pretraining op diverse chemische invoer (bulk + moleculen) verlaagt consistent de MAE ten opzichte van alleen-bulk of geen-SSL baselines.
• Kwaliteit van Representaties: Lineair proppen op bevroren embeddings toont aan dat TriForces fundamentele informatie behoudt (kristalstelsel, meerderheidselement, coördinatiegetal) die standaard MLIPs verliezen. TriForces bereikt 96–100% nauwkeurigheid bij kristalstelsel- en meerderheidselementclassificatie, terwijl baselines worstelen (55–73%).
• Retrieval: Het kader maakt effectieve k-NN retrieval mogelijk waarbij de samenstellingsstroom uitblinkt in element-set recall en de structuurstroom uitblinkt in ruimtegroep recall, een vermogen dat ontbreekt in single-stream modellen.

Betekenis en Claims

Het artikel positioneert TriForces niet louter als een zelftoezichtmethode, maar als een architecturaal kader waarvan de representaties verder worden versterkt door SSL.

• Regime-afhankelijkheid: De auteurs stellen dat stream-factorisatie de dominante winst biedt in grootschalige toezichtsettings, terwijl SSL het waardevolst is voor transfer met weinig data, organisatie van representaties en retrieval-taken.
• Ontkoppeling: Door samenstelling en structuur te scheiden, adresseert TriForces de "breekbare transfer" en "moeilijk te hergebruiken" problemen van huidige foundation-modellen. Het stelt modellen in staat representaties te leren die georganiseerd zijn voor analyse (bijvoorbeeld proppen, retrieval) in aanvulling op voorspelling.
• Praktijktoepasbaarheid: Het kader is model-onafhankelijk en plug-in, waardoor het direct toepasbaar is op bestaande of nieuwe atomaire architecturen. De auteurs geven voorgereede checkpoints en code vrij om hergebruik in downstream atomaire modellering te faciliteren.

Het werk suggereert dat toekomstige atomaire foundation-modellen voorbij moeten gaan aan single-stream voorspellingsarchitecturen en moeten bewegen naar gefactoriseerde representaties die expliciet de distincte fysieke factoren van chemische systemen behouden.