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

Riepilogo Tecnico: Espansioni dell'Operatore di Fermi a Temperatura Finita Apprese tramite Macchina

Enunciato del Problema
I calcoli della struttura elettronica, in particolare all'interno della Teoria del Funzionale Densità di Kohn-Sham (KS-DFT), sono limitati dal punto di vista computazionale dal costo di scalatura cubica della diagonalizzazione della matrice Hamiltoniana per risolvere il problema agli autovalori. Sebbene gli schemi di espansione ricorsiva dell'operatore di Fermi, come il metodo di Seconda Ordine di Proiezione Spettrale (SP2), offrano un modo per calcolare direttamente la matrice densità senza diagonalizzazione, le implementazioni efficienti esistenti sono limitate alla temperatura elettronica zero. A temperatura zero, la matrice densità è idempotente (le occupazioni sono rigorosamente 0 o 1). Tuttavia, molti sistemi fisici — come i metalli o i sistemi a temperature elettroniche elevate — richiedono occupazioni orbitali frazionarie per modellare accuratamente stati degeneri o lo sfocamento termico.

I precedenti tentativi di generalizzare SP2 a temperature finite hanno coinvolto l'intercettazione della ricorsione per introdurre uno sfocamento termico. Tuttavia, queste espansioni troncate sono intrinsecamente approssimate, non riescono a riprodurre la funzione di Fermi esatta, in particolare vicino al potenziale chimico dove l'accuratezza è critica. Metodi alternativi come le espansioni di Chebyshev o gli approssimanti di Padé richiedono o ordini polinomiali proibitivamente alti per sopprimere le oscillazioni di Gibbs o comportano un significativo sovraccarico computazionale dovuto alla risoluzione ripetuta di sistemi lineari.

Metodologia
Gli autori propongono un quadro che generalizza il metodo ricorsivo SP2 a temperature finite mappando la sua struttura algebrica su architetture di reti neurali profonde (DNN). L'idea fondamentale è che gli aggiornamenti ricorsivi SP2 assomiglino ai livelli di una rete neurale. Trattando i coefficienti di espansione come pesi e bias addestrabili, gli autori costruiscono modelli di apprendimento automatico capaci di approssimare la funzione di distribuzione di Fermi con occupazioni frazionarie a temperature arbitrarie.

I componenti metodologici chiave includono:

1. Architetture di Reti Neurali:

   • MLSP2 (SP2 Appresa tramite Macchina): Una generalizzazione di SP2 in cui le regole di aggiornamento quadratico ($X^2$ o $2X-X^2$) sono sostituite da polinomi quadratici apprendibili ($ax^2 + bx + c$) con un termine accumulatore. Ciò permette al modello di approssimare lo sfocamento termico esatto della funzione di Fermi piuttosto che una funzione a gradino troncata.
   • Max-SP2: Un'architettura più espressiva che incorpora "connessioni di salto" (skip connections), dove ogni livello è il quadrato di una combinazione lineare di tutti i livelli precedenti.
   • Skip-SP2: Una versione compressa di Max-SP2 che utilizza una memoria finita dei livelli recenti e degli accumulatori per bilanciare l'espressività e l'uso della memoria.

2. Approssimazione dell'Entropia:
   Gli autori sviluppano anche uno schema ricorsivo per approssimare la funzione di entropia elettronica, $s(x)$, necessaria per calcolare l'energia libera elettronica. Utilizzano un prodotto scalato della funzione di Fermi e del suo complemento, $f(x)(1-f(x))$, come ipotesi iniziale, che viene poi raffinata tramite un'espansione quadratica ricorsiva addestrata per corrispondere alla seconda derivata della vera entropia al potenziale chimico.

3. Addestramento e Ottimizzazione:
   I modelli sono addestrati su input scalari all'interno dell'intervallo unitario $[0, 1]$ piuttosto che su matrici complete, utilizzando l'algoritmo di Levenberg–Marquardt con accelerazione geodetica. I dati di addestramento sono campionati con un peso proporzionale alla derivata della funzione di Fermi per minimizzare l'errore massimo vicino al potenziale chimico.

4. Riscaldamento Affine e Trasferibilità:
   Un'innovazione critica è l'uso del riscaldamento affine per eliminare la necessità di riaddestramento quando i parametri di simulazione cambiano. Normalizzando l'Hamiltoniana ($H'$), il potenziale chimico ($\mu'$) e l'inverso della temperatura ($\beta'$), un singolo modello addestrato su parametri specifici $(\beta_0, \mu_0)$ può essere applicato a un'ampia "regione di validità" di altri parametri. Ciò è ottenuto riscaldando l'Hamiltoniana di input per corrispondere alle condizioni di addestramento, permettendo di utilizzare lo stesso insieme di pesi attraverso temperature e potenziali chimici variabili durante una simulazione.

5. Implementazione Hardware:
   Gli algoritmi sono ottimizzati per le moderne GPU e hardware per l'intelligenza artificiale (specificamente i Tensor Core NVIDIA). Gli autori sfruttano l'aritmetica a precisione mista (FP16/FP32) per eseguire operazioni di quadratura di matrici in modo efficiente, sfruttando la simmetria dell'Hamiltoniana per ridurre il numero di moltiplicazioni e trasferimenti di dati richiesti.

Risultati Chiave

• Accuratezza: I modelli MLSP2 raggiungono errori dell'ordine di $10^{-7}$ per l'approssimazione della funzione di Fermi, superando significativamente i metodi SP2 troncati (che hanno errori intorno a $10^{-2}$) e corrispondendo alla precisione della diagonalizzazione in doppia precisione in molti regimi.
• Prestazioni: Su una GPU NVIDIA RTX 6000 Ada, l'approccio MLSP2 dimostra un aumento di velocità di 16 volte rispetto alla diagonalizzazione in doppia precisione (utilizzando cuSOLVER) per dimensioni di matrici intermedie e un aumento di 9 volte per matrici più grandi. Anche rispetto alla diagonalizzazione in precisione singola, MLSP2 offre un aumento di velocità da 2x a 5x mantenendo stabilità e accuratezza superiori.
• Scalabilità: Il metodo si basa esclusivamente su kernel di moltiplicazione matrice-matrice altamente ottimizzati, evitando la diagonalizzazione esplicita. Il numero di livelli richiesto per raggiungere un'accuratezza target scala logaritmicamente con l'inverso della temperatura ($\beta$), permettendo un calcolo efficiente anche a basse temperature.

Significato e Affermazioni
Il documento afferma che questo approccio fornisce una soluzione robusta e generalizzabile per i calcoli della struttura elettronica a temperatura finita che evita i colli di bottiglia computazionali della diagonalizzazione. Generalizzando SP2 attraverso l'apprendimento automatico, gli autori abilitano il calcolo delle matrici densità per sistemi con occupazioni frazionarie a una frazione del costo dei metodi tradizionali.

Il significato risiede nella capacità di eseguire simulazioni dinamiche a temperatura finita (come la dinamica molecolare quantistica) in cui il potenziale chimico e la temperatura elettronica fluttuano tra i passi temporali. La strategia di riscaldamento affine garantisce che un singolo modello pre-addestrato possa essere riutilizzato per tutta la durata di una simulazione senza riaddestramento, rendendo il metodo pratico per applicazioni su larga scala. Inoltre, l'approccio è specificamente adattato per sfruttare le caratteristiche prestazionali dell'hardware AI moderno (Tensor Core), offrendo aumenti di velocità sostanziali rispetto alle routine di diagonalizzazione ottimizzate dai fornitori, mantenendo al contempo un'alta accuratezza numerica.