Introduction

Computing electronic properties of the ground-state of molecules and materials is a central task in chemistry and materials science. Even within the common mean-field formulation known as Kohn-Sham density functional theory (DFT)1, the computations can be prohibitively expensive for many applications, for example, in molecules with many atoms, or when many geometric configurations must be considered. To reduce this cost, machine learning techniques have risen to the fore. These substitute physics-based calculations with a machine learning model, turning the electronic structure problem into a regression task to map a chemical description to the electronic state and properties2,3.

The most common machine learning models for electronic structure target the electronic energy. In one approach, known as machine learning force fields, the inputs are the molecular geometry and element types (and sometimes atomic energies)4,5,6,7,8,9,10,11,12,13,14,15. Another common class of models requires first carrying out a simpler molecular electronic structure calculation (e.g., semi-empirical tight-binding density functional theory) to generate an input to a model that predicts the energy of a more sophisticated electronic structure method on the same molecule16,17,18. In both approaches, much effort has focused on improving the featurization and model architectures, for example, to explicitly incorporate the equivariance and invariance that arise from physical symmetries.13. These ML models have achieved considerable success, greatly reducing the time to obtain quantum mechanical quality energies. However, the energy is only a single facet of the electronic structure, and machine learning other properties of the quantum state has been comparatively less explored19.

Within density functional theory, due to the Hohenberg-Kohn theorem20, the ground-state electron density formally yields all properties of the associated quantum state. Machine learning the electron density has thus attracted recent attention17,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42. Like their counterparts in force fields, ML density models have been explored with different sources of input data, ranging from ones that just take in the geometry and element types21,22,23,24,25,27,28,29,30,31,32,33,35,36,37,38,39,40,41,42 to ones that first require a quantum mechanical calculation17,34.

We will demonstrate that a powerful way to learn an accurate molecular electron density is to start from its crudest physical approximation, the superposition of atomic densities, represented on a spatial grid. In fact the atomic densities can be chosen to be very crude, and serve mainly to denote the size and location of the atoms, but importantly, once the choice is made, they do not require recalculation for any molecule. After featurizing the input superposed atomic density on a coarse grid, our model learns the accurate molecular electron density on a high-resolution grid. This learning task is the 3D analog of image super-resolution, a well-established domain in computer vision43,44, where crude, lossy images at low resolution are enhanced to accurate, high-resolution images. We therefore adopt a standard image processing approach, using a simple convolutional residual neural net (ResNet)45 that does not enforce any physical symmetries. We find that we achieve state-of-the-art accuracy in density prediction, improving on all previous models. Furthermore, because we do not use explicit element types, our architecture generalizes, and can be easily fine-tuned, across chemical space, which we demonstrate in molecules including elements we do not train against. From the density, we compute energies and orbitals through a single diagonalization of the Kohn-Sham Hamiltonian based on our predicted density. We find that even though we do not train against these quantities, we can obtain accurate predictions from this additional diagonalization step. We propose that the strategy in this work not only provides an important step forward in the density prediction problem but also opens up a distinct class of methods to predict real-space electronic structure properties based on well-established image processing techniques. We conclude with a discussion of the future of such models, including the current limitations and potential future solutions.

Results

Density learning on the test set

In Table 1, we summarize the accuracy of our density model trained on the QM9 dataset, which contains 134K small neutral organic molecules composed of the CHONF elements46. As is usual, the dataset is split into training (used to optimize the model parameters), validation (used to determine the end-point of training), and test (used for independent assessment) datasets, and we follow the same number of training/validation/testing points used in Ref. 38.

Table 1 Density prediction mean absolute error on QM9 test set
Full size table

