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

Technische Samenvatting: Bayesiaanse Neurale Netwerken versus Deep Ensembles voor Onzekerheidskwantificatie in Machine Learning Interatomaire Potentialen

Probleemstelling
Machine Learning Interatomaire Potentialen (MLIPs) zijn essentieel geworden voor atomistische modellering, omdat ze een computationeel efficiënt alternatief bieden voor ab initio-methoden zoals Density Functional Theory (DFT). De standaard diepe leerarchitecturen zijn echter deterministisch en missen inherente mechanismen voor het kwantificeren van voorspellende onzekerheid. Deze beperking is cruciaal in de materiaalkunde, waar modellen frequent worden toegepast op out-of-distribution data. Zonder betrouwbare onzekerheidsschattingen lopen simulaties het risico misleidende resultaten te produceren, en missen actieve leerstrategieën voor modelverfijning de nodige sturing. Hoewel Deep Ensembles (DE) en Bayesiaanse Neurale Netwerken (BNNs) gevestigde strategieën zijn om dit aan te pakken, vereist hun vergelijkende prestatie in de specifieke context van MLIPs — met name met betrekking tot gezamenlijke energie- en krachtvoorspellingen over variërende dataregimes — een systematische evaluatie.

Methodologie
De auteurs implementeerden en vergeleken twee primaire benaderingen voor onzekerheidskwantificatie binnen het ænet-PyTorch framework:

1. Deep Ensembles (DE): Volgens de formulering van Lakshminarayanan et al. werden vijf onafhankelijk getrainde neurale netwerken geïnitialiseerd met verschillende willekeurige seeds. Onzekerheid werd gekwantificeerd als de steekproefstandaarddeviatie van de ensemblevoorspellingen.
2. Variational Bayesian Neural Networks (VBNN): De auteurs implementeerden BNNs met behulp van Variational Inference (VI) via het Pyro probabilistische programmeerframework. Twee specifieke variational guides werden getest:
   • Local Reparametrization Trick (LRT): Gebruikmakend van een mean-field benadering met een volledig gefactoriseerde Gaussische posterior, geoptimaliseerd met het Bayes by Backprop algoritme met LRT om de gradiëntvariantie te verminderen.
   • Radial Guide (RAD): Een non-mean-field benadering die radiale transformaties introduceert om globale correlaties tussen gewichten te vangen, geoptimaliseerd via standaard Monte Carlo sampling van de Evidence Lower Bound (ELBO).

Experimentele Opzet
De studie evalueerde deze methoden op twee verschillende datasets:

• TiO2 Dataset: 7.815 structuren die verschillende fasen van titaniumdioxide vertegenwoordigen. Modellen werden getraind op zowel de volledige dataset (High-Data regime) als een subset van 20% (Low-Data regime). Zowel energie-alleen als gezamenlijke energie-kracht training werd uitgevoerd.
• QM7 Dataset: 7.165 kleine organische moleculen om de generaliseerbaarheid over chemisch diverse systemen te testen. Modellen werden getraind in het high-data regime op enkel atomisatie-energieën.

Evaluatiemetrieken
Prestaties werden beoordeeld met een uitgebreid pakket aan metrieken:

• Nauwkeurigheid: Mean Absolute Error (MAE), Root Mean Squared Error (RMSE) en Negative Log-Likelihood (NLL).
• Kwaliteit van Onzekerheid: Root Mean Squared Calibration Error (RMSCE), Sharpness (gemiddelde voorspellende onzekerheid), Spearman rangcorrelatie ($r_s$) tussen onzekerheid en fout, en een "Overlap Score" die de uitlijning tussen hoge onzekerheidsvoorspellingen en hoge foutregio's meet.

Belangrijkste Resultaten

1. Voorspellende Nauwkeurigheid:

   • In het High-Data regime (TiO2) presteerde DE consequent beter dan VBNNs, waarbij het de laagste MAE (6 meV/atoom) en RMSE (9 meV/atom) behaalde, evenals de meest negatieve NLL (-4.2), wat wijst op superieure nauwkeurigheid en kalibratie.
   • In het Low-Data regime behield DE zijn superioriteit in nauwkeurigheid (MAE 35 meV/atom). Echter, LRT vertoonde concurrerende nauwkeurigheid (MAE 54 meV/atom) en een lagere RMSE (259 meV/atom) dan DE, wat duidt op een betere robuustheid tegen uitschieters. RAD presteerde slecht in low-data settings.
   • Op de QM7 dataset vond een omkering plaats: BNNs (LRT en RAD) behaalden lagere energie RMSE (13.0–13.3 meV/atom) dan DE (15.2 meV/atom), wat suggereert dat de variational prior zorgt voor begunstigende regularisatie in chemisch diverse ruimtes.

