Machine-learned, finite temperature Fermi-operator expansions suitable for GPUs and AI-hardware

Dit artikel presenteert een door machine learning verbeterde, recursieve Fermi-operator-ontwikkelingsmethode bij eindige temperatuur die is gebaseerd op het SP2-schema, elektronenstructuurberekeningen koppelt aan diepe neurale netwerkarctitecturen en door gebruik te maken van geoptimaliseerde matrix-matrixvermenigvuldigingen en affiene herschaling om expliciete diagonalisatie en modelhertraining te vermijden, snelheidswinsten van de orde van grootte op GPUs en AI-hardware mogelijk maakt.

De kern

Het Grote Geheel: Een Snellere Manier om Atomen te Simuleren

Stel je voor dat je probeert te voorspellen hoe een menigte mensen (elektronen) zich zal verplaatsen en zal interageren in een kamer (een materiaal). In de wereld van de kwantumfysica is dit ongelooflijk moeilijk. Om het exacte antwoord te krijgen, moet je meestal een enorm, complex raadsel oplossen dat "diagonalisatie" wordt genoemd.

Denk aan diagonalisatie als het proberen om een miljoen boeken te sorteren door elke enkele pagina van elk boek te lezen om de juiste volgorde te vinden. Het is accuraat, maar het kost veel tijd, vooral naarmate de kamer groter wordt.

De auteurs van dit artikel hebben een shortcut gebouwd. In plaats van elke pagina te lezen, hebben ze een "slimme gok"-machine gemaakt die leert hoe het boeken bijna direct moet sorteren. Ze noemen dit een Machine-learned Fermi-operator expansie.

Het Probleem: Hete versus Koude Menigten

In het verleden werkten deze shortcuts alleen goed wanneer de "menigte" zeer koud was (nul temperatuur). In een koude menigte staat iedereen stil in een zeer voorspelbare rij. De wiskunde is simpel: je zit óf in de rij, óf niet.

In de echte wereld zijn dingen echter vaak "heet". Wanneer elektronen heet worden, worden ze onrustig. Sommige mensen die in de rij stonden, stappen misschien uit, en sommige die wachtten, stappen erin. Dit creëert een "vage" grens waar mensen deels binnen en deels buiten zijn.

Eerdere shortcuts faalden hier omdat ze te stijf waren. Ze konden de "vage" aard van een hete menigte niet aan.

De Oplossing: Een Neuraal Netwerk Leren om te "Plotten"

De auteurs realiseerden zich dat de wiskunde die wordt gebruikt om de koude menigte te sorteren er precies hetzelfde uitziet als de structuur van een Deep Neural Network (het soort AI dat wordt gebruikt om gezichten te herkennen of gedichten te schrijven).

• De Oude Manier (SP2): Stel je een machine voor die een getal neemt en het ofwel kwadrateert ($x^2$) of een specifieke aftrekking uitvoert ($2x - x^2$). Het herhaalt dit keer op keer, "plottend" de getallen totdat ze 0 of 1 worden. Dit werkt geweldig voor koude menigten.
• De Nieuwe Manier (MLSP2): De auteurs namen deze machine en gaven het een "brein". In plaats van vaste regels te gebruiken, trainden ze de machine met Machine Learning. Ze leerden het om zijn eigen interne knoppen (coëfficiënten) aan te passen zodat het de "vage" hete menigte perfect kon hanteren.

Denk hierover na:

• Oude Machine: Een stijve stempel die alleen "Ja" of "Nee" print.
• Nieuwe Machine: Een flexibele 3D-printer die leert precies hoe hij het "Ja" en "Nee" moet vormen om een gladde, perfecte curve in het midden te creëren, afhankelijk van hoe heet de menigte is.

De Magische Truc: Eén Model Past bij Veel Temperaturen

Normaal gesproken moet je, als je de temperatuur van je simulatie verandert, je AI-model helemaal opnieuw trainen. Dat kost eeuwig.

De auteurs ontdekten een slimme truc genaamd Affine Rescaling.
Stel je een kaart van een stad voor. Als je wilt in- of uitzoomen, hoef je de hele stad niet opnieuw te tekenen; je strekt of verkleint gewoon de kaart.

De auteurs ontdekten dat ze hun AI-model slechts eenmaal konden trainen voor een specifiek "zoomniveau" (een specifieke temperatuur en chemisch potentiaal). Vervolgens konden ze, voor elke andere temperatuur binnen een bepaald bereik, de invoergegevens (de Hamilton-matrix) simpelweg "strekken" voordat ze deze aan het model gaven. Het model hoeft niets opnieuw te leren; het ziet de gegevens gewoon in een iets andere schaal en geeft het juiste antwoord.

Dit betekent dat ze simulaties kunnen uitvoeren waarbij de temperatuur constant verandert (zoals in een chemische reactie) zonder te stoppen om de AI opnieuw te trainen.

De Hardware: AI-chips Gebruiken voor Wetenschap