We consider two types of input densities for our density prediction models (see Methods for model details). The first is the superposition of spherical neutral atomic density (SAD) (similar to the atomic guess approach in NWChem47 but without building and diagonalizing a molecular Fock), based on pre-tabulated atomic natural orbitals expanded in a Gaussian atomic basis48,49,50. The other is of even lower (atrocious) quality, and is generated by modifying the SAD density by multiplying the Gaussian basis exponents of the atomic natural orbitals to be filled by electrons by two. This results in a highly spatially contracted density that is very poor in terms of Errρ but which still preserves qualitative information about the atom positions and the relative order of atom sizes.

As shown in Table 1, starting from the SAD guess, and using a spatial upscaling factor of 2 (i.e., the predicted density is output on a grid resolution that is uniformly improved from that of the input density by a factor of 2 along each axis) our model refines the density error of the input SAD density by two orders of magnitude. (As an immediate consequence of this, using the predicted density in self-consistent-field (SCF) iterations to obtain a converged DFT density required, on average, 5.5 fewer iterations than starting from the SAD guess, or roughly 35% of the iterations). The atrocious SAD density has a worse input density error by a factor of 2 from the SAD density (Errρ = 27.7%). However, using this input, the model achieves a 173-fold factor of improvement, and the predicted density saved 10.2 SCF cycles (49% of the iterations) on average compared to the atrocious SAD guess. The predicted error in the density is then almost the same as when using the standard SAD input (model prediction Errρ = 0.16% on the test set), testifying to the robustness of the model. If we spatially upscale the resolution by a factor of 4, the result is only slightly less accurate than upscaling by a factor of 2, with Errρ = 0.19%.

To place our results in a broader perspective, we compare to the density prediction error of some recent density learning approaches on the same datasets. ChargE3Net41 and DeepDFT38 are equivariant models that take the molecular structure and element types as inputs, while OrbNet-Equi17 uses input information from a molecular semi-empirical tight-binding DFT calculation. We reproduce their reported errors in Table 1 and we find that we achieve a better accuracy than all these prior models.

We further compare to the error from a direct numerical fit of the target density. For this, we use Gaussian density fitting, a common numerical representation in quantum chemistry which expands the density in an auxiliary Gaussian basis51. Here, we choose the auxiliary basis to be the default even-tempered Gaussian auxiliary basis in PYSCF for our DFT computational basis, and fit using both the Coulomb metric51 and the overlap metric52. Our density model outperforms even this direct numerical fit to the data by a factor of 2 in terms of density mean absolute error (MAE), although it falls behind according to some metrics, as discussed later.

Energies and properties from predicted densities

Within the framework of Kohn-Sham density functional theory, the total electronic energy and orbitals can be obtained from the exact electron density via a single diagonalization of the effective Kohn-Sham Hamiltonian (one-step DFT). Although this additional diagonalization means these quantities are not obtained within a pure data-driven approach, they provide a physically interpretable metric with which to understand the quality of a density prediction model, and one-step DFT metrics have been employed in several recent density prediction models28,38,41,53. It is important to note that the energy obtained from the one-step DFT process is not the same as the energy computed directly from the Kohn-Sham energy functional using the model-predicted density. The latter is challenging to compute due to the unknown kinetic energy functional in the Kohn-Sham framework. This discrepancy arises because the diagonalization step also updates the density, typically improving it (e.g., by two-fold on average on the QM9 test set).

We list our MAE errors in Table 2. The one-step DFT electronic energy error is small (well within chemical accuracy of 1 kcal/mol  ≈ 43 meV). The one-step DFT highest-occupied-molecular-orbital (HOMO) and lowest-unoccupied-molecular-orbital (LUMO) energies (ϵHOMO, ϵLUMO), and HOMO-LUMO gap (Δϵ) are less accurate than the total energy (although still within chemical accuracy). We find larger errors in the individual energy components: for the Coulomb energy MAE = 167 meV, the exchange-correlation energy MAE = 40 meV, and the Kohn-Sham kinetic energy MAE = 303 meV. This is because near the DFT electronic minimum, the energy error is quadratic in the density error (which is here small), while the relationship for other observables is linear53.

Table 2 Electronic property errors on QM9
Full size table

