Bayesian Neural Networks versus deep ensembles for uncertainty quantification in machine learning interatomic potentials

Sintesi Tecnica: Reti Neurali Bayesiane contro Deep Ensembles per la Quantificazione dell'Incertezza nei Potenziali Interatomici basati su Machine Learning

Definizione del Problema
I potenziali interatomici basati su Machine Learning (MLIP) sono diventati essenziali per la modellazione atomistica, offrendo un'alternativa computazionalmente efficiente ai metodi ab initio come la Teoria del Funzionale della Densità (DFT). Tuttavia, le architetture standard di deep learning sono deterministiche e mancano di meccanismi intrinseci per quantificare l'incertezza predittiva. Questa limitazione è critica nella scienza dei materiali, dove i modelli vengono frequentemente applicati a dati fuori distribuzione (out-of-distribution). Senza stime dell'incertezza affidabili, le simulazioni rischiano di produrre risultati fuorvianti e le strategie di apprendimento attivo per il raffinamento del modello mancano di guida. Sebbene i Deep Ensembles (DE) e le Reti Neurali Bayesiane (BNN) siano strategie consolidate per affrontare questo problema, la loro valutazione comparativa nel contesto specifico degli MLIP — in particolare riguardo alle previsioni congiunte di energia e forza attraverso diversi regimi di dati — richiede una valutazione sistematica.

Metodologia
Gli autori hanno implementato e confrontato due approcci primari di quantificazione dell'incertezza all'interno del framework ænet-PyTorch:

1. Deep Ensembles (DE): Seguendo la formulazione di Lakshminarayanan et al., cinque reti neurali addestrate indipendentemente sono state inizializzate con diversi semi casuali (random seed). L'incertezza è stata quantificata come la deviazione standard del campione delle previsioni dell'ensemble.
2. Reti Neurali Bayesiane Variazionali (VBNN): Gli autori hanno implementato le BNN utilizzando l'Inferenza Variazionale (VI) tramite il framework di programmazione probabilistica Pyro. Sono stati testati due specifici guide variazionali:
   • Local Reparametrization Trick (LRT): Utilizzando un'approssimazione mean-field con una distribuzione gaussiana posteriore completamente fattorizzata, ottimizzata tramite l'algoritmo Bayes by Backprop con LRT per ridurre la varianza del gradiente.
   • Radial Guide (RAD): Un approccio non mean-field che introduce trasformazioni radiali per catturare le correlazioni globali tra i pesi, ottimizzato tramite campionamento Monte Carlo standard dell'Evidence Lower Bound (ELBO).

Configurazione Sperimentale
Lo studio ha valutato questi metodi su due dataset distinti:

• Dataset TiO2: 7.815 strutture che rappresentano varie fasi del biossido di titanio. I modelli sono stati addestrati sia sull'intero dataset (regime High-Data) che su un sottoinsieme del 20% (regime Low-Data) per valutare l'efficienza dei dati. Sono state eseguite sia l'addestrazione solo su energia che quella congiunta energia-forza.
• Dataset QM7: 7.165 piccole molecole organiche utilizzate per testare la generalizzabilità attraverso sistemi chimicamente diversificati. I modelli sono stati addestrati nel regime high-data solo sulle energie di atomizzazione.

Metriche di Valutazione
Le prestazioni sono state valutate tramite una suite completa di metriche:

• Accuratezza: Errore Assoluto Medio (MAE), Errore Quadratico Medio (RMSE) e Log-Likelihood Negativa (NLL).
• Qualità dell'Incertezza: Root Mean Squared Calibration Error (RMSCE), Sharpness (incertezza predittiva media), correlazione di rango di Spearman ($r_s$) tra incertezza ed errore, e un "Overlap Score" che misura l'allineamento tra le previsioni ad alta incertezza e le regioni ad alto errore.

Risultati Chiave

1. Accuratezza Predittiva:

   • Nel regime High-Data (TiO2), i DE hanno superato costantemente le VBNN, raggiungendo il MAE più basso (6 meV/atomo) e l'RMSE più basso (9 meV/atomo), così come la NLL più negativa (-4,2), indicando una superiore accuratezza e calibrazione.
   • Nel regime Low-Data, i DE hanno mantenuto la loro superiorità in termini di accuratezza (MAE 35 meV/atomo). Tuttavia, l'LRT ha mostrato un'accuratezza competitiva (MAE 54 meV/atomo) e un RMSE inferiore (259 meV/atomo) rispetto ai DE, suggerendo una migliore robustezza agli outlier. Il RAD si è comportato male nei contesti low-data.
   • Sul dataset QM7, si è verificata un'inversione: le BNN (LRT e RAD) hanno ottenuto RMSE energetici inferiori (13,0–13,3 meV/atomo) rispetto ai DE (15,2 meV/atomo), suggerendo che il prior variazionale fornisce una regolarizzazione benefica per spazi chimicamente diversificati.

