Quantized SO(3)-Equivariant Graph Neural Networks for Efficient Molecular Property Prediction

Deze paper introduceert een kwantiseringsframework voor SO(3)-equivariante grafische neurale netwerken dat, door middel van innovatieve technieken zoals magnitude-richting-decoupling en takgescheiden training, de inferentie op randapparaten versnelt en de modelgrootte verkleint zonder in te leveren op nauwkeurigheid of fysische symmetrie.

De kern

Stel je voor dat je een superintelligente robot hebt die moleculen (de bouwstenen van alles om ons heen) kan begrijpen. Deze robot is een 3D-kaartlezer: hij kijkt niet alleen naar welke atomen er zijn, maar ook naar hoe ze in de ruimte staan. Als je een molecuul in je hand draait, moet de robot precies hetzelfde antwoord geven, alleen dan "gedraaid". Dit noemen wetenschappers SO(3)-equivariantie. Het is alsof je een kompas hebt: als je het draait, wijst de naald nog steeds naar het noorden, maar dan in de nieuwe richting.

Het probleem? Deze robot is enorm zwaar en traag. Hij is als een vrachtwagen die nodig is om een postzegel te bezorgen. Hij past niet op je telefoon of in een klein lab-apparaatje.

De auteurs van dit paper hebben een oplossing bedacht: ze hebben de robot ingeklemd (gequantizeerd) zonder dat hij zijn intelligentie of zijn kompas verliest. Ze noemen dit "Quantized SO(3)-Equivariant Graph Neural Networks".

Hier is hoe ze dat gedaan hebben, vertaald in alledaagse taal:

1. Het probleem: De "Klompige" Robot

Normaal gesproken werkt deze robot met heel precieze getallen (zoals 32-bit). Dat is als het meten van een afstand tot op een miljardste van een millimeter. Dat is nodig voor de robot om te begrijpen hoe een molecuul eruitziet, maar het kost veel energie en ruimte. Als je dit gewoon "afkapt" naar een lagere precisie (8-bit, zoals een simpele meetlat), gaat de robot in de war.

• Het risico: Hij vergeet niet alleen de grootte van een atoom, maar ook de richting. Het is alsof je een pijl tekent, maar door de ruwe meetlat verandert de pijl van richting. Dan wijst hij naar het noorden, terwijl hij naar het oosten moet wijzen. De robot is dan niet meer betrouwbaar.

2. De Oplossing: Drie Slimme Trucs

De auteurs hebben drie nieuwe trucs bedacht om de robot kleiner te maken zonder hem gek te maken:

Truc 1: Splitsen in "Grootte" en "Richting" (MDDQ)

Stel je voor dat je een pijl hebt. Deze pijl heeft twee eigenschappen: hoe lang hij is (grootte) en waar hij naartoe wijst (richting).

• De oude manier: Je probeerde de hele pijl in één keer te verkleinen. Dat ging vaak mis; de pijl werd korter én draaide een beetje.
• De nieuwe manier: De auteurs zeggen: "Laten we de lengte en de richting van elkaar scheiden."
  • Ze meten de lengte apart en ronden die af.
  • Ze meten de richting apart en houden die heel zuiver.
  • Analogie: Het is alsof je een foto van een vliegtuig maakt. Je verkleint de foto (grootte), maar je zorgt dat de neus van het vliegtuig nog steeds precies naar voren wijst (richting). Zelfs als de foto wazig is, weet je nog steeds waar het vliegtuig naartoe gaat.

Truc 2: Twee verschillende teams (Branch-Separated Training)

De robot heeft twee soorten hersenen:

1. De "Statische" hersenen: Deze kijken naar dingen die niet veranderen als je draait (zoals de totale energie).
2. De "Dynamische" hersenen: Deze kijken naar vectoren (richtingen) die wel veranderen als je draait.

• De fout: Veel mensen behandelen beide hersenen hetzelfde.
• De oplossing: De auteurs zeggen: "Behandel ze verschillend!" De statische hersenen mogen ruwer worden (simpeler), maar de dynamische hersenen krijgen een speciale, zorgzame behandeling. Ze trainen ze ook in twee stappen: eerst de statische hersenen, en pas later de dynamische hersenen, zodat de robot niet in paniek raakt tijdens het leren.

