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

Per oltre sessanta anni, gli scienziati si sono affidati a un potente insieme di regole per comprendere come gli atomi si uniscano per formare i materiali che ci circondano. Questo quadro, noto come teoria del funzionale della densità, tratta la nuvola elettronica che circonda un atomo non come una collezione di particelle individuali, ma come una densità continua e fluida. Questo approccio permette ai ricercatori di prevedere le proprietà di tutto, dai nuovi medicinali alle batterie avanzate. Tuttavia, c'è un problema. Per ottenere risultati accurati, la teoria richiede un calcolo complesso e ripetitivo che consiste nel risolvere il comportamento di ogni singolo elettrone nel sistema. Man mano che il sistema cresce, il tempo richiesto per questi calcoli esplode, crescendo così velocemente che simulare un pezzo di metallo di dimensioni modeste o una grande molecola biologica diventa impossibile anche sui supercomputer più potenti. Questo collo di bottiglia ha mantenuto molte importanti questioni scientifiche fuori portata, costringendo i ricercatori a scegliere tra accuratezza e scala.

Un team di ricercatori del California Institute of Technology ha trovato un modo per rompere questa barriera, insegnando a un computer come saltare la parte più costosa del calcolo senza perdere accuratezza. Invece di cercare di prevedere direttamente la risposta finale o di apprendere una difficile scorciatoia matematica che spesso fallisce, hanno insegnato a un'intelligenza artificiale di eseguire il passaggio specifico e ripetuto che avviene in ogni singolo ciclo di calcolo. Concentrandosi su questo passaggio intermedio, hanno creato un metodo in grado di simulare sistemi con decine di miglia di atomi su una singola scheda grafica, un compito che precedentemente richiedeva migliaia di processori. Il risultato è un nuovo modo per modellare la materia che è quasi veloce quanto le approssimazioni più semplici, ma mantiene l'alta precisione dei metodi più rigorosi.

Il cuore del problema risiede in come la teoria gestisce gli elettroni. Per trovare lo stato stabile di un materiale, il computer deve regolare ripetutamente una mappa della densità elettronica finché questa non smette di cambiare. Nel metodo tradizionale, questa regolazione richiede la risoluzione di un enorme problema di autovalori, che consiste essenzialmente nel trovare i pattern specifici del movimento elettronico per ogni atomo nel sistema. Questa operazione è così computazionalmente pesante che il suo costo cresce in modo cubico con il numero di elettroni; raddoppiare la dimensione del sistema rende il calcolo otto volte più difficile. Per decenni, gli scienziati hanno cercato di aggirare questo problema eliminando la necessità di tracciare interamente i pattern elettronici individuali, una strategia nota come teoria del funzionale della densità priva di orbitali (orbital-free). Tuttavia, i precedenti tentativi di apprendere questa scorciatoia utilizzando il machine learning hanno incontrato difficoltà. Alcuni hanno cercato di apprendere la relazione tra la densità elettronica e l'energia, ma questa relazione è matematicamente instabile e soggetta a errori. Altri hanno cercato di prevedere direttamente la densità elettronica finale dalla disposizione degli atomi, ma questi modelli fallivano quando venivano interrogati su sistemi più grandi o chimicamente diversi da quelli su cui erano stati addestrati.

I ricercatori hanno capito che la chiave era smettere di cercare di prevedere la destinazione finale e invece imparare il viaggio. In ogni passaggio del calcolo tradizionale, il computer prende una mappa delle forze che agiscono sugli elettroni e produce una nuova mappa di dove è probabile che si trovino gli elettroni. Questo passaggio in avanti è matematicamente stabile e ben definito. Il team ha addestrato una rete neurale specializzata, un tipo di intelligenza artificiale progettata per comprendere i pattern spaziali, per apprendere esattamente questa trasformazione. Hanno fornito alla rete migliaia di esempi di queste mappe di forze e delle corrispondenti mappe di densità elettronica che esse producevano, tratti da un insieme diversificato di 8.504 molecole e materiali solidi, inclusi composti organici, isolanti e metalli. La rete ha imparato a prevedere la nuova densità elettronica direttamente dalle forze, sostituendo efficacemente il calcolo lento e pesante con una previsione rapida e appresa.

