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

Resumen Técnico: Redes Neuronales Bayesianas frente a Ensembles Profundos para la Cuantificación de la Incertidumbre en Potenciales Interatómicos de Aprendizaje Automático

Planteamiento del Problema
Los Potenciales Interatómicos de Aprendizaje Automático (MLIP, por sus siglas en inglés) se han vuelto esenciales para el modelado atomístico, ofreciendo una alternativa computacionalmente eficiente a los métodos ab initio como la Teoría del Funcional de la Densidad (DFT). Sin embargo, las arquitecturas estándar de aprendizaje profundo son deterministas y carecen de mecanismos inherentes para cuantificar la incertidumbre predictiva. Esta limitación es crítica en la ciencia de materiales, donde los modelos se aplican frecuentemente a datos fuera de la distribución (out-of-distribution). Sin estimaciones de incertidumbre fiables, las simulaciones corren el riesgo de producir resultados engañosos, y las estrategias de aprendizaje activo para el refinamiento de modelos carecen de guía. Si bien los Ensembles Profundos (DE) y las Redes Neuronales Bayesianas (BNN) son estrategias establecidas para abordar esto, su rendimiento comparativo en el contexto específico de los MLIP —particularmente en la predicción conjunta de energía y fuerza a través de diversos regímenes de datos— requiere una evaluación sistemática.

Metodología
Los autores implementaron y compararon dos enfoques principales de cuantificación de la incertidumbre dentro del marco ænet-PyTorch:

1. Ensembles Profundos (DE): Siguiendo la formulación de Lakshminarayanan et al., se entrenaron cinco redes neuronales independientes con diferentes semillas aleatorias. La incertidumbre se cuantificó como la desviación estándar de la muestra de las predicciones del ensemble.
2. Redes Neuronales Bayesianas Variacionales (VBNN): Los autores implementaron BNN utilizando Inferencia Variacional (VI) a través del marco de programación probabilística Pyro. Se probaron dos guías variacionales específicas:
   • Truco de Reparametrización Local (LRT): Utilizando una aproximación de campo medio con una posterior gaussiana totalmente factorizada, optimizada mediante el algoritmo Bayes by Backprop con LRT para reducir la varianza del gradiente.
   • Guía Radial (RAD): Un enfoque de no campo medio que introduce transformaciones radiales para capturar correlaciones globales entre pesos, optimizado mediante el muestreo de Monte Carlo estándar del Límite Inferior de la Evidencia (ELBO).

Configuración Experimental
El estudio evaluó estos métodos en dos conjuntos de datos distintos:

• Conjunto de datos de TiO2: 7,815 estructuras que representan varias fases de dióxido de titanio. Los modelos fueron entrenados tanto en el conjunto de datos completo (régimen de Altos Datos) como en un subconjunto del 20% (régimen de Bajos Datos) para evaluar la eficiencia de los datos. Tanto el entrenamiento de energía únicamente como el de energía-fuerza conjunta se realizaron.
• Conjunto de datos QM7: 7,165 moléculas orgánicas pequeñas utilizadas para probar la generalizabilidad en sistemas químicamente diversos. Los modelos fueron entrenados en el régimen de altos datos solo con energías de atomización.

Métricas de Evaluación
El rendimiento se evaluó mediante un conjunto integral de métricas:

• Precisión: Error Absoluto Medio (MAE), Raíz del Error Cuadrático Medio (RMSE) y Log-Verosimilitud Negativa (NLL).
• Calidad de la Incertidumbre: Raíz del Error Cuadrático Medio de Calibración (RMSCE), Agudeza (Sharpness, la incertidumbre predictiva promedio), correlación de rango de Spearman ($r_s$) entre la incertidumbre y el error, y una "Puntuación de Traslape" (Overlap Score) que mide la alineación entre las predicciones de alta incertidumbre y las regiones de alto error.

Resultados Clave

