Electronic Structure Prediction of Multi-million Atom Systems Through Uncertainty Quantification Enabled Transfer Learning

TL;DR

This paper introduces a transfer learning approach with Bayesian neural networks to predict electron densities in large multi-million atom systems efficiently, with quantified uncertainty, surpassing previous computational limitations.

Contribution

It presents a novel transfer learning framework combined with uncertainty quantification for scalable, accurate electron density predictions across diverse and large-scale systems.

Findings

Lower data generation costs for training

Accurate predictions for systems with defects and alloys

Scalable to multi-million atom systems

Abstract

The ground state electron density -- obtainable using Kohn-Sham Density Functional Theory (KS-DFT) simulations -- contains a wealth of material information, making its prediction via machine learning (ML) models attractive. However, the computational expense of KS-DFT scales cubically with system size which tends to stymie training data generation, making it difficult to develop quantifiably accurate ML models that are applicable across many scales and system configurations. Here, we address this fundamental challenge by employing transfer learning to leverage the multi-scale nature of the training data, while comprehensively sampling system configurations using thermalization. Our ML models are less reliant on heuristics, and being based on Bayesian neural networks, enable uncertainty quantification. We show that our models incur significantly lower data generation costs while allowing confident -- and when verifiable, accurate -- predictions for a wide variety of bulk systems well beyond training, including systems with defects, different alloy compositions, and at unprecedented, multi-million-atom scales. Moreover, such predictions can be carried out using only modest computational resources.