Equivariant Masked Position Prediction for Efficient Molecular Representation
Authors: Junyi An
Chao Qu
Yun-Fei Shi
XinHao Liu
Qianwei Tang
Fenglei Cao
Yuan Qi
Venue: ICLR 2025 Poster (2025)
Keywords: Self-supervised Learning; Graph Neural Network; Molecular Property Prediction;
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Graph neural networks (GNNs) have shown considerable promise in computational chemistry. However, the limited availability of molecular data raises concerns regarding GNNs' ability to effectively capture the fundamental principles of physics and chemistry, which constrains their generalization capabilities. To address this challenge, we introduce a novel self-supervised approach termed Equivariant Masked Position Prediction (EMPP), grounded in intramolecular potential and force theory. Unlike conventional attribute masking techniques, EMPP formulates a nuanced position prediction task that is more well-defined and enhances the learning of quantum mechanical features. EMPP also bypasses the approximation of the Gaussian mixture distribution commonly used in denoising methods, allowing for more accurate acquisition of physical properties. Experimental results indicate that EMPP significantly enhances performance of advanced molecular architectures, surpassing state-of-the-art self-supervised approaches. Our code is released in https://github.com/ajy112/EMPP.

Published as a conference paper at ICLR 2025
EQUIVARIANT MASKED POSITION PREDICTION FOR EF- FICIENT MOLECULAR REPRESENTATION
Junyi An!; Chao Qu² Yunfei Shi', Xinhao Liu', Qianwei Tang', Fenglei Cao Yuan Qi Shanghai Academy of Artificial Intelligence for Science
PINFYY TECH (Shanghai) Co., Ltd.
Sccooo of Computer Science, Fudan University
*State Key Laboratory for Novel Software Technology, Nanjing University
5 Artificial Intelligence Innovation and Incubation (AI") Institute, Fudan University
ABSTRACT
Graph neural networks (GNNs) have shown considerable promise in computational chemistry. However, the limited availability of molecular data raises concerns regarding GNNs' ability to effectively capture the fundamental principles of physics and chemistry, which constrains their generalization capabilities. To address this challenge, we introduce novel self-supervised approach termed Equivariant Masked Position Prediction (EMPP), grounded in intramolecular potential and force theory. Unlike conventional attribute masking techniques, EMPP formulates a nuanced position prediction task that more well-defined and enhances the learning of quantum mechanical features. EMPP also bypasses the approximation of the Gaussian mixture distribution commonly used in denoising methods, allowing for more accurate acquisition of physical properties. Experimental results indicate that EMPP significantly enhances performance of advanced molecular architectures, surpassing state-of-the-ar self-supervised approaches.
INTRODUCTION
Graph neural networks (GNNs) have found widespread application in computational chemistry. However, unlike other fields such as natural language processing (NLP), the limited availability of molecular data hampers the development of GNNs in this domain. For example, one of the largest molecular dataset, OC20 (Chanussot et al.) 2021), contains only 1.38 million samples, and collecting more molecular data with ab initio calculations is both challenging and expensive. To address this limitation, molecular self-supervised learning has gained increasing attention. This approach enables molecular GNNs to learn more general physical and chemical knowledge, enhancing performance in various computational chemistry tasks, such drug discovery (Hasselgren & Oprea 2024) and catalyst design (Chanussot et al. 2021).
Current self-supervised methods for molecular learning contain two mainstream categories: masking and denoising. Masking methods (Hu et al. 2020; Hou et al.) 2022; Inae et al.) 2023) adapt the concept of masked token prediction from natural language processing (NLP) to graph learning, where graph information, such as node attribute, is masked instead of token. However, there are two major limitations: underdetermined reconstruction and lack of deep quantum mechanical (QM) insight. As illustrated in Figure a (i) reconstructing attribute of the masked iodine atom can yield multiple possible solutions, and as the number of masked atoms increases, the number of solutions will increase rapidly, making it difficult for training data to cover all possibilities; (ii) the attribute of the masked carbon atom can be inferred from the 2D geometric principles of the benzene ring, causing the model to overlook essential atomic interactions needed for learning QM properties (Messiah) 2014). In contrast, denoising methods (Zaidi et al. 2023 Feng et al. 2023) are physics-informed and facilitate self-supervised learning of QM information equilibrium structures. Their core idea involves adding noise to atomic positions and predicting them (see Figure [b)). In this process,
" 'Equal contribution: Junyi An and Chao Qu. Correspondence Junyi An: anjunyi sais com. com. "Corresponding authors: Fenglei Cao and Yuan Qi 'Our code is released in https//gituhb.com/ay wttp://///imy111//111/11 12/FMPP

---

Published as a conference paper at ICLR 2025
Masking
Denoising
(Mus)
Gaussian Mixture Approximation (Denoising)
[d)
on eloa aadill
Gravie
farce
Noise
forcx
Pooliio
Amriass
seguin force If arggan (IGhN())) GII
Ms-CI,
Figure 1: (a, b, c) Comparison of three molecular seff-supervised methods using real halobenzenes (Ph-X) as example: (a) Masking atomic attributes, such as atomic number, and reconstructing them; (b) Adding noise to atomic positions and predicting the noise; (c) Completely masking positions and inferring them based neighboring structures. (d) Principle of the denoising methods: they utilize Gaussian mixture distributions to approximate the local minima of the PES, allowing the noise terms to estimate the forces, i.e. derivatives of the PES. The two local minima in the figure correspond to two equilibrium atoms, each requiring a different standard deviation a for approximation. As the shape of the PES is unknown, determining a requires empirical approach.
the local minima of the potential energy surface (PES) are approximated using Gaussian mixture distributions (Ziddi al. 2023), making noise prediction equivalent to learning the forces, i.e., the derivatives of the PES. However, the actual PES presents diverse and unknown local minima shapes (see Figure d)), making it difficult accurately parameterize the Gaussian mixture distribution. Our ablation studies in Section 5.4 further illustrate that model performance is sensitive to the standard deviation a of the Gaussian mixture. We propose addressing these limitations through a position prediction process, which allows the model to effectively learn critical QM features.
In this paper, we introduce Equivariant Masked Position Prediction (EMPP), a novel training method designed to enhance molecular representations in GNNs. In EMPP, randomly mask an atom's 3D position while retaining its other attributes, such as atomic number, ensuring that position prediction remains well-posed problem compared to attribute masking methods. Furthermore, the masked atom's position determined by quantum mechanics within the neighboring structure, which our method designed learn. It is important denoising methods. As illustrated in Figure [c), EMPP completely removes the node of the masked atom and the embeddings of unmasked atoms predict its position. Consequently, EMPP can bypass the approximation of Gaussian mixture distributions, resulting in a more deterministic position prediction process, as detailed in Section 3.2 Additionally, since EMPP significantly alters the original molecular graph by removing nodes, it can generate a vast number of diverse data during training. Given dataset of M molecules with an average of N atoms each, EMPP can produce O(MN) well-posed examples.
note that EMPP fundamentally different from
EMPP is a versatile training paradigm. enables self-supervised training of models on unlabeled equilibrium data to capture general knowledge, while also serving as auxiliary task to connect known quantum properties with atomic positions. Experiments demonstrate that EMPP enhances the generalization of advanced equivariant GNNs (Thökee & Fabritis 2022 Liao & Smidt 2023) across various molecular tasks, regardless of whether extra data is used for pre-training. Additionally, EMPP achieves state-of of-thh--rtt performance compared to previous masking and denoising methods. Our key contributions can be summarized as follows: we propose a molecular learning paradigm that utilizes position prediction address the challenges posed by previous masking and denoising approaches, paving a novel path for molecular learning.
PRELIMINARIES
In this section, we briefly review the mathematical background, which includes molecular property prediction, equivariance, spherical harmonics, and More detailed introductions deferred to Appendix A We list the notations frequently used in the following. We denote the unit sphere as S? where the spherical coordinates (0, are the polar angle and the azimuth angle, respectively. The
2

---

Published as a conference paper at ICLR 2025
symbol IR represents the set of real numbers, while R. represents the rotation matrix for 3D vectors. We use (0(3) to denote the special orthogonal group, i.e. the 3D rotation group.
2.1 MOLECULAR PROPERTY PREDICTION
Molecular property prediction aims to construct projection from the molecular 3D structure to the molecular properties. In the following, we use the {1,..., 118} denote the atomic number, or z to denote richer atomic attributes, including atomic numbers, chemical environments, and other features. We use the term p to denote atomic 3I) position. In N-nodes molecule, properties can be divided into global properties y E R and node-wise properties y E TRA. In detail, given a 3D) molecular S {(z,, p;)li E {1, {l...., ]) the properties can be predicted by a GNN:
y y= PR and f GNN(S),
(I)
where f represents the node embeddings in the final GNN layer, and (RED(.) is the prediction head.
2.2 EQUIVARIANCE
Given any transformation parameter g E G, a function yo p: :: -> y is called equivariant to g if satisfies:
(2) where (g) y - and T()): : X X' denote the corresponding transformations over y and X, respectively. Invariance is special case of equivariance where (g) is identity transformation. In this paper, we mainly focus on the SO(3) eguivariance and invariance, since it is closely related to the interactions between atoms in molecule In other words. The backbone and prediction head should adhere to equation 2
[=]] ((T()]]=)
2.3 SPHERICAL. HARMONICS AND STEERABLE VECTOR
Spherical harmonics, class of functions defined over the sphere S², form an orthonormal basis and have some special algebraic properties widely used in equivariant models (Kondor et al. 2018 Cohen et al. 2018). In this paper, we use the real-valued spherical harmonics denoted as {Y S + RJ. where and m denote degree and order, respectively. It known that any square-integrable function defined over 5 can be expressed in spherical harmonic basis via
1(0 Σ
Σ Y..
where 11th is the Fourier coefficient. For any vector F with orientation (0, Φ), we define (F/IIFF)) = (O, v) (0,, v); 1=141 (0, v)....: Y/ v/) vector with 27 + elements. Furthermore, we define the spherical harmonics representation for any direction:
F
sh'
forms (1 + 1)² vector. Equivariant models (Zitnick et al. 2022 Liao & Smidt 2023 An et al. 2024) typically set maximum degree Lmax and construct node embeddings with C channels, resulting in an embedding size of (Imaa + 1)2 x C. The full mathematical form of spherical harmonics can be found in Appendix A.
commonly used property of the spherical harmonics is that for any R E SO(3). we have Y' (RF)- (F) where D' (R) is a (21 1) 1)) (21 + 1) matrix known as a Wigner-D matrix with degree 1. Therefore, R and D'(R) corresponds to T(2) and (4),,, respectively in equation Following the convention in (Chami et al. 2019 Brandstetter al., 2022), we say Y (FF) is steerable by Wigner-D matrix of the same degree I. The 21+ 1-dmmnnnno vector space on which a Wigner-D matrix of degree act is termed type-l steerable vector space, denoted by the superscript (l) this paper.
Equivariant operations. To ensure the equivariance of the entire model, each operation must maintain equivariance. The e3nn library (Geiger al. (022) offers common equivariant operations, including SO(3) linear transformations, SO(3) normalizations, gate activations, and (CG) tensor products. Passaro & Zitnick (2023) further extends certain nonlinear equivariant opera- tions to the sphere. In our work, we leverage existing equivariant models as backbones and employ aforementioned equivariant operations to construct our position prediction block.
Invariance of translation trivially satisfied by taking the relative positions inputs.
3

---

Published as a conference paper at ICLR 2025
Properties:
Self-Supervised Label: Fu Pr Pk
Fu
Pk
Position:
KI--Div
Attributes:
Encoding
Embedding
Masking Position
Backbone GNN
Prediction
Figure 2: The overall framework of EMPP. The masked position can be recounstructed by the GNNs output features of the neighboring nodes, with the position determined by the predicted directions and radius from those nodes.
METHODOLOGY
In this section, we first revisit the existing molecular self-supervised methods and outline their primary limitations. Subsequently, we present approach, Equivariant Masked Position Prediction (EMPP), and claborate on its implementation details.
3.1 REVISITING THE VANILLA MASK METHOD MOLECULE LEARNING
Similar to NLP, molecular self-supervised methods aim to learn the underlying chemical and physical mechanisms molecular systems, such as valence bond theory (Shaik & Hiberty 2007) and force fields (Ponder& & Case 2003). To this end, AttrMask (Hu et al. (220)) pioneers method that randomly masks atoms and predicts their attributes. More formally, we assume the i, atoms are masked and the modified molecule denoted as S = {(Z1, P1), .... (M. p.), .... (M. p,), (ZN, C.N.N) pN)) } where M denotes the mask vector (like the [MASK] token in BERT (Devlin et al. 2019)). The attributes the masked atoms are predicted by a GNN as follows:
Ji..... = PRED(..,...). fi....,fy GNN(S),
(5)
where fi.... represents the GNN output features of the masked atoms, and PRED(.) typically refers neural network. The objective is minimize the discrepancy between the predicted attributes J...... and the actual attributes z..j.... forming a self-supe learning. Two key limitations emerge: the iosedesesssss of attribute prediction and the inability capture deep quantum mechanical features. These limitations, mentioned in the introduction, can be observed in Figure [a). Additionally, denoising methods (Zaidi et al.] 2023 Feng et al., 2023) noise the atomic positions, where the modified molecule is denoted as = (Z1,
(Z,, P. + EE),,... (Z,, Py + E2),.... ..., (Z,, PN))- The GNNs required to produce equivariant features and the PRED(.) is used to predict noises C1 c..... These methods assume that Boltzmann distribution (i.e. exponent of PES) around equilibrium positions can be approximated by Gaussian mixture distribution (Zaidi al. 2023). However, the assumed distribution can not always approximate the true distribution. In fact, the shapes of PES in local minima are diverse and unknown in advance (Messiah, 2014). making challenging to define the parameters of the Gaussian mixture distribution.
3.2 EQUIVARIANT MASKED POSITION PREDICTION (EMPP)
The overall framework of EMPP is depicted in Figure We begin the process by masking the position of an atom, causing the corresponding node in the graph to vanish. Next, we utilize equivariant backbones to generate equivariant node embeddings for the masked molecule. These embeddings are then input into a position prediction module, which outputs distributions for directions and radius. Both of them determine the predicted position. By aligning the true and predicted positions, EMPP enables the GNNs to capture atomic interactions within 3D molecular structure.
3.2.
MASK POSITION AND NEIGHBOUR ENCODING
To address the limitations of previous works, EMPP only masks the atomic position and predicts using its atomic attributes and the representations of unmasked atoms. Assuming the i-th atom is
4

