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

Seit mehr als sechzig Jahren verlassen sich Wissenschaftler auf ein mächtiges Regelwerk, um zu verstehen, wie Atome zusammenhaften, um die Materialien um uns herum zu bilden. Dieser Rahmen, bekannt als Dichtefunktionaltheorie, betrachtet die Elektronenwolke, die ein Atom umgibt, nicht als eine Sammlung einzelner Teilchen, sondern als eine glatte, kontinuierliche Dichte. Dieser Ansatz ermöglicht es Forschern, die Eigenschaften von allem vorherzusagen, von neuen Medikamenten bis hin zu fortschrittlichen Batterien. Es gibt jedoch einen Haken. Um genaue Ergebnisse zu erhalten, erfordert die Theorie eine komplexe, repetitive Berechnung, die beinhaltet, das Verhalten jedes einzelnen Elektrons im System zu lösen. Wenn das System größer wird, explodiert die für diese Berechnungen benötigte Zeit, sodass die Simulation eines moderat großen Metallstücks oder eines großen biologischen Moleküls selbst auf den leistungsstärksten Supercomputern unmöglich wird. Dieser Engpass hat viele wichtige wissenschaftliche Fragen unerreichbar gehalten und Forscher gezwungen, zwischen Genauigkeit und Skalierbarkeit zu wählen.

Einem Team von Forschern am California Institute of Technology ist es gelungen, diese Barriere zu durchbrechen, indem sie einem Computer beigebracht haben, den teuersten Teil der Berechnung zu überspringen, ohne dabei an Genauigkeit zu verlieren. Anstatt zu versuchen, das Endergebnis direkt vorherzusagen oder eine schwierige mathematische Abkürzung zu erlernen, die oft scheitert, haben sie eine künstliche Intelligenz darauf trainiert, den spezifischen, wiederholten Schritt auszuführen, der in jedem einzelnen Berechnungszyklus vorkommt. Indem sie sich auf diesen Zwischenschritt konzentrierten, entwickelten sie eine Methode, mit der Systeme mit zehntausenden Atomen auf einer einzigen Grafikkarte simuliert werden können – eine Aufgabe, die zuvor tausende Prozessoren erfordert hätte. Das Ergebnis ist eine neue Art, Materie zu modellieren, die fast so schnell ist wie die einfachsten Näherungsverfahren, aber die hohe Präzision der strengsten Methoden beibehält.

Der Kern des Problems liegt darin, wie die Theorie mit Elektronen umgeht. Um den stabilen Zustand eines Materials zu finden, muss der Computer eine Karte der Elektronendichte wiederholt anpassen, bis sie sich nicht mehr verändert. In der traditionellen Methode erfordert diese Anpassung das Lösen eines massiven Eigenwertproblems, was im Wesentlichen das Finden der spezifischen Bewegungsmuster der Elektronen für jedes Atom im System ist. Diese Operation ist so rechenintensiv, dass ihre Kosten kubisch mit der Anzahl der Elektronen steigen; verdoppelt man die Größe des Systems, wird die Berechnung achtmal schwieriger. Jahrzehntelang haben Wissenschaftler versucht, dies zu umgehen, indem sie die Notwendigkeit, einzelne Elektronenmuster zu verfolgen, vollständig eliminierten – eine Strategie, die als orbitalfreie Dichtefunktionaltheorie bekannt ist. Bisherige Versuche, diese Abkürzung mittels maschinellem Lernen zu erlernen, waren jedoch problematisch. Einige versuchten, die Beziehung zwischen der Elektronendichte und der Energie zu erlernen, doch diese Beziehung ist mathematisch instabil und fehleranfällig. Andere versuchten, die endgültige Elektronendichte direkt aus der Anordnung der Atome vorherzusagen, doch diese Modelle versagten, wenn sie mit Systemen konfrontiert wurden, die größer oder chemisch anders waren als die, mit denen sie trainiert worden waren.

Die Forscher erkannten, dass der Schlüssel darin lag, nicht das Endziel vorherzusagen, sondern die Reise zu erlernen. In jedem Schritt der traditionellen Berechnung erstellt der Computer eine Karte der Kräfte, die auf die Elektronen wirken, und produziert daraus eine neue Karte dessen, wo die Elektronen wahrscheinlich zu finden sind. Dieser Vorwärtsschritt ist mathematisch stabil und wohldefiniert. Das Team trainierte ein spezialisiertes neuronales Netzwerk, eine Art künstliche Intelligenz, die darauf ausgelegt ist, räumliche Muster zu verstehen, um genau diese Transformation zu erlernen. Sie speisten das Netzwerk mit tausenden Beispielen dieser Kraftkarten und der entsprechenden Elektronendichtekarten, die sie produzierten, extrahiert aus einem vielfältigen Satz von 8.504 Molekülen und Festkörpern, einschließlich organischer Verbindungen, Isolatoren und Metalle. Das Netzwerk lernte, die neue Elektronendichte direkt aus den Kräften vorherzusagen, wodurch die langsame, schwere Berechnung durch eine schnelle, gelernte Vorhersage ersetzt wurde.

