Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

Each language version is independently generated for its own context, not a direct translation.

🧪 De Grote Moleculaire Ontmaskering: Wat zit er echt in AI's hoofd?

Stel je voor dat je een supergeavanceerde keukenrobot hebt die perfect kan voorspellen hoe een gerecht smaakt. Maar als je vraagt: "Hoe weet deze robot precies wat er in het gerecht zit en hoe het eruitziet?", dan houdt hij zijn mond dicht.

Wetenschappers Joshua Steier en zijn team wilden dit geheim onthullen. Ze keken naar moderne AI-modellen die moleculen (de bouwstenen van alles) analyseren. Deze modellen moeten twee dingen tegelijk begrijpen:

De Ingrediënten: Wat zit erin? (Bijvoorbeeld: 6 koolstofatomen, 12 waterstofatomen).
De Vorm: Hoe staan die atomen precies in de ruimte? (Is het een rechte lijn of een geknikte bocht?)

Het probleem is dat deze twee dingen vaak door elkaar lopen. Als je een model vraagt om iets te voorspellen, is het vaak te makkelijk om te zeggen: "Oh, het heeft 6 koolstofatomen, dus het moet deze waarde zijn." De AI neemt dan een "korte weg" en negeert de vorm.

De onderzoekers wilden weten: Zit de informatie over de vorm (de geometrie) echt los van de ingrediënten in het brein van de AI?

🛠️ De Oplossing: De "Ingrediënten-Verwijderaar" (CPD)

Om dit te testen, bedachten ze een slimme truc genaamd CPD (Compositional Probe Decomposition).

Stel je voor dat je een smoothie hebt gemaakt van fruit (de ingrediënten) en ijsblokjes (de vorm). Je wilt weten of je de ijsblokjes nog kunt voelen als je de fruitsmaak eruit haalt.

Ze nemen de AI's "gedachten" (de interne data).
Ze gebruiken een wiskundige filter (een soort zeef) om alle informatie over de ingrediënten eruit te halen.
Dan kijken ze wat er overblijft: de "geometrische rest".

Vervolgens proberen ze met een simpele test (een lineaire probe) te voorspellen of de AI nog steeds de vorm van het molecuul kan "zien" in die rest.

Het verrassende resultaat:
Sommige AI-modellen zijn als een goed georganiseerde bibliotheek. Als je de boeken over de auteursnamen (ingrediënten) verwijdert, staan de boeken over de verhaallijnen (vorm) nog steeds perfect op de planken. Je kunt ze makkelijk vinden.
Andere modellen zijn als een rommelige zolder. Als je de auteursnamen verwijdert, liggen de verhaallijnen verward onder een hoop oud papier. Je kunt ze niet vinden, tenzij je heel hard gaat graven (met complexe, niet-lineaire methoden).

🏆 De Drie Factoren die het Verschil Maken

De onderzoekers keken naar 10 verschillende AI-modellen en ontdekten een groot verschil in hoe goed ze hun "geometrische bibliotheek" hadden opgezet. Dit hangt af van drie dingen:

1. De Doelstelling (Het Belangrijkste!) 🎯
Dit is de grootste verrassing. Het maakt niet uit hoe slim de architectuur van de robot is; het maakt er wel toe waarvoor hij is getraind.

Analogie: Stel je voor dat je een chef-kok traint om alleen soep te maken (energie). Hij leert dan vooral welke groenten erin zitten. Als je hem later vraagt om de vorm van de groenten te beschrijven, kan hij dat niet goed.
Maar als je hem traint om specifieke smaken te maken die afhangen van hoe de groenten liggen (zoals de HOMO-LUMO kloof, een elektronische eigenschap), dan leert hij de vorm van de groenten heel goed.
Conclusie: Modellen die getraind zijn op de juiste taak (vorm-gevoelige taken) zijn 6,6 keer beter in het loskoppelen van vorm en ingrediënten dan modellen die alleen op energie zijn getraind.

2. De Architectuur (De Bouwstijl) 🏗️
Sommige modellen zijn gebouwd met speciale "symmetrie-regels" (ze weten dat draaien en kantelen geen verschil moet maken).