Truc 3: Een Stabilisator voor de Aandacht (Robust Attention)

De robot gebruikt een "aandachtssysteem" om te beslissen welke atomen belangrijk zijn voor elkaar.

• Het probleem: Bij lage precisie kunnen kleine rekenfouten ervoor zorgen dat de robot denkt dat een heel klein atoom belangrijker is dan een groot atoom, puur door een rekenfout.
• De oplossing: Ze voegen een normeerder toe. Het is alsof je bij een vergadering zegt: "Iedereen mag even spreken, maar we meten niet hoe hard iemand schreeuwt (grootte), maar alleen wat hij zegt (richting/inhoud)." Hierdoor wordt het systeem veel stabieler en maakt hij minder fouten, zelfs als de rekenmachine niet heel precies is.

3. Het Resultaat: De "Sportauto"

Door deze drie trucs te combineren, hebben ze het volgende bereikt:

• Snelheid: De robot is nu 2,4 tot 2,7 keer sneller. Hij doet wat hij voorheen in een seconde deed, nu in een fractie van een seconde.
• Grootte: Het model is 4 keer kleiner. Het past nu op een telefoon of een klein chipje in een laboratorium.
• Kwaliteit: En het allerbelangrijkste: hij is niet dommer geworden. Hij maakt bijna dezelfde fouten als de zware, dure versie. Hij draait nog steeds perfect mee als je het molecuul draait.

Waarom is dit belangrijk?

Vroeger moest je een molecuul analyseren op een enorme, dure supercomputer in een koele serverruimte. Met deze nieuwe techniek kun je straks een mobiel chemielab hebben.

• Stel je voor dat je een onbekend stofje op je telefoon scant en direct ziet of het giftig is of hoe het als medicijn werkt.
• Of een sensor in een fabriek die direct controleert of een chemische reactie veilig verloopt, zonder internetverbinding.

Kortom: Ze hebben de "vrachtwagen" omgebouwd tot een snelle, wendbare sportauto, zonder dat hij minder kracht heeft. Ze hebben de symmetrie (de wetten van de natuurkunde) gered, zelfs in een kleine, snelle verpakking.

Technische samenvatting

Probleemstelling

