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

Résumé Technique : Réseaux de Neurones Bayésiens versus Ensembles Profonds pour la Quantification de l'Incertitude dans les Potentiels Interatomiques par Apprentissage Automatique

Énoncé du Problème
Les potentiels interatomiques par apprentissage automatique (MLIP) sont devenus essentiels pour la modélisation atomistique, offrant une alternative calculatoirement efficace aux méthodes ab initio telles que la théorie de la fonctionnelle de la densité (DFT). Cependant, les architectures standards de deep learning sont déterministes et manquent de mécanismes intrinsèques pour quantifier l'incertitude prédictive. Cette limitation est critique en science des matériaux, où les modèles sont fréquemment appliqués à des données hors distribution. Sans estimations d'incertitude fiables, les simulations risquent de produire des résultats trompeurs, et les stratégies d'apprentissage actif pour le raffinement des modèles manquent de directives. Bien que les Ensembles Profonds (DE) et les Réseaux de Neurones Bayésiens (BNN) soient des stratégies établies pour répondre à cela, leur performance comparative dans le contexte spécifique des MLIP — particulièrement concernant les prédictions conjointes d'énergie et de force à travers différents régimes de données — nécessite une évaluation systématique.

Méthodologie
Les auteurs ont implémenté et comparé deux approches principales de quantification de l'incertitude au sein du framework ænet-PyTorch :

1. Ensembles Profonds (DE) : Suivant la formulation de Lakshminarayanan et al., cinq réseaux de neurones indépendamment entraînés ont été initialisés avec différentes graines aléatoires. L'incertitude a été quantifiée comme l'écart type de l'échantillon des prédictions de l'ensemble.
2. Réseaux de Neurones Bayésiens Variationnels (VBNN) : Les auteurs ont implémenté des BNN utilisant l'inférence variationnelle (VI) via le framework de programmation probabiliste Pyro. Deux guides variationnels spécifiques ont été testés :
   • Astuce de Reparamétrage Local (LRT) : Utilisant une approximation de champ moyen avec une distribution postérieure gaussienne entièrement factorisée, optimisée via l'algorithme Bayes by Backprop avec LRT pour réduire la variance du gradient.
   • Guide Radial (RAD) : Une approche non de champ moyen introduisant des transformations radiales pour capturer les corrélations globales entre les poids, optimisée via l'échantillonnage de Monte Carlo standard de la borne inférieure de l'évidence (ELBO).

Configuration Expérimentale
L'étude a évalué ces méthodes sur deux ensembles de données distincts :

• Jeu de données TiO2 : 7 815 structures représentant diverses phases de dioxyde de titane. Les modèles ont été entraînés à la fois sur l'ensemble complet du jeu de données (régime de données élevées) et sur un sous-ensemble de 20 % (régime de données faibles) pour évaluer l'efficacité des données. Un entraînement conjoint énergie-force ainsi qu'un entraînement sur l'énergie seule ont également été effectués.
• Jeu de données QM7 : 7 165 petites molécules organiques utilisées pour tester la généralisabilité à travers des systèmes chimiquement divers. Les modèles ont été entraînés en régime de données élevées sur les énergies d'atomisation uniquement.

Métriques d'Évaluation
La performance a été évaluée par une suite complète de métriques :

• Précision : Erreur Absolue Moyenne (MAE), Erreur Quadratique Moyenne (RMSE) et Log-Vraisemblance Négative (NLL).
• Qualité de l'Incertitude : Erreur de Calibration Quadratique Moyenne (RMSCE), Netteté (Sharpness - incertitude prédictive moyenne), corrélation de rang de Spearman ($r_s$) entre l'incertitude et l'erreur, et un « Score de Recouvrement » (Overlap Score) mesurant l'alignement entre les prédictions à haute incertitude et les régions à haute erreur.

Résultats Clés

1. Précision Prédictive :

   • Dans le régime de données élevées (TiO2), les DE ont systématiquement surpassé les VBNN, atteignant la MAE la plus faible (6 meV/atome) et la RMSE la plus faible (9 meV/atome), ainsi que la NLL la plus négative (-4,2), indiquant une précision et une calibration supérieures.
   • Dans le régime de données faibles, les DE ont maintenu leur supériorité en termes de précision (MAE 35 meV/atome). Cependant, le LRT a montré une précision compétitive (MAE 54 meV/atome) et une RMSE plus faible (259 meV/atome) que les DE, suggérant une meilleure robustesse aux valeurs aberrantes. Le RAD a mal performé dans les contextes de données faibles.
   • Sur le jeu de données QM7, une inversion s'est produite : les BNN (LRT et RAD) ont obtenu des RMSE d'énergie plus faibles (13,0–13,3 meV/atome) que les DE (15,2 meV/atome), suggérant que le prior variationnel fournit une régularisation bénéfique pour les espaces chimiquement divers.

