Electron Density-enhanced Molecular Geometry Learning

1 Introduction

Machine learning force fields (MLFF) is a computationally efficient and low-cost method for learning the interactions between atoms in molecular systems, bringing revolutionary advances to molecular dynamic simulations (MD) in many fields such as physics, chemistry and biology, and materials science [Chmiela et al., 2017; Xiang et al., 2024b;

∗Corresponding author (xzeng@hnu.edu.cn)

Wang et al., 2024b]. Recent MLFF methods use geometric deep learning that represent atoms in molecular systems as nodes in a geometric graph and take into account physical symmetries have been shown to be effective in learning molecular force fields (MFF) [Liu et al., 2022; Wang et al., 2024a]. However, previous studies focused on mining physical quantities at the atomic level (such as coordinates, multi-body interactions, etc.) [Batzner et al., 2022; Liao and Smidt, 2023; Wang et al., 2024a], ignoring information at the electronic level.

Electron density (ED) is a core quantum mechanical property of the distribution of electrons within a molecule and is crucial for accurately predicting the quantum chemical properties of MFF [Sunshine et al., 2023; Skogh et al., 2024]. The application of ED faces two major challenges:

Challenge 1: High computational complexity of ED. Unlike the number of atoms in a molecule, which is usually on the scale of hundreds, the ED relies on a continuous spatial distribution and as the resolution increases, the number of data points may reach millions or even tens of millions (See Appendix A for details). As shown in Figure 1(a), there are two direct ways to represent ED: point cloud [Guo et al., 2020] and voxel [Gong et al., 2023]. As shown in Figure 1(b), we empirically show the limitations of point clouds and voxels as ED representations in energy prediction in force fields, efficiency of GPU memory and training (See Appendix B for details). In particular, point clouds and voxels are directly related to the resolution of the ED, so the computational efficiency will decrease as the resolution increases. The limitations above motivate us to propose a novel multi-view RGBD image (the right subfigure of Figure 1(a)) for accurate and efficient ED representation [Xiang et al., 2024a], which is independent of ED resolution and compresses the continuous ED signal in space into pixels. Compared with point clouds and voxels, the proposed images improve energy prediction ability, GPU memory efficiency, and training efficiency by 38.4%, 42.1%, and 4.8%, respectively.

Challenge 2: Expensive ED acquisition. The acquisition of ED data mainly relies on two types of technologies: experimental measurements and theoretical computations. Experimental measurements, such as X-ray crystallography [Nienaber et al., 2000] or neutron diffraction [Goncharenko and Loubeyre, 2005], require high-precision instruments, sophisticated experimental setups, substantial time and technical expertise, making them resource-intensive. Theoretical computations, such as density functional theory (DFT) [Kohn and Sham, 1996], require a lot of computational time to obtain high-quality ED [Hegde and Bowen, 2017; Lee and Kim, 2024], which relies on high-performance computing clusters, resulting in significant cost. Given the limitations above and to improve data efficiency, as shown in Figure 1(c), we transform the ED-enhanced geometry learning problem into teaching excellent students (geometry) using an ED-aware teacher (image). Specifically, we first use DFT to obtain 2 million high-quality ED data and use them to train an ED representation learning model (called ImageED). Subsequently, we transfer the knowledge in ImageED to an EDaware teacher that takes images without ED information as input, and use an ED-aware teacher to distill excellent geometry students (called EDG). This scheme allows us to obtain ED data only once and no more ED data is needed at any other time, which greatly improves the computational efficiency. As shown in Figure 1(d), student models (SchNet [Sch¨utt et al., 2017], SphereNet [Liu et al., 2022] and ViSNet [Wang et al., 2024a]) equipped with EDG achieve significant performance improvements.

We summarize the main contributions as follows:

• To the best of our knowledge, we are the first to exploit ED images to enhance molecular geometry learning.
• We propose an efficient multi-view ED images with RGB-D channels and design an ED representation learning method, called ImageED, to automatically extract ED-related features from images.
• We propose an ED-enhanced molecular geometry representation learning framework, called EDG, which equips with ED-aware teacher to improve the performance of a large number of geometry models.
• We show that our method achieves significantly better performance on 12 datasets from QM9 and 10 datasets from rMD17 and can substantially improve the performance of existing geometry representation models.

