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

Resumo Técnico: Redes Neurais Bayesianas Versus Ensembles Profundos para Quantificação de Incerteza em Potenciais Interatômicos de Aprendizado de Máquina

Declaração do Problema
Os Potenciais Interatômicos de Aprendizado de Máquina (MLIPs) tornaram-se essenciais para a modelagem atomística, oferecendo uma alternativa computacionalmente eficiente aos métodos ab initio, como a Teoria do Funcional da Densidade (DFT). No entanto, as arquiteturas padrão de aprendizado profundo são determinísticas e carecem de mecanismos inerentes para quantificar a incerteza preditiva. Essa limitação é crítica na ciência dos materiais, onde os modelos são frequentemente aplicados a dados fora da distribuição (out-of-distribution). Sem estimativas de incerteza confiáveis, as simulações correm o risco de produzir resultados enganosos, e as estratégias de aprendizado ativo para o refinamento do modelo carecem de orientação. Embora os Ensembles Profundos (DE) e as Redes Neurais Bayesianas (BNNs) sejam estratégias estabelecidas para abordar isso, seu desempenho comparativo no contexto específico de MLIPs — particularmente em relação às previsões conjuntas de energia e força através de diferentes regimes de dados — requer uma avaliação sistemática.

Metodologia
Os autores implementaram e compararam duas abordagens primárias de quantificação de incerteza dentro do framework ænet-PyTorch:

1. Ensembles Profundos (DE): Seguindo a formulação de Lakshminarayanan et al., cinco redes neurais treinadas independentemente foram inicializadas com diferentes sementes aleatórias. A incerteza foi quantificada como o desvio padrão amostral das previsões do ensemble.
2. Redes Neurais Bayesianas Variacionais (VBNN): Os autores implementaram BNNs usando Inferência Variacional (VI) por meio do framework de programação probabilística Pyro. Dois guias variacionais específicos foram testados:
   • Truque de Reparametrização Local (LRT): Utilizando uma aproximação de campo médio com uma posterior Gaussiana totalmente fatorizada, otimizada usando o algoritmo Bayes by Backprop com LRT para reduzir a variância do gradiente.
   • Guia Radial (RAD): Uma abordagem de não-campo médio introduzindo transformações radiais para capturar correlações globais entre os pesos, otimizada via amostragem de Monte Carlo padrão do Limite Inferior da Evidência (ELBO).

Configuração Experimental
O estudo avaliou esses métodos em dois conjuntos de dados distintos:

• Conjunto de Dados TiO2: 7.815 estruturas representando várias fases de dióxido de titânio. Os modelos foram treinados tanto no conjunto completo (regime de Alta Quantidade de Dados) quanto em um subconjunto de 20% (regime de Baixa Quantidade de Dados). Tanto o treinamento de energia apenas quanto o de energia-força conjunta foram realizados.
• Conjunto de Dados QM7: 7.165 pequenas moléculas orgânicas usadas para testar a generalização em sistemas quimicamente diversos. Os modelos foram treinados no regime de alta quantidade de dados apenas para energias de atomização.

Métricas de Avaliação
O desempenho foi avaliado por meio de um conjunto abrangente de métricas:

• Acurácia: Erro Médio Absoluto (MAE), Raiz do Erro Quadrático Médio (RMSE) e Log-Verossimilhança Negativa (NLL).
• Qualidade da Incerteza: Erro de Calibração de Raiz do Erro Quadrático Médio (RMSCE), Sharpness (incerteza preditiva média), correlação de ranking de Spearman ($r_s$) entre incerteza e erro, e um "Score de Sobreposição" (Overlap Score) medindo o alinhamento entre previsões de alta incerteza e regiões de alto erro.

Resultos Principais

1. Acurácia Preditiva:

   • No regime de Alta Quantidade de Dados (TiO2), o DE superou consistentemente as VBNNs, alcançando o menor MAE (6 meV/átomo) e RMSE (9 meV/átomo), bem como o NLL mais negativo (-4.2), indicando superioridade em acurácia e calibração.
   • No regime de Baixa Quantidade de Dados, o DE manteve sua superioridade em acurácia (MAE 35 meV/átomo). No entanto, o LRT mostrou acurácia competitiva (MAE 54 meV/átomo) e um RMSE menor (259 meV/átomo) que o DE, sugerindo melhor robustez a valores discrepantes (outliers). O RAD teve um desempenho ruim em cenários de baixa quantidade de dados.
   • No conjunto de dados QM7, ocorreu uma reversão: as BNNs (LRT e RAD) alcançaram RMSE de energia mais baixos (13.0–13.3 meV/átomo) do que o DE (15.2 meV/átomo), sugerindo que a prior variacional fornece uma regularização benéfica para espaços quimicamente diversos.