2. Quantification de l'Incertitude :

   • Calibration : Les DE ont généralement fourni les incertitudes les mieux calibrées (plus faible RMSCE) dans les expériences TiO2. Dans le régime de données faibles, le RAD a étonnamment maintenu une excellente calibration (RMSCE 0,05), comparable à sa performance en régime de données élevées, tandis que le LRT et les DE ont vu leur calibration se dégrader.
   • Netteté (Sharpness) : Les DE ont produit les prédictions les plus nettes (les plus confiantes). Le LRT a eu tendance à produire des intervalles d'incertitude trop larges (valeurs de netteté élevées), particulièrement dans les régimes de données faibles.
   • Corrélation et Recouvrement : Les DE ont démontré la plus forte corrélation monotone ($r_s$) entre l'incertitude prédite et l'erreur réelle dans le régime de données faibles (0,45), surpassant le LRT (0,28) et le RAD (0,19). Les DE ont également obtenu le score de recouvrement le plus élevé (55 %), indiquant qu'ils identifient le plus efficacement les régions à haute erreur.

3. Prédictions de Force :

   • Dans le régime de données faibles, les VBNN (LRT et RAD) ont atteint des RMSE de force significativement plus faibles que les DE, probablement en raison de la régularisation fournie par le prior. Cependant, la calibration de la force est restée médiocre pour tous les modèles, seul le DE montrant une calibration raisonnable dans le régime de données faibles.

4. Efficacité Computationnelle :

   • Entraînement : Les DE ont un coût par époque plus faible mais nécessitent l'entraînement de cinq modèles indépendants. Le temps total d'entraînement pour les DE était comparable à celui du LRT en régime de données élevées, mais légèrement supérieur en régime de données faibles.
   • Inférence : Les DE sont nettement plus rapides lors de l'inférence (0,03 s) par rapport aux VBNN (0,64–0,71 s), car les BNN nécessitent 20 échantillons de Monte Carlo par prédiction, incluant la différenciation automatique pour les forces.

Signification et Conclusions
L'article conclut que bien que les Ensembles Profonds soient théoriquement moins fondés que les BNN, ils surpassent empiriquement les BNN variationnels dans la plupart des métriques concernant la précision, la calibration et l'efficacité d'inférence, particulièrement lorsque l'espace chimique est relativement étroit (TiO2). Les DE sont recommandés comme le choix privilégié lorsque une quantification de l'incertitude fiable est la priorité et qu'une simplicité de mise en œuvre est souhaitée.

Cependant, l'étude souligne que les VBNN, spécifiquement la variante LRT, offrent un compromis compétitif et peuvent être supérieurs dans les régimes de données rares ou chimiquement divers (comme vu pour QM7) où le prior variationnel agit comme un régularisateur efficace. Les auteurs notent que les BNN sont plus sensibles à l'initialisation aléatoire et nécessitent plusieurs cycles d'entraînement pour garantir la robustesse. Ce travail fournit des directives pratiques pour la sélection de MLIP conscients de l'incertitude, suggérant que le choix entre DE et VBNN doit dépendre du régime de données spécifique et de la diversité chimique du système cible.

Contributions

• Introduction d'une implémentation de VBNN accélérée par GPU pour les MLIP au sein du framework ænet-PyTorch, supportant l'entraînement conjoint énergie-force.
• Une comparaison systématique de DE, LRT et RAD à travers des régimes de données hautes et faibles sur des jeux de données tant inorganiques (TiO2) qu'organiques (QM7).
• Démonstration que, bien que les DE soient généralement supérieurs, la performance relative des BNN peut s'inverser dans des ensembles de données chimiquement divers, remettant en question l'universalité de la supériorité des ensembles dans tous les contextes.
• Fourniture d'une nouvelle métrique de « Score de Recouvrement » pour évaluer l'utilité des estimations d'incertitude dans les scénarios d'apprentissage actif.