2 RelatedWork

Molecular Geometry Representation Learning. Geometry deep learning, which studies the interactions between atoms in molecular systems, is the key to the success of machine learning force fields (MLFF) [Liu et al., 2022; Wang et al., 2024a]. Recently, the main approaches have been to incorporate physical constraints such as roto-translational invariance of the geometry into the model architecture [Zaidi et al., 2023; Wang et al., 2024b], making the output features of the model invariant to the roto-translation of the molecule. Equivariant neural network (ENN) [Satorras et al., 2021] is the most representative one and has been greatly developed in geometric representation learning. A simple way to achieve rotational-translational invariance is to construct invariant features based on the geometric conformation of the molecule, such as inter-atomic distances [Fuchs et al., 2020], angles [Liu et al., 2022], molecular descriptors [Todeschini and Consonni, 2009], etc. Besides these, there are many ENNs designed for invariance, such as models for modeling inter-atomic interactions [Sch¨utt et al., 2017; Satorras et al., 2021; Liao and Smidt, 2023] and models for modeling multi-body interactions [Wang et al., 2024b; Wang et al., 2024a]. Our approach is agnostic to the model architecture, which can enhance any geometric representation learning model from a novel electronic perspective.

Electron Density Representation Learning. Existing electron density (ED) representation methods can be mainly divided into two categories: point cloud-based and voxelbased methods. The former considers all density values in the ED as a collection of points. For example, PointNet [Qi et al., 2017] is used to classify the symmetry of inorganic compounds [Kim et al., 2024]. The latter treats each point in the ED as a voxel. For example, 3D convolutional neural network (CNN) [Liu et al., 2015] is used for the prediction of molecular exchange energy [Gong et al., 2023] and discovery of guests of host molecules [Parrilla-Guti´errez et al., 2024], respectively. 3D-UNet [C¸ ic¸ek et al., 2016] is used to segment reactive sites in molecules and classify substances [Singh et al., 2024]. Different from previous methods, we propose a novel multi-view RGB-D image to represent ED and design a representation learning method to extract features.

3 Our Method

3.1 Preliminaries

Background. Electron density (ED) is a key bridge to understand the prediction of energy and forces in molecular force fields (MFF). ED is not only the core output of quantum mechanical calculations, but also provides a solid theoretical basis for constructing high-precision MFF and understanding complex interactions between molecules. This demonstrates the significance of the proposed method to introduce ED into the learning framework of geometry representation. We provide more background details in the Appendix C.

Notation and Problem Formulation. The molecular geometry and the corresponding ground-truth labels with t prediction tasks are {Gi = (Vi, Ei)}ni=1 and {yi}ni=1 ∈ Rk respectively, where V ∈ Rnvi ×(3+dvi ) (nvi , 3, dvi are the number of atoms, the coordinates of the atoms, and the feature dimensions of the atoms) and E ∈ Rnvi ×nvi ×dei (dei is the feature dimensions of bonds). The corresponding multi-view ED image and structural image are U ∈ RV×4×H×W and S ∈ RV×3×H×W respectively, where V , H, W represent the number of views, the height, and the width of images, respectively. 3 and 4 represent RGB channels and RGB-D channels, respectively. This paper mainly defines three problems: (1) Pre-train a masked autoencoder (MAE) [He et al., 2022] architecture consisting of an ED encoder fEDE and an ED decoder fEDD to learn useful representations F U ∈ RdU (dU is the dimension of features) from the ED image U; (2) Pre-train an ED-aware teacher fS and ED predictor fEDP so that it can complete the mapping from structural images S to structural features F S ∈ RdS and then to ED-related features F S→U ∈ RdU ; (3) Distill a strong geometry student fG using ED-aware teacher, ED predictor and mapper fM. Among them, the geometry student fG takes the molecular geometry as input and extracts the corresponding geometry features F G ∈ RdG and the mapper is used to convert the geometry features into features recognized by ED predictor F G→S ∈ RdS for distillation.