2. Onzekerheidskwantificatie:

   • Kalibratie: DE leverde over het algemeen de best gekalibreerde onzekerheden (laagste RMSCE) in de TiO2 experimenten. In het low-data regime behield RAD verrassend genoeg een uitstekende kalibratie (RMSCE 0.05), vergelijkbaar met de high-data prestaties, terwijl LRT en DE een gedegradeerde kalibratie vertoonden.
   • Sharpness: DE produceerde de scherpste (meest zelfverzekerde) voorspellingen. LRT had de neiging om overly brede onzekerheidsintervallen te produceren (hoge sharpness waarden), vooral in low-data regimes.
   • Correlatie en Overlap: DE demonstreerde de sterkste monotone correlatie ($r_s$) tussen voorspelde onzekerheid en werkelijke fout in het low-data regime (0.45), waarmee het superieur was aan zowel LRT (0.28) als RAD (0.19). DE behaalde ook de hoogste Overlap Score (55%), wat aangeeft dat het hoog-fout regio's het meest effectief identificeert.

3. Krachtvoorspellingen:

   • In het low-data regime behaalden VBNNs (LRT en RAD) significant lagere kracht RMSE dan DE, waarschijnlijk door de regularisatie die de prior biedt. Echter, de krachtkalibratie bleef slecht voor alle modellen, waarbij alleen DE een redelijke kalibratie vertoonde in het low-data regime.

4. Computationele Efficiëntie:

   • Training: DE heeft een lagere kosten per epoch maar vereist het trainen van vijf onafhankelijke modellen. De totale trainingstijd voor DE was vergelijkbaar met LRT in het high-data regime, maar iets hoger in het low-data regime.
   • Inference: DE is aanzienlijk sneller tijdens inference (0.03 s) vergeleken met VBNNs (0.64–0.71 s), aangezien BNNs 20 Monte Carlo samples per voorspelling vereisen, inclusief automatische differentiatie voor krachten.

Betekenis en Conclusies
De paper concludeert dat hoewel Deep Ensembles theoretisch minder gegrond zijn dan BNNs, ze empirisch beter presteren dan variational BNNs op de meeste metrieken wat betreft nauwkeurigheid, kalibratie en inference-efficiëntie, vooral wanneer de chemische ruimte relatief smal is (TiO2). DE wordt aanbevolen als de voorkeurskeuze wanneer betrouwbare onzekerheidskwantificatie de prioriteit is en implementatiesimpliciteit gewenst is.

De studie benadrukt echter dat VBNNs, specifief de LRT variant, een competitieve trade-off bieden en superieur kunnen zijn in data-arme of chemisch diverse regimes (zoals gezien bij QM7) waar de variational prior fungeert als een effectieve regularisator. De auteurs merken op dat BNNs gevoeliger zijn voor willekeurige initialisatie en meerdere trainingsruns vereisen om robuustheid te garanderen. Het werk biedt praktische richtlijnen voor het selecteren van onzekerheidsbewuste MLIPs, waarbij gesuggereerd wordt dat de keuze tussen DE en VBNNs afhangt van het specifieke dataregime en de chemische diversiteit van het doelsysteem.

Bijdragen

• Introductie van een GPU-versnelde VBNN implementatie voor MLIPs binnen het ænet-PyTorch framework, met ondersteuning voor gezamenlijke energie-kracht training.
• Een systematische vergelijking van DE, LRT en RAD over high- en low-data regimes op zowel anorganische (TiO2) als organische (QM7) datasets.
• Demonstratie dat hoewel DE over het algemeen superieur is, de relatieve prestaties van BNNs kunnen omkeren in chemisch diverse datasets, wat de universaliteit van de superioriteit van ensembles in alle contexten uitdaagt.
• Levering van een nieuwe "Overlap Score" metriek om de bruikbaarheid van onzekerheidsschattingen voor actieve leerscenario's te evalueren.