Despite the use of one-step DFT errors as a metric in density learning28,38,41,53, reported numbers for the multiple quantities in Table 2 for the QM9 dataset are scarce. One comparable example is given by the DeepDFT model, which reported the distribution of energy errors on this dataset. We show the comparison with the distribution in Fig. 1 and we see that our errors are on average smaller (note the logarithmic axes), and we do not see large outliers. Removing the outliers of DeepDFT (defined as density MAE  > 0.6%) cannot change its prediction accuracy more than 0.01% due to the relatively few number of outliers (15 data points, according to Fig. 2 of Ref. 38 compared to 10,000 total testing data). While our ResNet model achieves higher accuracy regarding the density MAE than a direct fitting to the SCF-converged density (Table 1), it does not in terms of energy errors (Table 2). One might prioritize energy-related metrics during training if energy errors are deemed more practically significant.

Fig. 1: Energy error distribution from a one-step DFT procedure based on the predicted electron density.
figure 1

The equivariant DeepDFT result was extracted from Ref. 38. Source data are provided as a Source Data file.

Full size image
Fig. 2: Zero-shot density prediction mean absolute error of the QM9 model and timings on a water cluster dataset.
figure 2

Representative molecular configurations of a three-water cluster and a thirty-water cluster are shown. The ResNet timing was performed using one A100 GPU, and the DFT timing was performed on one compute node with 28 cores [CPU model: Intel(R) Xeon(R) E5-2680 v4 @ 2.40GHz]. The computational cost of ResNet grows linearly with the system size due to the increasing box volume, where we maintain a constant margin around the actual extent of the atoms. Source data is provided as a Source Data file.

Full size image

In addition to the above quantities, we can also assess the quality of prediction of properties that are directly available from the density without diagonalization, such as the multipole moments (non-traceless, using the center of electron charge density as the origin). The L2 norm of our dipole, quadrupole and octupole errors are 68 mD, 0.14 a.u., and 1.0 a.u. respectively. Another model that reports these quantities on the QM9 data set is OrbNet17. Our results can be compared to the reported 185 mD, 0.856 a.u., and 8.68 a.u. errors for OrbNet17.

Transferability on unseen geometries and elements

We now assess the transferability of our model. As a first task, we consider the performance of our ResNet model trained on QM9 on a different dataset of the isomers C7O2H10. Unlike QM9, which consists only of equilibrium molecular geometries, the C7O2H10 dataset contains diverse conformers sampled at high temperature, and thus the geometries and associated electron densities extend outside of the training space of QM9. Specifically, the C7O2H10 dataset comprises 5000 conformations for each of the 113 isomers of C7O2H10. As these conformations were sampled at 1 fs intervals from molecular dynamics (MD) simulations and geometries at adjacent time steps are very correlated, we evenly downsampled by a factor of 100 for each isomer trajectory (keeping the 0-th, 99-th, ... geometries), resulting in a total of 5650 less correlated data points.

As shown in Table 3, our ResNet model trained on QM9 (denoted zero-shot) moderately decreases in accuracy when evaluated on the C7O2H10 dataset, illustrating a qualitative difference between equilibrium and non-equilibrium structure densities. This hypothesis is supported by the high accuracy we observe when evaluating our model densities only for the initial conformation of each of the 113 isomer MD trajectories (Errρ = 0.13%), which correspond to local energy minima. We also assess the zero-shot eqDeepDFT (trained on QM9) model. This gives a zero-shot density error of 1.8% on the full diluted dataset and 0.29% on the first conformation of each isomer. The zero-shot error of ResNet is not only smaller than that of eqDeepDFT, but the increase from the QM9 test error is also smaller than that of eqDeepDFT. This shows that our density model not only fits better but also generalizes better.

Table 3 Transferability check of the ResNet model trained on QM9
Full size table