1. Precisión Predictiva:

   • En el régimen de Altos Datos (TiO2), DE superó consistentemente a las VBNN, logrando el MAE más bajo (6 meV/átomo) y el RMSE más bajo (9 meV/átomo), así como el NLL más negativo (-4.2), lo que indica una precisión y calibración superiores.
   • En el régimen de Bajos Datos, DE mantuvo su superioridad en precisión (MAE 35 meV/átomo). Sin embargo, LRT mostró una precisión competitiva (MAE 54 meV/átomo) y un RMSE menor (259 meV/átomo) que DE, lo que sugiere una mejor robustez ante valores atípicos. RAD funcionó deficientemente en entornos de bajos datos.
   • En el conjunto de datos QM7, ocurrió una reversión: las BNN (LRT y RAD) lograron un RMSE de energía más bajo (13.0–13.3 meV/átomo) que DE (15.2 meV/átomo), sugiriendo que el prior variacional proporciona una regularización beneficiosa para espacios químicamente diversos.

2. Cuantificación de la Incertidumbre:

   • Calibración: DE proporcionó generalmente las incertidumbres mejor calibradas (menor RMSCE) en los experimentos de TiO2. En el régimen de bajos datos, RAD sorprendentemente mantuvo una excelente calibración (RMSCE 0.05), comparable a su desempeño en el régimen de altos datos, mientras que LRT y DE vieron degradada su calibración.
   • Agudeza (Sharpness): DE produjo las predicciones más agudas (más confiadas). LRT tendió a producir intervalos de incertidumbre excesivamente amplios (altos valores de agudeza), particularmente en regímenes de bajos datos.
   • Correlación y Traslape: DE demostró la correlación monotónica más fuerte ($r_s$) entre la incertidumbre predicha y el error real en el régimen de bajos datos (0.45), superando tanto a LRT (0.28) como a RAD (0.19). DE también logró la Puntuación de Traslape más alta (55%), indicando que identifica más eficazmente las regiones de alto error.

3. Predicciones de Fuerza:

   • En el régimen de bajos datos, las VBNN (LRT y RAD) lograron un RMSE de fuerza significativamente menor que DE, probablemente debido a la regularización proporcionada por el prior. Sin embargo, la calibración de la fuerza fue deficiente en todos los modelos, siendo solo DE el que mostró una calibración razonable en el régimen de bajos datos.

4. Eficiencia Computacional:

   • Entrenamiento: DE tiene un costo por época menor, pero requiere el entrenamiento de cinco modelos independientes. El tiempo total de entrenamiento para DE fue comparable al de LRT en el régimen de altos datos, pero ligeramente superior en el de bajos datos.
   • Inferencia: DE es significativamente más rápido en la inferencia (0.03 s) en comparación con las VBNN (0.64–0.71 s), ya que las BNN requieren 20 muestras de Monte Carlo por predicción, incluyendo la diferenciación automática para las fuerzas.

Significado y Conclusiones
El artículo concluye que, aunque los Ensembles Profundos son teóricamente menos fundamentados que las BNN, superan empíricamente a las BNN variacionales en la mayoría de las métricas respecto a precisión, calibración y eficiencia de inferencia, particularmente cuando el espacio químico es relativamente estrecho (TiO2). Se recomienda DE como la opción preferida cuando la prioridad es la cuantificación de la incertidumbre fiable y se desea simplicidad de implementación.

Sin embargo, el estudio destaca que las VBNN, específicamente la variante LRT, ofrecen un intercambio competitivo y pueden ser superiores en regímenes de pocos datos o químicamente diversos (como se vio en QM7), donde el prior variacional actúa como un regularizador efectivo. Los autores señalan que las BNN son más sensibles a la inicialización aleatoria y requieren múltiples ejecciones de entrenamiento para asegurar la robustez. El trabajo proporciona orientación práctica para la selección de MLIP conscientes de la incertidumbre, sugiriendo que la elección entre DE y VBNN debe depender del régimen de datos específico y la diversidad química del sistema objetivo.

Contribuciones

• Introducción de una implementación de VBNN acelerada por GPU para MLIP dentro del marco ænet-PyTorch, con soporte para el entrenamiento conjunto de energía-fuerza.
• Una comparación sistemática de DE, LRT y RAD a través de regímenes de altos y bajos datos en conjuntos de datos tanto inorgánicos (TiO2) como orgánicos (QM7).
• Demostración de que, si bien DE es generalmente superior, el rendimiento relativo de las BNN puede revertirse en conjuntos de datos químicamente diversos, desafiando la universalidad de la superioridad de los ensembles en todos los contextos.
• Provisión de una nueva métrica de "Puntuación de Traslape" para evaluar la utilidad de las estimaciones de incertidumbre en escenarios de aprendizaje activo.