Het artikel benadrukt dat deze methode specifiek is gebouwd voor moderne computerchips, met name GPUs (Graphics Processing Units) en Tensor Cores (chips die zijn ontworpen voor AI).

• De Analogie: Traditionele diagonalisatie is als een meester-timmerman die elk meubelstuk met de hand uithouwt. Het is nauwkeurig maar traag.
• De Nieuwe Methode: Dit is als het gebruik van een hoge-snelheids 3D-printer. Het maakt gebruik van de specifieke architectuur van AI-chips om enorme berekeningen (matrixvermenigvuldigingen) ongelooflijk snel uit te voeren.

De auteurs testten dit op een Nvidia RTX 6000 Ada GPU. Ze ontdekten dat hun methode 9 tot 16 keer sneller was dan de standaard, sterk geoptimaliseerde methoden die wetenschappers vandaag de dag gebruiken, terwijl ze toch een hoge nauwkeurigheid behielden.

Samenvatting van Resultaten

1. Snelheid: Ze bereikten een enorme snelheidswinst (tot 16x) bij het berekenen van hoe elektronen zich gedragen in materialen, vooral op moderne AI-hardware.
2. Nauwkeurigheid: Ze kunnen "hete" elektronen (fractieve bezetting) modelleren met extreme precisie, iets wat eerdere shortcuts niet goed konden doen.
3. Efficiëntie: Door het model één keer te trainen en wiskundige trucs te gebruiken om invoer te herschalen, vermijden ze de noodzaak om het model elke keer opnieuw te trainen wanneer de temperatuur in een simulatie verandert.
4. Geen "Magische" Diagonalisatie: Ze vermijden volledig de trage, zware wiskunde van diagonalisatie en vertrouwen in plaats daarvan op snelle, herhaalde vermenigvuldigingsstappen die AI-chips graag uitvoeren.

Kortom, de auteurs hebben een traag, stijf wiskundig proces omgezet in een snel, flexibel, door AI aangedreven hulpmiddel dat ongelooflijk efficiënt draait op moderne computerchips, waardoor wetenschappers complexe materialen veel sneller kunnen simuleren dan voorheen.

Technische samenvatting

Technische Samenvatting: Machine-geleerde, eindige-temperatuur Fermi-operator-expansies

Probleemstelling
Berekeningen van elektronische structuur, met name binnen Kohn-Sham Dichtefunctietheorie (KS-DFT), worden computatiekundig beperkt door de kubische schaling van de kosten voor het diagonaliseren van de Hamiltoniaan-matrix om het eigenwaardeprobleem op te lossen. Hoewel recursieve Fermi-operator-expansieschema's, zoals de Second-Order Spectral Projection (SP2)-methode, een manier bieden om de dichtheidsmatrix direct te berekenen zonder diagonalisatie, zijn bestaande efficiënte implementaties beperkt tot een elektronische temperatuur van nul. Bij een temperatuur van nul is de dichtheidsmatrix idempotent (bezettingen zijn strikt 0 of 1). Echter, vereisen veel fysische systemen—zoals metalen of systemen bij verhoogde elektronische temperaturen—fractionele orbitaalbezettingen om gedegenereerde eigenstates of thermische vervaging nauwkeurig te modelleren.

Eerdere pogingen om SP2 te generaliseren naar eindige temperaturen hielden in dat de recursie werd afgekapt om thermische vervaging in te voeren. Deze afgekapt expansies zijn echter inherent benaderend en falen in het reproduceren van de exacte Fermi-functie, met name in de buurt van het chemische potentiaal waar nauwkeurigheid kritiek is. Alternatieve methoden zoals Chebyshev-expansies of Padé-benaderingen vereisen óf verboden hoge polynoomordes om Gibbs-oscillaties te onderdrukken óf brengen aanzienlijke computatiekosten met zich mee door herhaalde oplossingen van lineaire systemen.

Methodologie
De auteurs stellen een raamwerk voor dat de recursieve SP2-methode generaliseert naar eindige temperaturen door de algebraïsche structuur ervan af te beelden op deep neural network (DNN)-architecturen. Het kerninzicht is dat de recursieve SP2-bijwerkingen lijken op de lagen van een neurale netwerk. Door de expansiecoëfficiënten te behandelen als trainbare gewichten en bias-termen, construeren de auteurs machine learning-modellen die in staat zijn de Fermi-verdelingsfunctie te benaderen met fractionele bezettingen bij willekeurige temperaturen.

Belangrijke methodologische componenten zijn:

1. Neurale Netwerk Architecturen:

   • MLSP2 (Machine-Learned SP2): Een generalisatie van SP2 waarbij de kwadratische bijwerkregels ($X^2$ of $2X-X^2$) worden vervangen door leerbare kwadratische polynomen ($ax^2 + bx + c$) met een accumulator-term. Dit stelt het model in staat de exacte thermische vervaging van de Fermi-functie te benaderen in plaats van een afgekapt stapfunctie.
   • Max-SP2: Een expressievere architectuur die "skip connections" incorporateert, waarbij elke laag het kwadraat is van een lineaire combinatie van alle voorgaande lagen.
   • Skip-SP2: Een gecomprimeerde versie van Max-SP2 die een eindig geheugen van recente lagen en accumulators gebruikt om expressiviteit en geheugengebruik in evenwicht te brengen.