3D Graph Neural Networks (GNN's) die equivariant zijn voor 3D-rotaties (de groep SO(3)), zoals NequIP en So3krates, hebben de staat van de kunst bereikt in het voorspellen van moleculaire eigenschappen (zoals energie en krachten). Deze modellen respecteren de fysieke symmetrieën van moleculen, wat essentieel is voor nauwkeurigheid.

Echter, het inzetten van deze modellen op edge-apparaten (zoals mobiele telefoons of lab-on-chip sensoren) is zeer uitdagend vanwege:

1. Hoge rekenkosten: De tensor-operaties en aandachtmechanismen (attention) die nodig zijn voor SO(3)-equivariantie zijn computationally zwaar.
2. Kwetsbaarheid voor kwantisatie: Het toepassen van standaard lage-bit kwantisatie (bijv. 8-bit) op deze modellen leidt vaak tot een drastische daling in nauwkeurigheid en het verbreken van de rotatie-symmetrie. Standaard methoden verstoren de richting en grootte van vectoriële features, wat fysisch onzin oplevert.
3. Uniforme behandeling: Bestaande kwantisatiemethoden behandelen scalare (invariante) en vectoriële (equivariante) features vaak hetzelfde, terwijl deze fundamenteel verschillende verdelingen en rollen hebben.

Methodologie

De auteurs stellen een equivariantie-bewust kwantisatiekader voor dat specifiek is ontworpen voor transformer-achtige SO(3)-GNN's (gebaseerd op So3krates). Het kader bestaat uit drie kerninnovaties:

1. Magnitude-Direction Decoupled Quantization (MDDQ):

   • In plaats van vectorcomponenten direct te kwantiseren (wat de richting kan verstoren), worden equivariante vectorfeatures ontkoppeld in hun grootte (norm, $r$) en richting (eenheidsvector, $\hat{h}$).
   • De grootte wordt gekwantiseerd met een scalair kwantiseerder, en de richting wordt gekwantiseerd en vervolgens genormaliseerd tot een eenheidsvector.
   • Dit zorgt ervoor dat rotatie-informatie behouden blijft, zelfs bij lage precisie, en voorkomt dat kleine vectoren naar nul "instorten".

2. Branch-Separated Quantization-Aware Training (QAT):

   • Het model wordt opgesplitst in twee takken: een invariante tak (scalars, $\ell=0$) en een equivariante tak (vectoren, $\ell=1$).
   • Deze takken worden verschillend behandeld tijdens het trainen met kwantisatiebewustheid:
     • De scalare tak gebruikt standaard 8-bit kwantisatie.
     • De vectoriële tak gebruikt de MDDQ-methode.
   • Er wordt een gestructureerde trainingsschema gebruikt (warm-up): eerst wordt alleen de scalare tak gekwantiseerd, waarna de vectoriële tak wordt geactiveerd. Dit voorkomt destabilisatie van de geometrische features in de vroege trainingsfasen.

3. Robuuste Attention Normalisatie:

   • Aandachtscorrecties (dot-products) zijn gevoelig voor kwantisatieruis.
   • De auteurs introduceren een $\ell_2$-normalisatie van de query- en key-vectoren voordat de dot-product wordt berekend.
   • Dit beperkt de waarden van de attention-scores tot het interval $[-1, 1]$ en zorgt ervoor dat de aandacht voornamelijk afhangt van de richting van de vectoren en niet van hun grootte, wat de stabiliteit van INT8-berekeningen aanzienlijk verbetert.

4. Equivariantiebehoudende Loss (LEE Regularisatie):

   • Om te garanderen dat het gekwantiseerde model nog steeds correct reageert op rotaties, wordt een Local Equivariance Error (LEE) term toegevoegd aan de loss-functie tijdens het trainen.
   • Deze term straft afwijkingen af waarbij het voorspellen van een geroteerde input niet overeenkomt met het roteren van de voorspelling van de originele input.

Belangrijkste Bijdragen

• Eerste werk dat zich specifiek richt op de kwantisatie van SO(3)-equivariante GNN's.
• Ontwikkeling van MDDQ, een methode die de geometrische integriteit van vectorfeatures behoudt onder kwantisatie.
• Een branch-specifieke QAT-strategie die rekening houdt met de verschillende aard van scalare en vectoriële features.
• Demonstratie dat een 8-bit model prestaties kan leveren die vergelijkbaar zijn met full-precision (FP32) modellen, terwijl de efficiëntie drastisch wordt verbeterd.

Resultaten

De methode is geëvalueerd op de QM9 en rMD17 datasets:

• Nauwkeurigheid: Het 8-bit model bereikte een energie-MAE van 8.9 meV (op QM9) en een kracht-MAE van 22.6 meV/Å (op rMD17). Dit is slechts ~4.7% en ~6.6% hoger dan het FP32-baseline model, wat aanzienlijk beter is dan naieve kwantisatiemethoden (zoals PTQ of Degree-Quant) die grote nauwkeurigheidsverliezen lieten zien.
• Equivariantie: De Local Equivariance Error (LEE) bleef laag (~2 meV/Å), wat aantoont dat de rotatie-symmetrie behouden blijft.
• Efficiëntie:
  • Snelheid: 2.37x tot 2.73x snellere inferentie op CPU.
  • Geheugen: Ongeveer 4x kleinere modelgrootte.
• Agressieve kwantisatie (W4A8): Zelfs met 4-bit gewichten en 8-bit activaties behaalde het model betere resultaten dan het FP32-baseline op sommige taken, wat suggereert dat QAT fungeert als een effectieve regularisator.

Betekenis

Dit onderzoek opent de deur voor het praktische inzetten van symmetrie-bewuste AI-modellen op resource-beperkte apparaten. Het maakt het mogelijk om moleculaire eigenschappen in real-time te voorspellen op mobiele apparaten of geïntegreerde sensoren, zonder in te leveren op de fysieke correctheid (symmetrie) of nauwkeurigheid. De voorgestelde technieken bieden een fundamentele basis voor het comprimeren van andere symmetrie-bewarende modellen in de natuurwetenschappen en materialkunde.