Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

Al meer dan zestig jaar vertrouwen wetenschappers op een krachtige set regels om te begrijpen hoe atomen aan elkaar blijven plakken om de materialen om ons heen te vormen. Dit kader, bekend als dichtheidsfunctionaaltheorie, behandelt de elektronenwolk rond een atoom niet als een verzameling individuele deeltjes, maar als een gladde, continue dichtheid. Deze benadering stelt onderzoekers in staat om de eigenschappen te voorspellen van alles, van nieuwe medicijnen tot geavanceerde batterijen. Er is echter een addertje onder het gras. Om nauwkeurige resultaten te krijgen, vereist de theorie een complexe, repetitieve berekening die inhoudt dat het gedrag van elke individuele elektron in het systeem moet worden opgelost. Naarmate het systeem groter wordt, explodeert de tijd die nodig is voor deze berekeningen, waardoor het zo snel gaat dat het simuleren van een bescheiden stuk metaal of een groot biologisch molecuul onmogelijk wordt, zelfs op de krachtigste supercomputers. Deze flessenhals heeft veel belangrijke wetenschappelijke vragen buiten bereik gehouden, waardoor onderzoekers gedwongen werden te kiezen tussen nauwkeurigheid en schaal.

Een team van onderzoekers aan het California Institute of Technology heeft een manier gevonden om deze barrière te doorbreken door een computer te leren het meest kostbare deel van de berekening over te slaan zonder aan nauwkeurigheid in te boeten. In plaats van te proberen het eindantwoord direct te voorspellen of een moeilijke wiskundige afkorting te leren die vaak faalt, hebben ze een kunstmatige intelligentie geleerd om de specifieke, herhaalde stap uit te voeren die in elke enkele berekeningscyclus voorkomt. Door zich te concentreren op deze tussenstap, hebben ze een methode gecreëerd die systemen met tienduizenden atomen kan simuleren op een enkele grafische kaart, een taak die voorheen duizenden processoren vereiste. Het resultaat is een nieuwe manier om materie te modelleren die bijna net zo snel is als de eenvoudigste benaderingen, maar de hoge precisie behoudt van de meest rigoureuze methoden.

De kern van het probleem ligt in de manier waarop de theorie met elektronen omgaat. Om de stabiele toestand van een materiaal te vinden, moet de computer herhaaldelijk een kaart van de elektrondichtheid aanpassen totdat deze niet meer verandert. In de traditionele methode vereist deze aanpassing het oplossen van een massaal eigenwaardeprobleem, wat in essentie het vinden is van de specifieke patronen van elektronbeweging voor elk atoom in het systeem. Deze operatie is zo rekenintensief dat de kosten kubisch groeien met het aantal elektronen; het verdubbelen van de grootte van het systeem maakt de berekening acht keer zwaarder. Decennialang hebben wetenschappers geprobeerd dit te omzeilen door de noodzaak om individuele elektronpatronen te volgen volledig weg te nemen, een strategie die bekend staat als orbitaalvrije dichtheidsfunctionaaltheorie. Echter, eerdere pogingen om deze afkorting met behulp van machine learning te leren, zijn gestruikeld. Sommigen probeerden de relatie tussen de elektrondichtheid en de energie te leren, maar deze relatie is wiskundig instabiel en gevoelig voor fouten. Anderen probeerden de uiteindelijke elektrondichtheid direct te voorspellen vanuit de rangschikking van atomen, maar deze modellen faalden wanneer ze werden gevraagd om systemen te voorspellen die groter of chemisch anders waren dan de systemen waarop ze waren getraind.

De onderzoekers realiseerden zich dat de sleutel lag in het stoppen met het proberen te voorspellen van de eindbestemming en in plaats daarvan de reis leren kennen. In elke stap van de traditionele berekening neemt de computer een kaart van de krachten die op de elektronen werken en produceert een nieuwe kaart van waar de elektronen waarschijnlijk te vinden zijn. Deze voorwaartse stap is wiskundig stabiel en goed gedefinieerd. Het team trainde een gespecialiseerd neuraal netwerk, een type kunstmatige intelligentie ontworpen om ruimtelijke patronen te begrijpen, om precies deze transformatie te leren. Ze voedden het netwerk met duizenden voorbeelden van deze krachtkaarten en de bijbehorende elektrondichtheidskaarten die zij produceerden, afkomstig van een diverse set van 8.504 moleculen en vaste stoffen, waaronder organische verbindingen, isolatoren en metalen. Het netwerk leerde de nieuwe elektrondichtheid direct vanuit de krachten te voorspellen, waardoor de trage, zware berekening effectief werd vervangen door een snelle, geleerde voorspelling.