3.2 Overview of the Method

Here, we propose the Electron Density-enhnaced molecular Geometry representation learning framework (called EDG). The overview of EDG is illustrated in Figure 2, which is divided into 4 main modules: (a) Given 2 million molecular conformations, necessary DFT data are generated and further processed together with the conformations into multiview ED images with RGB-D channels (Section 3.3); (b) ImageED receives the 2 million multi-view ED image as input and utilizes two pre-training tasks to learn ED-related knowledge (Section 3.4); (c) We design an ED-aware teacher, which is optimized by minimizing the difference between the predicted ED features from the structural images and the true ED features from the ImageED on 2 million molecules (Section 3.5). (d) The ED-aware teacher is used to distill a strong geometry student by using ED predictor and mapper (Section 3.6). We summarize the main processed in Appendix D.

3.3 Generation of RGB-DElectron Density Images

As shown in Figure 2(a), we first obtain 2 million molecular conformations from PCQM4Mv2 [Hu et al., 2021] and use density functional theory (DFT) with the basis set of 6-31G**/+G** and the exchange-correlation functional of B3LYP to generate DFT data for these molecular conformations [Sud, 2016]. For each molecule, the DFT data includes an electrostatic potential (ESP) file and an ED file stored in the form of a three-dimensional grid. Next, we describe the main details of ED image generation. The structural loader uses the command load {conformation file} in PyMol [DeLano and others, 2002] to load the structural information from the molecular conformation file. ED loader uses the commands load {ED file}, ED; load {ESP file}, ESP; ramp new legend, ESP, [-0.08, 0, 0.08], [red, white, blue]; isosurface surface, ED, 0.05; set surface color, legend, surface in PyMol to load the ED information, which uses a red-white-blue distribution range to describe the electrostatic potential distribution of ED and red, white, and blue represent positive, neutral, and negative regions, respectively. Finally, multi-view joint render uses commands set transparency, 0.4; turn {axis}, {angle}; png {path}, width={width}, height={height} in PyMol to render the ED information and structural information into multi-view RGB-D ED images U ∈ R6×4×224×224 . Specifically, {axis} and {angle} represent the rotation angle degrees along axis and we set ({axis}, {angle}) to (x, 0), (x, 180), (x, 90), (x, -90), (y, 90), (y, -90), which means generating images from 6 different views. {width} = 224 and {height} = 224 represent the width and height of the rendered image and {path} represents the path where the image is saved. We describe more details of ED image rendering in Appendix E.

3.4 EDRepresentation Learning with ImageED

The proposed ED images have two properties: (1) RGB and D channels use color to represent the distribution of ED and depth to represent the spatial layout, respectively, which indicates that each pixel has a clear physical meaning; (2) the distribution of ED is continuous, which means that ED can be predicted by the context. Therefore, we propose a novel ED representation learning framework (called ImageED) with mask prediction and restoration prediction tasks to learn pixel-level local and contextual ED information from 2 million ED images. ImageED is an encoder-decoder architecture, which is built following the ViT-Base/16 [Dosovitskiy et al., 2020] of MAE [He et al., 2022]. For a given batch (n molecules) of ED images u ∈ Rn×c×V×H×W (where V, c,H,W = 6, 4, 224, 224), we first use a view-agnostic patch embed layer with a patch size of np = 16 to transform u into a pile of multi-view tokens mt.

We assume nt = H/np = W/np and mt = {mti,j |i, j ∈ {0, 1, ..., nt}} ∈ Rn×V×n2t ×(n2p×c) . Next, we add the positional embeddings and expand along the view to get tokens t =∈ Rn×(V×n2t )×(n2p×c) . We shuffle the order of tokens and randomly mask 25% of the tokens to obtain the masked tokens tm ∈ Rn×(0.25×V×n2t )×(n2p×c) and unmasked tokens tu ∈ Rn×(0.75×V×n2t )×(n2p×c) respectively. We input the tu into the ED encoder fEDE to get the encoded tokens:

where dU represents the feature dimension of hu . Afterwards, we use additional mask tokens ˆhm that is initialized by 0 to the encoded tokens to obtain all encoded tokens h:

where π represents arrange them in the order of the images and pos represents the positional embeddings. The final predicted masked tokens ˆ tm and unmasked tokens ˆ tu can be obtained by input h into the ED Decoder fEDD.

In order to optimize the fEDE and fEDD in ImageED, we define a mask prediction task LMP :

where sim() represents the Euclidean distance. However, we find that using only mask prediction tasks will limit the understanding of local features. We further introduce the restoration prediction task LRP :

Finally, the overall loss function of ImageED is formulated:

where λMP and λRP is the balance coefficient and we set them to 1. After pre-training ImageED with 2 million molecules, we use the fEDE to extract ED features.

3.5 Pretraining of ED-aware Teacher

Since the generation of ED data requires a lot of computing resources, it is resource-intensive and impractical to calculate ED data for every downstream task. Therefore, we hope to use an easily accessible intermediary to replace ED to generate ED features. We assume that the molecular conformation, ED data generated by DFT, and ED features are s, x, h, respectively. The acquisition of ED features follows this path: s DFT --→ x ImageED -----→ h, which shows that when s is known, x and h can be obtained. Therefore, according to the probability chain rule, the joint probability distribution p(h, x|s) can be decomposed into: p(h, x|s) = p(h|x, s) · p(x|s). According to the Markov hypothesis [Markov, 1960], we can get: p(h|x, s) = p(h|x). In order to simplify p(h, x|s) to a probability distribution that is independent of x, we further marginalize x (integrating with respect to x) and get:

Therefore, we can approximate p(h|x) · p(x|s) by directly learning a p(h|s). Here, we propose an ED-aware teacher fS and ED predictor fEDP as p(h|s) to learn the mapping of ED features directly from molecular conformations. We choose multi-view structural images as the input of the fS (See Appendix F for specific reasons). Specifically, structural render uses command template turn {axis}, {angle}; png {path} in PyMol to render molecular structures from 2 million conformations into multi-view images S. Considering computational efficiency, we generate 4 views here and ({axis}, {angle}) is set to (x, 0), (x, 180), (y, 180), (z, 180). fS uses ResNet18 [He et al., 2016] with a view-wise average pooling and fEDP is a Multilayer Perceptron (MLP) with Linear Layer→Softplus Activator→Linear Layer. Given a batch (n molecules) of structural images s from S and ED images u from U, we obtain structural features F S and ED features F S→U :

Next, we freeze the ED encoder that accepts ED images u as input and use a token-wise average pooling to convert the token features output by the ED encoder into ED features F U . Finally, we take F U as the ground-truth and train the EDaware teacher and ED predictor to learn the mapping from structural features to ED features on 2 million molecules. The loss function Lalign is defined as:

where L1 represent L1 distance. With the ED-aware teacher, the costly ED image can be replaced by a cheaper structural image, significantly reducing DFT-related costs.

3.6 ED-enhancedMolecular Geometry Learning

In the training stage of downstream tasks, we first convert the molecular conformations in the dataset into geometric data and multi-view structural images using structural loader and structural render. Subsequently, geometry data and images are input into the geometry student fG and frozen EDaware teacher to extract features F G and F S , respectively. Please note that fG can be any geometry-based model, such as SchNet [Sch¨utt et al., 2017], EGNN [Satorras et al., 2021], etc. Next, a mapper fM is used to map the geometry features into the structural space to obtain fM(F G ). Subsequently, the frozen ED predictor accepts fM(F G ) and F S as input and obtains predicted ED features:

To distill the ED knowledge from the teacher model to the student model, we define a consistency loss LED:

where SL1 represents smooth L1 distance [Girshick, 2015]. In order to obtain task-related labels, we define a task predictor fT , which accepts geometry features F G to output taskrelated logits ˆ y = fT (F G ). The task-related loss function is defined as:

The final loss of EDG is formulated:

where λ is the balanced coefficient. In the inference phase, the prediction result is obtained by sequentially inputting the geometry data into the student network fG and the task predictor fT . Therefore, images are only involved in the training of the model and are not needed during inference, which further improve the efficiency.

4 Experiments and Results

4.1 Experimental Settings

Datasets and evaluation protocol. To pre-train ImageED, the ED-aware teacher, and the ED predictor, we select the first 2 millions unlabeled molecular conformations and their DFT-computed ED data from the EDBench database [Xiang et al., 2025], generated by Psi4 software [Turney et al., 2012] with a grid spacing of 0.4. In evaluation stage, we select 12 widely used tasks related to quantum mechanic properties from QM9 [Ramakrishnan et al., 2014] and 10 common tasks related to energy/force from revised MD17 (rMD17) [Christensen and Von Lilienfeld, 2020]. It is worth noting that for the force prediction, we first predict the molecular energy and use the gradient of each node position as the force, that is, force=-torch.autograd.grad(outputs=energy, inputs=positions) in PyTorch [Paszke et al., 2019]. The dataset split follows Geom3D [Liu et al., 2024], i.e., using 110K for training, 10K for validation, and 11K for testing in QM9 and 950 for training, 50 for validation, and 1000 for testing in rMD17. We use mean absolute error (MAE) as evaluation metric.

Baselines. To verify the effectiveness of EDG, we select many geometry-based models with different architectures, such as SchNet [Sch¨utt et al., 2017], EGNN [Satorras et al., 2021], Equiformer [Liao and Smidt, 2023], SphereNet [Liu et al., 2022], ViSNet [Wang et al., 2024a], as geometry students to verify the generalizability of EDG. Following [Liu et al., 2022; Wang et al., 2024a; Liu et al., 2024], we ensure that each baseline is fully trained. For example, SchNet, EGNN, and SphereNet are trained for 1,000 epochs with a learning rate of 5e-4; Equiformer is trained for 300 epochs with a learning rate of 5e-4; and ViSNet is trained for 3,000 epochs with a learning rate of 0.0002. The batch size of SchNet, EGNN, SphereNet, and Equiformer is set to 128 in QM9 and 1 in rMD17, the batch size of ViSNet is set to 4 in rMD17.

Implementation details. The encoder and decoder of ImageED are built based on ViT-Base/16. In pre-training of ImageED on 2 million ED molecules, we use a learning rate of 1.5e-4, a batch size of 64, a mask ratio of 0.25, λMP and λRP of 1 for 20 epochs on 8 GeForce RTX 4090 (See Appendix G for more details). In pre-training of EDaware teacher on 2 million molecules, we divide 2% as validation set and the rest as training set. We use a learning rate of 5e-3 and a batch size of 128 to train the ED-aware teacher and ED predictor for about 280k steps (See Appendix H for more details). In distillation stage of EDG, we select hyper-parameters λ from 1e-4 and 5e-4 to 1.0 with a 10x increasing in steps. Following [Wang et al., 2024a; Liu et al., 2024], we run the experiments with exactly the same parameter settings as the baselines and report test scores corresponding to the best validation performance. The mapper and task predictor consists of a simple linear layer.

4.2 Main Results

We first evaluate the performance of EDG on the 12 quantum properties from QM9 with 4 baselines (SchNet, EGNN, Equiformer, SphereNet) and Table 1 shows the main results. We find baselines equipped with EDG achieve the best performance. We observe that regardless of the architecture, the baselines after equipping EDG achieve consistent performance improvement with a relative performance increase ranging from 2.2% to 6.4% in average MAE performance. Except for the property U in Equiformer, all other performance are enhanced.