2. Quantificação de Incerteza:

   • Calibração: O DE geralmente forneceu as incertezas mais bem calibradas (menor RMSCE) nos experimentos de TiO2. No regime de baixa quantidade de dados, o RAD surpreendentemente manteve uma excelente calibração (RMSCE 0.05), comparável ao seu desempenho de alta quantidade de dados, enquanto o LRT e o DE viram sua calibração degradar.
   • Sharpness: O DE produziu as previsões mais nítidas (sharpest/mais confiantes). O LRT tendeu a produzir intervalos de incerteza excessivamente amplos (valores de sharpness altos), particularmente em regimes de baixa quantidade de dados.
   • Correlação e Sobreposição: O DE demonstrou a mais forte correlação monotônica ($r_s$) entre incerteza prevista e erro real no regime de baixa quantidade de dados (0.45), superando tanto o LRT (0.28) quanto o RAD (0.19). O DE também alcançou o maior Score de Sobreposição (55%), indicando que identifica mais efetivamente regiões de alto erro.

3. Previsões de Força:

   • No regime de baixa quantidade de dados, as VBNNs (LRT e RAD) alcançaram significativamente menor RMSE de força do que o DE, provavelmente devido à regularização fornecida pela prior. No entanto, a calibração de força permaneceu ruim em todos os modelos, sendo o DE o único a mostrar uma calibração razoável no regime de baixa quantidade de dados.

4. Eficiência Computacional:

   • Treinamento: O DE possui um custo por época menor, mas requer o treinamento de cinco modelos independentes. O tempo total de treinamento para o DE foi comparável ao do LRT no regime de alta quantidade de dados, mas ligeiramente superior no regime de baixa quantidade de dados.
   • Inferência: O DE é significativamente mais rápido na inferência (0.03 s) comparado às VBNNs (0.64–0.71 s), pois as BNNs requerem 20 amostras de Monte Carlo por previsão, incluindo diferenciação automática para forças.

Significância e Conclusões
O artigo conclui que, embora os Ensembles Profundos sejam teoricamente menos fundamentados que as BNNs, eles superam empiricamente as BNNs variacionais na maioria das métricas relativas a acurácia, calibração e eficiência de inferência, particularmente quando o espaço químico é relativamente estreito (TiO2). O DE é recomendado como a escolha preferencial quando a quantificação de incerteza confiável é a prioridade e a simplicidade de implementação é desejada.

Contudo, o estudo destaca que as VBNNs, especificamente a variante LRT, oferecem um compromisso competitivo e podem ser superiores em regimes de escassez de dados ou quimicamente diversos (como visto no QM7), onde a prior variacional atua como um regularizador eficaz. Os autores observam que as BNNs são mais sensíveis à inicialização aleatória e requerem múltiplas execções de treinamento para garantir robustez. O trabalho fornece orientação prática para a seleção de MLIPs conscientes da incerteza, sugerindo que a escolha entre DE e VBNNs deve depender do regime de dados específico e da diversidade química do sistema alvo.

Contribuições

• Introdução de uma implementação de VBNN acelerada por GPU para MLIPs dentro do framework ænet-PyTorch, suportando treinamento conjunto de energia-força.
• Uma comparação sistemática de DE, LRT e RAD através de regimes de alta e baixa quantidade de dados em ambos os conjuntos de dados inorgânicos (TiO2) e orgânicos (QM7).
• Demonstração de que, embora o DE seja geralmente superior, o desempenho relativo das BNNs pode reverter em conjuntos de dados quimicamente diversos, desafiando a universalidade da superioridade dos ensembles em todos os contextos.
• Provisão de uma nova métrica de "Score de Sobreposição" (Overlap Score) para avaliar a utilidade das estimativas de incerteza em cenários de aprendizado ativo.