Despite the decrease in accuracy of our predicted densities when fed into a one-step diagonalization, our zero-shot ResNet model still predicts the total energy within chemical accuracy of 1 kcal/mol (24.5 meV = 0.56 kcal/mol). Since we did not train on energies and the density learning was trained without conformational diversity of this isomer dataset, this illustrates the surprising robustness of the model.

To obtain better performance, we further fine-tune our ResNet model after training on QM9, using a small dataset comprising the first two conformations from the downsampled dataset, including the one at the energy minimum—from each of the 113 trajectories. During training, we monitored the validation loss on the third confirmation of each trajectory. Since the conformations in the dataset are ordered by the MD simulation time, the use of the first 3 conformations for training and validation limits our model’s fine-tuning to future conformations.

In Table 3, we show the accuracy of the fine-tuned model on the last 40 conformations of each downsampled trajectory. As such, the middle data points (4th-9th) of each trajectory are never used, and this gap between the training/validation set and test set helps minimize the correlation between the training and test data, creating a more stringent test. We note that using the first two conformations of each isomer corresponds to only 4% of our downsampled C7O2H10 dataset (and 0.04% of the original dataset). Typically, such limited training data (113 × 2 = 226 data points) does not suffice to train a deep neural network, thus, a fine-tuning approach is essential. Training on this small dataset alone yields the from-scratch ResNet model, which is less accurate than the predictions from the zero-shot ResNet model. Fine-tuning leads to a (more than) two-fold decrease of the density and energy error on the C7O2H10 set and a performance close to that achieved by ResNet in the QM9 testing set.

The isomer test evaluates the transferability of our model to moderately non-equilibrium geometries of a single gas-phase molecule. To assess the transferability of our model, which is trained on single-molecule data to aggregates and condensed-phase settings, we use a water cluster dataset54. We tested the zero-shot prediction of the QM9 model on the energy-minimized geometries of water clusters ranging from 3 to 30 molecules. As shown in Fig. 2, our model trained on gas-phase single molecule data extrapolates reasonably well to the larger water clusters, with a computational cost that is a small fraction of that of the DFT calculations.

Comparing to literature models, the equivariant e3nn model trained on 7-water clusters achieved an Errρ of around 0.65% on 30-water clusters, and including clusters up to 14 waters in the training data brought the error down to  ~0.4% (see Fig. 3 of Ref. 36). To test if similar behavior exists for our model, we fine-tuned the QM9 model on a few conformations containing up to 9 or 13 waters (see Methods for training and dataset details) and achieved a density error Errρ of 0.17% and 0.12%, respectively on the 30-water clusters. For a different comparison, we also trained from scratch on 4635 conformations of up to 13 waters and obtained Errρ = 0.11% on 30-water clusters.

Fig. 3: Density prediction mean absolute error and potential energy surfaces of a single H2 and N2 molecule as a function of bond distance.
figure 3

Source data is provided as a Source Data file.

Full size image

The water cluster dataset further provides an opportunity to assess the quality of electrostatics associated with our predicted density. The Coulomb energy Ecoul spans a large range in this dataset (its standard deviation on the 15-30 water data is 1.6 × 103 kcal/mol/molecule). Despite this, the Coulomb energy for the same 15–30 water data using the from-scratch model density prediction, followed by a single Fourier transform, gave an MAE of only 1.1 kcal/mol/molecule. As is well known, capturing long-range interactions is a challenge for commonly employed ML approaches based on short-range potentials, and a hybrid approach55 combining a density prediction model for the long-range electrostatic interactions with an ML potential for the remaining interactions, is a promising possibility.

As another test for model transferability on unseen geometries, we examined the density MAE of the QM9 model on H2 and N2 molecules. The potential energy surface was scanned every 0.01 Å, and the density error as well as the one-step DFT energy were computed at the same resolution. H-H and N-N bonds are poorly represented in the training data, and we observe that Errρ is larger than in the transferability test on the C7O2H10 dataset, but the one-step DFT energy remains accurate. Thus, the predicted equilibrium bond lengths are accurate (exact within the 0.01 Å spacing used in the PES scan). The harmonic frequencies at the minimum are 4153 cm−1 (one-step DFT) and 4172 cm−1 (reference) for H2; 2425 cm−1 (one-step DFT) and 2330 cm−1 (reference) for N2, computed from finite difference with δx = 0.01 Å.