Ciò che rende unico questo approccio è il modo in cui gestisce la complessità di diversi sistemi. I ricercatori hanno progettato la rete affinché fosse invariante rispetto al dominio, il che significa che può applicare le stesse regole apprese a una minuscola molecola e a un enorme cristallo senza dover essere riaddestrata per ogni nuova dimensione. Hanno inoltre garantito che la rete rispettasse le simmetrie fondamentali dello spazio, in modo che comprenda che ruotare o traslare una molecola non cambia le sue proprietà fisiche. Quando testata su sistemi che non aveva mai visto, il modello si è dimostrato straordinariamente robusto. Su un set di molecole simili a farmaci contenenti elementi e dimensioni non presenti nei dati di addestramento, il nuovo metodo ha mantenuto un'alta accuratezza, mentre i precedenti modelli di predizione diretta vedevano i loro errori salire alle stelle man mano che le molecole diventavano più grandi. La natura autoconsistente del metodo permette di correggere i propri errori ad ogni passaggio, proprio come un navigatore che controlla frequentemente la propria posizione invece di indovinare l'intero percorso in una volta sola.

La potenza di questo metodo è diventata evidente quando il team lo ha applicato a sistemi periodici, come i cristalli, e a difetti estesi nei metalli. Nei calcoli standard per i cristalli, il computer deve ripetere il calcolo pesante per molti diversi punti nello spazio matematico che descrive il pattern ripetitivo del cristallo. Il nuovo metodo salta completamente questa ripetizione, facendo evolvere la densità elettronica sulla cella unitaria senza dover campionare questi punti individualmente. Ciò ha permesso ai ricercatori di far convergere i calcoli per sistemi metallici e semiconduttori con un singolo modello, riproducendo gli spettri elettronici e le proprietà strutturali con la stessa accuratezza del metodo tradizionale. Inoltre, il modello ha mostrato una sorprendente capacità di trasferirsi tra diversi tipi di approssimazioni teoriche. Un modello addestrato su un tipo di calcolo poteva essere utilizzato con un'approssimazione diversa e più accurata per le forze senza alcun riaddestramento, semplicemente scambiando le forze in ingresso.

La prova suprema è arrivata quando i ricercatori hanno applicato il metodo a una dislocazione di magnesio, un tipo di difetto in un cristallo metallico che è fondamentale per comprendere come i materiali si deformano e si rompono. Questi difetti creano campi elastici a lungo raggio che richiedono simulazioni di migliaia di atomi per essere modellati accuratamente. I precedenti tentativi di simulare un tale sistema con alta precisione richiedevano un supercomputer con migliaia di unità di elaborazione grafica e impiegavano giorni per essere eseguiti. Il nuovo metodo, eseguito su una singola moderna scheda grafica, ha completato con successo il calcolo per un sistema contenente 8.250 atomi di magnesio e 82.500 elettroni di valenza. Il tempo necessario per completare il calcolo cresceva in modo quasi lineare con la dimensione del sistema, confermando che il collo di bottiglia computazionale era stato rimosso. Le densità elettroniche risultanti erano indistinguibili dal riferimento ad alta precisione, permettendo ai ricercatori di risolvere le sottili differenze di energia tra diversi tipi di difetti che determinano la resistenza del metallo.

Questo lavoro dimostra che il machine learning può estendere la portata dei calcoli della struttura elettronica non sostituendo l'intera teoria fisica, ma accelerando la sua operazione più costosa e ripetuta. Imparando la mappa diretta dalle forze alla densità, i ricercatori hanno creato uno strumento che è stabile, trasferibile tra diversi tipi di materia e capace di gestire sistemi a una scala che prima era inaccessibile. Il metodo fornisce un diagnostico diretto per la propria affidabilità: se il calcolo non converge, segnala che il sistema si trova al di fuori dell'intervallo di fiducia del modello, una caratteristica che i modelli di predizione diretta non possiedono. Sebbene l'attuale modello predica la densità elettronica, l'energia totale e altre proprietà possono ancora essere ottenute con un singolo passaggio finale, preservando l'accuratezza della teoria completa. Questo approccio apre la porta alla simulazione di materiali complessi, sistemi biologici e reazioni chimiche a un livello di dettaglio e scala che un tempo era dominio dell'impossibilità teorica.