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

Dieser Beitrag stellt eine durch maschinelles Lernen verbesserte, bei endlichen Temperaturen angewandte rekursive Fermi-Operator-Entwicklungsmethode vor, die auf dem SP2-Schema basiert, elektronische Strukturberechnungen auf tiefe neuronale Netzarchitekturen abbildet und durch Nutzung optimierter Matrix-Matrix-Multiplikationen sowie affiner Skalierung, um eine explizite Diagonalisierung und ein erneutes Modelltraining zu vermeiden, eine Beschleunigung um Größenordnungen auf GPUs und KI-Hardware ermöglicht.

Das Wesentliche

Das große Ganze: Ein schnellerer Weg, Atome zu simulieren

Stellen Sie sich vor, Sie versuchen vorherzusagen, wie sich eine Menschenmenge (Elektronen) in einem Raum (ein Material) bewegt und interagiert. In der Welt der Quantenphysik ist dies unglaublich schwierig. Um die exakte Antwort zu erhalten, müssen Sie normalerweise ein riesiges, komplexes Rätsel lösen, das „Diagonalisierung" heißt.

Denken Sie an Diagonalisierung wie den Versuch, eine Million Bücher zu sortieren, indem Sie jede einzelne Seite jedes Buches lesen, um die richtige Reihenfolge zu finden. Es ist genau, aber es dauert lange, besonders wenn der Raum größer wird.

Die Autoren dieses Papers haben einen Abkürzungsweg entwickelt. Anstatt jede Seite zu lesen, haben sie eine „intelligente Schätzmachine" erstellt, die lernt, die Bücher fast augenblicklich zu sortieren. Sie nennen dies eine maschinell gelernte Fermi-Operator-Entwicklung.

Das Problem: Heiße vs. kalte Menschenmengen

In der Vergangenheit funktionierten diese Abkürzungen nur dann gut, wenn die „Menschenmenge" sehr kalt war (Nulltemperatur). In einer kalten Menge steht jeder in einer sehr vorhersehbaren Reihe still. Die Mathematik ist einfach: Man ist entweder in der Reihe oder man ist es nicht.

In der realen Welt sind Dinge jedoch oft „heiß". Wenn Elektronen heiß werden, werden sie unruhig. Manche Leute, die in der Reihe standen, treten vielleicht aus, und manche, die warteten, treten ein. Dies erzeugt eine „verschwommene" Grenze, bei der Menschen teilweise drin und teilweise draußen sind.

Frühere Abkürzungen scheiterten hier, weil sie zu starr waren. Sie konnten die „Verschwommenheit" einer heißen Menge nicht bewältigen.

Die Lösung: Einem neuronalen Netz beibringen, zu „quetschen"

Die Autoren erkannten, dass die Mathematik, die verwendet wird, um die kalte Menge zu sortieren, exakt der Struktur eines tiefen neuronalen Netzwerks (der Art von KI, die Gesichter erkennt oder Gedichte schreibt) entspricht.

• Der alte Weg (SP2): Stellen Sie sich eine Maschine vor, die eine Zahl nimmt und entweder quadriert ($x^2$) oder eine spezifische Subtraktion durchführt ($2x - x^2$). Sie wiederholt dies immer wieder und „quetscht" die Zahlen, bis sie entweder 0 oder 1 werden. Dies funktioniert hervorragend für kalte Mengen.
• Der neue Weg (MLSP2): Die Autoren nahmen diese Maschine und gaben ihr ein „Gehirn". Anstatt feste Regeln zu verwenden, trainierten sie die Maschine mittels Maschinellem Lernen. Sie lehrten sie, ihre eigenen internen Regler (Koeffizienten) anzupassen, damit sie die „verschwommene" heiße Menge perfekt bewältigen konnte.

Stellen Sie es sich so vor:

• Alte Maschine: Ein starrer Stempel, der nur „Ja" oder „Nein" druckt.
• Neue Maschine: Ein flexibler 3D-Drucker, der lernt, genau wie er das „Ja" und „Nein" formen muss, um je nach Hitze der Menge eine glatte, perfekte Kurve dazwischen zu erzeugen.

Der magische Trick: Ein Modell passt zu vielen Temperaturen