Wat deze aanpak uniek maakt, is hoe deze omgaat met de complexiteit van verschillende systemen. De onderzoekers ontwierpen het netwerk om domein-invariant te zijn, wat betekent dat het dezelfde geleerde regels kan toepassen op een klein molecuul en een massief kristal zonder dat het voor elk nieuw formaat opnieuw getraind hoeft te worden. Ze zorgden er ook voor dat het netwerk de fundamentele symmetrieën van de ruimte respecteert, zodat het begrijpt dat het roteren of verschuiven van een molecuul de fysieke eigenschappen niet verandert. Wanneer het werd getest op systemen die het nog nooit had gezien, bleek het model opmerkelijk robuust. Op een set van medicijnachtige moleculen die elementen en groottes bevatten die niet aanwezig waren in de trainingsdata, behield de nieuwe methode een hoge nauwkeurigheid, terwijl eerdere directe voorspellingsmodellen zagen dat hun fouten exponentieel toenamen naarmate de moleculen groter werden. De zelfconsistente aard van de methode stelt het systeem in staat om zijn eigen fouten bij elke stap te corrigeren, vergelijkbaar met een navigator die regelmatig zijn positie controleert in plaats van de hele route in één keer te gokken.

De kracht van deze methode werd duidelijk toen het team het toepaste op periodieke systemen, zoals kristallen, en op uitgebreide defecten in metalen. In standaardberekeningen voor kristallen moet de computer de zware berekening herhalen voor vele verschillende punten in de wiskundige ruimte die het herhalende patroon van het kristal beschrijft. De nieuwe methode slaat deze herhaling volledig over door de elektrondichtheid op de eenheidscel te laten evolueren zonder deze punten individueel te hoeven bemonsteren. Hierdoor konden de onderzoekers de berekeningen voor metallische systemen en halfgeleiders laten convergeren met één enkel model, waarbij de elektronische spectra en structurele eigenschappen met dezelfde nauwkeurigheid als de traditionele methode werden gereproduceerd. Bovendien toonde het model een verrassend vermogen om te transfereren tussen verschillende theoretische benaderingen. Een model dat getraind was op één type berekening kon worden gebruikt met een andere, nauwkeurigere benadering voor de krachten zonder dat er hertraining nodig was, simpelweg door de invoerkrachten te verwisselen.

De ultieme test kwam toen de onderzoekers de methode toepasten op een magnesiumdislocatie, een type defect in een metaalkristal dat cruciaal is voor het begrijpen van hoe materialen vervormen en breken. Deze defecten creëren lang reikende elastische velden die simulaties van duizenden atomen vereisen om accuraat te modelleren. Eerdere pogingen om een dergelijk systeem met hoge precisie te simuleren vereisten een supercomputer met duizenden grafische processoren en namen dagen om te draaien. De nieuwe methode, draaiend op een enkele moderne grafische kaart, slaagde erin de berekening voor een systeem met 8.250 magnesiumatomen en 82.500 valentie-elektronen te laten convergeren. De tijd die nodig was om de berekening te voltooien, groeide bijna lineair met de grootte van het systeem, wat bevestigde dat de computationele flessenhals was verwijderd. De resulterende elektrondichtheden waren ononderscheidbaar van de hoog-precieze referentie, waardoor de onderzoekers de subtiele energieverschillen tussen verschillende soorten defecten konden oplossen die de sterkte van het metaal bepalen.

Dit werk demonstreert dat machine learning de reikwijdte van elektronische structuurberekeningen kan uitbreiden, niet door de volledige fysieke theorie te vervangen, maar door de meest dure, herhaalde operatie te versnellen. Door de voorwaartse kaart van krachten naar dichtheid te leren, hebben de onderzoekers een instrument gecreëerd dat stabiel is, overdraagbaar is naar verschillende soorten materie en in staat is om systemen op een schaal te hanteren die voorheen ontoegankelijk was. De methode biedt een directe diagnose voor de eigen betrouwbaarheid; als de berekening niet convergeert, signaleert dit dat het systeem buiten het vertrouwde bereik van het model valt, een kenmerk dat directe voorspellingsmodellen missen. Hoewel het huidige model de elektrondichtheid voorspelt, kunnen de totale energie en andere eigenschappen nog steeds worden verkregen met één laatste stap, waardoor de nauwkeurigheid van de volledige theorie behouden blijft. Deze aanpak opent de deur naar het simuleren van complexe materialen, biologische systemen en chemische reacties op een niveau van detail en schaal die ooit het domein van theoretische onmogelijkheid waren.