Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

Questo lavoro presenta un modello fondazionale grafico per l'atomistica addestrato su scala exascale che, sfruttando l'architettura HydraGNN e l'ottimizzazione iperparametrica su supercomputer come Frontier, permette di setacciare in pochi secondi oltre un miliardo di strutture atomiche, accelerando drasticamente la scoperta di materiali rispetto ai metodi computazionali tradizionali.

L'essenziale

Immagina di dover trovare l'ingrediente segreto per un nuovo super-farmaco o un materiale che salva il pianeta, ma invece di avere una ricetta, devi provare a cucinare un miliardo di piatti diversi per vedere quale funziona.

Fino a poco tempo fa, questo compito sarebbe stato impossibile. I metodi scientifici tradizionali (chiamati "calcoli quantistici") sono come cuochi che preparano un solo piatto alla volta, ma ci vogliono anni per cucinarlo e assaggiarlo. Se dovessi provarne un miliardo, ci vorrebbero secoli.

Questo documento racconta come un team di ricercatori ha costruito un "super-cuoco digitale" capace di assaggiare un miliardo di piatti in 50 secondi.

Ecco come hanno fatto, spiegato in modo semplice:

1. Il problema: Troppi ingredienti, troppe ricette

I ricercatori avevano 16 "libri di ricette" (dataset) diversi. Alcuni parlavano di molecole organiche (come quelle del corpo umano), altri di metalli, altri di catalizzatori.

• Il problema: I libri erano di dimensioni diverse (alcuni avevano poche pagine, altri milioni) e usavano unità di misura diverse (alcuni erano precisi al millimetro, altri solo al centimetro).
• La sfida: Se mescoli tutto in una pentola gigante senza ordine, le ricette piccole vengono sommerse da quelle grandi e il risultato è un pasticcio. Inoltre, alcuni libri erano scritti in dialetti diversi (fidelità diverse).

2. La soluzione: Il "Chef Poliglotta" (Il Modello Fondamentale)

Hanno creato un'intelligenza artificiale chiamata HydraGNN (un modello a "testa multipla").

• L'analogia: Immagina uno chef che ha un'unica grande conoscenza di base su come funzionano gli ingredienti (le interazioni atomiche), ma ha 16 cuffie diverse (una per ogni libro di ricette).
• Come funziona: Mentre impara, lo chef ascolta tutti i libri contemporaneamente. Le cuffie specifiche adattano le istruzioni generali alle sfumature di ogni singolo libro. In questo modo, lo chef non si confonde e impara a riconoscere sia le molecole delicate che i metalli pesanti, anche se i dati sono sbilanciati.

3. La potenza: La cucina su scala "Esascale"

Per addestrare questo chef, hanno usato i computer più potenti del mondo (come Frontier, Aurora e Perlmutter), che sono come cucine con 16.000 fornelli elettrici che lavorano tutti insieme.

• Hanno fatto un "gara di cucina" (ottimizzazione degli iperparametri) per trovare la combinazione perfetta di ingredienti e tecniche.
• Hanno scoperto che un tipo specifico di architettura (chiamata PaiNN) era il più veloce e preciso, capace di imparare senza impazzire.

4. Il trucco magico: La "Cucina Intelligente"

Non basta avere molti fornelli; bisogna anche non sprecare tempo. Hanno inventato dei trucchi per velocizzare il processo:

• Riutilizzo dell'enciclopedia: Invece di leggere l'intero libro di ricette 16 volte (una per ogni cuffia), lo chef legge il libro una sola volta e salva le note, poi le usa per tutte le 16 versioni.
• Salto dei rami: Se una cuffia dice "questo piatto non è importante", lo chef la ignora e non spreca tempo a cucinarlo.
• Fusione: Invece di calcolare i sapori uno alla volta, li calcola tutti insieme in un unico movimento fluido.

5. Il risultato: Un miliardo di assaggi in 50 secondi

Grazie a questi trucchi, il modello è riuscito a:

• Assaggiare 1,1 miliardi di strutture atomiche in soli 50 secondi.
• Fare un lavoro che, con i metodi vecchi, avrebbe richiesto 6,7 anni di calcolo continuo su un supercomputer.
• Essere così preciso che, una volta addestrato, può essere usato per scoprire nuovi materiali per batterie, farmaci o energie pulite in pochi minuti invece che in decenni.