---

Published as a conference paper at ICLR 2025
masked, the modified molecule is denoted as $ ((11.1)).... (z.).... (ZN.PN) According to physics, the force at equilibrium is zero, with atomic force primarily governed by atomic interactions, such as Coulombic forces (Messiah 2014). Thus, EMPP aims to find a position in equilibrium structures that satisfies the condition Žj=1zæ forcej; = O, where forcey, represents the inter- atomic force and is function of the unknown masked position p: = (ar, y. 2), given that the atomic attributes are fixed. Intuitively, there exists a unique optimal position under the force condition, since the number of unknown variables matches the number of equations. Organic chemistry further supports that the position of atoms is uniquely determined in most cases. For instance, in Figure I when any carbon atom is masked, its position can still be uniquely determined. Similarly, for the iodine atom, its position can be uniquely identified once the structure of its neighboring atoms is known. Further details and examples can be found in Appendix [.I Additionally, EMPP employs equivariant representation capable of storing vector features, enabling the accurate description of the true forcey. Therefore, EMPP can be regarded as a nearly well-posed method, effectively avoiding the approximation of Gaussian mixture distributions to learn forces or other quantum features.
P1),
In the feedforward process of backbone GNN, the node of masked atom is completely removed. We encode the attributes z; of the masked atom into the embedding of unmasked atoms. Note that traditional backbone GNNs use an embedding layer to project atomic attributes into the embedding space: we employ the same module augmented with a two-layer multi-layer pereeptron (MLP) to embedding the attributes z;:
e, -MLP (LLBBEB))(2.) Next, we aggregate e; with the embeddings of the unmasked nodes:
(6)
V j (1, - 1,i + I,..,,N N),
(7)
where v denotes the node embedding. We apply equation 7 at each layer of the backbone GNN. Notably, e; is invariant under SO(3) transformation. If the node embedding v, is a spherical harmonic representation, we aggregate e; with the its type-0 vector (0))
Masking multiple atoms in EMPP increases the complexity of solution space, may make the position prediction ill-posed To mitigate this issue, we mask different atoms sequentially within a single molecule. First, we randomly generate the indices of the masked atoms as Mask = { {Maski, . Masky,.. Maskn. We then create multiple masked molecules, each con- Each masked molecule independently
taining only one masked atom, denoted as S2, predicts the position of its masked atom, and we average the resulting loss values. The objective function for training multiple masked molecules is given by:
Lmalti = n Σ singte (PMasks Pred (GN)()
where Pred(-) refers to the position prediction block, and Ceingle (-) represents the loss function for each individual masked atom. Details of them will be provided in the following sections. Since the nodes of masked atoms in EMPP are completely deleted, there is a distinct difference between (Š1, $2,...,.)). In contrast, previous methods do not delete nodes but perturb the features of the nodes. Intuitively, EMPP has stronger capability to generate vast array of diverse data.
3.2.2 POSITION PREDICTION
Equivarance for predicted position. After masking i-th atom and encoding its attributes into the backbone GNN, we obtain the node representations of the masked molecule. Notably, we retain the neighbor set of the masked atom, denoted as N₁. and use the embeddings of the neighboring nodes to predict the position. For an unmasked node with embedding fę (where E AN)) and its position Pæ: the probability distribution of the masked atom's 's position is given by pk- (F(ff)) with the predicted position denoted as p; = P& + F. This distribution Pk must satisfy both normalization and non-negativity constraints: Jt. P& (F)dV = 1 and pk (F) > 0 where dV = dædydz is the volume element and Days represents all space in Cartesian coordinates. Another challenge ensuring that the distribution remains SO(3) caanssfrmationcqquivveeo
adhering to symmetry principles:
Pk (Fl = Pı (RF Df.)) = pk (RF (NN(
In organic chemistry, local system composed of two more atoms often exhibits symmetry, allowing for multiple possible positions (See examples in Appendix B.
5

---

Published as a conference paper at ICLR 2025
where R epresents any rotation matrix, and D is the Wigner-D D matrix corresponding to the same transformation parameters. The term R. {Z, Rp:)li E {1,... 1, i 1, .....)) denotes the rotated molecular structure. Note that equation disregards translation, as its effect neutralized by
the relative position prediction Pmosk P
Distribution prediction. Recent molecular models have demonstrated that high-degred equivariant features can effectively describe the PES and forces field (Passaro & Zitnick 2023 Liao al.) 2024). Therefore, we can directly predict the distribution of masked positions using equivariant node features. In our prediction module, we use each neighboring atom independently predict a positional distribution, which is then aggregated. In practice, we first pass the node embeddings through an equivariant two-layer MLP. The mathematical process is expressed as:
fg -AAFAAARRIIIAAAAA
(10)
f.- -ELINEAR (f.c @ EMBED(Z))).
(11)
where ELINEAR(.) denotes the SO(3) linear layer, Gate(.) denotes gate activation using SiLU and & denotes the CG tensor product. In the second layer of the MLP, we integrate the features with the intrinsic attributes of the masked atom. We transform position prediction to Spherical coordinates and decompose it into two parts: the radius distribution, p prodius(r) (r), and the directional distribution, p ddrrcction (0, 6). For the radius distribution, we uniformly partition the space into 128 intervals between 0.9 Aand 5 A, covering the key organic molecular interaction distances. As distance is a
rotation-invariant feature, we predict the radius distribution directly using the type-0 vector:
(12) where SoftMax(() is used to normalize the distribution and is the temperature coefficient. For the directional distribution, we apply a Fourier transform to project the node representation onto the spherical surface, allowing us to capture features from all directions:
pk
radius SoftMax (LINEAR( (() )/T),
Lmna Σ
F(0,
"y mm (0,
(13)
The Fourier transform spherical harmonics is described in detail in Appendix A.1.2 Afterward, we grid the spherical surface to obtain a finite feature matrix denoted as Fgrid. Specifically, fz E 1) xCC contains C' channels. After gridding, Ferid RSC retains information from all channels, where S represents the grid sampling rate] We then apply shared MLP to project channel dimension into a 1D space, followed by the Softmax function to produce a normalized directional distribution.
(14) It is important to note that operations within the channel dimension preserve overall equivariance. As a result, equation satisfies the condition outlined in equation | Multi-channel features are utilized to retain high-frequency components. While EMPP's method of masking atoms and predicting coordinates using spherical harmonic projection shares similarities with Symphony (Daigavane et al.. 2024), their primary focuses diverge significantly: Symphony tackles molecular generation, whereas EMPP addresses a self-supervised task. This difference in objective leads to substantial variations in implementation, such masking strategies and the design of the multi-channel spherical harmonic projection module. Detailed comparisons are provided in Appendix B.5
Pk
direction SoftMax (MP((FFr 1/r),
3.2.3 LOSS FUNCTION
Radius loss. During training. the model is optimized by the distribution of true radius ||Fa||| and direction Fu/ /lFF ||, where F, p. P Intuitively, the deterministic vector F& should correspond to a Dirac delta distribution denoted as 6 - Fix). )- In practice, however, the model can learn the accurate F. long as the defined distribution can uniquely represent F.g. The Dirac delta distribution can reduce training stability and lacks the ability to transfer to other conformations. To address this, we use a Gaussian distribution as a surrogate:
dk radius (r) ~ ~
(15)
*The grid sampling rate must follow the Nyquist rate criterion: S (2Tman 1
"A common method embedding distance (Thomas al. 2018 Zitnick et al. 2022

---

Published as a conference paper at ICLR 2025
where a is typically set to 0.5 A. The Bullback-Leible divergence (KL-div) is utilized to compute the loss:
radius
Cradius
Σ
qk
llg Pk
(16)
N
pradius
KENN
Direction loss. Similarly, we choose soft directional distribution instead of a Dirac delta function to represent the ground truth for direction:
W. (0, = exp( Σ sh
fm
Fish
(0, Ø)),
1 direction (0, Φ) ~ W (0, 9k
(17)
FsR
where sh(.)) denotes the spherical harmonics representation the direction, and W = Jnow we (0, d))0dá serves to normalize the distribution. The true distribution is then gridded, and the KI. divergence between Pkc
direction and the true distribution is computed:
lk direction
ddrrection
[ REN,
Σ
4k
Airection log direction
(18)
The total loss for single masked atom, as defined in equation Edirection
g is given by Loingle fradiue
3.3 APPLICATION OF EMPP
Pre-training without annotation. The previously described EMPP can learn quantum knowledge from equilibrium molecules without quantum property labels. Therefore, EMPP can be used to pre-train GNNs to learn transferable knowledge. For example, the PCQM4Mv2 dataset (Nakata & Shimazaki 2017) contains vast collection of equilibrium molecules but provides labels for only one chemical property, the gap between HOMO and LUMO. We can use EMPP to pre-train GNNs on PCQM4Mv2 and use the pre-train models to predict additional chemical properties in other datasets. Auxiliary task for property prediction. When training a model for a specific property prediction, calculate two losses: the prediction loss for the targeted property and the EMPP loss defined in equation 8 These losses are then combined for gradient descent. In this case, the goal EMPP is not to learn the equilibrium position of the masked atom, but rather learn the position corresponding the known property value, which may be in a non-equilibrium state. For example, when the force on the masked atom is non-zero (indicating non-equilibrium) EMPP can capture the relationship between the true position and this force. Similarly, for other quantum labels, implicit relationships exist with the masked atom's 's position. In practice, we encode the label into the GNN. If the label is global invariant property like energy, the encoding can be written as:
(0) + LINEAR
(ynervvu)))))))
(19)
where GAUSS(.) denotes a Gaussian block. Specifically, for node-wise equivariant labels such as force y forces* we map it into the spherical harmonic representations and add it to embeddings:
V& + Vk + ELINEAR (sh(y) forces)))
(20)
Then, we use the GNN output embeddings to predict position. In this scenario, we treat EMPP as an auxiliary task for property prediction. By modeling the relationship between the target property and position, EMPP further enhances the generalization ability of the property prediction model.
4 RELATED WORK
Graph self-supervised methods. Leveraging the inherent graph structure of molecules, many graph- based self-supervised methods have the potential to train molecular models that capture transferable knowledge. For example, Hu et al. (2020) proposed graph context prediction and attribute masking methods to enhance molecular property prediction. GraphMAE (Hou et al., (022) pre-trained molecular models using a generative decoder to reconstruct atomic and bond attributes. D-SLA (Kim et al. 2022) applied contrastive learning based on graph edit distance, improving predictions of molecular biochemical activities. Additionally, graph motifs -ootifs-induced subgraphs that describe recurrence and significance-Bae have increasingly been utilized to construct self self-super learning frameworks for molecules (Rong et al.] 22000] [hhang et al.) [021] [nae et al.) 2023). facilitating the learning of multi-scale molecular information. These pre-training methods primarily focus graph
7