We next evaluate the ResNet model transferability on unseen elements and molecular charges. These are non-trivial tasks for models that explicitly require the element types as input, and (indeed, without any modification) are impossible for certain models. However, our network can be applied without any modification in both these cases.

As a first test, we create a training and validation set for QM9 that deliberately omits the N atom, and train a model to this limited data. The model yields a density error of 1.4% when tested on the remaining nitrogen-containing QM9 molecules. This error is significantly higher than what we saw in the previous transferability tests on the isomer and water cluster dataset. However, the 1.4% error still represents a substantial improvement over the initial SAD guess, which has a density error of 13.6%. Further analysis reveals that the density error specifically associated with nitrogen is 3.3%, compared to a 7.8% error in the SAD guess around the nitrogen atoms. (The density error around the atoms was evaluated using grid points within the atomic radii as taken from Ref. 56). This indicates that although the pre-training itself is not sufficient for a quantitatively accurate prediction of the unseen nitrogen element, it enhances the model’s ability to predict densities for the unseen element by learning about the density of other elements.

As a second, even harder test, we consider a large biological system, namely, a model of the GTP hydrolysis reaction in microtubules57 (denoted here the MT dataset). This consists of, on average, 211 atoms and a wide range of conformations involving close interactions between GTP, water, and protein amino acid residues, and bond-breaking during GTP hydrolysis. From the density prediction perspective, the system is especially challenging as it carries a total charge of -4 and also includes two new elements compared to the QM9 dataset, namely phosphorus and magnesium.

The previous transferability task omitted nitrogen in the QM9 dataset but trained on data containing oxygen and fluorine elements. Thus testing on molecules with the unseen nitrogen element tests the interpolation ability of the model between two adjacent elements in the periodic table. In contrast, moving from the QM9 dataset to the MT dataset involves extrapolating elemental and charge state behavior. As such, very accurate predictions are not anticipated in this scenario. Indeed, we find using the zero-shot ResNet model a relatively large average error (2.1%) with a large density error around the phosphorus and magnesium elements (7.8% and 33%, respectively). A visual check confirms that the density errors are predominantly around the unseen elements (Fig. 4).

Fig. 4: Density prediction on a microtubule-mediated GTP hydrolysis model.
figure 4

The first two rows show the electron density isosurface at 0.1 e/Bohr3. The third row shows the density prediction error (ρpred − ρref) isosurfaces at 0.003 e/Bohr3 (transparent red) and −0.003 e/Bohr3 (transparent blue). Coloring of atoms: oxygen (red), nitrogen (blue), carbon (silver), hydrogen (white), phosphorus (gold), and magnesium (pink).

Full size image

To improve on the zero-shot model, we can use fine-tuning. Again, we limit ourselves to a small training set, containing a few initial conformations of the MD trajectories. Despite using limited training data (110 data points, comprising only 8% of the total data), the fine-tuned ResNet model achieves a 0.33% accuracy in density prediction on the test data (approximately 88% of the entire MT dataset), again approaching the accuracy of our state-of-the-art model on the simpler QM9 training dataset. In comparison, training from scratch with the same limited data results in a model that is much less accurate than the fine-tuned one. This highlights that pre-training our density model on a standard dataset such as QM9, in conjunction with a modest amount of fine-tuning, achieves quantitative accuracy even in the very challenging case of unseen systems with new charge states and unusual elements.

Discussion

We have demonstrated that a simple machine learning strategy, inspired by image super-resolution methods, is a promising approach for electron density prediction. The model is based on a convolutional residual network, and the input is a crude guess of the density (superposition of atomic densities; SAD) represented on a coarse 3D real-space grid. We showed that even very poor input atomic guesses make virtually no difference to the accuracy, thus the input density mainly encodes the atomic positions and atom types (the latter is implicitly represented by the sizes of the atomic densities).

