A Function-Centric Graph Neural Network Approach for Predicting Electron Densities

Electronic structure predictions are relevant for a wide range of applications, from drug discovery to materials science. Since the cost of purely quantum mechanical methods can be prohibitive, machine learning surrogates are used to predict the results of these calculations. This work introduces the Basis Overlap Architecture (BOA), an equivariant graph neural network architecture based on a novel message passing scheme that utilizes the overlap matrix of the basis functions used to represent the predicted ground state electron density. BOA is evaluated on QM9 and MD density datasets, surpassing the previous state of the art in predicting accurate electron densities. Excellent generalization to larger molecules of up to nearly 200 atoms is demonstrated using a model trained only on QM9 molecules of at most 9 heavy atoms. INTRODUCTION Accurate electronic structure predictions are crucial for the development of new catalysts, improved batteries, or more specific drugs. Today's gold standard for reasonably sized systems is Kohn-Sham density functional theory (KS-DFT). It accounts for a significant fraction of worldwide supercomputing time, and three of its cornerstone methods are amongst the ten most cited publications of all time and fields (Van Noorden, 2025) . Still, its computational cost prohibits routine use on large systems, or in very high throughput scenarios. In response, machine learning surrogates are developed to either circumvent or speed up KS-DFT calculations. These methods range from property prediction (directly predicting observables from molecular geometry (