---

Published as a conference paper at ICLR 2025
characteristics, while neglecting the intrinsic quantum mechanisms within molecules. As result, they are limited in their ability to predict molecular quantum mechanical properties.
3D molecular representation. Given the strong correlation between the quantum characteristics of molecules and their 3D structures, recent molecular models have increasingly focused on 3D representations (Liao & Smidt, 2023 Passaro & Zitnick 2023 Liao al. 2024). As result, self- supervised techniques have also evolved to operate in 3D space. For instance, Unimols (Zhou et al.] 2023 Lu et al., (2023) masked atomic properties and restored them using 3D molecular models, while denoising methods (Zaidi et al. 2023 Feng et al. 2023) introduced a series of physics-informed pre- training paradigms. We provide a detailed discussion of these methods in Section 3.1 and highlight their main limitation: accurately defining the parameters of Gaussian mixture distributions can be challenging. In contrast, EMPP learns quantum mechanical features through a position prediction process, effectively bypassing the difficulties associated with denoising methods.
Spherical harmonics projection. High-degree spherical harmonic representations with grid projec- tions (Liao & Smidt 2023; Passaro & Zitnick (023) have demonstrated significant capabilities in spatial description. Building on this, Symphony (Daigavane et al. 2024) introduced a neighbor- based spatial position prediction method, developing framework for molecular generation. While EMPP shares some technical similarities with Symphony, their distinct tasks (molecular generation versus self- superviee learning) lead to differing implementation priorities. Symphony prioritizes flexibility in position prediction to ensure the sampling of diverse molecules, whereas EMPP emphasizes accuracy and well-posed position prediction facilitate learning of genuine physical interactions.
5 EXPERIMENTS
In this section, we present experiments to evaluate the effectiveness of EMPP across several 3D molecular benchmarks. Since EMPP can be applied in both unlabeled and labeled scenarios, we evaluate it in two settings: (i) self-supervised tasks for learning transferable molecular knowledge, and (ii) auxiliary tasks for enhancing the prediction of supervised molecular properties.
5.1 DATASETS AND CONFIGURATIONS
Datasets. We evaluate quantum property prediction using the QM9 (Ramakrishnan et al. 2014) and MD17 (Chmiela et al. 2017) datasets. QM9 comprises 134,000 stable small organic molecules made up of C, H. O, N. and atoms, with one conformation per molecule, and includes labels for 12 quantum properties. MD17 contains molecular dynamics trajectories for 8 small organic molecules, providing between 150,000 and nearly million conformations per molecule, along with corresponding total energy and force labels. Notably, MD17 features significant number of non cquiiiiiiii molccules. Additionally, we utilize the PCQM4Mv2 (Nakata & Shimazaki] (017)) dataset pre-train GNN backbones, which consists of 3.4 million organic molecules, each with one equilibrium conformation, and is widely used for pre-training.
Baselines. Our baselines include state-of-the-ee self- supervised methods for 3D molecular structures, such as AttrMask (Hu et al.) 2020), OP-TorchMD--NE (Zaidi et al., 2023). 3D- -MMG (Jiao et al. 2023), SE(3)-DDM (Liu al. 2022 Transformer-M (Luo al., 2022). ) and Frad (Feng et al. 2023). These methods pre-train GNNs on the PCQM4My2 dataset hefore predicting molecular properties in QM9 and MD17, leveraging additional molecular data. In contrast, EMPP can operate without extra data, so we also compare it to molecular models trained solely on QM9 or MD17, including SchNet (Schütt al. 2018). PaiNN (Schüttt et al., 2021 DimeNet++ (Gasteiger al. 2020). TorchMD-NBi
(Doerr et al. 2021), SEGNN (Brandstetter et al.) (022). and Equiformer (Liao & Smidi 2023).
Detailed configurations for EMPP and the aforementioned baselines can be found in Appendix C.I
5.2 RESULTS WITHOUT PRE-TRAIE TRAININGG
QM9. We first evaluate EMPP without incorporating additional data, treating it as an auxiliary task for QM9 and MD17. The total loss comprises equation & combined with the MAE loss from the original property prediction. In these experiments, EMPP is implemented using the Equiformer backbone (Liao & Smidi (023) due to its high-degree equivariant representation, which effectively captures interatomic features. As shown in Table I two key observations emerge. First, EMPP enhances prediction accuracy across all QM9 tasks, achieving the best results in 11 tasks without additional data. Moreover, multi-masking show better performance, demonstrating that the data generated by EMPP is diverse and reliable. Second, while EMPP is designed to learn interaction
8

---

Published as a conference paper at ICLR 2025
Table Results on QM9 property prediction without pre-trainging on extra molecular data. f denotes
using different data partitions. In this experiment, EMPP uses the Equiformer backbone. n-Mask denotes masking atoms in each molecule during training. The masking strategy follow equation X
Task
Ac meV
ЕhолO meV
ELUMO meV
Cs
G
H
R
Un
ZPVE
Units
bohu
califmool K)
meV
meV
bohr²
meV
meV
meV
SchNet
.235
63
.033
.033
14
.073
14
1.70
DineNet++
.044
33
25
20
.030
.023
8
7
.331
6
1.21
PaiNN
.045
46
28
20
012
024
7.35
5.98
.066
5.83
5,85
1.28
TorchMD-NET
.059
36
026
7.62
6.15
SEGNN"
.660
42
21
.033
.031
.660
15
Equiformer
.046
30
15
14
.011
.023
7.63
6.63
.251
6.74
6.59
1.26
EMPP (I-Mask)
.041
27
14
13
.0108
.021
6.89
5.38
.189
6.05
5.88
1.20
EMPP (3-Mask)
.039
26
12
.0996
.019
6.32
5.02
.154
5.72
5.25
1.18
Table 2: Results on MD17 testing without pre-trainging on molecular data. Energy and force are in units of meV and m/ A. The "energy only" and "force only" are based on "1 Mask strategy.
Aspiria
Benzene
Fihanol
Malotaldeiiyde
Naphthalene
Salieclic noid
Tolucne
Umaeil
Methods
energy
Eurves
energy
forces
energy
forces
energy
forves
energy
funces
energy
forces
energy
forves
energy
Eorces
SchNet
13.4
25.6
6.1
DimeNer
8.8
2.8
16.6
5.8
5.0
PoiNN
6.9
14.7
2.7
13.
4.9
8.5
4.5
1orchMD-NNET
5.3
11.0
2,5
8.5
2.3
4.7
3.3
7.3
3,7
2.6
4.0
5.6
3.2
2.9
4.1
4.1
Equaformer
5.3
6.6
2,5
8.1
2.2
2.9
3.2
5.4
4.4
2,0
4,3
3.9
3.7
2.1
4.3
3.4
EMPP Maskk)
5.1
6.4
2,3
1.5
2.2
2.7
3.0
5.2
4.1
1.9
4,3
3.7
3.5
2.1
4.1
3.4
EMPP Maskk)
5.0
6.2
7.3
2.0
3.0
5.0
3.7
4,1
36
3,3
2.0
4,0
3,0
EMPP (energy cally)
4.8
4.0
(fence only)
4,3
3.7
4.7
between adjacent atoms, also leads to improvements in properties not directly related to atomic interaction, such as нOММ and ELUMO (Pope 1999), further confirming its effectiveness.
MD17. The MD17 dataset provides both global labels (energy) and node-wise labels (forces), along with numerous non-equilibrimm conformations that present new challenges for models and self-supervised methods. To address this, we encode both the global energy and the forces of the masked atom into the backbone GNNs. EMPP predicts atom positions where the predicted forces align with the ground truth. The results, presented in Table show that EMPP achieves the best performance on most tasks. Additionally, we conducted two experiments where either energy or force alone is encoded into the backbone, referred to as "energy only" and "force only". Notably, when only energy is encoded, Table iablleiiniicats indicates that energy prediction can be further improved. the cases of Aspirin, Naphthalene, and Toluene, the "energy only" approach with single-atom masking even surpasses the results obtained by three-atoms masking while encoding both energy and forces. These findings suggest that encoding single label simplify EMPP's learning process, thereby enhancing its effectiveness for predicting specific properties. We further investigate the performance of EMPP without pre-training on the GEOM-Drug dataset (Axelrod & Komez-Bombarell 2022). which contains abundant non-equilibrium data, Results can be found in the Appendix C.2.
5.3 RESULTS WITH PRE-TNAININ ON PCQM4MV2
In this section, we assess the effectiveness of EMPP as a pre-training task. To maintain consistency with state-of-the-as self-supervised methods (aidi al. 2023 Feng al. 2023), we TorchMD- Net (ET) (Thölke & Fabritiis| 2022) as the backbone model. To enhance the fine granularity of the spherical representation after Fourier transformation, we extend the node representation of TorchMD- Net to Lmaz = 3 while retaining all other core operations (see details in Appendix B.3). During the pre-training phase, EMPP encodes only the atomic numbers of the masked atoms into the backbone. In the fine-tuning stage, EMPP encodes properties following the same approach outlined in Section 5.2
As shown in Table 3 the attribute masking method produces poorer results for quantum property prediction, consistent with our earlier observations: recovering attributes based on simple features limits the model's ability to capture deep quantum characteristics. In contrast, denoising methods such as TorchorhMD- and Frad deliver competitive performance due their physically interpretable paradigms. Notably, our EMPP surpasses denoising methods in nine tasks, achieving the best results in seven tasks. This success attributed to EMPP's more precise paradigm, which employs a precise paradigm to learn interactions instead of approximation of mixture distribution.