Normalerweise müssen Sie Ihr KI-Modell von Grund auf neu trainieren, wenn Sie die Temperatur Ihrer Simulation ändern. Das dauert ewig.

Die Autoren entdeckten einen cleveren Trick namens Affine Reskalierung.
Stellen Sie sich eine Stadtkarte vor. Wenn Sie hinein- oder herauszoomen wollen, müssen Sie nicht die ganze Stadt neu zeichnen; Sie dehnen oder verkleinern einfach die Karte.

Die Autoren stellten fest, dass sie ihr KI-Modell nur einmal für einen bestimmten „Zoom-Level" (eine bestimmte Temperatur und ein bestimmtes chemisches Potential) trainieren konnten. Dann konnten sie für jede andere Temperatur innerhalb eines bestimmten Bereichs die Eingabedaten (die Hamilton-Matrix) einfach „dehnen", bevor sie sie dem Modell zuführten. Das Modell muss nichts neu lernen; es sieht die Daten nur in einem leicht anderen Maßstab und liefert die korrekte Antwort.

Das bedeutet, dass sie Simulationen durchführen können, bei denen sich die Temperatur ständig ändert (wie in einer chemischen Reaktion), ohne das KI-Modell neu trainieren zu müssen.

Die Hardware: KI-Chips für die Wissenschaft nutzen

Das Paper hebt hervor, dass diese Methode speziell für moderne Computerchips entwickelt wurde, insbesondere für GPUs (Grafikprozessoren) und Tensor Cores (Chips, die für KI entwickelt wurden).

• Die Analogie: Die traditionelle Diagonalisierung ist wie ein Meisterhandwerker, der jedes Möbelstück von Hand schnitzt. Es ist präzise, aber langsam.
• Die neue Methode: Dies ist wie die Verwendung eines Hochgeschwindigkeits-3D-Druckers. Sie nutzt die spezifische Architektur von KI-Chips, um massive Berechnungen (Matrixmultiplikationen) unglaublich schnell durchzuführen.

Die Autoren testeten dies auf einer Nvidia RTX 6000 Ada GPU. Sie stellten fest, dass ihre Methode 9 bis 16 Mal schneller war als die Standardmethoden, die heute von Wissenschaftlern verwendet werden und die hochgradig optimiert sind, während sie gleichzeitig eine hohe Genauigkeit beibehielten.

Zusammenfassung der Ergebnisse

1. Geschwindigkeit: Sie erzielten eine massive Beschleunigung (bis zu 16-fach) bei der Berechnung, wie sich Elektronen in Materialien verhalten, insbesondere auf moderner KI-Hardware.
2. Genauigkeit: Sie können „heiße" Elektronen (fraktionale Besetzung) mit extremer Präzision modellieren, etwas, das frühere Abkürzungen nicht gut bewältigen konnten.
3. Effizienz: Indem sie das Modell einmal trainieren und mathematische Tricks zur Reskalierung der Eingaben verwenden, vermeiden sie die Notwendigkeit, das Modell bei jeder Temperaturänderung in einer Simulation neu zu trainieren.
4. Keine „magische" Diagonalisierung: Sie umgehen die langsame, schwere Mathematik der Diagonalisierung vollständig und verlassen sich stattdessen auf schnelle, wiederholte Multiplikationsschritte, die KI-Chips gerne ausführen.

Kurz gesagt: Die Autoren haben einen langsamen, starren mathematischen Prozess in ein schnelles, flexibles, KI-gestütztes Werkzeug verwandelt, das auf modernen Computerchips unglaublich effizient läuft und es Wissenschaftlern ermöglicht, komplexe Materialien viel schneller als zuvor zu simulieren.

Technische Zusammenfassung

Technische Zusammenfassung: Maschinengelernte Fermi-Operator-Entwicklungen bei endlicher Temperatur