2. Quantificazione dell'Incertezza:

   • Calibrazione: I DE hanno generalmente fornito le incertezze meglio calibrate (RMSCE più basso) negli esperimenti su TiO2. Nel regime low-data, il RAD ha sorprendentemente mantenuto un'eccellente calibrazione (RMSCE 0,05), paragonabile alle sue prestazioni nel regime high-data, mentre LRT e DE hanno visto una degradazione della calibrazione.
   • Sharpness: I DE hanno prodotto le previsioni più nitide (sharpest/più sicure). L'LRT tendeva a produrre intervalli di incertezza eccessivamente ampi (alti valori di sharpness), particolarmente nei regimi low-data.
   • Correlazione e Overlap: I DE hanno dimostrato la più forte correlazione monotonica ($r_s$) tra incertezza predetta ed errore effettivo nel regime low-data (0,45), superando sia l'LRT (0,28) che il RAD (0,19). I DE hanno anche ottenuto l'Overlap Score più elevato (55%), indicando che identificano più efficacementamente le regioni ad alto errore.

3. Previsioni della Forza:

   • Nel regime low-data, le VBNN (LRT e RAD) hanno raggiunto un RMSE della forza significativamente inferiore rispetto ai DE, probabilmente a causa della regolarizzazione fornita dal prior. Tuttavia, la calibrazione della forza è rimasta scarsa per tutti i modelli, con i DE che sono stati gli unici a mostrare una calibrazione ragionevole nel regime low-data.

4. Efficienza Computazionale:

   • Addestramento: I DE hanno un costo per epoca inferiore, ma richiedono l'addestramento di cinque modelli indipendenti. Il tempo totale di addestramento per i DE è stato comparabile a quello di LRT nel regime high-data, ma leggermente superiore nel regime low-data.
   • Inferenza: I DE sono significativamente più veloci in fase di inferenza (0,03 s) rispetto alle VBNN (0,64–0,71 s), poiché le BNN richiedono 20 campioni Monte Carlo per ogni previsione, inclusa la differenziazione automatica per le forze.

Significato e Conclusioni
Il documento conclude che, sebbene i Deep Ensembles siano teoricamente meno fondati rispetto alle BNN, essi superano empiricamente le BNN variazionali in quasi tutte le metriche riguardanti accuratezza, calibrazione ed efficienza di inferenza, specialmente quando lo spazio chimico è relativamente ristretto (TiO2). I DE sono raccomandati come la scelta preferibile quando la priorità è una quantificazione dell'incertezza affidabile e si desidera la semplicità di implementazione.

Tuttavia, lo studio evidenzia che le VBNN, specificamente la variante LRT, offrono un compromesso competitivo e possono essere superiori in regimi con scarsezza di dati o alta diversità chimica (come visto in QM7), dove il prior variazionale agisce come un efficace regolarizzatore. Gli autori osservano che le BNN sono più sensibili all'inizializzazione casuale e richiedono molteplici sessioni di addestramento per garantire la robustezza. Il lavoro fornisce una guida pratica per la selezione di MLIP capaci di gestire l'incertezza, suggerendo che la scelta tra DE e VBNN debba dipendere dal regime di dati specifico e dalla diversità chimica del sistema target.

Contributi

• Introduzione di un'implementazione di VBNN accelerata su GPU per MLIP all'interno del framework ænet-PyTorch, che supporta l'addestramento congiunto energia-forza.
• Un confronto sistematico tra DE, LRT e RAD attraverso regimi high-data e low-data su entrambi i dataset (inorganico TiO2 e organico QM7).
• Dimostrazione del fatto che, sebbene i DE siano generalmente superiori, le prestazioni relative delle BNN possono invertirsi in dataset chimicamente diversificati, sfidando l'universalità della superiorità degli ensemble in tutti i contesti.
• Fornitura di una nuova metrica, l'"Overlap Score", per valutare l'utilità delle stime dell'incertezza in scenari di apprendimento attivo.