In sintesi

Hanno trasformato la ricerca scientifica da un'attività lenta e costosa (come cercare un ago in un pagliaio a mano) in un processo veloce e massivo (come usare un aspirapolvere industriale).

Questo modello non è solo un "grande numero"; è uno strumento pratico che permette agli scienziati di esplorare spazi chimici immensi che prima erano inaccessibili, accelerando la scoperta di soluzioni per i problemi reali del mondo. È come se avessimo dato agli scienziati un telescopio che vede non solo le stelle vicine, ma l'intero universo in un battito di ciglia.

Sommario tecnico

Titolo: Modelli di Base Grafici Multi-Task su Scala Exascale per Dati Atomistici Squilibrati e Multi-Fiducia

1. Il Problema

La scoperta di materiali, chimica e design molecolare è ostacolata dall'elevato costo computazionale dei metodi di first-principles (come la Teoria del Funzionale Densità - DFT), che li rendono impraticabili per lo screening su larga scala di spazi chimici vasti. Sebbene i modelli di apprendimento automatico (MLIPs - Machine Learning Interatomic Potentials) e i modelli di base grafici atomistici (GFMs) offrano un'alternativa più veloce, esistono sfide critiche irrisolte:

• Dati Eterogenei e Squilibrati: Nessun singolo dataset o teoria elettronica è affidabile uniformemente per sistemi organici, inorganici e ibridi. Addestrare su dati multi-sorgente e multi-fiducia (diverse approssimazioni teoriche, diversi set di basi) introduce rumore nelle etichette, riferimenti energetici incoerenti e squilibri nell'ottimizzazione.
• Scalabilità e Selezione del Modello: Non esiste un'architettura di rete neurale (backbone) universalmente superiore. La selezione del modello richiede un'ottimizzazione su larga scala che bilanci accuratezza, tempo di addestramento e costi computazionali.
• Adattamento ai Dati Scarsi: Molti compiti scientifici rilevanti dispongono di dati etichettati insufficienti per l'addestramento from scratch.
• Screening su Bilioni: La necessità di esplorare spazi chimici vasti richiede la capacità di valutare miliardi di strutture in tempi brevi, superando i limiti dei metodi tradizionali.

2. Metodologia

Il lavoro presenta un workflow exascale integrato basato su HydraGNN, progettato per addestrare modelli di base grafici su 16 dataset aperti di principi primi, comprendenti oltre 544 milioni di strutture e 85+ elementi.

• Architettura Multi-Task (MTL):

  • Viene utilizzata un'architettura con livelli di passaggio messaggi (message-passing) condivisi per apprendere caratteristiche di interazione atomica trasferibili, mentre testine di output specifiche per dataset assorbono le caratteristiche dipendenti dal dataset (offset di fiducia, rumore).
  • Questo approccio riduce l'interferenza dei gradienti tra dataset e stabilizza l'addestramento su corpora fortemente sbilanciati.
  • Un MLP condizionato alla composizione apprende dinamicamente i pesi per fondere le previsioni delle 16 testine in base alla composizione chimica dell'input, permettendo l'interpolazione su sistemi chimici non visti durante l'addestramento.

• Pipeline di Dati e Infrastruttura:

  • Utilizzo di ADIOS2 e DDStore per la gestione dei dati: i dati vengono stagionati su NVMe locale (burst buffer) e distribuiti in modo sharded per evitare colli di bottiglia nel file system condiviso durante l'addestramento su migliaia di nodi.
  • Addestramento distribuito su 16.384 GPU (Frontier, OLCF).

• Selezione del Modello su Larga Scala (HPO):

  • Esecuzione di sei campagne di Ottimizzazione degli Ipertparametri (HPO) su larga scala utilizzando DeepHyper su Frontier.
  • Sono stati esplorati 6 backbone di reti neurali a passaggio messaggi (MPNN): EGNN, SchNet, DimeNet, MACE, PaiNN e PNAEq.
  • I modelli migliori sono stati promossi a sessioni di addestramento continuo su 2.048 nodi per valutare la convergenza reale e il rapporto costo/accuratezza.