---

Published as a conference paper at ICLR 2025
Table Results on QM9 property prediction with pre-trainging on extra molecular data. EMPP use the TorchMD-Net backbone. Bold and underline indicate the best result, and the second best result.
Task
A
ENOMO meV
C2
G
R²
Un meV
ZPVE
Units
bohr
meV
meV
cul/(mol K)
meV
meV
bohr
meV
meV
AttrMask
072
50.0
31.3
37.8
020
0.062
11.2
11.4
.423
10.8
10.7
1.90
Transformer-M
041
27.4
17.5
16.2
.037
0.022
9.63
.075
9.41
9.37
1.18
SE(3)-DM
046
40.2
23.5
19,5
.015
.02
7.65
7.09
.122
6.99
6.92
1.31
3D-EMGP
.057
37.1
21.3
18.2
.020
.06
9.30
8.70
092
8.60
1.38
0517
31.8
17.7
14.3
.200
6.45
.449
6.57
1.71
Frad
0.037
27.8
15.3
13.7
.010
.200
6.19
5.55
.342
5.62
5.33
1,42
EMPP (3-Mask)
.035
25,8
13.7
13.4
.014
.019
6.45
5.73
.241
5.34
5.08
1.27
Table 4: Ablation study on key hyperparameters. Bold indicates the default configuration.
Sampling Rate
еномо
Loss Weight
еномо
20
.050
20.4
0.1
.046
15.4
502
.043
18.5
0.5
.045
15.2
100²
.041
14.2
.041
14.2
150a
.041
14.7
.044
15.8
5.4 ABLATION STUDY
In this section, we address several key questions: (i) the relationship between parameters of Gaussian mixtures and performance in denoising method. (ii) the impact of key hyper-parameters in EMPP.
Additional ablation studies are provided in Appendix C.2
To assess the impact of Gaussian mixture dis- tributions in denoising, we applied the DP- TorchMD-NNt (Zaidi et al., (023) to QM9 (HOMO, LUMO) as an auxiliary task, with the results displayed in Figure 3 By varying the standard deviation a of the Gaussian mix- ture, which controls the curvature, we observed changes in the model's generalization perfor- mance. The experiment with lent to training without denoising. From Figure 3 we derive two key insights: (i) Without utiliz- ing external data, the denoising method demon- strates limited potential to improve the model's generalization ability, yielding results that are similar to or worse than the baseline without denoising. (ii) The performance of denoising is highly sensitive to the standard deviation, with different values leading significantly diver- gent outcomes. In practical applications, because the curvature near local minima is often unknown, the optimal standard deviation typically requires empirical tuning, making challenging to account for all possible local minima.
70
HOMO
JUMO
60
MAE
is equiva-
0.001
0.005
0.01
0.02
0.05
0.5
o(A)
Figure 3: The curve performance varying with the standard deviation a.
Another ablation experiment rate and loss weight. From the Table we draw two conclusions: (i) A high sampling rate over the sphere beneficial better predicting all possible positions, thereby enhaneing generalization ability. (ii) When EMPP is used as an auxiliary task, it can be of the same importance as the original task.
verify the key hyperparameters in EMPP, including the sampling
CONCLUSION
We identified key limitations in mainstream molecular self-supervised learning: attribute masking introduces ill-posedness and fails to capture quantum features, while denoising struggles with accurately modeling Gaussian mixture in unknown distributions. To address these issues, we propose EMPP, which predicts masked atom positions through neighbor structures, bypassing Gaussian mixture approximation and turning the task into well-posed one. Our experiments show EMPP achieves state-of- the-art results in quantum property prediction. While we employ equivariant representations in EMPP extending to higher-order representation (1 (>> 3) remains open question.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We are thankful to the anonymous reviewers for their helpful comments. The computations in this research were performed using the CFFF platform of Fudan University. This research was supported by the National Natural Science Foundation of China under Grant 82394432 and 92249302. The corresponding authors are Fenglei Cao and Yuan Qi.
REFERENCES
Junyi An, Chao Qu, Zhipeng Zhou, Fenglei Cao, Xu Yinghui, Yuan Qi, and Furao Shen. Hybrid directional graph neural network for molecules. In The Twelfth International Conference Learning Representations, 2024.
Simon Axelrod and Rafael Gomez-Bombarelli Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling. Geometric and physical quantities improve (3) equivariant message passing. The Tenth International Conference on Learning Representations, 2022.
Ines Chami, Zhitao Ying, Christopher Ré, and Jure Leskovec. Hyperbolic graph convolutional neural networks. Advances neural information processing systems, 32, 2019.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras- Doraasoo Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 11(10):055-6600 6072, 2021.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 36))e1603015, 2017.
Taco S Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical cnns. arXiv preprint arXiv: 2018.
Ameya Daigavane, Song Eun Kim. Mario Geiger, and Tess Smidt. Symphony: Symmetry-equivarien point-centered spherical harmonics for 3d molecule generation. In The Twelfth International Conference on Learning Representations (ICLR), 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long and Short Papers), 2019.
Stefan Doerr, Maciej Majewski, Adrià Pére7, Andreas Kramer, Cecilia Clementi, Frank Noe, Toni Giorgino, and Gianni De Fabritis. Torchmd: A deep learning framework for molecular simulations. Journal of chemical theory and computation, 17(4):2355-254) 2021.
Shikun Feng, Yuyan Ni, Yanyan Lan, Zhi-Ming Ma, and Wei- Ma. Fractional denoising for 3d molecular pre-training. In International Conference on Machine Learning, pp. 9938 19889966,125 2023.
Johannes Gasteiger, Shankari Giri, Johannes T Margraf, and Stephan Günnnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115. 2020.
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kostiantyn Lapchevskyi, Maurice Weiler, Michal Tyszkiewicz, Simon Batzner, Dylan Madis- etti, Martin Uhrin, Jes Frellsen, Nuri Jung, Sophia Sanborn, Mingjian Wen, Josh Rackers. Marcel Rød, and Michael Bailey. Fuclidean neural networks: eănn, April 2022. URL
https::/doo www/oor................o 5281 /tw///ezooo..m 6459381
Catrin Hasselgren and Tudor Oprea. Artificial intelligence for drug discovery: Are we there yet?
Annual Review of Pharmacology and Toxicology, 64(1):527-550, 2024.
11

---

Published as a conference paper at ICLR 2025
Zhenyu Hou, Xiao Liu, Yukuo Cen, Yuxiao Dong, Hongxia Yang, Chunjie Wang, and Jie Tang. Graphmae: supervied masked graph autoencoders. In Proceedings of the 28th ACM SICKDD Conference Knowledge Discovery and Data Mining, pp. 000559440 604, 2022.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. In International Conference Learning Representations. ICLR 2020, 2020.
Eric Inae, Gang Liu, and Meng Jiang. Motif-aware attribute masking for molecular graph pre-training, arXiv preprint urXiv:2309.0444 04499, 2023.
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. Energy-motivann equivariant pre- training for 3d molecular graphs. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 37, pp.
8104, 2023.
Dongki Kim, Jinheon Back, and Sung Ju Hwang. Graph self-supervised learning with accurate discrepancy learning. Advances in Neural Information Processing Systems, 35:14085-14099 2022.
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch Clebsch-gordan nets: a fully fourier space spherical convolutional neural network. Advances in Neural Information Processing Systems, 31, 2018.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023.
Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to bigher-degree representations. In The Twelfth International Conference on Learning Representations, 2024.
Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with se (3)-invariant denoising distance matching. arXiv preprint arXiv:2206.13602. 2022.
Shuqi Lu, Zhifeng Gao, Di He, Linfeng Zhang, and Guolin Ke. Highly accurate quantum chemical property prediction with uni-mol+, arXiv preprint arXiv:2303. 16982, 2023.
Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2d & 3d molecular data. In The Eleventh International Conference on Learning Representations, 2022.
Albert Messiah. Quantum mechanics. Courier Corporation, 2014.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project: large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 57(6): 1300-1308, 2017.
Saro Passaro and Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant gnns. volume 202 of Proceedings of Machine Learning Research, pp. 27420-27438. PMLR, 2023.
Jay W Ponder and David A Case. Force fields for protein simulations. Advances in protein chemistry, 66:27-85, 2003.
M Pope.
Electronic processes in organic crystals and polymers, 1999.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. Advances in neural information processing systems, 33: 1355999-12577. 2020.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-9388, 2021.
12