From a computer vision perspective, the above process mimics the well-studied image super-resolution problem. From a physical perspective, the model may be viewed as indirectly learning the Hohenberg-Kohn (HK) map from the nuclear electrostatic potential vext(r) to the ground-state density, where the nuclear potential is encoded in the features of the input density. The latter may be seen as in the spirit of Bright-Wilson’s argument58,59 that the nuclear potential of the Hamiltonian may be identified from the features of the density. It therefore provides an alternative target for ML approaches to the construction of the HK map from vext(r) to ρ(r) via minimization of ML density functionals60,61,62,63,64,65,66,67,68,69,70.

Following the density prediction, a one-step DFT calculation enables the extraction of various DFT electronic properties, such as the total energies and molecular orbitals. We found that compared to other recent density prediction models that have used this one-step diagonalization metric, we obtained smaller average errors in the energies with fewer outliers. Our model also displayed significant transferability to previously unseen molecular geometries but only moderate transferability for interpolated chemical elements (with atomic numbers in between elements seen during training). Although we do not achieve quantitative accuracy in the zero-shot predictions for new elements, our model is easily fine-tuned on new data with unseen chemical elements, because element types are not part of the input. After fine-tuning our pre-trained models on limited data corresponding to the new geometries and elements, we found that we could achieve accurate predictions in all the cases we examined. Such transferability and fine-tunability suggests our approach is well-suited to an ensemble-based active learning setup71,72 on large unlabeled datasets. In particular, given its transferability, we expect a significant reduction in the need for expensive DFT data labeling.

Our model features a real-space convolution that necessitates a uniform grid. Such a requirement prevents us from directly working with angular-radial grids (such as atom-centered grids) that are more economical for small gas-phase molecules when all-electron calculations are conducted. Currently, the model predicts the valence density as the result of employing pseudo-potentials in the periodic DFT calculations used for data generation. The prediction of all-electron densities may be facilitated by a hybrid approach that employs models that predict density fitting coefficients for the core part of the density, while maintaining the ResNet approach for the valence part.

The success of our real-space approach to learning densities also suggests applications to other objects that are naturally expressed on real-space grids. In addition, while the electron density provides access to the electronic observables such as the energy after a single diagonalization, as performed in this work, it is desirable to use a data-driven approach to bypass this expensive computational step. These directions will be explored in the future.

Methods

Density learning and the superposition of atomic densities

The main objective of image super-resolution is to establish a mapping from a low-resolution and potentially inaccurate image (e.g., generated by lossy compression), denoted x, to its high-resolution counterpart, X. We view density learning through the same lens. Namely, given an inaccurate and low-resolution molecular density, we wish to predict the accurate molecular density, here defined as the self-consistent pseudopotential DFT density, on a finer grid. The density learning essentially performs two tasks: it transforms the inaccurate density to an accurate one, and it upscales the spatial resolution of the density. The latter is similar to interpolation and can be achieved without machine-learning techniques. The former is a more complex map, and is where data-driven approaches are most useful.

Since we wish to learn accurate densities for large molecules, it is essential for our input data to be cheap to generate. This suggests we should avoid approaches where we must first perform a low-level molecular quantum mechanical calculation, as these will be prohibitive in large problems. On the other hand, the list of atomic types and positions alone seems more parsimonious than necessary. We will define our input data to be the physically motivated approximation of the superposition of (spherical) neutral atomic densities, normalized to the total number of electrons, expressed on a low-resolution uniform grid. The spherical neutral atomic densities are obtained from filling electrons into atomic natural orbitals according to the Aufbau occupations, which need to be tabulated once for the periodic table, while the cost of assembling the SAD guess on the molecular grid requires no additional quantum calculation, is linear in the molecular size, and scalable to large molecules. The SAD guess preserves important physical constraints (e.g., it is positive) and provides qualitative information about the size and location of the atoms, but otherwise does not provide an accurate density, as quantified below.