• Precisione e Inference:

  • Caratterizzazione sistematica degli effetti della precisione (BF16, FP32, FP64) sia durante l'addestramento che l'inference.
  • Ottimizzazione dell'inference tramite: riuso dell'encoder, skipping dei rami, gradienti fusi (fused gradient) e compilazione (torch.compile), raggiungendo un'accelerazione fino a 33x.

3. Contributi Chiave

1. Primo Workflow Exascale Multi-Task: Addestramento congiunto su un corpus eterogeneo di 544M+ strutture provenienti da 16 dataset diversi, risolvendo il problema dell'ottimizzazione su dati squilibrati e multi-fiducia.
2. Selezione Architetturale Guidata dai Dati: Dimostrazione che la selezione del modello non può basarsi solo sulla perdita di validazione, ma deve considerare il "tempo-to-epoch" e l'efficienza delle risorse. Il modello PaiNN è emerso come leader, completando 100 epoche entro il budget fisso di tempo.
3. Portabilità Cross-System: Il workflow è stato eseguito con successo su tre supercomputer leader (Frontier, Aurora, Perlmutter) mantenendo scalabilità e fedeltà senza riprogettazione del codice.
4. Screening su Bilioni di Strutture: Capacità di valutare 1,1 miliardi di strutture atomistiche in 50 secondi su Frontier, un compito che richiederebbe circa 6,7 anni di calcolo first-principles continuo.
5. Adattamento Efficiente (Fine-Tuning): Dimostrazione che il pre-addestramento su scala exascale permette un adattamento efficace a compiti downstream con dati scarsi, superando di un ordine di grandezza i modelli addestrati da zero.

4. Risultati

• Performance di Addestramento:
  • Addestramento su 16.384 GPU con scalabilità forte e debole quasi lineare su Frontier, Aurora e Perlmutter.
  • Il modello finale (PaiNN) ha circa 12,1 milioni di parametri e occupa ~92 MB di memoria, rendendolo leggero per il deployment.
• Accuratezza nel Fine-Tuning:
  • Per compiti legati alla superficie di energia potenziale (PES), il fine-tuning con backbone "scongelato" ha ridotto l'errore di validazione fino a 10 volte rispetto all'addestramento da zero o con backbone congelato.
  • Per compiti non-PES (es. classificazione metalli), i benefici sono stati minori, indicando che il modello apprende principalmente strutture fisiche energetiche.
• Throughput di Inference:
  • Su Frontier (9.300 nodi), il throughput sostenuto è di 21,8 milioni di strutture al secondo (293 strutture/s/GPU).
  • L'ottimizzazione dell'inference ha permesso un'accelerazione di 33x rispetto alla baseline, mantenendo la precisione FP64 necessaria per la stabilità scientifica.
• Trade-off Precisione:
  • FP64 garantisce la massima accuratezza.
  • FP32 introduce un errore medio assoluto (MAE) di ~0,021 eV.
  • BF16 introduce un errore di ~0,080 eV, utile solo per screening preliminari ad altissimo throughput.

5. Significato e Implicazioni

Questo lavoro rappresenta un salto qualitativo nella scoperta di materiali:

• Democratizzazione dello Spazio Chimico: Trasforma lo screening di spazi chimici vasti da un problema combinatorio intrattabile a un flusso di lavoro pratico e veloce.
• Validazione Scientifica: Dimostra che i modelli di base addestrati su scala exascale non sono solo "grandi modelli", ma strumenti scientifici pratici che catturano strutture fisiche riutilizzabili, permettendo l'apprendimento accurato anche in regimi a dati scarsi.
• Futuro dell'Computing Scientifico: Stabilisce un paradigma per l'integrazione di infrastrutture dati (ADIOS2/DDStore), ottimizzazione multi-task e hardware eterogeneo, fornendo un modello replicabile per future applicazioni ML scientifiche su scala exascale.
• Impatto Industriale: La capacità di valutare miliardi di candidati in secondi accelera drasticamente la scoperta di nuovi materiali per catalizzatori, batterie e farmaci, riducendo i costi di R&D e i tempi di sviluppo.

In sintesi, il paper dimostra come la combinazione di architetture multi-task avanzate, infrastrutture dati scalabili e ottimizzazione su scala exascale possa superare i limiti attuali della simulazione computazionale dei materiali.