2. Entropiebenadering:
   De auteurs ontwikkelen ook een recursief schema om de elektronische entropiefunctie, $s(x)$, te benaderen, wat noodzakelijk is voor het berekenen van de elektronische vrije energie. Zij gebruiken een geschaald product van de Fermi-functie en haar complement, $f(x)(1-f(x))$, als een eerste schatting, die vervolgens wordt verfijnd via een recursieve kwadratische expansie die is getraind om de tweede afgeleide van de ware entropie bij het chemische potentiaal te matchen.

3. Training en Optimalisatie:
   Modellen worden getraind op scalair input binnen het eenheidsinterval $[0, 1]$ in plaats van volledige matrices, gebruikmakend van het Levenberg–Marquardt-algoritme met geodetische versnelling. De trainingsdata wordt bemonsterd met een weging die evenredig is met de afgeleide van de Fermi-functie om de maximale fout in de buurt van het chemische potentiaal te minimaliseren.

4. Affiene Herschaling en Overdraagbaarheid:
   Een kritieke innovatie is het gebruik van affiene herschaling om de noodzaak van hertraining te elimineren wanneer simulatieparameters veranderen. Door de Hamiltoniaan ($H'$), het chemische potentiaal ($\mu'$) en de inverse temperatuur ($\beta'$) te normaliseren, kan een enkel model dat is getraind bij specifieke parameters $(\beta_0, \mu_0)$ worden toegepast op een breed "geldigheidsgebied" van andere parameters. Dit wordt bereikt door de invoer-Hamiltoniaan te herschalen om de trainingscondities te matchen, waardoor dezelfde set gewichten kan worden gebruikt tijdens een simulatie over variërende temperaturen en chemische potentialen.

5. Hardware-implementatie:
   De algoritmen zijn geoptimaliseerd voor moderne GPU's en AI-hardware (specifiek NVIDIA Tensor Cores). De auteurs maken gebruik van gemengde precisie-aritmetiek (FP16/FP32) om matrixkwadrateringsoperaties efficiënt uit te voeren, waarbij ze de symmetrie van de Hamiltoniaan benutten om het aantal benodigde vermenigvuldigingen en gegevensoverdrachten te verminderen.

Belangrijkste Resultaten

• Nauwkeurigheid: De MLSP2-modellen bereiken fouten in de orde van grootte van $10^{-7}$ voor de Fermi-functie-benadering, wat aanzienlijk beter presteert dan afgekapt SP2-methoden (die fouten rond $10^{-2}$ hebben) en de precisie van diagonalisatie in dubbele precisie in veel regimes evenaart.
• Prestaties: Op een NVIDIA RTX 6000 Ada GPU toont de MLSP2-aanpak een 16-voudige snelheidswinst ten opzichte van diagonalisatie in dubbele precisie (met cuSOLVER) voor intermediaire matrixgroottes en een 9-voudige snelheidswinst voor grotere matrices. Zelfs in vergelijking met diagonalisatie in enkele precisie biedt MLSP2 een 2x tot 5x snelheidswinst terwijl het superieure stabiliteit en nauwkeurigheid behoudt.
• Schaalbaarheid: De methode vertrouwt uitsluitend op sterk geoptimaliseerde kernroutines voor matrix-matrix vermenigvuldiging en vermijdt expliciete diagonalisatie. Het aantal lagen dat nodig is om een doelnaauwkeurigheid te bereiken, schaalt logaritmisch met de inverse temperatuur ($\beta$), wat efficiënte berekening mogelijk maakt zelfs bij lage temperaturen.

Betekenis en Claims
Het artikel claimt dat deze aanpak een robuuste, generaliseerbare oplossing biedt voor berekeningen van elektronische structuur bij eindige temperaturen die de computatieknelpunten van diagonalisatie vermijdt. Door SP2 te generaliseren via machine learning, stellen de auteurs de berekening van dichtheidsmatrijzen voor systemen met fractionele bezettingen mogelijk voor een fractie van de kosten van traditionele methoden.

De betekenis ligt in het vermogen om dynamische simulaties bij eindige temperaturen uit te voeren (zoals kwantum-moleculair dynamica) waarbij het chemische potentiaal en de elektronische temperatuur fluctueren tussen tijdstappen. De strategie voor affiene herschaling zorgt ervoor dat een enkel voorgetraind model gedurende een simulatie kan worden hergebruikt zonder hertraining, waardoor de methode praktisch wordt voor grootschalige toepassingen. Bovendien is de aanpak specifiek afgestemd om de prestatiekenmerken van moderne AI-hardware (Tensor Cores) te benutten, wat aanzienlijke snelheidswinsten biedt ten opzichte van door leveranciers geoptimaliseerde diagonalisatieroutines, terwijl hoge numerieke nauwkeurigheid wordt behouden.