Problemstellung
Elektronenstrukturrechnungen, insbesondere innerhalb der Kohn-Sham-Dichtefunktionaltheorie (KS-DFT), sind rechnerisch durch die kubische Skalierung der Kosten für die Diagonalisierung der Hamilton-Matrix zur Lösung des Eigenwertproblems begrenzt. Während rekursive Fermi-Operator-Entwicklungsschemata, wie die Second-Order Spectral Projection (SP2)-Methode, einen Weg bieten, die Dichtematrix direkt ohne Diagonalisierung zu berechnen, beschränken sich bestehende effiziente Implementierungen auf die elektronische Temperatur von Null. Bei Temperatur Null ist die Dichtematrix idempotent (Besetzungen sind strikt 0 oder 1). Viele physikalische Systeme – wie Metalle oder Systeme bei erhöhten elektronischen Temperaturen – erfordern jedoch fraktionale Orbitalbesetzungen, um entartete Eigenzustände oder thermische Verschmierung genau zu modellieren.

Frühere Versuche, SP2 auf endliche Temperaturen zu verallgemeinern, beinhalteten das Abschneiden der Rekursion, um thermische Verschmierung einzuführen. Diese abgeschnittenen Entwicklungen sind jedoch inhärent approximativ und scheitern daran, die exakte Fermi-Funktion wiederherzustellen, insbesondere in der Nähe des chemischen Potentials, wo Genauigkeit kritisch ist. Alternative Methoden wie Chebyshev-Entwicklungen oder Padé-Approximanten erfordern entweder prohibitiv hohe Polynomordnungen, um Gibbs-Oszillationen zu unterdrücken, oder verursachen erhebliche rechnerische Mehrkosten durch wiederholte Lösungen linearer Gleichungssysteme.

Methodik
Die Autoren schlagen einen Rahmen vor, der die rekursive SP2-Methode auf endliche Temperaturen verallgemeinert, indem sie deren algebraische Struktur auf Deep-Neural-Network-Architekturen (DNN) abbildet. Der Kerngedanke besteht darin, dass die rekursiven SP2-Aktualisierungen den Schichten eines neuronalen Netzwerks ähneln. Indem die Expansionskoeffizienten als trainierbare Gewichte und Bias-Werte behandelt werden, konstruieren die Autoren maschinelle Lernmodelle, die in der Lage sind, die Fermi-Verteilungsfunktion mit fraktionalen Besetzungen bei beliebigen Temperaturen zu approximieren.

Zu den wichtigsten methodischen Komponenten gehören:

1. Neuronale Netzwerk-Architekturen:

   • MLSP2 (Machine-Learned SP2): Eine Verallgemeinerung von SP2, bei der die quadratischen Aktualisierungsregeln ($X^2$ oder $2X-X^2$) durch lernfähige quadratische Polynome ($ax^2 + bx + c$) mit einem Akkumulator-Term ersetzt werden. Dies ermöglicht dem Modell, die exakte thermische Verschmierung der Fermi-Funktion anstelle einer abgeschnittenen Stufenfunktion zu approximieren.
   • Max-SP2: Eine ausdrucksstärkere Architektur, die „Skip Connections" (Überbrückungen) integriert, wobei jede Schicht das Quadrat einer linearen Kombination aller vorherigen Schichten ist.
   • Skip-SP2: Eine komprimierte Version von Max-SP2, die ein begrenztes Gedächtnis der jüngsten Schichten und Akkumulatoren verwendet, um Ausdruckskraft und Speichernutzung auszubalancieren.

2. Entropie-Approximation:
   Die Autoren entwickeln zudem ein rekursives Schema zur Approximation der elektronischen Entropiefunktion $s(x)$, die für die Berechnung der elektronischen freien Energie erforderlich ist. Sie nutzen ein skaliertes Produkt der Fermi-Funktion und ihres Komplements, $f(x)(1-f(x))$, als Anfangsschätzung, das anschließend durch eine rekursive quadratische Entwicklung verfeinert wird, die so trainiert ist, dass sie die zweite Ableitung der wahren Entropie am chemischen Potential nachbildet.

3. Training und Optimierung:
   Modelle werden auf skalaren Eingaben im Einheitsintervall $[0, 1]$ trainiert, anstatt auf vollständigen Matrizen, unter Verwendung des Levenberg–Marquardt-Algorithmus mit geodätischer Beschleunigung. Die Trainingsdaten werden mit einer Gewichtung proportional zur Ableitung der Fermi-Funktion abgetastet, um den maximalen Fehler in der Nähe des chemischen Potentials zu minimieren.