---

Published as a conference paper at ICLR 2025
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
Sason Shaik and Philippe C Hiberty, A chemist's guide valence bond theory. John Wiley & Sons, 2007.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. International Conference on Learning Representations, 2022.
Nathaniel
Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translationeequivarian neural networks for 3d point clouds. Neurips 2018, 2018.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Tch. Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. In The Eleventh International Conference on Learning Representations, 2023.
7aixi Zhang, Qi Liu, Hao Wang, Chengqiang Lu, and Chee-Kong Lee. Motif-based graph self- supervised learning for molecular property prediction. Advances in Neural Information Processing Systems, 34: 148700--88882 2021.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023.
C Lawrence Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Mubammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Neurips 2022, 2022.
13

---

Published as a conference paper at ICLR 2025
APPENDIX
A The Mathematics
14
A.1
The Mathematics of Spherical Harmonics
14
A.2 Equivariant Operation
16
B
Model Details
18
B.1
The Position Distribution in Organic Chemistry
18
B.2
Details of Prediction Module
18
B.3
TorchMD-Net with Higher-order Spherical Harmonic Representation
19
B.4
Prediction with Neighbors
19
B.S
Difterences from Symphony
19
C
Details of Experiments and Supplementary Experiments
20
Implementation Details
20
Supplementary Experiments
21
A THE MATHEMATICS
A.1 THE MATHEMATICS OF SPHERICAL HARMONICS
A.1.1 THE PROPERTIES OF SPHERICAL HARMONICS
The spherical harmonics Y," (0, o) are the angular portion of the solution to Laplace's equation in spherical coordinates where azimuthal symmetry is not present. Some must be taken in identifying the notational convention being used. In this entry, is taken as the polar (colatitudinal) coordinate with 0 in [0, n]. and Φ as the azimuthal (Iongitudinal) coordinate with á in [0, 2T).
Spherical harmonics satisfy the spherical harmonic differential equation, which is given by the angular part of Laplace's equation in spherical coordinates. If we define the solution of Laplace's equation as F = 2 (()), the equation can be transformed as:
(()) d sin 0 d0 Here we omit the derivation process and just show the result. The (complex-value) spherical harmonies are defined by:
de
(
d
sin
0
+ [(I 1)0(0)
(21)
do
sin
do²
21 + (1 mm)!
(0,
=
pt (cos "m,
(22)
4m (1 + m)!
where pl. (cos 0) is an associated Legendre polynomial. Spherical harmonics are integral basis, which satisfy:
T 14 (0, & (0, o) (0, Φ)
sin Ad0do
(1 11(21 +
4
1
1/
4
1
1/
(23)
4m
m1
m2
m3
14

---

Published as a conference paper at ICLR 2025
where h
12
Is a Wigner 3j-symbol (which is related to the Clebsch-Gordan coeffi-
m.
m2
mg
cients). We list few spherical harmonics which are:
1 (0 = 2V π
1
(0,00) =
sin de -i6
2
2m
1
3
(0)
2
V
COS 0
-1
3
(0,
V
sin
Ae
2m
1
15
(0,00
4
V
sin
20
21
(24)
1
15
(0,0)
sin cos -i
2
2π
1
5
(0
(3 cos 1)
4
V
π
-1
15
(7 = 2
sin O cos 0e
V
2m
1
15
(0,
V
sin
20
2m
In this work, the real-value spherical harmonics rather than the complex-value one.
A.1.2 FOURIER TRANSFORMATION OVER $
is well known that the spherical harmonic Y",, form a complete set orthonormal functions and thus form orthonormal basis of the Hilbert space of square-integrable function. On the unit sphere 5², any square-integrable function f can thus be expanded as a linear combination of these:
0(0,))=> =
f y. (0,
(25)
I=0 m=-!
The coefficient fh, can be obtained by the Fourier transformation over S², which
f'h / F(F)
[o 2x I d0 sin ff(0, ") (0,
(26)
Usually we define a vector fl - 14,, 1. We now investigate how the fourier coefficients transforms if we rotate the input signal. More precisely, we want calculate the coefficient f/2 of the signal f(Rr), where R E SO((3) a rotation matrix.
fi] to denote the Fourier coefficients with degree
Using the fact =D')))D((R)Y(/ and equation 25
we know
S(RF) =
1 Y
(RF)
Σ
Σ
Σ
D
mome Yt me (P).
I= m==
I= m m
m72
Therefore f D'' and it is steerable.
A.1.3 THE RELATIONSHIP BETWEEN SPHERICAL HARMONICS AND WIGNER-D D MATRIX
A rotation R sending the I' to F' can be regarded linear combination of spherical harmonics that are set to the same degree. The coefficients of linear combination represent the complex conjugate of an element of the Wigner D-matrix. The rotational behavior of the spherical harmonics is perhaps
15

---

Published as a conference paper at ICLR 2025
their quintessential feature from the viewpoint of group theory. The spherical harmonics Y,,, provide a basis set functions for the irreducible representation of the group SO(3) with dimension (21 + 1).
The Wigner-D matrix can be constructed by spherical harmonics. Consider transformation Y (f') Ym (.77%, where r denote the x-orientation. a, 3, denotes the items of Euler angle. Therefore, Y, (T) is invariant with respect to rotation angle y. Based on this viewpoint, the Wigner-D matrix with shape (2l + 1) (21 1) can be defined by:
D (Ra,g.) = 21 + IY, m (T).
(27)
In this case, the orientations are encoded in spherical harmonics and their Wigner-D matrices, which are utilized in our cross module.
A.2 EQUIVARIANT OPERATION
A.2.1 EQUIVARIANCE OF CLEBSCH- GORDAN TENSOR PRODUCT
The Sebschh-Gordan Tensor Product shows a strict equivariance for different group representations, which make the mixture representations transformed equivariant based on Wigner-D matrices. We to denote the element of Wigner-D D matrix. The Clebsch-Gordan coefficient satisfies:
use mm my
Cilla.ma Emn my (fr.m (1a,m5) D'g my (9) mimma (g)
((o.mm)) (lo.m')
(28)
Em D mome (g) (2.m1))(22.m)
Therefore, the spherical harmonics can be combined equivariantly by CG Tensor Product:
la
CG Emn Dda (g) Emm D'a mamb my.
(g) my mo
mams ml) Σ mm, me (1)m1) Emg Dtt Diimms (g) mi Emg D Dimms (g)
m4
(29)
Eng Dlo
mamb (0) Σ m1 ma
C (1,m))(lmmm
(10.m6)
yh y.2. mi m/
Em, Dlo Co
(ylh, yrl,
mame
m/
equation? 29
represents a relationship between scalar. If we transform the scalar to vector or matrix, equal
equation 29
(D = @ v)'.
(30)
The tensor CG product mixes two representations to a new representation under special rule. For example, 1.ww type-0 vectors will only generate type-0 representations; 2.1ype-11 and type-l2 can generate type-l1 + la vector at most. Note that some widely-used products are related to tensor product: scalar product (1, 0, l2 1, 1), dot product (4, = 1, l2 = 1, 0) and cross product (4 = 1, l/ = 1, = 1). However, for each element with > O, there are multi mathematical operation for the connection with weights. The relation between number of operations and degree quadratic. Thus, as degree increases, the amount of computation increases significantly, making calculation of the CG tensor product slow for higher order irreps. This statement be proven by the implementation of e3nn
A.2.2 LEARNABLE PARAMETERS TENSOR PRODUCT
We utilize the e3nn library Geiger et al. 2022) to implement the corresponding tensor product. is crucial to emphasize that the formulation of CG tensor product is devoid of any learnable parameters, as CG coefficients remain constant. In the context of e3nn, learnable parameters are introduced into each path, represented as w(u't @ v²2). Importantly, these learnable parameters will not destory the equivariance of each path. However, they are limited in capturing directional information. In equivariant models, the original CG tensor product primarily captures directional information. We have previously mentioned our replacement the CG tensor product with learnable modules. is worth noting that our focus lies on the CG coefficients rather than the learnable parameters in the e3nn implementation.
A.2.3 GATE ACTIVATION AND NORMALIZATION
The gate activation and normalization used in HDGNNN are implement by eônn code framework.
16

---

Published as a conference paper at ICLR 2025
Gate Activation. In equivariant models, the gate activation combines two sets of group represen- tations. The first set consists of scalar irreps (1 = 0), which are passed through standard activation functions such sigmoid, ReLU and SiLU. The second set comprises higher-order irreps ((1 > 0)), which are multiplied by an additional set of scalar irreps that are introduced solely for the purpose of the activation layer. These scalar irreps also passed through activation functions.
The gate activation allows for the controlled integration of different types of irreps in the network. The scalar irreps capture global and local patterns, while the higher-order irreps capture more complex relationships and interactions. By combining these irreps in gate-like manner, the gate activation enables the model selectively amplify or suppress information flow based on the importance of different irreps for given task.
Normalization. Normalization is technique commonly used neural networks normalize the activations within each layer. It helps stabilize and accelerate the training process by reducing the internal covariate shift, which refers to the change in the distribution of layer inputs during training.
The normalization process involves computing the mean and variance across the channels. In equivariant normalization, the variance computed using the root mean square value the 1.2nnorm of each type-l vector. Additionally, this normalization removes the mean term. The normalized activations are then passed through learnable affine transformation without learnable bias, which enables the network to adjust the mean and variance based on the specific task requirements.
Figure 4: Determinacy of atomic positions in organic molecules. The relationship between atoms and colors is (Carbon, gray; Hydrogen, white: Oxygen, red; Nitrogen, blue). When the system is fixed, most atoms have a uniquely determined position. The translucent represents the atom of interest, and the arrow indicates its unique position.
Figure 5: Atoms with uncertain positions exist. In organic molecules, the positions of atoms have multiple possibilities, corresponding to several local minima of the potential energy surface. We represent the two possible positions of the atom with arrows.
17

---

Published as a conference paper at ICLR 2025
B MODEL DETAILS
B. THE POSITION DISTRIBUTION IN ORGANIC CHEMISTRY
the realm organic chemistry, the majority atomic positions within molecules are unequivocally determined, reflecting the precise architectural principles that govern molecular structure. This certainty stems from the well-defined rules of covalent bonding, hybridization, and the tetrahedral geometry that (C, H. O, N, F, S, ...) atoms often exhibit, which together constrain the spatial arrangement of atoms in most organic compounds. As illustrated in the Figure [ for given molecule, the majority of its constituent atoms possess uniquely determined positions. Any other arrangements would result in an implausible molecular structure.
However, there exists minority subset of scenarios where atomic positions may be ambiguous due the existence of multiple potential locations, corresponding to various local minima on the potential typically associated with conformational flexibility,
energy surface. These sites of uncertainty such as in molecules with large rotational freedom about single bonds, or in cases where there are multiple stable conformations or rotameric states. We take Figure [ as example, the hydroxyl group of the molecule can rotate, allowing for two possible positions of the hydrogen atom.
Furthermore, when considering the positional possibilities of atoms simultaneously. the combina- torial complexity of their potential configurations increases significantly. This multiplicative effect consequence of the interplay between intramolecular forces and the three-dimensiona nature of molecular space, which allows for vast array of spatial orientations and relative positions. As illustrated the Figure masking multiple atoms introduces a multitude of possibilities, rendering the position prediction of EMPP ill ill-posed problem.
Figure 6: The Possible positions of multiple atoms are masked. When calculating the forces acting on each atom, the influence of another masked atom must also be taken into account.
In summary, while most atomic positions in organic molecules are rigidly defined, there are instances where positional ambiguity arises, and this complexity is magnified when multiple atoms are consid- ered. This interplay between certainty and uncertainty in atomic positioning is fascinating aspect of organic chemistry that has profound implications for our understanding molecular structure and behavior.
B.2 DETAILS OF PREDICTION MODULE
Our position prediction module consists of multiple equivariant operations, designed based the structure the backbone. For the Equiformer, the output feature f, is represented as "128x0++64x1e+32X22
(irreducible representations in e3nn (Geiger al. 2022 ). We first standard-
ize the number of channels using a linear layer, resulting "64x0e+64x eexeece+xxxeee Through a gating
mechanism and second linear layer, f₃ is mapped to representation of "32x0e+32x 32x3222"
18

---

Published as a conference paper at ICLR 2025
Following Fourier transformation and gridding, we apply shared MLP the spherical surface, which is equivalent to 1D convolution. Since nonlinear transformations on the sphere do not disrupt equivariance, the shared MLP first projects the channels into 16-dimensional space and then applies the SiLU activation function. Finally, the 16-dimensional vector is projected back down to one dimension, yielding non-normalized distribution.
B.3 TORCHMD-NT WITH HIGHER- ORDER SPHERICAL HARMONIC REPRESENTATION
In the position prediction module, the EMPP must map the equivariant representations onto distribution on the spherical surface. A greater diversity of frequencies in the frequency domain representation leads a finer time-domain representation post-rransformatio. EMPP is built upon Equiformer and TorchMD-Net while Equiformer incorporates high-order spherical harmonics representations (analogous to higher frequencies), TorchMD-Ner includes only invariant features of I = 0 and three-dimensional equívariant features of 1. Consequently, the output after Fourier transformation may not effectively capture the fine- fine- grained positional distribution. To address this, convert the TorchMD-NM Net embeddings into high-order spherical harmonics representations. First, we project the embeddings onto representations function. Next, we expand the number of attention coefficients from TorchMD--Net, applying them to each degree representation (whereas the original model only applied these coefficients to the I representation). When calculating the inner product, the inner product of high-order spherical harmonics remains rotationally invariant, ensuring that the modified TorchMD-Net maintains overall equivariance. We apply this variant in our experiments. Note that we do not change the core operations in TorchMD-ND
= [0, 1, 2, using the spherical harmonics
B.4 PREDICTION WITH NEIGHBORS
In equations equation and equation 18 we restrict the prediction of masked positions to neighboring nodes within the cutoff radius. This restriction arises from the fact that, in GNNs, atoms beyond the cutoff do not directly interact with the masked atom; instead, they influence the PES indirectly by transmitting features to the embeddings of neighboring nodes through multi-layce forward propagation. Including points outside the cutoff during loss calculation would compromise the inductive bias inherent in GNNs. Furthermore, our ablation study demonstrates that incorporating atoms beyond the cutoff can diminish the model's generalization performance.
B.5 DIFFERENCES FROM SYMPHONY
Symphony (Daigavane al. 2024) proposes molecular generation framework based equivariant
representations. It constructs a fragment sequence by iteratively masking atoms, with each subsequent fragment masking one additional atom. Symphony then uses spherical harmonic projections of high- degree equivariant representations to encode the spatial distribution of atomic positions, enabling it to sample new atoms based on the predicted distribution. EMPP uses the similar technique define position prediction module where the relative positions are projected onto spatial distribution based on high-degree equivariant representations. However, EMPP and Symphony differ in several key aspects:
Masking Strategy: Symphony, geared towards molecular generation, masks all information about an atom, relying on the model's prediction of a "focus" atom to sample subsequent atoms. EMPP, conversely, masks only the position of single atom, retaining the ground truth of its atomic type and neighbor information. EMPP's position inference leverages all available real information from the molecule. These distinct strategies reflect their respective goals: Symphony prioritizes generating diverse molecules, while EMPP focuses on learning precise interaction patterns.
Prediction Horizon:
Symphony predicts the position of an atom without considering the positions of subsequent atoms. EMPP, however, uses all available positional information except for the masked atom's position. This difference again stems from their objectives: Symphony emphasizes diversity and fragment baageent-bbsedd generation, while EMPP prioritizes accuracy by minimizing interfering factors in position prediction to capture genuine physical interactions.
19

---

Published as a conference paper at ICLR 2025
Symphony normalizes multi-channel features using the ag- gregation log C exp(.)). While parameter-efficien, this approach may make the chanmet
Spatial Projection Module:
resulting distribution sensitive the spherical harmonic representation. fa. particularly after gridding. EMPP employs a neural network (detailed in equation 14) to aggregate multi-channel features the sphere. Ablation experiments (Appendix C.2 demonstrate the greater stability of the distribution generated by EMPP's neural network.
Symphony uses a Dirac delta distribution as the label for spatial pro- jection to encourage sampling of physically plausible positions. EMPP, not designed for sampling, relaxes this constraint and adopts smoother distribution to improve generalization performance.
Label Definition:
Table 5: Hyper-parametent for QM9 experiments without pre-train.
Hyper- paraeeerrseeeerrs
Value description
Optimizer
Adam
Leamnngg scheduling
Cosine learning rate with linear warmup
Warmup epuchs
5
Maximum learning rate
5 10-1, 1.5 x 10-4
Batch size
128, 64
Number of epochs
300, 600
Weight decay
5 x 10-3, 0
Dropout rate
0.0, 0.2
Cutoff radius (A)
5
EMPP
The maximum value of degree
2
Number of Sampling
100a
Intermediate dimensions Output dimensions
[64, 0). (64, (64.
[32, 0), (32, (32,
Encoding dimensions
(128, 0), (64, (32,
MLP on the spherical representation
Linear (3,, 16) SSLU
Linear 1)
Loss weights
Temperature
0.1
C DETAILS OF EXPERIMENTS AND SUPPLEMENTARY EXPERIMENTS
C.1 IMPLEMENTATION DETAILS
First, we introduce the hyperparameter configuration of EMPP when used as an auxiliary task, with the basic training configurations based the Equiformer setup. The configurations for QM9 and MD17 7 are recorded in Table 5 and T respectively. Note that there are two sets of configurations for QM9, the one with a longer epoch or smaller batch size is for the four tasks of G, H, U and Uo tasks.
In the experiments on PCQM4MV2, we only use EPMM for pre-training. The experimental config- urations are shown in Table The validation set of PCQM4MV2 is used to evaluate the accuracy of position prediction. Most of our training configurations are the same as those of the Denoising method (Zaidi et al., 2023). During the fine-tuning phase QM9, we use the training parameters from Denoising method and the model papeeererreeters in Table 7 The loss weight for EMPP is set to 1.
Hapee-parameters of Baselines. In Table I and Table 2, the results of baselines is from (Liao & Smidi 2023) In Table 3 the results of baselines is from (Feng et al. 2023). Additionally, the result of Attr is from (Luo et al. 2022).
In the ablation study, verified the relationship between the denoising method and the standard deviation. The experimental setup for this part follows Work A. The difference that we did not introduce PCQ-based pre-training model, and the loss weight for denoising was set 1.
20

---

Published as a conference paper at ICLR 2025
Table 6: ayper-parameters for MD17 dataset without pre-train.
Value or description
Optimizer
AdamW
Leurning rate scheduling
Cosine learning rate with linear warmup
Warmup epochs
10
Maximum learning rate
5 x 10
Batch size
Number epochs
2000
Weight decay
x 10
Dropout rate
0.0
Weight for energy loss
Weight for force loss
80
Cutoff radius (A)
EMPP
The maximum value degree Number of Sampling Intermediate dimensions
1002
[64, (64, (32.3)
Output dimensions
(32, (32.)).3321)) (32 (32.
Encoding dimensions
(128, 0), (64, (64,) (32.3)
ML.P on the spherical representation
Lineor (32, 16) S:LU [1near(16,
Loss weights Temperature
50
0.1
C.2 SUPPLEMENTARY EXPERIMENTS
GEOM-Drug, GEOM-Drug (Axelrod & Gomez- Bombarelli 2022 is dataset containing diverse collection of non-equilibrium molecular data, featuring a broader range of atomic types and molecular sizes. In our experiments, we use Equiformer as the backbone model. The training configuration follows the setup used QM9 experiments, with the number training epochs set to 300. The absolute energy of each conformation is used as the target label.
Table Hyper-parameters for PCQM4MV2 dataset.
Hyper-parameters
Value or description
Optimizer
AdamW
Leurning rate scheduling
Cosine learning rate with linear warmup
Warmup steps
10000
Maximum learning rate
x 10
Batch size
70
Number epochs
20
Weight decay
0.0
Cutoff radius (A)
EMPP
The maximum value degree
Number of Sampling Intermediate dimensions
100
[64, (64, (64.3)
Output dimensions
(32, 0), (32, 1). (32, (32.3)
Encoding dimensions
(128, 0), (128, (128, 2), (128,3)
MI.P on the spherical representation
Lineor (32, 16) S:LU Lánear 1)
Temperature
0.1
Table 8: Ablation experiments of multiple atom masking on QM9 tasks.
Task
Ac
EНоно meV
EL mcV
Units
bohr
meV
Mask atom
.041
27
13
Mask
.046
31
16
16
Mask 3 atom
.052
33
20
19
For data preparation, we randomly sample 200,000 molecules from GEOM-Drug as the training set
and 10,000 molecules as the validation set. To ensure the reliability of the validation results, SMILES strings appearing in the validation set are excluded from the training set, recognizing that a single SMILES can represent multiple conformational data points.
21

---

Published as a conference paper at ICLR 2025
The experimental results, summarized in Table 13 demonstrate that EMPP achieves substantial performance improvements on GEOM-Drug, reducing the energy MAE by 48%. From this, we draw two key conclusions: (a) EMPP is effective for more complex organic molecular systems. (b) Despite the prevalence of non-equilibrium structures in GROM-Drug, HMPP remains highly effective as an auxiliary task.
This experiment highlights EMPP's ability to overcome the limitations of denoising methods, which typically restricted approximating equilibrium structures.
Table 9: Ablation experiments of spherical operations on QM9 tasks.
Task
a bohr²
Ae
еномо meV
ÊLUMO
Units
meV V
meV
Shared MLP
.041
27
14
13
log (2 channet exp(.))
.043
28
14
14
Table 10: Ablation study label distribution.
Method
Éномо
Without self-supervision
.046
15.4
Baseline
.041
14.2
Directly predict relative positions
.044
15.1
Dirac delta for radius
.042
14.6
Dirac delta for direction
.044
14.6
Dirac delta for both
.044
14.8
Dirac delta distribution vs. Gaussian distribution. We have mentioned in Section 3.2.3 that whether Dirac delta distribution and Gaussian distribution are a one-to-one projection ground truth Fu, they are all theoretically correct. Here, we conduct experiments to investigate their difference. In detail, we set the equation to 0.0001 to make the radius distribution similar to Dirac delta distribution (The ideal Dirac delta distribution is hard to sample). Similarly, we add the temperature coefficient to equation| [7 and set it to 0.0001 to make the spherical distribution closed to Dirac delta distribution. Additionally, We also conducted set of experiments to directly predict relative positions Tik-
The results shown in Table 10 where "Baseline" denotes using original setting. found that the Dirac delta distribution can also achieve improvements compared to distribution in equation [6 and equation]] 17 However, the sharp distribution will reduce the stability of training: when repeat the experiments, the baseline model converges to the optimal result every time, however, EMPP based on the Dirac distribution requires more than five repetitions of the experiment find the results presented in Table IO, Finally, there is another observation from Table 10 directly predict the relative positions limits the effectiveness of EMPP
Transferability We also assessed the transferability of two distributions. We pre-trained on PCQM4MV2 and transferred QM9. We use the TorchMD-Net model and the "baseline" employed EMPP's 1-mask strategy. "Dirac delta for both" involved changing the distributions in both pre- training and fine-tuning to Dirac delta distribution, similar to the experiments Table TO From it can be seen that relaxed distribution contributes to the transferability of EMPP.
Table 11
Table 11: Ablation study label transferability. The backbone is TorchMD-Net.
Method
еномо
Without self-supervision and pre-train
.059
20.0
Baseline
.037
14.3
Dirac delta for both
.049
17.3
Training time consumption. EMPP is self-supervised approach that generates new data indirectly, which can lead to increased training resource consumption. We have compared the computational time required for calculating the EMPP loss with other methods in Table T2] The findings indicate
22

---

Published as a conference paper at ICLR 2025
that the time taken to compute EMPP is similar to that of calculating losses for property prediction denoising tasks. Consequently, employing n-mask EMPP as an auxiliary task is expected to roughly multiply the training time by factor of n. Based on the experimental results presented in Tables [ 2 we find that the 1-mask strategy is generally sufficient for most scenarios: doubles the training time but also yields substantial performance enhancements. When computational resources and time are not constraints, we suggest opting for higher in the mask strategy for even better outcomes.
Table 12:
Ablation study on training time consumption.
Method
Samples per second
The cost of each iterations (ms)
Property prediction
291.57
439
Denoising
290.25
441
EMPP
281.94
454
Masking multiple atoms in a molecule. EPMM based on crucial theoretical premise: when single atom is masked, its position is well-posed in most cases, but when two atoms are masked, the position prediction becomes ill-posed. We empirically validate this assertion through experiments. We construct a variant in which multiple atoms are masked in each molecule, and the remaining atoms' embeddings are used to predict the positions all atoms simultancously. From Table 8 We found that when multiple atoms were masked, the performance actually declined, and only when single atom was masked did EMPP surpass the baseline. This experiment confirmed our concerns regarding the ill posed nature of masking multiple atoms. Furthermore, in practical applications, recommend the approach of repeatedly masking a single atom, which is decribed in equation g
Table 13: Ablation study on hyper-parametent in label distribution.
Method
a
НHоМО
Without self-supervision
.046
15.4
Baseline (7 0.1/o 0.5)
.041
14.2
0.1/o 0.3)
.041
14.2
(T = 0.1/a 0.1)
.041
14.3
(T 0.1/o 0.05)
.042
14.2
(T = 0.1/a = 0.01)
.042
14.3
(T 0.5/a 0.5)
.041
14.2
T 0.3/o = 0.5)
.041
14.1
(T 0.05/O 0.5)
.042
14.3
T 0.01/a 0 = 0.5)
.042
14.4
mentioned that operations the sphere do affect equivariance. In order to map the infommaion of all channels to ID, we employed a shared fully connected layer. Previous methods have used aggregation log log(Eetonnee! exp(-)) for mapping to ID. It can be observed from the that the fully connected layer introduces more parameters, but it also enhances performance. Moreover, parameter-fre methods may encounter numerical overflow issues during the training process.
Operations Sphere. In Section 3.2.2
Table 14: The difference between prediction based on neighbor structure and global molecular graph.
Task
a bohr
EHOMO meV
LUMO meV
Units
EMPP (Neighbor)
.041
14
EMPP (Global)
.045
118
16
EMPP (Global + Dist Weights)
.044
15
The impact of tareee-psrameters in label distribution. т in equation 14 and or in equation are two important hyperparameters, and we conducted ablation studies evaluate them. This ablation
23

---

Published as a conference paper at ICLR 2025
study not based pre-training and used Equiformer to train QM9 with EMPP auxiliary task. From Table [3] we found that when we changed these two hyperpurameters, there was no significant change performance, which reflects that EMPP is not sensitive to the distribution of labels.
Table 15: Results GEOM-Drug property prediction without pre- trainriing
Method
Energy MAE (kcal/mol)
Equiformer
0.07517
Equiformer+EMPP
0.03912
Prediction with all the unmasked atoms. In EMPP, we use the embeddings of neighboring nodes to predict the position of the masked atom. However, in reality, all atoms interact with the masked atom, even if the interactions are minimal. Here, we have designed two variants of EMPP. The first variant uses all atoms predict the position ("Global"),, and the second variant uses all atoms to predict and adds distance-related weights ("Gllobal Dist Weights"). ). In other words, the second method weights the loss generated by the first variant, with greater weight given to atoms that are closer to the masked atom. The weighting function is vr = max(0, [) where C is hyper-parameter set to 8 Á.
From Table 14 we observe that utilizing all unmasked atoms to predict positions diminishes the model's generalization capability. We believe this may be due to inaccurate long-range interactions af- fecting the model's learning. Although incorporating distance weights improves the global prediction effect, it still fails to match the performance of predicting based on neighboring atoms.
The impact of molecules of different sizes on EMPP. EMPP improves its generalization perfor- mance by masking atomic positions and then restoring them. Intuitively, as the number of atoms in molecule increases, more atoms be masked, allowing EMPP to generate more diverse data. To evaluate the impact of EMPP on molecules with different sizes, we conducted experiments by categorizing the QM9 training data into three groups based the number of atoms: (0-16), (17-19), and (20+). These categories contain roughly equal amounts of data. In each experiment, we computed the EMPP loss using only molecules from of these categories. As shown in Table T6 the results demonstrate that EMPP consistently improves performance across different molecular sizes, with larger molecules experiencing more sigmificant performance gains.
Table 16: Ablation study on the impact of molecules of different sizes.
Method
Baseline
.041
14.2
EMPP on (0-16)
.044
14.9
EMPP on (17-19)
.043
14.4
EMPP (20+)
.043
14.5
24


Reviews:
[{"review_id": "Jw1zv5luxY", "reviewer": "Reviewer_cKne", "summary": "This paper introduces the Equivariant Masked Position Prediction (EMPP) model, a self-supervised learning approach for molecular representation aimed at improving molecular property prediction in graph neural networks (GNNs). EMPP differs from conventional methods by predicting the relative position of masked atoms based on surrounding atomic structures rather than masking node attributes. This avoids the Gaussian mixture approximation used in denoising methods, potentially enhancing accuracy. EMPP is presented as a versatile technique that can function both as a pre-training task on unlabeled data and as an auxiliary task for labeled property prediction, achieving strong results across molecular benchmarks.", "questions": "- **Physical Assumptions in Modeled Distributions**: It is interesting to understand the foundational physical or chemical assumptions behind EMPP’s use of relative position modeling for capturing atomic interactions. Denoising methods typically rely on a Boltzmann distribution near local minima, with harmonic potential assumptions. By contrast, EMPP’s approach seems to focus on relative distances and directional distributions between masked and neighboring atoms. This reviewer would like clarification on the theoretical basis for this choice. Specifically, do the authors view this approach as a distinct modeling assumption, or is it comparable to the physical basis provided by Boltzmann distributions in denoising methods?\n\n- **Comparison with Denoising Methods**: Given that both EMPP and denoising methods model unknown data distributions, EMPP’s key distinction is its focus on relative, rather than absolute, atomic positions. The authors claim that one of the reasons for the performance improvement over the existing method is that they avoided the erroneous gaussian approximation (Figure 1-(d)). Then, can we expect the same level of performance improvement even if we do not adopt the method of modeling the distance and angular parts separately in EMPP? Is it possible to adopt a method that directly estimates the positional difference vector? Is it possible that the good performance of the model is due to modeling the distance and angular part separately? Please conduct ablation experiments on this.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- **Originality**: EMPP’s approach of modeling relative atomic positions marks a novel shift from traditional self-supervised molecular methods, which typically rely on attribute masking or denoising with Gaussian mixtures. By avoiding Gaussian approximations, EMPP proposes a fresh method that could more precisely capture quantum mechanical properties without prior assumptions on potential energy surface shape.\n- **Quality**: The methodology and experiment sections are comprehensive, offering detailed descriptions of the model architecture and training procedures. EMPP’s application across QM9 and MD17 benchmarks highlights its effectiveness and consistency in molecular property prediction.\n- **Clarity**: The paper is clearly structured, and the motivations for EMPP are presented well. The provided comparisons, particularly with denoising methods, effectively illustrate the advantages of EMPP.\n- **Significance**: EMPP’s adaptability for both pre-training and auxiliary tasks in GNNs is particularly valuable, as it could streamline the use of self-supervised learning in various molecular and materials science applications.", "weaknesses": "- **Unclear Evaluation in Pre-training**: EMPP is proposed for both pre-training and auxiliary loss tasks, yet performance declines are observed in some properties (e.g., HOMO, LUMO, G, H) when used as a pre-training task (see Table 1 vs. Table 3). This suggests a possible limitation in generalizability when used for pre-training. I guess this is due to dataset constraints, as QM9 may not properly diverse molecular dataset to evaluate the most recent methods.\n\n- **Too simple benchmark** : The molecules in the benchmark test sets, particularly QM9 and MD17, are relatively small or simple compounds, which may not sufficiently assess EMPP’s performance on more larger and realistic molecular structures. This makes it challenging to evaluate EMPP's generalization capabilities for applications requiring more intricate molecular structures.\n\nTo address this, the reviewer's suggestion:\n\n**Conduct experiments on the Geom-Drugs dataset**: the dataset includes larger and more complex molecular structures than QM9. This dataset can provide a more realistic test of EMPP's ability to generalize to more diverse molecular systems. Especially, please perform evaluations with and without pre-training to compare how EMPP's pre-training performance translates to more complex molecules in this dataset. This will help determine if the performance drop observed in QM9 pre-training persists or changes with different datasets.\n\n**Analyze Performance Scaling with Molecular Complexity and Size**:\nExtend experiments across molecules of varying complexity and size to observe how EMPP scales. This may involve assessing EMPP’s accuracy, efficiency, and convergence with increasing molecular complexity.\nPresent scaling analysis results, ideally quantifying any trends where EMPP’s accuracy vary with molecular size. This will provide insights into EMPP's robustness and applicability to more complex molecular systems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qlQCAPe6ZQ", "reviewer": "Reviewer_Koif", "summary": "This paper presents a novel self-supervised framework for molecular representation learning, named Equivariant Masked Position Prediction (EMPP). EMPP pretrains the model by directly predicting the masked positions of atoms and addresses limitations in existing self-supervised methods (attribute masking and denoising). Experimental results show EMPP achieves competitive performance. EMPP can also enhance supervised learning when used as an auxiliary task.", "questions": "In Equation 8, does the averaging over multiple masked atoms effectively act as training for a longer time? How does it impact the efficiency of the model? Could you please provide empirical results comparing training times and computational efficiency between single and multiple masked atoms? Additionally, discussing potential trade-offs between performance gains and computational costs would be appreciated.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The paper is well-organized. The method has clear explanations.\n- The proposed EMPP approach is simple yet effective, achieving state-of-the-art on benchmark datasets.", "weaknesses": "In Section 5.2, the comparison does not seem entirely fair. EMPP serves as an auxiliary task and relies on the Equiformer backbone, yet it is compared with other architectures that lack this component.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "15zkQl6sxG", "reviewer": "Reviewer_DiCi", "summary": "This paper introduces a self-supervised learning approach called Equivariant Masked Position Prediction (EMPP) for molecular graph neural networks. The main idea behind is to mask the 3D position of atoms while keeping their other attributes, then predict these positions using the surrounding molecular structure. The authors intend to address the limitations of other self-supervised approaches. Results show comparable/superior performance to other state-of-the-art approaches for quantum-mechanical property prediction.", "questions": "1) What is the computational overhead of calculating distributions over the 100² spherical grid points compared to simpler approaches? Perhaps providing training time and memory requirements compared to baseline methods could also help.\n\n2) How sensitive are the results to the sampling rate on the sphere?\n\n3) Providing baseline results using the modified ET from TorchMD-Net could help distill how significant is the improvement due to the inclusion of the masking task. (By the way, the citation of the ET architecture is wrong).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The core idea of predicting 3D atomic positions while keeping atomic attributes seems a reasonable evolution of existing molecular self-supervised approaches, taking into account the claimed limitations of previous approaches. The paper's strongest aspect is its clear presentation. The problem motivation and methodology are well articulated. The experimental validation, while having some gaps, provides adequate evidence on standard benchmarks. These strengths suggest incremental results.", "weaknesses": "The performance improvements on MD17 are relatively modest compared to the baseline Equiformer. This raises questions about whether the added complexity of EMPP is justified by the gains. Furthermore, for non-pretaining comparisons, there are models that substantially outperform the proposed method, based also on higher-order tensors.\n\nAlso, the computational overhead of EMPP isn't adequately addressed. The method requires calculating distributions over a spherical grid (100² points according to Table 7) for each masked atom, which could be computationally expensive. An analysis of training time and memory requirements compared to simpler approaches is missing.\n\nI feel that altering the Equivariant Transformer architecture while testing the validity of the method simultaneously makes it difficult to distill where the improvements come from.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Tm7fm3G4NB", "reviewer": "Reviewer_RrYB", "summary": "This paper proposed EMPP, a new approach for pre-training GNNs for 3D molecules with equivariant mask positional prediction. Instead of directly denoising the position from the approximation of Gaussian mixtures, the authors proposed to model the distribution over positions with equivariant GNNs. Downstream tasks on QM9 and MD17 demonstrated improvements in the performance after such a pre-training.", "questions": "1. What is the training objective in the end-to-end training for the QM9/MD17 dataset? See Weakness 1.\n2. Why can the proposed framework address the Gaussian mixture issue? In your pre-training objectives, Gaussian assumptions (or other softmax-based assumptions) were also explicitly assumed to make the ground truth distribution target. See Weakness 2.\n3. What is the major innovation of the proposed framework as a \"force prediction\" model? Previous work has clearly indicated the close relationship between force field prediction and the denoising process of atom positions. Equivariant GNNs were also used in these works. In this sense, the approach proposed in this work is essentially the same as existing denoising frameworks. See Weakness 3.\n4. How do you address the issue of inconsistent pre-training objectives if Gaussians are assumed independently for radii and angles? See Weakness 4.\n5. What is the impact of $\\tau$ in Eq.13? What is the impact of $\\sigma$ in Eq.14? According to the ablation study of the denoising baselines, these values that control the \"sharpness\" of the distributions have a large impact on the final performance. However, they were not fully tested in the paper with the proposed approach.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The idea of predicting distributions over positions instead of directly predicting the atom positions is novel. Previous work, as was demonstrated in the paper, predominantly relied on the direct prediction of the mask atoms.\n- Experimental results on QM9 and MD17 demonstrated the improvement of pre-training compared to the baseline models, in both end-to-end training and pre-training settings.\n- An anonymous code link is provided for better reproducibility.", "weaknesses": "1. The delivery of this paper is poor, with noticeable inconsistencies regarding the proposed method, irrelevant materials and experiments.\n    - **Irrelevant related work**. The proposed method has nothing to do with language models (not even the language modeling for chemical sequences like SMILES). However, the authors started with LLM, which was misleading and confusing.\n    - **Irrelevant experiments**. The ablation study in Section 5.4 demonstrated the results for denoising pre-training models, but not for the proposed method in the paper. In contrast, the ablation study in Appendix C was designed for the proposed model. It is unclear why the authors put the irrelevant ablation studies in the main text.\n    - **Objective in pre-training**. In Figure 1 and Section 3.2.1, the authors seemed to indicate their proposed approach is a force-prediction network for pre-training. However, Figure 2 and Section 3.2.2 indicated that the displacement vector $r_{ij}$ was instead regressed as the learning target.\n    - **Objective in finetuning**. In Section 3.3, the energy and force prediction seemed to be the auxiliary input for pre-training. However, in the QM9/MD17 setting, they are the prediction targets. It is unclear how the model was trained end-to-end without the proposed pre-training objective.\n\n2. **The claim of Gaussian mixture issues is questionable**. The authors claimed that the proposed approach can \"effectively avoid\nthe approximation of Gaussian mixture\" (Section 3.2.1). However, the loss functions in Eq.14 also explicitly rely on the Gaussian mixture assumption to derive the KL loss. The direction loss in Eq.16 assumes an alternative form based on the softmax function. The only difference is the Gaussian is defined for the radii and angles instead of Euclidean coordinates.\n\n3. **The innovation in force prediction is lacking**. It is unclear about the major motivation behind the proposed method of \"force prediction\". Previous work (e.g., [1], also cited in the paper, [2], [3]) has clearly indicated the relation between the potential energy surface and the force field prediction with equivariant GNNs. In this sense, the proposed approach is exactly the same as other denoising pre-training models. Furthermore, the indicated issues with the Gaussian mixture of variable variance landscape can be effectively addressed with iterative sampling methods (i.e., diffusion models like [4]) and annealed Langevin dynamics. Alternatively, [2] introduced invariant noise-scale prediction which can handle different noise scales.\n\n4. **The local equivariant prediction of the position distribution is not consistent as a pre-training objective** (unless the Gaussian variance is 0). This can be easily demonstrated by the fact that the radii and direction angles are constructed with respect to different neighboring atoms such that the expanded ground truth distribution roughly follows the shape of a small region of a spherical shell. For neighboring atoms at different positions, the ground truth is inconsistent unless the variance is 0 (i.e., a Dirac distribution).\n\n\n[1] Zaidi, Sheheryar, et al. \"Pre-training via denoising for molecular property prediction.\" arXiv preprint arXiv:2206.00133 (2022).\n\n[2] Jiao, Rui, et al. \"Energy-motivated equivariant pretraining for 3d molecular graphs.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 7. 2023.\n\n[3] Jiao, Rui, et al. \"Equivariant Pretrained Transformer for Unified Geometric Learning on Multi-Domain 3D Molecules.\" arXiv preprint arXiv:2402.12714 (2024).\n\n[4] Hoogeboom, Emiel, et al. \"Equivariant diffusion for molecule generation in 3d.\" International conference on machine learning. PMLR, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)