In order to verify the effectiveness of EDG in more tasks, we further evaluate on 10 energy/force prediction tasks from rMD17 with 3 baselines (SchNet, SphereNet, ViSNet). Table 2 and Table 3 show the prediction performance on energy and force respectively. We find the same conclusion as on the QM9 benchmark, that is, EDG improves the performance of all baselines with average MAE performance improvements ranging from 8.1% to 33.7% on energy and 1.5% to 5.3% on force. It is worth noting that EDG has a larger improvement on energy prediction than on force. This is because ED can capture the global energy distribution, while force, as an energy gradient, depends on local atomic interactions, making the improvement brought by ED is not as obvious as energy. In any case, the performance improvements on energy and force prove the effectiveness of EDG. In addition, we also visualize absolute value of the difference between the predicted energy ypred and the ground-truth ytrue for all trajectories in the test set by showing the absolute value of the difference between them in Figure 3, which shows that EDG outperforms the baselines in energy prediction for almost all trajectories.

4.3 Hyperparameters Analysis

λ in Formula 13 is a parameter used to control the strength of distilling knowledge from the ED-aware teacher into the geometry students and a larger value will force the student to learn more knowledge from the teacher. Figure 4 shows the line figures of the performance of EGNN and ViSNet with different LED on QM9 and rMD17, respectively. Overall, we find that EDG can improve the performance of baselines to varying degrees with LED. For example, on the aspirin task, the performance gain of EDG fluctuates from 8.2% to 16.2% with the adjustment of LED. In addition, we also find several patterns: on the α and U tasks, as LED increases, the performance decreases overall; on the ∇E, ELUMO, aspirin, and malonaldehyde tasks, the performance curve changes with LED in a U-shaped manner. These findings suggest that by tuning the appropriate LED, EDG can better enhance the baselines.

4.4 Results of ED images on Energy-related Tasks

Here, we describe the advantages of the proposed ED image on energy-related tasks. We sample 10,000 molecules from 2

million DFT data and predict the energy of the molecular system given the ED information. We use exactly the same experimental settings and hyperparameters and randomly split the dataset into training/validation/test sets with an 8:1:1 ratio for evaluation. More settings see Appendix B. For each ED representation, we select the corresponding popular encoder to extract features, such as point cloud-based PointNet [Qi et al., 2017], voxel-based ResNet3D [Hara et al., 2018] and image-based ResNet18 [He et al., 2016]. As shown in Table 4, we find the proposed ED iamge achieves the best performance on 6 energy-related tasks with a relative performance gain ranging from 7.8% to 71.6%, which demonstrates the effectiveness of image as a representation of ED and that 2D images are easier to learn compared to 3D representations.

Table 4: RMSE (Root Mean Squared Error) performance of different ED representations on 6 energy prediction tasks. Point (point cloud), voxel, and image use PointNet, ResNet3D, and ResNet18 as encoders, respectively. E1-E6 represent DF-RKS Final Energy, Nuclear Repulsion Energy, One-Electron Energy, Two-Electron Energy, DFT Exchange-Correlation Energy, and Total Energy, respectively. ∆ represents the relative performance gain of the image compared to the best other results.

4.5 Visualization of ImageED

As shown in Figure 5, we find that ImageED can generate ED images well compared original images, which indicates that ImageED can learn ED-related knowledge well. In addition, we find that simply applying the masked prediction task (ImageED w/o LRP ) will limit the understanding of ImageED in local pixels, which shows the importance of restoration prediction task in ImageED. We show more examples in the Appendix I.

5 Conclusion

In this work, we propose a novel ED-enhanced molecular Geometry representation learning framework (called EDG), which is the first attempt to exploit ED images to improve the performance of geometry-based methods. We propose an efficient ED representation learning, called ImageED, to extract ED knowledge from images and further transfer the knowledge in ImageED to an ED-aware teacher to save the cost of DFT. By exploiting ED-aware teacher, EDG can significantly improve the performance of geometry-based methods without any architectural modifications on a large number of quantum chemical benchmarks. In addition, we experimentally show that using ED images can more accurately predict energy-related prediction tasks while saving memory and computational costs, enabling direct use of ED images in broader tasks like drug discovery and materials science.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (grant nos. U22A2037, 62425204, 62122025, 62450002, 62432011), Grants of Ningbo 2023CX050011.