4. Affine Neuskalierung und Übertragbarkeit:
   Eine entscheidende Innovation ist die Verwendung einer affinen Neuskalierung, um die Notwendigkeit eines erneuten Trainings zu eliminieren, wenn sich Simulationsparameter ändern. Durch die Normalisierung der Hamilton-Matrix ($H'$), des chemischen Potentials ($\mu'$) und der inversen Temperatur ($\beta'$) kann ein einzelnes Modell, das bei spezifischen Parametern $(\beta_0, \mu_0)$ trainiert wurde, auf einen breiten „Gültigkeitsbereich" anderer Parameter angewendet werden. Dies wird erreicht, indem die Eingabe-Hamilton-Matrix so neu skaliert wird, dass sie den Trainingsbedingungen entspricht, wodurch derselbe Gewichtungssatz während einer Simulation über verschiedene Temperaturen und chemische Potentiale hinweg verwendet werden kann.

5. Hardware-Implementierung:
   Die Algorithmen sind für moderne GPUs und KI-Hardware (insbesondere NVIDIA Tensor Cores) optimiert. Die Autoren nutzen gemischte Genauigkeitsarithmetik (FP16/FP32), um Matrixquadratur-Operationen effizient durchzuführen, wobei sie die Symmetrie der Hamilton-Matrix ausnutzen, um die Anzahl der erforderlichen Multiplikationen und Datentransfers zu reduzieren.

Hauptergebnisse

• Genauigkeit: Die MLSP2-Modelle erreichen Fehler in der Größenordnung von $10^{-7}$ für die Approximation der Fermi-Funktion und übertreffen damit abgeschnittene SP2-Methoden (die Fehler um $10^{-2}$ aufweisen) erheblich. In vielen Bereichen entsprechen sie der Präzision der Diagonalisierung mit doppelter Genauigkeit.
• Leistung: Auf einer NVIDIA RTX 6000 Ada GPU zeigt der MLSP2-Ansatz für mittlere Matrixgrößen eine 16-fache Beschleunigung gegenüber der Diagonalisierung mit doppelter Genauigkeit (unter Verwendung von cuSOLVER) und eine 9-fache Beschleunigung für größere Matrizen. Selbst im Vergleich zur Diagonalisierung mit einfacher Genauigkeit bietet MLSP2 eine 2- bis 5-fache Beschleunigung bei gleichzeitiger Aufrechterhaltung überlegener Stabilität und Genauigkeit.
• Skalierbarkeit: Die Methode stützt sich ausschließlich auf hochoptimierte Kernel für Matrix-Matrix-Multiplikation und vermeidet eine explizite Diagonalisierung. Die Anzahl der Schichten, die erforderlich ist, um eine Zielgenauigkeit zu erreichen, skaliert logarithmisch mit der inversen Temperatur ($\beta$), was eine effiziente Berechnung auch bei niedrigen Temperaturen ermöglicht.

Bedeutung und Behauptungen
Die Arbeit behauptet, dass dieser Ansatz eine robuste, verallgemeinerbare Lösung für elektronenstrukturrechnungen bei endlicher Temperatur bietet, die die rechnerischen Engpässe der Diagonalisierung umgeht. Durch die Verallgemeinerung von SP2 mittels maschinellen Lernens ermöglichen die Autoren die Berechnung von Dichtematrizen für Systeme mit fraktionalen Besetzungen zu einem Bruchteil der Kosten traditioneller Methoden.

Die Bedeutung liegt in der Fähigkeit, dynamische Simulationen bei endlicher Temperatur (wie Quantenmolekulardynamik) durchzuführen, bei denen das chemische Potential und die elektronische Temperatur zwischen den Zeitschritten schwanken. Die Strategie der affinen Neuskalierung stellt sicher, dass ein einziges vortrainiertes Modell während einer gesamten Simulation wiederverwendet werden kann, ohne dass ein erneutes Training erforderlich ist, was die Methode für groß angelegte Anwendungen praktikabel macht. Darüber hinaus ist der Ansatz speziell darauf zugeschnitten, die Leistungsmerkmale moderner KI-Hardware (Tensor Cores) auszunutzen und bietet erhebliche Beschleunigungen gegenüber herstelleroptimierten Diagonalisierungsroutinen bei gleichzeitiger Aufrechterhaltung hoher numerischer Genauigkeit.