Was diesen Ansatz einzigartig macht, ist die Art und Weise, wie er mit der Komplexität verschiedener Systeme umgeht. Die Forscher entwarfen das Netzwerk so, dass es domäneninvariant ist, was bedeutet, dass es dieselben gelernten Regeln auf ein winziges Molekül und einen massiven Kristall anwenden kann, ohne für jede neue Größe neu trainiert werden zu müssen. Sie stellten zudem sicher, dass das Netzwerk die fundamentalen Symmetrien des Raums respektiert, sodass es versteht, dass das Rotieren oder Verschieben eines Moleküls dessen physikalische Eigenschaften nicht verändert. Bei Tests an Systemen, die es zuvor noch nie gesehen hatte, erwies sich das Modell als bemerkenswert robust. Bei einem Satz drogenähnlicher Moleküle, die Elemente und Größen enthielten, die nicht in den Trainingsdaten vorhanden waren, behielt die neue Methode eine hohe Genauigkeit bei, während die Fehler früherer direkter Vorhersagemodelle mit zunehmender Größe der Moleküle drastisch anstiegen. Die selbstkonsistente Natur der Methode ermöglicht es ihr, ihre eigenen Fehler bei jedem Schritt zu korrigieren, ganz ähnlich wie ein Navigator, der seine Position häufig überprüft, anstatt die gesamte Route auf einmal zu erraten.

Die Leistungsfähigkeit dieser Methode wurde deutlich, als das Team sie auf periodische Systeme, wie Kristalle, und auf erweiterte Defekte in Metallen anwandte. Bei Standardberechnungen für Kristalle muss der Computer die schwere Berechnung für viele verschiedene Punkte im mathematischen Raum wiederholen, der das sich wiederholende Muster des Kristalls beschreibt. Die neue Methode überspringt diese Wiederholung vollständig, indem sie die Elektronendichte auf der Elementarzelle entwickelt, ohne diese Punkte einzeln abtasten zu müssen. Dies ermöglichte es den Forschern, Berechnungen für metallische Systeme und Halbleiter mit einem einzigen Modell zu konvergieren und dabei die elektronischen Spektren und strukturellen Eigenschaften mit der gleichen Genauigkeit wie die traditionelle Methode zu reproduzieren. Darüber hinaus zeigte das Modell eine überraschende Fähigkeit zum Transfer zwischen verschiedenen theoretischen Näherungsverfahren. Ein Modell, das auf einer Art von Berechnung trainiert wurde, konnte mit einer anderen, genaueren Näherung für die Kräfte verwendet werden, ohne dass ein erneutes Training erforderlich war, indem man einfach die Eingangsfrüchte austauscht.

Der ultimative Test erfolgte, als die Forscher die Methode auf eine Magnesiumversetzung anwandten, einen Typ von Defekt in einem Metallkristall, der entscheidend dafür ist, zu verstehen, wie Materialien sich verformen und brechen. Diese Defekte erzeugen elastische Fernfelder, die Simulationen von tausenden Atomen erfordern, um sie genau zu modellieren. Frühere Versuche, ein solches System mit hoher Präzision zu simulieren, erforderten einen Supercomputer mit tausenden Grafikkarten und dauerten Tage. Die neue Methode, die auf einer einzigen modernen Grafikkarte läuft, konnte die Berechnung für ein System mit 8.250 Magnesiumatomen und 82.500 Valenzelektronen erfolgreich konvergieren. Die Zeit, die für den Abschluss der Berechnung benötigt wurde, wuchs fast linear mit der Größe des Systems, was bestätigte, dass der rechnerische Engpass beseitigt wurde. Die resultierenden Elektronendichten waren von der hochpräzisen Referenz nicht zu unterscheiden, was es den Forschern ermöglichte, die subtilen Energiedifferenzen zwischen verschiedenen Arten von Defekten aufzulösen, die die Festigkeit des Metalls bestimmen.

Diese Arbeit zeigt, dass maschinelles Lernen die Reichweite von Berechnungen zur Elektronenstruktur erweitern kann, indem es nicht die gesamte physikalische Theorie ersetzt, sondern ihren teuersten, wiederholten Betrieb beschleunigt. Durch das Erlernen der Vorwärtsabbildung von Kräften zu Dichte haben die Forscher ein Werkzeug geschaffen, das stabil, über verschiedene Arten von Materie transferierbar und in der Lage ist, Systeme in einem Maßstab zu handhaben, der zuvor unzugänglich war. Die Methode bietet ein direktes Diagnoseinstrument für ihre eigene Zuverlässigkeit: Wenn die Berechnung nicht konvergiert, signalisiert dies, dass sich das System außerhalb des vertrauenswürdigen Bereichs des Modells befindet – eine Eigenschaft, die direkten Vorhersagemodellen fehlt. Während das aktuelle Modell die Elektronendichte vorhersagt, können die Gesamtenergie und andere Eigenschaften immer noch mit einem einzigen abschließenden Schritt ermittelt werden, wodurch die Genauigkeit der vollständigen Theorie gewahrt bleibt. Dieser Ansatz öffnet die Tür zur Simulation komplexer Materialien, biologischer Systeme und chemischer Reaktionen auf einem Detailgrad und einer Skala, die einst im Bereich des theoretisch Unmöglichen lagen.