Analogie: Het is alsof je een auto bouwt met wielen die altijd recht blijven staan, ongeacht hoe je het stuur draait.
Dit helpt, maar alleen als de auto ook de juiste bestemming heeft (zie punt 1). Een supergeavanceerde auto (MACE) die naar de verkeerde bestemming rijdt, doet het slechter dan een simpele auto (SchNet) die wel de juiste route volgt.

3. De Diversiteit van de Data (De Reizigers) 🌍
Als een model getraind is op een enorme, diverse dataset (duizenden verschillende soorten moleculen), leert het beter dan als het alleen op één klein soort molecuul is getraind.

Analogie: Een kok die in 10 verschillende landen heeft gewerkt, kent de vorm van groenten beter dan een kok die alleen in één restaurant heeft gewerkt. Maar zelfs deze wereldkok kan niet helemaal concurreren met de kok die specifiek getraind is op jouw favoriete gerecht.

🧠 De Geheime Gangen in het Brein (MACE)

Bij één specifiek model (MACE) ontdekten ze iets fascinerends. Het model heeft verschillende "kanalen" of gangen in zijn brein:

De Scalar-gang (L=0): Hier worden eigenschappen bewaard die geen richting hebben (zoals de energie-kloof).
De Vector-gang (L=1): Hier worden eigenschappen bewaard die een richting hebben (zoals een magnetisch veld of dipoolmoment).

Het model heeft geleerd om deze informatie schoon te scheiden. Het is alsof het model twee verschillende postvakken heeft: één voor brieven en één voor pakketten. Bij een ander model (ViSNet) zitten de brieven en pakketten door elkaar in één grote doos.

⚠️ Een Waarschuwing: De "Truc" van de Niet-Lineaire Test

De onderzoekers ontdekten ook een valkuil. Als je probeert te meten wat er overblijft na het verwijderen van de ingrediënten, en je gebruikt een zeer complexe test (zoals een "Gradient Boosted Tree"), dan haalt die test de ingrediënten er weer uit!

Analogie: Het is alsof je een raam dichtdoekt, maar de test kijkt door de kieren en zegt: "Ik zie nog steeds de zon!"
De onderzoekers waarschuwen: gebruik alleen simpele, lineaire tests. Die zijn eerlijk en zeggen je wat er écht overblijft.

📝 Samenvatting voor de Leek

Het doel is koning: Als je een AI wilt gebruiken om de vorm van moleculen te begrijpen, kies dan een model dat getraind is op taken die gevoelig zijn voor vorm. De bouwstijl van het model is minder belangrijk dan wat het geleerd heeft.
Schoon scheiden: De beste modellen houden de "wat" (ingrediënten) en de "hoe" (vorm) netjes gescheiden. Dit maakt het makkelijker om nieuwe dingen te voorspellen.
Voorzichtig met meten: Gebruik geen te complexe meetinstrumenten om te kijken wat een AI "weet", anders krijg je nep-resultaten.

Kortom: Wat een AI leert, is belangrijker dan hoe het eruitziet. Als je een model traint op de juiste manier, ontstaat er vanzelf een helder, goed georganiseerd inzicht in de wereld van de moleculen.

Each language version is independently generated for its own context, not a direct translation.

Hier is een gedetailleerde technische samenvatting van het artikel "Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement" in het Nederlands.

Probleemstelling

In de moleculaire machine learning (ML) worden foundation-modellen (zoals MACE, SchNet, PaiNN) gebruikt om eigenschappen van moleculen te voorspellen. Een fundamentele vraag blijft echter open: hoe zijn de interne representaties van deze modellen georganiseerd? Specifiek: scheiden deze modellen de informatie over de samenstelling (welke atomen en in welke verhoudingen) en de geometrie (hoe die atomen in de ruimte zijn gerangschikt) van elkaar?