DFT data generation

Our training data consisted of DFT densities computed using periodic boundary conditions in the Gaussian-Plane-Wave approach73,74,75 with the Perdew-Burke-Ernzerhof (PBE) functional76, the Goedecker-Teter-Hutter (GTH) pseudopotential77, and the GTH-TZV2P basis set73, computed on a fine uniform grid (plane wave cutoff = 400 Ry), as implemented in the PYSCF74,75 package. The relation between the coarse grid for the input density and the predicted fine uniform grid is a factor of 2 or 4 in each dimension as specified below (an overall factor of 8 or 64 in the number of density points; note, the high-resolution grid is fixed, so the coarse-grid is scaled).

For all the datasets, the box size was determined based on the closest even-number multiple of the grid spacing while ensuring a minimum of 4 Å of space around the molecule in every dimension. This box size occasionally led to SCF convergence issues for some long conjugated chains in QM9, and a 5 Å margin was used instead for these cases. A SCF convergence threshold of 10−11 Hartree was set for the QM9, C7O2H10, and water cluster calculations. For the water cluster dataset, we consider three subsets of the full data optimized by the TTM2.1-F potential (https://sites.uw.edu/wdbase/database-of-water-clusters/): (1) 1904 structures that are within 1 kJ/mol of the minimum (denoted as Wmin/n) (2) the first few structures of each of the n-water clusters within 1 kcal/mol of the minimum (in the order of geometries listed in the xyz file) such that n times the number of structures does not exceed 1500 (denoted as W1500/n; 692 geometries in total), and (3) structures of each of the n-water clusters within 1 kcal/mol of the minimum such that n times the number of geometries does not exceed 50,000 (denoted as W50,000/n; 5137 geometries in total). The GTP hydrolysis in microtubules (MT) dataset contains 1442 configurations extracted from 11 metadynamics runs57. The total simulation time for the metadynamics trajectories was about 750 ps. Configurations spaced every 12 hours of wall time, corresponding to a 0.5-1 ps simulation interval, were included in the dataset. For the MT dataset, the SCF convergence criteria were 10−7 Hartree for the energy and 10−5 Hartree for the orbital gradient. To aid SCF convergence, we used an electronic smearing with an electronic temperature of 0.01 Hartree.

Model details

The model is a convolutional residual neural network, adapted from the ResNet part of the SRGAN superresolution model78 for electron density prediction. Our modifications were (1) changing the 2D convolution to a 3D one, (2) substituting batch normalization with instance normalization, (3) appending a ReLU layer to the network output to guarantee positive density values, and (4) normalizing the ReLU output to match the correct electron number, i.e., ρ ← Neρ/∫ρdr. All convolutions maintained a stride of 1, with padding applied to preserve the original size. Circular padding was employed to respect the periodic boundary conditions. The resulting architecture is depicted in Fig. 5. Note that, as we do not target human perception of the visual quality of the image, we do not use an adversarial network for our loss function (as used in image-processing tasks to produce images that are visually realistic to humans78) but instead use a standard mean absolute error loss

$${{{{\rm{Err}}}}}_{\rho }=\frac{1}{{N}_{e}}\int{{{\rm{d}}}}{{{\bf{r}}}}| {\rho }^{{{{\rm{pred}}}}}({{{\bf{r}}}})-{\rho }^{{{{\rm{ref}}}}}({{{\bf{r}}}})| .$$
(1)

The reported MAE over a dataset is computed as the average of the MAE of individual molecules.

Fig. 5: Model architecture.
figure 5

K indicates the kernel size used in the convolution.

Full size image

Supplementary Table 1 summarizes the hyperparameters and density errors for all models in this study. Each model underwent 50 epochs of training using the density MAE as the training loss, with the best model chosen based on the lowest validation error. The first epoch used a warm-up with a 105 times reduced learning rate, followed by a cosine scheduler from the second epoch starting with a learning rate specified in Supplementary Table 1. We employed the Adam optimizer79 with β1 = 0.9 and β2 = 0.99. Training data was augmented with a random rotation along the x, y, or z axes. A batch size of 1, combined with instance normalization and data augmentation, proved sufficient for regularization across all cases, eliminating the need for weight decay. The data splitting for QM9 was random, and the number of training, testing, and validation points followed Ref. 38. The QM9 model was tested on the Wmin/3-30 datasets to produce Fig. 2. Fine-tuning on water clusters was performed on W1500/3-9 (59 geometries), and the validation error on W1500/10 was monitored. Another fine-tuning trial was performed on W1500/3-13 (226 geometries) and validated on W1500/14. Training from scratch on water clusters was performed on W50,000/3-13 (4635 geometries) and validated on W50,000/14. As mentioned in the main text, the original C7O2H10 dataset, comprising 5000 conformations for each of the 113 isomers, was downsampled by a factor of 100 to yield less correlated data. The first two conformations of each isomer in the diluted dataset served as the training data, with the third used for validation. The entire (diluted) dataset was used for testing if the model was not trained on the dataset, and otherwise the final 40 conformations were used for testing. For the MT dataset, the first 10 conformations from each of the 11 metadynamics trajectories were used for training, the 11th for validation, and the 17th to the last for testing. Fine-tuning on C7O2H10 and MT followed exactly the same training protocol, but the models were initialized from the best QM9-trained model.

Recent learning approaches for molecular systems have favored the incorporation of physical symmetries into the model, such as equivariance under the coordinate transformations corresponding to translations, rotations, and reflections. Equivariant models for density prediction have also been formulated, for example, in Refs. 17,24,27,28,35,36,38,40,41. Our approach satisfies a transformation property that the predictions are unchanged if the molecule is rotated and the grid is simultaneously rotated with it. This is because the input SAD density is invariant under the simultaneous coordinate transformations of the molecule and grid points, and the model uses only the values of the density (not the coordinates of the grid points) thus, the predictions are unchanged. However, the set of density values is not invariant under molecular symmetry transformations if the grid points are not also transformed, i.e., the set \(\{\langle {{{\bf{r}}}}| \rho ({{{\bf{M}}}})\rangle \}\ne \{\langle {{{\bf{r}}}}| \rho (\hat{R}{{{\bf{M}}}})\rangle \}\). A simple way to improve this is to augment the density data with that from rotated molecules on the same fixed grid, and we thus performed this data augmentation in our model training.

Computation of electronic properties

Some electronic properties were obtained from a one-step DFT procedure, which requires one Fock build and one diagonalization. Explicitly, a Fock matrix was built using the predicted electron density ρ(r):

$${f}_{\mu \nu }={h}_{\mu \nu }^{{{{\rm{core}}}}}+{J}_{\mu \nu }[\rho ]+\int{{{\rm{d}}}}{{{\bf{r}}}}{\phi }_{\mu }({{{\bf{r}}}}){v}_{{{{\rm{XC}}}}}[\rho ]({{{\bf{r}}}}){\phi }_{\nu }({{{\bf{r}}}})$$
(2)

with μ and ν indexing atom orbitals, \({h}^{{{{\rm{core}}}}}\), J, and vXC adopting their usual quantum chemical meanings of the core Hamiltonian, Coulomb matrix, and exchange-correlation potential. The one-body reduced density matrix γμν and orbital energies were obtained from diagonalizing the Fock matrix. The total energy and the Kohn-Sham kinetic energy were evaluated using the resulting γμν. The dipole moment, the Coulomb energy, and the exchange-correlation energy were obtained directly from the predicted density without the need for one step of DFT.

Reporting summary

Further information on research design is available in the Nature Portfolio Reporting Summary linked to this article.