Technische Samenvatting: Compacte SO(3)-equivariante atomaire fundamentele modellen via structurele pruning

Probleemstelling

SO(3)-equivariante grafische neurale netwerken (GNN's) hebben zich gevestigd als het dominante paradigma voor atomaire fundamentele modellen, waarbij ze hoge nauwkeurigheid en data-efficiëntie bereiken door rotatiesymmetrie expliciet te coderen. Deze modellen vertrouwen echter op tensorproductoperaties van hogere orde om geometrische kenmerken te koppelen, wat resulteert in rekenkosten die oplopen tot $O(L_{max}^6)$, waarbij $L_{max}$ de maximale orde van de irreducibele representaties is. Dit creëert een aanzienlijke afweging tussen modelnauwkeurigheid en inferentie-efficiëntie.

Bestaande strategieën voor modelcompressie – specifiek Kennisdistillatie (KD) en Quantisatie – hebben beperkingen in dit domein. KD vereist vaak het trainen van een studentmodel vanaf nul en hergebruikt geen leerkrachtgewichten direct, terwijl quantisatie de bitprecisie verlaagt maar het aantal parameters en de tensorproductpaden onveranderd laat. Bovendien kunnen standaard structurele pruning-methoden die zijn ontwikkeld voor Large Language Models (LLM's) niet direct worden toegepast op SO(3)-equivariante modellen. Het naïef verwijderen van kenmerkdimensies breekt de rotatiesymmetrie van het tensorproductpad, waardoor de fundamentele equivariantie-eigenschap die vereist is voor fysieke consistentie wordt geschonden. Eerdere pogingen tot pruning van atomaire modellen waren beperkt tot grofkorrelige blokreductie (bijvoorbeeld het verwijderen van hele interactielagen), wat het receptieve veld van het model verkleint en bijzonder schadelijk is voor ondiepe architecturen.

Methodologie

De auteurs stellen een structureel pruningkader voor dat specifiek is ontworpen voor SO(3)-equivariante atomaire fundamentele modellen. De methode werkt op voorgeöefende checkpoints om kleinere, dichte modellen te produceren die SO(3)-equivariantie en hoge nauwkeurigheid behouden.

1. Equivariante Pruninggranulariteit

De kerninnovatie ligt in de definitie van de pruningeenheid. In SO(3)-equivariante modellen worden kenmerken ontbonden in irreducibele representaties van orde $l \in [0, L_{max}]$. Een kenmerkvector voor een specifiek kanaal $k$ en orde $l$ bestaat uit $2l+1$ componenten geïndexeerd door $m \in [-l, l]$.

• Beperking: De rotatieoperatie in SO(3) koppelt alle $m$-componenten binnen een specifieke orde $l$. Het verwijderen van individuele componenten (elementsgewijze pruning) breekt de equivariantie.
• Oplossing: De methode behandelt de deelruimte die wordt opgespannen door de $2l+1$ componenten van een specifieke orde $l$ en kanaal $k$ als een atomaire eenheid. Pruning wordt uitgevoerd op het tuple $(k, l)$, weergegeven door een binaire masker $z^{(t)}_{k,l}$. Dit zorgt ervoor dat als een kenmerkblok behouden blijft, al zijn rotatiecomponenten worden behouden, waardoor SO(3)-equivariantie strikt wordt gehandhaafd.

2. Fysiek-informeerde Relevantiecriterium

Om te bepalen welke $(k, l)$-blokken behouden moeten blijven, stellen de auteurs een relevantiescore voor die gebaseerd is op de gevoeligheid van het potentieel-energieoppervlak.

• Dubbel doel: MLIP's moeten zowel potentiële energie ($E$) als krachten ($F_i = -\nabla_{r_i} E$) voorspellen.
• Maatstaf: De relevantiescore $I^{(t)}_{k,l}$ is afgeleid van een Taylor-ontwikkeling van eerste orde van de energieverandering bij het verwijderen van een kenmerkblok. Het aggregeert het product van de kenmerkgrootte en de gradiënt van de energie met betrekking tot dat kenmerk over alle atomen en representatiedimensies:
  $$I^{(t)}_{k,l} = \mathbb{E}_{R,Z \sim \mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h^{(t)}_{j,klm}} \cdot h^{(t)}_{j,klm} \right|$$
• Voordeel: Dit criterium is SO(3)-invariant en veroorzaakt verwaarloosbare overhead, omdat de gradiënten $\frac{\partial E}{\partial h}$ al worden berekend tijdens standaard krachtinferentie.

3. Structurele Uitlijning en Pijplijn

Het kader volgt een pijplijn in vier fasen:

1. Kalibratie: Een kleine subset van pre-training data wordt door het grote model gevoerd om relevantiescores te berekenen.
2. Doelgestuurde Pruning: Een doelarchitectuur wordt gedefinieerd. De best gerangschikte $(k, l)$-blokken worden behouden op basis van relevantiescores.
3. Structurele Uitlijning: De modelstructuur wordt fysiek gewijzigd om overeen te komen met de gereduceerde dimensies. Dit omvat:
   • Tensorproductherconfiguratie: Het snijden van radiale gewichten en Clebsch-Gordan-paden om uit te lijnen met behouden kanalen.
   • Gewichtssnijden: Het fysiek verwijderen van rijen en kolommen uit lineaire transformatiegewichten om een dichte, kleinere gewichtstensor te creëren.
4. Opnieuw trainen & Fine-tunen: Het gepreunteerde model ondergaat lichtgewicht opnieuw trainen op een subset van pre-training data om te herstellen van benaderingsfouten, gevolgd door fine-tunen op specifieke downstream-taken.

Belangrijkste Bijdragen

1. Equivariante Pruningsmethode: Introductie van een structurele pruningmethode voor SO(3)-equivariante tensors van hogere orde over kanaal- en ordedimensies die strikt SO(3)-equivariantie behoudt door $(k, l)$-blokken als atomaire eenheden te behandelen.
2. Pijplijn in vier fasen: Een complete pijplijn (Kalibratie, Pruning, Opnieuw trainen, Fine-tunen) voor het comprimeren van fundamentele modellen.
3. Generaliseerbaarheid: Demonstratie dat de methode van toepassing is op meerdere architecturen, waaronder MACE, SevenNet en eSCN.
4. Combineerbaarheid: Validatie dat structurele pruning kan worden gecombineerd met quantisatie en kennisdistillatie voor cumulatieve efficiëntiewinsten.

Experimentele Resultaten

De methode werd voornamelijk geëvalueerd op MACE-MP (anorganische materialen) en MACE-OFF (organische moleculen), met generaliseerbaarheidstests op SevenNet en eSCN.

Nauwkeurigheid versus Trainen vanaf Nul

• Superioriteit: Gepreunteerde modellen presteren consequent beter dan modellen van dezelfde grootte die vanaf nul zijn getraind.
  • Op MACE-MP bereikten gepreunteerde modellen een lagere Force MAE en betere prestaties buiten de verdeling op de sAlex-dataset in vergelijking met basismodellen die vanaf nul waren getraind.
  • Op Matbench Discovery presteerde een gepreunteerd MACE-MP-model (met slechts 25% van de pre-training data) beter dan het officiële vanaf nul getrainde kleine model op 7 van de 9 maatstaven.
  • Op MACE-OFF kwamen gepreunteerde modellen overeen met of overtroffen basismodellen die vanaf nul waren getraind over diverse moleculaire subsets, zelfs bij gebruik van slechts 50% van de pre-training data.
• Downstream-prestaties: Het fine-tunen van gepreunteerde modellen op acht anorganische downstream-datasets verlaagde de energie- en krachtfouten met respectievelijk 70,1% en 34,4% in vergelijking met het trainen van taakspecifieke modellen vanaf nul.

Efficiëntiewinsten

• Parameterreductie: Gecomprimeerde modellen bevatten 1,5× tot 4× minder parameters.
• Trainingskosten: De volledige pijplijn vereist 2,5× tot 4× minder pre-training rekenkracht dan het trainen van een klein model vanaf nul (bijvoorbeeld 181 GPU-uren versus 750 GPU-uren voor een specifieke configuratie).
• Inferentiesnelheid: Gepreunteerde modellen bereiken een 2,7× tot 3,0× versnelling van de inferentie en tot 5,7× reductie in piekgeheugengebruik.

Combinatie met Andere Methoden

• Quantisatie: Het combineren van pruning met FP32-quantisatie leverde een 9,6× verbetering in doorvoer op (relatief ten opzichte van het basis FP64-model).
• Kennisdistillatie: Het combineren van pruning met Hessiaan-gebaseerde KD verfijnde de modelprestaties verder, wat aantoont dat de twee methoden complementair zijn.

Betekenis en Claims

Het artikel beweert dat structurele pruning een natuurlijke en effectieve weg biedt om de kloof tussen nauwkeurigheid en efficiëntie in atomaire fundamentele modellen te overbruggen. Door gewichten direct te erven van grote, voorgeöefende checkpoints en deze te comprimeren terwijl fysieke symmetrieën worden behouden, produceert de methode compacte modellen die:

1. Nauwkeuriger zijn dan modellen van dezelfde grootte die vanaf nul zijn getraind.
2. Aanzienlijk goedkoper zijn om te trainen en te implementeren.
3. Fysisch consistent zijn, waarbij de strikte SO(3)-equivariantie wordt behouden die vereist is voor betrouwbare moleculardynamica-simulaties (gevalideerd via MD-simulaties van vloeibaar water die een stabiele energiedrift en accurate radiale verdelingsfuncties tonen).

De auteurs benadrukken dat het kernprincipe – het evalueren van het belang van gekoppelde irreducibele representaties – breed toepasbaar is op elke equivariante architectuur die kenmerken weergeeft als sommen van SO(3)-irreducibele representaties. Het werk suggereert dat pruning practitioners in staat stelt om hoogwaardige fundamentele modellen te implementeren in omgevingen met beperkte middelen, zonder de transferabiliteit op te offeren die tijdens grootschalige pre-training is geleerd.