Als een model deze factoren "verstrengeld" (entangled) encodeert, moeten downstream-taakjes deze factoren opnieuw ontrafelen, wat inefficiënt is. Echter, traditionele "probing"-methodes (het trainen van een simpele predictor op interne activaties) zijn misleidend in de moleculaire context. Omdat samenstelling en geometrie sterk gecorreleerd zijn, kan een probe op ruwe representaties hoge scores halen door simpelweg op samenstelling te vertrouwen, in plaats van echte geometrische informatie te vinden. Bovendien bleek dat niet-lineaire probes (zoals Gradient Boosted Trees) kunstmatig hoge scores halen op residuen waar samenstelling uit is verwijderd, doordat ze de verwijderde signalen via niet-lineaire interacties reconstrueren.

Methodologie: Compositional Probe Decomposition (CPD)

De auteurs introduceren Compositional Probe Decomposition (CPD), een methode om lineair toegankelijke geometrische informatie te isoleren.

Projectie van Samenstelling: Voor elke molecule wordt een samenstellingsvector $Z$ (elementfracties en genormaliseerd atoomaantal) gedefinieerd. Er wordt een Ordinary Least Squares (OLS) regressie uitgevoerd om de lineaire relatie tussen de representaties $X$ en $Z$ te modelleren.
Residuen: De lineaire component van de samenstelling wordt afgetrokken van de representaties, waardoor een geometrisch residu $X_{geom}$ overblijft:
$X_{geom} = X - Z\hat{\beta}$
Dit residu bevat alle informatie die lineair orthogonaal is aan de samenstelling (topologie, bindingsconnectiviteit, conformatie).
Fold-wise Implementatie: Om data-lekkage te voorkomen, wordt de projectie binnen elke cross-validatie-fold alleen op de trainingsdata gefit.
Probing: Er wordt een Ridge-regressie (lineaire probe) getraind op $X_{geom}$ om een doelstelling (bijv. HOMO-LUMO gap) te voorspellen. De verkregen $R^2$ ( $R^2_{geom}$ ) meet hoeveel geometrische informatie lineair toegankelijk is.
Validatie: De methode wordt gevalideerd met een benchmark voor structurele isomeren (moleculen met identieke samenstelling maar verschillende geometrie). Hier moet de samenstellingscomponent willekeurig presteren (50%), terwijl de geometrische residuen hoge classificatie-accuraatheid moeten tonen.

Belangrijkste Bijdragen

CPD als gevalideerde methode: Een nieuwe standaard voor het meten van lineaire disentanglement in moleculaire modellen, inclusief een bewijs dat niet-lineaire probes (GBT) systematisch overgeïnfundeerde scores geven op residuen.
Het "Linear Accessibility Gradient": Een kwantificering van hoe goed modellen geometrische informatie scheiden, over tien modellen uit vijf architectuurfamilies.
Dominantie van Taak-Alignement: Het aantonen dat de trainingsdoelstelling (task alignment) de belangrijkste factor is, meer dan de architectuur of equivariantie.
Informatie-routing in MACE: Het ontdekken dat MACE informatie routeert via irreducibele representaties (L=0 scalair vs. L=1 vector), waarbij scalair kanalen scalare eigenschappen coderen en vector kanalen vector-eigenschappen.

Resultaten

1. De Lineaire Toegankelijkheidsgradiënt

Over tien modellen op de QM9-dataset varieert de $R^2_{geom}$ voor de HOMO-LUMO gap met een factor van 6,6x (van 0,081 tot 0,533). Drie factoren verklaren deze gradiënt:

Factor 1: Taak-Alignement (Dominant): Modellen die getraind zijn op de HOMO-LUMO gap (een eigenschap die sterk afhankelijk is van geometrie) presteren aanzienlijk beter ( $R^2_{geom} \approx 0,44 - 0,53$ $R_{g eo m}^{2} \approx 0, 44 - 0, 53$ ) dan modellen die alleen op energie zijn getraind ( $R^2_{geom} \approx 0,08 - 0,31$ $R_{g eo m}^{2} \approx 0, 08 - 0, 31$ ).
- Voorbeeld: PaiNN getraind op HOMO-LUMO scoort 0,533; dezelfde architectuur (PaiNN) getraind op energie scoort 0,310. Het verschil ( $\Delta \approx 0,22$ ) is groter dan het verschil tussen verschillende architecturen.
Factor 2: Equivariantie (Conditioneel): Equivariante modellen (zoals MACE, PaiNN) presteren niet automatisch beter. Een equivariant model getraind op de "verkeerde" taak (energie) kan slechter presteren dan een invariant model (SchNet) getraind op dezelfde taak. Equivariantie versterkt de prestaties alleen als het trainingsdoelstelling ook geometrisch gevoelig is.
Factor 3: Data-diversiteit (Compenserend): Modellen die op grote, diverse datasets zijn voorgeïmplementeerd (zoals MACE pretrained op MPTraj) presteren beter dan modellen die alleen op QM9 zijn getraind, maar halen de prestaties van taak-gealigneerde modellen niet volledig.

2. Informatie-routing in MACE

Binnen de MACE-architectuur wordt informatie gesorteerd op symmetrie:

L=0 (Scalair) kanalen: Domineren voor scalare eigenschappen zoals de HOMO-LUMO gap ( $R^2 = 0,76$ ).
L=1 (Vector) kanalen: Domineren voor vector-eigenschappen zoals het dipoolmoment ( $R^2 = 0,59$ ).
Dit patroon is afwezig in ViSNet, wat suggereert dat dit specifiek is voor architecturen die tensor-producten gebruiken om irreducibele representaties expliciet te behouden tot de output.

3. Niet-lineaire Probe Inflatie

Een kritieke bevinding is dat Gradient Boosted Trees (GBT) op de residuen ( $X_{geom}$ ) hoge $R^2$ -scores halen (0,68–0,95) voor een puur samenstellingsdoel (gemiddelde atoommassa). Dit bewijst dat GBTs de verwijderde samenstellingsinformatie via niet-lineaire interacties kunnen reconstrueren. Lineaire probes (Ridge) geven daarentegen correcte scores ( $\approx 0$ ). Dit betekent dat niet-lineaire probes onbetrouwbaar zijn voor het meten van wat er overblijft na het verwijderen van een concept.

4. Validatie met Isomeren

Op een benchmark van structurele isomeren (identieke formule, verschillende vorm) scoorde de samenstellingscomponent willekeurig (52,5%), terwijl de geometrische residuen van PaiNN een classificatie-accuraatheid van 94,6% bereikten. Dit bevestigt dat CPD samenstelling succesvol verwijdert terwijl geometrische informatie behouden blijft.

Significantie en Implicaties

Voor de Praktijk: Bij het selecteren van een voorgeïmplementerd model voor een downstream-taak, is de trainingsdoelstelling belangrijker dan de architectuur. Voor geometrie-gevoelige taken (zoals elektronische eigenschappen) is een model dat specifiek op die taak is getraind superieur, zelfs als het architectonisch "simpeler" is dan een equivariant foundation model dat op energie is getraind.
Representatie-organiserende principes: De studie toont aan dat "disentanglement" geen universele eigenschap is van equivariante netwerken, maar sterk afhankelijk is van de supervisie (de loss functie). Supervisie domineert architecturale inductieve bias.
Methodologische Waarschuwing: De bevinding over GBT-inflatie is breed toepasbaar. Het waarschuwt onderzoekers in NLP, computer vision en ML dat het gebruik van niet-lineaire probes op concept-geëraseerde representaties leidt tot valse positieven over wat een model heeft geleerd.
Sample Efficiency: Modellen met goed lineair ontkoppelde representaties (zoals PaiNN getraind op HOMO-LUMO) vereisen aanzienlijk minder data om geometrische signalen te extraheren dan modellen met verstrengelde representaties.

Kortom, deze paper biedt een robuust raamwerk om te begrijpen hoe moleculaire foundation-modellen informatie structureren, en benadrukt dat de keuze van de trainingsdoelstelling de cruciale factor is voor het creëren van bruikbare, lineair toegankelijke geometrische representaties.