Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Authors: Shuqi Lu
Xiaohong Ji
Bohang Zhang
Lin Yao
Siyuan Liu
Zhifeng Gao
Linfeng Zhang
Guolin Ke
Venue: ICML 2025 poster (2025)
Keywords: Molecular pretrained representation
Molecular Property
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: we propose a novel framework, SpaceFormer, which incorporates additional 3D space beyond atoms to enhance molecular pretrained representation ability

Abstract: Molecular pretrained representations (MPR) has emerged as a powerful approach for addressing the challenge of limited supervised data in applications such as drug discovery and material design. 
While early MPR methods relied on 1D sequences and 2D graphs, recent advancements have incorporated 3D conformational information to capture rich atomic interactions. However, these prior models treat molecules merely as discrete atom sets, overlooking the space surrounding them. We argue from a physical perspective that only modeling these discrete points is insufficient. We first present a simple yet insightful observation: naively adding randomly sampled virtual points beyond atoms can surprisingly enhance MPR performance. In light of this, 
we propose a principled framework that incorporates the entire 3D space spanned by molecules. We implement the framework via a novel Transformer-based architecture, dubbed SpaceFormer, with three key components:
(1)grid-based space discretization; (2)grid sampling/merging; and (3)efficient 3D positional encoding.
Extensive experiments show that SpaceFormer significantly outperforms previous 3D MPR models across various downstream tasks with limited data, validating the benefit of leveraging the additional 3D space beyond atoms in MPR models.

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Shuqi Lu
Xiaohong Ji
Bohang Zhang?
Lin Yao
Siyuan Liu
Zhifeng Gao
Linfeng Zhang' Guolin Ke
Abstract
Rong et al., 2020). The high cost of acquiring labeled data from wet-lab experiments or simulations limits the perfor- mance of models trained on small datasets. MPR addresses this by leveraging unsupervised pretraining on large-scale unlabeled data to learn molecular representations, which can be fine-tuned on limited labeled data for improved per- formance. The evolution of this field has progressed from ID sequences (Xu et al., 2017; Wang et al., 2019; Heller et al., 2015) and 2D graphs (Hu et al., 2019; Rong et al., 2020; Li et al., 2021; Wang et al., 2022b) to 3D conforma- tions (Stärk et al., 2022; Zhou et al., 2023), incorporating increasingly rich physical information and achieving supe- rior performance. all these prior 3D MPR models, atoms play a central role. More specifically, these models take the types and 3D positions of atoms (or atom tuples) inputs and focus modeling atomic interactions within 3D space, typically using graph neural networks or transformers (Zhou et al., 2023; Feng et al., 2023; Wang et al., 2023; Cui et al., 2024; Yang et al., 2024).
Molecular pretrained representations (MPR) has emerged as a powerful approach for addressing the challenge of limited supervised data in ap- plications such as drug discovery and material design. While early MPR methods relied on ID sequences and 2D graphs, recent advancements have incorporated 3D conformational informa- tion to capture rich atomic interactions. However, these prior models treat molecules merely as dis- crete atom sets, overlooking the space surround- ing them. We argue from physical perspective that only modeling these discrete points is insuffi- cient. We first present simple yet insightful ob- servation: naively adding randomly sampled vir- tual points beyond atoms can surprisingly enhance MPR performance. In light of this, we propose a principled framework that incorporates the entire 3D space spanned by molecules. We implement the framework via a novel Transformer-base ar- chitecture, dubbed SpaceFormer, with three key components: (1) grid-based space discretization; (2) grid sampling/merging: and (3) efficient 3D positional encoding. Extensive experiments show that SpaceFormer significantly outperforms previ- ous 3D MPR models across various downstream tasks with limited data, validating the benefit of leveraging the additional 3D space beyond atoms in MPR models.
While this atom-based MPR approach appears straightfor- ward, we argue that it has an inherent limitation: does not model the surrounding space beyond atoms. On the surface, one may feel that the "empty" surrounding space is inconsequential as the molecular structure can be fully determined given only atom positions. However, such an argument overlooks critical physical principles. In the realm of microscopic physics, the space beyond atoms is not truly empty; it is permeated by electrons, various electromag- netic fields, and quantum phenomena (Atkins & Friedman, 2011: 1; Zee, 2010; Weinberg, 1995). Moreover, many com- putational simulation methods in physics require modeling the entire 3D space, not just the positions of atoms. For instance, electronic density distributions and potential fields, which closely related to certain molecular properties, are functions of the entire 3D space (Atkins & Friedman, 2011; Parr al., 1979; Szabo & Ostlund, 2012). This naturally raises the following question:
1. Introduction
Molecular pretrained representation (MPR) has been a key area of research for its crucial role in utilizing limited su- pervised data, particularly in real-world applications such as drug design and material discovery (Gilmer et al., 2017;
Eqqua contribution 'DP Technology, Beijing, China "Peking University, Beijing, China. Correspondence to: Guolin Ke <kegl@dp.tech..
Can modeling the 3D space beyond atoms improve molecular pretrained representation?
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
To test the hypothesis, we start with a simple experiment. Given the point set of a molecule, we introduce randomly sampled virtual points (VPs) within the molecular's 's sur-

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Previous 3D MRL Models
type
pos.
N
(4.1, 4.2)
Atom Extracting
(0.8, 3.5)
GNN, SE(3) Transformer,
(1.8, 2.2)
(1.2, 0.5)
SpaceFormer
cell type
cell pbs.
Gell type
in-cell pos
Null
(2.0, (2.0.6.0)
Null
(0, O.)
(4.1, 4.2)
(-0.4, -0.3
Nuil
3.5
In-cell pos. cal.
Null
(0,
Model
Nuil
(7.0, 1.0
Null
Gridding
Adaptive Grid Merging
Featurizing
(Cs, ci)
(c,, cy)
Linear
Multi-head Attention (FlashAttn)
Norm
Layer
Linear
3D Directional
PE with RoPE
Distance PE with RFF
Layer Norm
FFN
Linear
dg
Embedding
Figure 1. Overview of SpaceFormer. For clarity, we illustrate the model in 2D space. SpaceFormer begins by discretizing the 3D cuboid containing molecules into grid cells. Grid sampling/merginn strategy is then applied to reduce the number of input cells for efficiency. Note that exact atomic positions are still encoded after discretization owing to various positional encoding (PE). To efficiently encode pairwise positional relationships, SpaceFormer utilizes 3D Directional PE with RoPE and 3D Distance PE with Random Fourier Features.
rounding space and feed the augmented point set to exist- ing 3D MPR model. Surprisingly, observe a non-trivial performance gain after adding small number of VPs, even though these VPs are inherently noisy and theoretically carry no meaningful information. Notably, the improvement remains consistent and gradually increases as the number of VPs grows, eventually reaching plateau (see Fig. 2).
one is adaptive grid merging, where adjacent empty cells recursively merged into large cells (see Fig. 1). Both meth- ods work well in practice, reducing the number of cells by approximately an order of magnitude while still maintain- ing performance. Additionally, we design efficient versions of pairwise positional encoding based on RoPE (Su et al., 2024) and random Fourier features (Rahimi & Recht, 2007), which enjoy linear complexity relative to the number of cells. We term the resulting architecture SpaceFormer.
This finding is intriguing and further motivates to de- velop more principled approach that directly models the entire 3D space spanned by molecules. To this end, we pro- pose a fundamentally different MPR framework: we treat a molecule as a 3D "image" by discretizing the 3D space into small grid cells. By using a sufficiently small cell size, we ensure that each grid cell contains most one atom. The input feature for each cell can then be defined the position of the cell combined with the type and position of the atom it contains (if any). This grid-based representation can subsequently be processed by any model capable of handling sets, such as Transformers.
Pretraining approach. A key advantage of our proposed framework is that it can be integrated with powerful pre- training approach known as Masked Auto-Encoder (MAH) (He et al., 2022). Specifically, during pretraining, we can mask portion of grid cells for each molecule and let the model predict whether each masked cell contains atoms. If it is true, the model further predicts the atom type and its precise position within the cell. Remarkably, our analysis reveals a potential advantage of MAE pretraining compared with existing approaches such as denoising (Zaidi et al., 2022), showing that the model may learn additional knowl- edge via masked prediction.
Efficient implementation. One potential drawback of the above framework lies its efficiency. Indeed, the number of grid cells grows cubically as the cell length decreases, which limits the scalability. To address this issue, we im- plement two intuitive strategies that significantly reduce the number of cells. The first one is grid sampling, where a random subset of empty grid cells is selected; The second
We conduct extensive experiments to evaluate the effective- and efficiency of SpaceFormer. Across total of 15 diverse downstream tasks, SpaceFormer achieves the best performance on 10 tasks and ranks within the top 2 on 14 tasks. Ablation studies further confirm that each component
2

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
plays a critical role in enhancing either the performance or efficiency of SpaceFormer. Overall, these results underscore the importance of modeling the 3D space beyond atom po- sitions, and we believe our framework can serve as a new paradigm towards advancing the area of MPR.
points to align and fuse features from 2D camera images and 3D LiDAR data, effectively addressing the resolution disparity between the sensors. Various studies (Song et al., 2023a;b; Mahmoud et al., 2023) have advanced this direc- tion by utilizing virtual points as a bridge to align data from heterogeneous sensors. In this paper, virtual points are used in a quite different context, in the sense that it a starting point to motivate our grid-based SpaceFormer architecture.
2. Related Work
Molecular Pretrained Representation. Molecular pre- trained representation has explored various modalities, giv- ing rise to diverse range of methods that leverage different molecular input formats. Some approaches rely on ID se- quences, such as SMILES (Wang et al., 2019; Xu et al., 2017), while others focus on 2D topologies, including meth- ods like MoICLR (Wang et al., 2022b), MoIGNet (Li et al., 2021), ContextPred (Hu al., 2019), GROVER (Rong et al., 2020). Additionally, several works have enhanced 2D MPR models by integrating 3D information, seen in approaches like GEM (Fang et al., 2022), 3D-Infomax (Stärk et al., 2022), MoleBLEND (Yu et al., 2024), GraphMVP (Liu ct al., 2021), and Transformer-M (Luo et al., 2022).
Broadly speaking, the idea of adding virtual points can be generalized to the use of auxiliary tokens, a trick that has been adopted across various domains beyond 3D point clouds, such as vision and language. In Vision Transformers, Darcet et al. (2023) introduced additional "register" tokens into input sequence patches to mitigate artifacts. Similarly, in language modeling, Pfau et al. (2024) demonstrated that incorporating dot tokens ("...") as Chain-of-thougee prompts surprisingly enhances the performance of large language models. Our approach aligns partially with this design prin- ciple; However, the underlying mechanism is fundamentally distinct: the empty grid cells in SpaceFormer are guided by physical principles along with meaningful 3D positional en- codings, and enjoys MAE pretraining. In Sec. 4.3, we will demonstrate that these aspects contribute to the improved performance of our model.
Recently, starting with Noisy Nodes (Zaidi et al., 2022) and Uni-Mol (Zhou et al., 2023), pure 3D MPR models have demonstrated superior performance across various tasks. This has inspired a surge of subsequent works (Feng et al., 2023; Wang et al., 2023; Cui et al., 2024; Yang et al., 2024), which further explored the potential of 3D MPR. Besides, several other works also focus on 3D conformations, such as deep potential models (Schütt et al., 2017: Thomas et al., 2018; Gasteiger al., 2020; 2021; Liu al., 2022; Wang et al., 2022a; Jiao et al., 2023), protein folding (Jumper et al., 2021; Abramson et al., 2024), and 3D conformation generation (Shi et al., 2021: Zhu et al., 2022b; Xu et al., 2022; 2021). Yet, all of these works use 3D atomic posi- tions as inputs without modeling the empty space beyond atoms. Earlier 3D grid- methods like AtomNet (Wal- lach et al.. 2015) and SchNet (Schütt et al., 2017) both use 3D convolutional neural networks, with SchNet focus- ing on atom positions. In contrast, SpaceFormer employs transformer- architecture to capture interactions across global grid cells, extending beyond just atoms positions. Separately, research on potential energy surfaces (Wang et al., 2024; Wallach et al., 2015) prioritizes long-range and higher-order atomic interactions for force field appli- cations: SpaceFormer, however, is not optimized for such force-specific tasks, distinguishing its scope and objectives.
3. Method
3.1. Motivation: the effect of adding virtual points
As mentioned in the introduction, we posit that modeling the surrounding space beyond atoms can enhance 3D MPR due its ability to capture critical physical principles such as electromagnetic fields and quantum phenomena. To illus- trate this, we design the following motivating experiment. Given an input molecule represented as a set of atoms, we introduce virtual points (VPs) randomly sampled from the 3D space within the circumscribed cuboid of the atom set. This augmented point set is then fed into a state-of-the- 3D MPR model, Uni-Mol (Zhou et al., 2023). We follow exactly the same pretraining pipeline as described in Zhou et al. (2023), which involves randomly masking atom types and adding Gaussian noise to their positions, while training the model to predict the correct atom types and the noises. Our goal is to investigate how the model's 's performance varies with the number of random virtual points, where the baseline corresponds to no virtual points.
The result is illustrated in Fig. 2. One can clearly see that the validation loss gets improved with the help of virtual points. Specifically, adding just 10 virtual points already yields a noticeable performance gain. This improvement remains consistent and gradually increases as the number of virtual points (VPs) grows, eventually reaching a plateau. Such a result perhaps surprising as the additional virtual
The use of virtual points and auxiliary tokens. Virtual points has been extensively used in the area of 3D point cloud learning to serve as intermediate representations. For example, Wu et al. (2023); Yin et al. (2021) converted 2D camera images into 3D virtual points, which are then fused with 3D LiDAR points to create unified input representa- tion. Similarly, Zhu et al. (2022a) introduced sparse virtual
3

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
-5.5
Input format. Based on the discretization, the input to our framework is simply the set of grid cells, where each cell consists of two parts: its type and coordinates. Since each grid cell contains at most one atom, the cells can be categorized as either atom cells or non-atom cells. For each atom cell, its type and coordinates are defined as the atom type and its precise 3D coordinates. For each non-atom cell, we assign it with the special type NULL and define the coordinates to be its cell center. In the following, we will use 4 and E R to denote the type and coordinates of the i-th cell, respectively. The coordinates will be further pro- cessed in the following embedding layer as well as relative positional encoding in Transformer layers.
Uni Mal
5.6
Adding viruual points
Log Valid Loss
5.7
-5.8
-6.0
-6.1
50
100
150
200
Number of sampled points
Figure Log-scale valid loss of Uni-Mol when adding virtual points. The blue horizontal line represents the original Uni-Mol model, while the orange curve represents the Uni-Mol performance with respect to the number of virtual points.
Embedding. Given each cell represented as Cz), we process it with the following embedding layer. First, we compute the inner position of the atom (or cell center for non-atom cells) within the cell, defined as é; C2 mod e', followed by discretization e; = - where ct is hyper- parameter setting to very small value, i.e., ct 0.01A. In this way, the resulting e; is an integer vector with ele- ments in range 0, cm Next, we convert all discrete input
points is purely random and can be interpreted noise. In Sec. 4.3, we further demonstrate that the improved pretrain- ing loss translates to improved performance on downstream tasks, and notably, the improvement is not attributed to the (slightly) increased computational costs.
While the above experiment confirms that simply adding virtual points can already enhance 3D MPR, it further raises the question of whether more systematic approach might yield even better results. In the next section, we will explore this question by introducing a principled framework that directly models the entire 3D space.
features o; == (li, es) E N' into continuous feature repre- sentations by summing the corresponding embedding layers, i.e., 3; = Σi= Embed((a...). where Embed,(.) is the pa- rameterized embedding function that maps discrete inputs to continuous representations, and x; is the resulting input embedding for the i-th cell.
3.2. Overall of Our framework
Our approach is inspired by the standard input format in computer vision, where inputs are uniformly sampled from the continuous 2D space and represented as discrete pixel grids. This can be naturally applied to our setting by dis- cretizing 3D space into small, uniform grid cells. Compared to randomly sampling virtual points described in Sec. 3.1, the grid-based sampling is clearly better as it ensures uni- form coverage of the 3D space. To choose a proper grid size (resolution), we use a fundamental physical principle: two distinct atoms must maintain a minimum separation distance due to repulsive forces. Consequently, as long as the cell where d represents the minimum distance between any pair of atoms, the discretization will be accurate enough in the sense that each grid contains at most one atom. Note that d typically does not depend on specific tasks: for small organic molecules, is approxi- mately 0.96Å corresponding to the O-H bond length. We will fix dl = 0.49Å throughout this paper.
Architecture. Given the set of input features {x,}, we can process it using any model capable of handling sets, in particular, the Transformer architecture (Vaswani, 2017). We call the resulting model SpaceFormer. However, a naive application of the SpaceFormer suffers from significant com- putational costs, as its complexity scales quadratically with the number of grid cells, which can be substantial due to the discretization process. For instance, in widely used organic molecule datasets like ZINC (Sterling & Irwin, 2015), the average number of cells is approximately 8,000. Such high cost can be partially mitigated by using advanced attention implementation such as FlashAttention (Dao et al., 2022), which is adopted in this paper. Yet, it remains higher than that of previous atom-based methods.
length c' satisfies c' < .d V3*
In the subsequent subsections, we will propose concrete implementations that aim to significantly improves the effi- ciency of our proposed SpaceFormer.
Data augmentation. We employ two simple types of data augmentation. First, we randomly rotate the molecules prior to discretization to enhance the model's robustness to rotation. Second, after discretization, we randomly pad each face of the grid boundary by up to 2 grid cells, ensuring that the surrounding space of the atoms is further expanded to capture broader spatial context.
3.3. Grid sampling and merging
To reduce computational costs, a straightforward direction is to reduce the number of non-atom cells. We propose two intuitive strategies to achieve this goal, as presented below.
Grid sampling, In this strategy, we just randomly select a portion of non-atom grid cells and feed the selected cells
4

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
following, we propose two types of tinear-completty 3D positional encoding: the first directly encodes AB, which captures the raw directional information; the second encodes the geometric distance || 2. which is fundamental in- variant under coordinate system transformations.
6
2 Figure 3. Adaptive grid merging. Dark cells represent merged cells, and L denotes the number of merging steps.
L
1
3D directional PE with ROPE. Recently, Rotary Positional Encoding (RoPE) (Su et al., 2024) has gained popularity due its imner--omplexity in encoding relative positions, Here, we extend RoPE to 3D continuous space to encode directional information AB, capturing pairwise directional relationships across all three axes.
along with all atom cells into the model. Empirically, we find that this strategy already works well when sampling only 10%-20% non-atom cells.
The potential limitation of the above sampling strategy that does not incorporate funda- mental physical principles and involves tuning additional pypee-parameter. Specifically, in quantum physics, regions close to atoms typically exhibit higher electron density, and the density within these regions can vary significantly with position. Conversely, regions far from all atoms generally have electron densities approaching zero. Consequently, computational simulations often employ fine-grained sam- pling near atoms accurately capture these dynamic vari- ations, while coarse- grrssne sampling is used in distant regions to optimize computational efficiency (Parr et al., 1979). Drawing inspiration from this approach, we intro- duce an adaptive grid merging strategy, as illustrated in Fig. 3. Specifically, starting from the full grid cells, we group them by x 2 x blocks and merge each block of eight adjacent cells into "larger" cell if all of them are non-atom cells. This process can be executed recursively until convergence. After merging, the coordinate of each merged cell is still defined to be its cell center, while its type will further distinguish the cell size (or merging level). Empirically, we find that this strategy can reduce the number of grid cell by an order, while maintaining the performance without tuning any hyper-parameter.
Adaptive grid merging.
Recall that in RoPE, the Query and Key vector of each token is multiplied by 2D position-aware rotational matri- ces before performing dot-product attention, which has the following form:
(R(1)g) ((())) =q (R(i) R()) = =q. R(j - k)
where qi is the Query vector of the i-th token, ky is the Key vector of the j-th token, and R(i)) is the d x d rotational matrix with parameters 0, ... 0a/2 defined below:
cos i0, sin i0 sin i0;
cos i0
R(i)
cos i0 sin i0 sin ila cos itl a/2
To extend RoPE to 3D continuous space, we split each Query and Key vector into 3 blocks and apply rotational matrix separately for each block. Formally,
3.4. Efficient 3D relative positional encoding
(R(ca) R(c(2)
9:.1
R(caa)
Positional encoding (PE) critical in Transformer-basee models since the original Transformer architecture lacks in- ductive bias for its input structure. While absolute positional information is typically injected into the input, relative PE is equally important as it captures pairwise relationships between grid cells. However, naively integrating relative PE is computationally inefficient due to its quadratic mem- ory cost relative to the number of grid cells. This ineffi- ciency becomes particularly impractical when combined with FlashAttention.
R(ca2)
kj2
R(c,,s)
9:.3
R(c,,))
k3,
q. T R(c,, -c,, kkj))
In this way, the directional information of each axis is en- coded in the Transformer layer via
3D distance PE with RFF. The geometric distance is fundamental measure of pairwise interactions in quantum mechanics. To overcome the quadratic memory overhead associated with computing pairwise distances, we propose using random Fourier features (RFF) to approximate the Gaussian kernel (Rahimi Recht, 2007), which has the
To address this issue, we propose efficient positional encoding method tailored for continuous 3D coordinates. Given two points A, B in 3D space, denote their coordinates as CA and CB. respectively. The relative positional informa- tion can be defined as the vector AB- == CB CA- In the
5

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
following form:
(e.g.. Zaidi et al., 2022; Zhou et al., 2023). We argue that our MAE-based pretraining could be more advantageous, be- cause it involves a stronger training task that helps the model acquire additional knowledge during pretraining. Specifi- cally, the prediction task in MAE can be decomposed into two components: the first involves predicting whether an atom exists within a certain range in 3D space, while the sec- ond refines the exact coordinates relative to the grid center. One can see that the second task is exactly a form of "de- noising". whereas the first task is absent in denoising-based approaches. Moreover, unlike denoising, which perturbs certain atoms, we remove them entirely, making the pretrain- ing task even harder. We hypothesis that these differences contribute to the improved performance, as will be demon- strated in the following section.
exp (- |c0 - c-l12) T(()),
20
2(c,) = Vā 2 cos
c (b))
o
w E p2xd I), be Rd ~ u(0 2T)"),
where controls the bandwidth of the Gaussian kernel, each element of w is i..d. sampled from standard normal distribution, and b is sampled from uniform distribution over [0, )T.
The random Fourier features are then combined with the Query and Key after applying the RoPE:
i = gg =T. k2 = (c,)))
4. Experiments
Here, the concatenation ensures that the effect of RoPE and RFF are separated as q, k, = q. k2 + 2(c,)" z(e,).
To evaluate our proposed SpaceFormer, we first conduct unsupervised pretraining on large-scale unlabeled data, and then fine-tune on various tasks with limited labeled data. Moreover, we perform extensive ablation studies to assess the contribution of its key components. Finally, we compare the effectiveness and efficiency of Space with atom- based MPR models that incorporate virtual points.
3.5. Pretraining
A key advantage of our proposed framework is that it ad- mits a powerful and efficient pretraining approach known as Masked Auto-Encoder (MAE) (He et al., 2022). In MAF, we randomly mask a portion of (either atom or non-atom) grid cells and only feed the remaining cells into Space- Former to obtain the representation of these cells. Then, we use another decoder architecture that takes the repre- sentations along with all masked cells as inputs to predict their type and coordinates. The MAE approach is highly efficient because (1) the encoder operates only on unmasked cells, and (2) the decoder, which takes all cells as inputs, is significantly smaller than the encoder.
4.1. Settings
Pretraining settings. We use the same pretraining dataset as Zhou et al. (2023), which contains a total of 19 million molecules. Details of the pretraining settings are provided in Table 5 in the Appendix. For grid merging, we set the merging level to 3, which corresponds to convergence. This configuration results in a model with approximately 67.8M (encoder) parameters and requires about 50 hours of training using NVIDIA A A100 GPUs
Implementation details. While the above pretraining framework is straightforward, several implementation de- tails distinguish it from image pretraining, which we high- light below. First, we must handle the prediction of atom and non-atom cells differently. Note that the input coordi- nates of masked cells are defined by their center positions. Therefore, if cell is predicted to be an atom cell, the model must additionally predict the offset relative to the cell center. Moreover, since the numbers of atom and non-atom cells are highly imbalanced, we assign different loss weights to the prediction of each cell type to ensure their overall con- tributions are balanced. Second, caution is required when combining MAE with adaptive grid merging, as masking merged grids does not make sense (since it can be trivially predicted). Instead, the masking process is applied prior to adaptive grid merging, and all masked cells are excluded thereafter in the encoder.
Baseline models. Our primary baseline is the Uni-Mol (Zhou et al., 2023), recent 3D MPR model that achieved state-of-the-art performance across a wide range of molecu- lar property prediction tasks. We use the same pretraining dataset as Uni-Mol to enable a fair comparison. We also include Mol-AE (Yang et al., 2024), which extends Uni-Mol with (atom-based) MAE pretraining. For more compre- hensive comparison, we further include 3D Infomax (Stärk et al., 2022), a GNN based representation incorporating 3D information, and two 2D graeh MPR models: GROVER (Rong et al., 2020) and GEM (Fang et al., 2022).
Downstream tasks. We develop new benchmark to com- prehensively evaluate MPR models!, particularly focusing
Most prior works follow MoleculeNet (Wu et al., 2018) for downstream task evaluation. However. recent studies (Walters, 2023) identified several limitations of this dataset, including the presence of invalid structures, inconsistent chemical representa-
Comparison with prior pretraining pipelines. In the liter- ature, most prior MPR approaches are based on denoising
6

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Table 1. Performance on molecular computational property and experimental properties prediction tasks. The best results are highlighted in bold, and the second best results are underlined.
Task
#Samples
GROVER
GEM
3D Infomax
Uni-Mol
Mol-AE
SpaceFormer
HOMO (Hartree) +
20,000
0.0075
200-4
0.0068 # 7.0-5
0.0065 # 1.0e55
0.0052 * 2.0e55
0.0050 # 8.00-5
0.0042 #: 1.0c-5
LUMO (Hartree)
20,000
0.0086
8.0e4
0.0080
2.0e5
0.0070 + 1.0c4
0.0060 + 6.0c-5
0.0057 + 4,70-4
0.0040
2.0e-5
GAP (Hartree) +
20,000
0.0109
0.0107
1.9e-
0.0095 + 1.0c4
0,0081 4.0c-5
0.0080
8.00-5
0.0064
1.2c-4
E1-CC2 (eV)
21,722
0.0101
9.70-4
0.0090 1.3e-4
0.0089 # 2.0c-4
0.0067 + 40k-5
0.0070
= 6.00-5
0.0058
8.0c-5
E2-CC2 (
21,722
0.0129
=4.60-4
0.0102 # 2.3e-4
0.0091 + 3.0e-4
0.0080 + 4.0c-5
0.0080
# 400-5
0.0074 8.4c-
f1-CC2 4
21,722
0.0219-
3.5e-4
0.0170 4.3e-
0.0172 + 4.0c-4
0.0143 + 2.0c-4
0.0140 = 4.00-5
0,0142
3.7ca-
12-CC2 4
21,722
0.0401
0.0352 4e--
0.0364 + 9.0c-
0,0309 + 9.4.4.4
0.0307 1.30-3
0.0294
7.14
Dipmom (Debye) 4
8,678
0.0752 = 1.1e-3
0.0289 1.3a-3
0.0291 + 1.7o-3
0.0106 + 3.10-4
0.0113 = 4.7044
0,0083
5.0c-4
alp (eV) 4
8,678
0.1467 1.5p-2
0.0207 = 2.6-4
0.0526 + Ae-
0.0095 + 6.Ac-4
0.0103 # 1.30-4
0.0090
= 5.9c-4
D3 disp coTr (eV) 4
8,678
0.2516
2.2e
0.0077 = 6.6e 4
0.2285 + 7.5e
0.0047 1: 5.fc
0.0077 + 1.30-3
0.0053
1.2e
HLM L
3,087
0.4313
2
0.3319 3
0.3333 +: 9,0e-3
0.3100 4 M.2e-
0.3127
1.5e-2
0.3098
7.46
MME
2,642
0,4438
1.7e22
0.3132
1.4e-2
0.3275 + 3.6e-
0.2901 + 3.0e-3
0.2947
1.te22
0.2933
6.8e-3
Solu I
2,173
0,2660 = 1.8e22
0.2290
# 3.1e:3
0.4535 # 1.4e41
0.3388
# 7.1e-3
0.2320
# 1.2e-2
0.2205
1.No-3
BBBP 1
1,965
0.8672
20e-2
0.8863
+ 2.5e:2
0,8276
+ 2.7e-2
0,9066 + 2.3e-2
0.8995
# 4.60-3
0.8605
2.2e-2
BACE
1,513
0.9053 2 0.9851
2.8e:3
0.9612 + 7
0,9523 + 2.0-2
0.9605
1.3e-2
0.9657
3.1e22
on the setting of limited data. Our benchmark comprises two categories of tasks: molecular computational properties and molecular experimental properties. For computational properties, we sample a 20K subset from the huge dataset of GDB-17 (Ramakrishnan et al., 2014; Ruddigkeit et al., 2012) and select the electronic properties HOMO, LUMO and GAP. Additionally, we another 21k subset from the same dataset following (Ramakrishnan et al., 2015), select- ing the energy properties E1-CC2, E2-CC2, f1-CC2 and 12-CC2. Furthermore, we incorporate the dataset (Wahab et al., 2022) of cata-condensed polybenzenoid hydrocar- bons comprising from to 11 rings which contain 8k sam- ples and selected the mechanical and electronic properties within and between molecules, like Dipmom, alP and D3 Dispersion Corrections (D3 DD3_disp_corr) For experimental properties, we select the BBBP and BACE datasets from MoleculeNet, ensuring that all duplicate and structurally invalid molecules were excluded. Additionally, we employ the HLM, MDRI-MDCK ER (MME), and Solubility (Solu) datasets from the Biogen ADME dataset (Fang et al., 2023). A detailed description of these tasks is provided in Table 7 in the Appendix. In all tasks, datasets were split into training, validation, and test sets in an 8:1:1 ratio. We applied the -f--sribiitittiinn splitting methods, where the sets are divided based on scaffold similarity. This eventually results in 15 tasks, allowing for a thorough evaluation. We ensure that the syper-parametee search space is consistent across all tasks and baseline models (see Table 6 in the Appendix). For each set of
4.2. Results
Table present our experimental results. The results clearly demonstrate SpaceFormer's ss superior performance, which ranks the first in 10 out of 15 tasks and within the top two in 14 out of 15 tasks. In particular, SpaceFormer signifi- cantly outperforms all baselines in computational properties such as HOMO, LUMO, GAP, E1-CC2 and Dipmom tasks, where it surpasses the runner-up models by about 20%. This verifies the effectiveness of modeling 3D space beyond atom positions and justifies the design of SpaceFormer.
Space has limited effectiveness on the BBBP task, one reason can be the complexity and variability of the mea- surement environments for experimental properties, which can be easily influenced by external factors, like tempera- ture and blood-brain barrier biofilm status for BBBP. Such variability can affect the stability of measurement results and may not accurately. Additionally, as mentioned in the footnote of Section 4.1, the MoleculeNet dataset for the experimental task (including BBBP) has several limitations, including invalid structures and inconsistent chemical repre- sentations, which affect its ability to differentiate molecular pretraining models.
4.3. Ablation Studies
In this subsection, we conduct a series of experiments to evaluate the proposed components of SpaceFormer. We choose the HOMO, LUMO, E1-CC2, and E2-CC2 proper- ties throughout all ablation experiments.
we train each model 3 times using different random seeds and report the the aver- age result and standard deviation. The checkpoint with the best validation loss is selected for all experiments.
The effect of grid sampling/merging. As discussed in Sec. 3.3, we propose using either random sampling or adap- tive grid merging for non-atom cells to reduce training costs. Here, we evaluate the efficiency and performance of both approaches by varying the number of input cells. For adap- tive grid merging, we experiment with different merging levels L. A level of L = 0 indicates no merging (using
tions, and data curation errors. Morcover, Sun et al. (2022) pointed out that MoleculeNet fails to adequately distinguish the perfor- mance of different molecular pretraining models. Instead, this paper proposes new benchmark to address these problems.
7

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Table 2. Ablation studies on grid sample/merging strategy.
Uni-Mal Spacefarmer
Pre-train cost
600
Default setting
No.
Sample Ratio
HOMO
LUMOL EI-CC2) E2-CC2L
Pretraining
cost
#cells
avg.
400
1.0
0.0043
0.0041
0.0057
0.0073
656h
8.4K
0.0044
0.0041
0.0059
0.0073
122h
2.2K
0.0042
0.0043
0.0058
0.0075
56h
1.1K
200
0.0042
0.0040
0.0058
0.0074
50h
1.0K
0.27
0.0044
0.0042
0.0058
0.0077
122h
2.2K
0
0.13
0.0044
0.0044
0.0058
0.0075
56h
LIK
200
400
600
800
1000
1200
0.12
0.0044
0.0044
0.0058
0.0076
50h
LUK
Number of sampled cells/points
8
0.096
0.0044
0.0045
0.0058
0.0072
44h
800
Figure 4. Efficiency comparison of SpaceFormer and Uni-Mol with different numbers of cells/points.
0.048
0.0046
0.0047
0,0060
0.0075
28h
400
10
0.024
0.0047
0.0047
0.0062
0.0077
25h
200
Table 3. Ablation studies on 3D positional encoding.
4.4. Further analysis
No. ROPF
RFF
HOMO
LUMO + EI-CC2 4
E2-CC2
In this subsection, we perform a more in-depth analysis of SpaceFormer via a set of carefully designed experiments, as presented in Table 4. These experiments can offer deeper insights into the underlying source of SpaceFormer's ' perfor- mance gains and further elucidate our major contributions.
v
0.0042
0.0040
0.0058
0.0074
V
X
0.0044
0.0043
0,0059
0,0073
x
X
0,0047
0,0048
0.0066
0.0078
the full set of grid cells), while L indicates merging until convergence. For random sampling, we test different sampling ratios, ensuring that the number of sampled cells matches that of merged cells at different levels for a direct comparison of the two approaches. The results are presented in Table 2, where we can draw the following conclusions:
Base settings. We train a purely "atom-based" SpaceFFormer without using any non-atom cells and refer to this setting as S1. Then, we add two settings S2 and S3 according to Table 2 No. 10 and No. 4, respectively (we use "S" instead of "No." t to distinguish from previous tables).
(No. 1-4) As the number of merging levels increases, the model's efficiency improves significantly while the per- formance is relatively stable. In particular, the merging- until-convergence strategy achieves the best efficiency without compromising accuracy.
Regarding FLOPs. We first examine whether the perfor- mance improvement stems solely from increased compu- tational costs (FI.Oss). To achieve this, train a signif- icantly larger SpaceFormer model by scaling the model width by 4x from the baseline setting S1, resulting in set- ting S4. The model width is adjusted to ensure that the training speed becomes slightly slower than that of setting S3, As shown Table 4, although enlarging the model size leads to notable performance gains, the results of S4 still fall short of S3 on most downstream tasks. This indicates that the observed performance improvements are not merely consequence of higher computational costs.
2. (No. 5-10) As the sampling ratio decreases, the the model's efficiency improves significantly. However, when the number of sampled cells is less than IK, the performance begins to drop quickly. The best tradcoff between efficiency and accuracy happens when sam- pling 800-1K OO grid cells.
Can absed models achieve similar performance by simply using virtual points? We answer this question by conducting a set of experiments using the atom-based model Uni-Mol, while varying the number of virtual points as described in Sec. 3.1. This corresponds settings S6- S10 in Table 4. We can see that while the introduction of virtual points can improve the model performance on several tasks, there is still significant gap between UniMol and SpaceFormer, so our improvement should not be simply attributed by using virtual points. The next experiment will further highlight another difference: the use of MAE.
3. (No. 2-7) Overall, the adaptive merging strategy per- forms slightly better than the random sampling strategy, though the difference is not significant. However, the key advantage of adaptive merging is its ability adap- tively adjust the number of grid cells without the need to the tune the sampling ratio.
3D positional encoding. In this section, we evaluate the contributions of the two 3D positional encoding methods proposed in S.c.3.4. Specifically, we consider two settings: (I) using only RoPE (No. 2), and (2) removing both PE (No. 3). For the latter setting, positional information is incorporated by simply adding the linear projection of the 3D position c to the input embeddings x; The results are presented in Table 3. We can clearly see that both 3D positional encodings play a crucial role in enhancing the model performance.
MAE vs. denoising. As described in Sec. 3.5, MAE corresponds to a stronger pretraining task than denoising. To see the difference empirically, we design a variant of MAE where only atom cells are masked, so the model will not need to predict whether a masked cell contains atom.
8

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
biomolecular interactions with alphafold 3. Nature pp.
Table Comparison between SpaceFormer and Uni-Mol adding
1-3, 2024.
VPs. See Sec. 4.4 for details.
No.
Model
non-atom cells/VPs
HOMOL LUMO! E1-CC2) E2-CC2
Atkins, P. W. and Friedman, R. S. Molecular quantum
mechanics. Oxford University Press, USA, 2011.
SpaceFormer
0,0051
0.0052
0.0065
0.0077
SpaceFormer
200
0,0047
0.0047
0,0062
0.0077
Cui, T. Tang, C., Su, M., Zhang, S., Li, Y., Bai, L.,, Dong, Y. Gong, X., and Ouyang, W. Geometry-enhancee pre- training interatomic potentials. Nature Machine Intel- ligence, 6(4):428 (14::4284436,, 2024.
3
SpaceFormer
1000
0.0042
0.0040
0.0058
0.0074
Spucel'ornerr
0.0046
0.0045
0.0060
0.0073
Large
SpaceF-onner-
5
1000
0.0044
0.0041
0.0059
0.0075
denoising
Uni-Mol
0.0052
0.0060
0.0067
0.0080
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAt- tention: Fast and memory-efficiet exact attention with IO-awareness. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
7
Uni-Mol
10
0.0052
0.0058
0.0063
0.0078
8
Uni-Mol
50
0.0053
0.0056
0.0063
0.0079
9
Uni-Mol
100
0.0053
0.0057
0.0065
0.0080
10
Uni-Mol
200
0.0052
0.0058
0.0067
0.0080
Darcet, T., Oquab, M., Mairal, J., and Bojanowski, P.
Vision transformers need registers. arXiv preprint arXiv:2309.16533 2023.
This corresponds to setting S5 in Table 4. We can see that S3 outperforms S5 in all four downstream tasks, thus confirming the benefits of using MAE pretraining.
Fang, C., Wang, Y., Grater, R., Kapadnis, S., Black, C., Trapa, P., and Sciabola, S. Prospective validation of machine learning algorithms for absorption, distribution, metabolism, and excretion prediction: An industrial per- spective. Journal of Chemical Information and Modeling, 63(11):3263-33
Scalability with the number of cells. Finally, we show that owing to our proposed 3D PE, the model enjoys much better scalability than UniMol with the increase of cells. As illustrated in Fig. 4, the pretraining cost of UniMol scales quadratically to the number of sampled cells, while the curve of SpaceFormer tends to have a linear shape. This verifies the efficiency of our proposed limear-comllextty PE.
2023.
Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang. Geometry-enhance molecular representation learning for property prediction. Nature Machine Intelligence, 4(2):127-134, 2022.
5. Conclusion
In this paper, introduce SpaceFormer, novel MPR framework that incorporates the 3D space beyond atomic positions to enhance molecular representation. To effec- tively and efficiently process 3D space, SpaceFormer is composed of 3 key components: grid-based space discretiza- tion, grid sampling/merging and efficient 3D positional en- coding. The resulting framework can be easily integrated with MAE training. Extensive experiments validate the ef- fectiveness and efficiency of SpaceFormer across various tasks. In future research, a promising direction is to extend SpaceFormer to force field tasks and larger systems, such as proteins and complexes.
Feng, S., Ni, Y., Lan, Y., Ma, Z.-M., and Ma, W.-Y. Frac- tional denoising for 3d molecular pre-training In Interna- tional Conference Machine Learning, pp. 9938-999090 PMLR, 2023.
Gasteiger, J., GroB, J., and Günnemann, S. Directional message passing for molecular graphs. arXiv preprint MrXiv:2003.03123: 2020.
Gasteiger, J., Becker, F,, and Günneeaan, S. Gemnet: Uni- versal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34: 6790-6802, 2021.
Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chem- istry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
He, K., Chen, X., Xie. S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders are scalable vision learners. In Proceedings of the IEFE/CVF conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
References
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T.,
Heller, S. R., MeNaught, A., Pletnev, I., Stein, S., and Tchekhovskoi, D. Inchi, the iupac international chemical identifier. Journal of cheminformatics, 7:1-34, 2015.
Pritzel, A., Ronneberger, O.,. Willmore, L., Ballard, A. J., Bambrick, J., al. Accurate structure prediction of

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Hu, W., Liu, B., Gomes, J., Zitnik, M., Liang, P., Pande, V., and Leskovec, J. Strategies for pre-training graph neural networks. arXiv preprint 1903: 2019.
Ramakrishnan, R., Hartmann, M., Tapavicza, E., and Von Lilienfeld, O. A. Electronic spectra from tddft and machine learning in chemical space. The Journal of chem- ical physics, 143(8), 2015.
Jiao, R., Han, J., Huang, W., Rong, Y., and Liu, Y. Energy- motivated equivariant pretraining for 3d molecular graphs. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 37, pp. 8096-8106 2023.
Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W., and Huang, J. Self-supervised graph transformer on large- scale molecular data. Advances in neural information processing systems, 33:12559-12571 2020.
Jumper, J.. Evans, R., Pritzel, A., Green, T., Figurnov. M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zidek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-588) 2021.
Ruddigkeit, I.., Van Deursen, R., Blum, I. C., and Reymond, J.-L. Enumcration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(1 (2111):2864225) 2875, 2012.
Li, P., Wang, J., Qiao, Y., Chen, H., Yu, Y., Yao, X., Gao, P., Xie, G., and Song. S. An effective self supervise frame- work for learning expressive molecular global representa- tions to drug discovery. Briefings in Bioinformatics, 22 (6):bbab109 2021.
Schütt, K., Kindermans, P.-J., Sauceda Felix, H. E., Chmiela, S. Tkatchenko, A., and Müller, K.-R. Schnet: A continuous- filter convolutional neural network for model- ing quantum interactions. Advances in neural information processing systems, 30, 2017.
Liu, S., Wang. H., Liu, W., Lasenby, J., Guo. H., and Tang, J. Pre-training molecular graph representation with 3d geometry. arXiv preprint arXiv:2110.077288 2021.
Shi, C., Luo, S., Xu, M., and Tang, J. Learning gradient fields for molecular conformation generation. In Interna- tional conference on machine learning, pp. 9558-9568. PMLR, 2021.
Liu, Y., Wang. L.,, Liu, M., Lin, Y., Zhang, X., Oztekin, B., and Ji, S. Spherical message passing for 3d molec- ular graphs. In International Conference on Learning Representations (ICLR), 2022.
Song, Z., Wei, H., Bai, L.,, Yang, L., and Jia, C. Graphalign: Enhancing accurate feature alignment by graph matching for multi-modal 3d object detection. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 3358-3369, 2023a.
I Luo,5 S., Chen, T., Xu, Y., Zheng, S., Liu, T.-Y,, Wang, I.., and He, D. One transformer can understand both 2d & 3d molecular data, In The Eleventh International Conference on Learning Representations, 2022.
Song, Z., Wei, H., Jia, C., Xia, Y., Li, X., and Zhang, C. Vp-net: Voxels as points for 3-d object detection. IEEE Transactions on Geoscience and Remote Sensing, 61:1- 12, 2023b.
Mahmoud, A., Hu, S., and Waslander, S. L. Dense voxel fusion for 3d object detection. In Proceedings of the IEFF/CVF winter conference on applications of computer vision, pp. 663-672, 2023.
Stärk, H., Beaini, D., Corso, G., Tossou, P., Dallago, C., Günnemann, S., and Liò, P. 3d infomax improves gnns for molecular property prediction. International Con- ference on Machine Learning, pp. 20479-20502. PMLR, 2022.
Parr, R. G., Gadre, S. R., and Bartolotti, J. Local density functional theory of atoms and molecules. Proceedings of the National Academy of Sciences, 76(6):2522-25266 1979.
Sterling, T. and Irwin, J. Zinc 15-ligand discovery for everyone. Journal of chemical information and modeling, 55(11):2324-25) 2015.
Pfau, J., Merrill, W. and Bowman, S. R. Let's think dot by dot: Hidden computation in transformer language models. arXiv preprint arxiv:2404. 15758 2024.
Su, J., Ahmed, M., Lu. Y., Pan. S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:1 127063, 2024.
Rahimi, A. and Recht, B. Random features for large-scale kernel machines. Advances in neural information pro- cessing systems, 20, 2007.
Sun, R., Dai, H., and Yu, A. W. Does gnn pretraining help molecular representation? Advances in Neural Informa- tion Processing Systems, 35: 1501669991209, 2022.
Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Szabo, A. and Ostlund, N. S. Modern quantum chem- istry: introduction to advanced electronic structure the- ory. Courier Corporation, 2012.
10

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Thomas, N., Smidt, T., Kearnes, S., Yang, L.,, Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and cansslatio-equivarian neural networks for 3d point clouds. arXiv preprint arXiv:1802.08.. 2018.
Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Ge- niesse, C., Pappu, A. S., Leswing, K., and Pande, V. Moleculenet: a benchmark for molecular machine learn- ing. Chemical science, 9()):513-530 2018.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Xu, M., Luo, S., Bengio, Y. Peng, J., and Tang, J. Learning neural generative dynamics for molecular conformation generation. arXiv preprint 11Xx::2102.102:000 2021.
Wahab, A., Pfuderer, L., Paenurk, E., and Pershoni-Poranne R. The compas project: A computational database of polycyclic aromatic systems. phase I: cata-condensed polybenzenoid hydrocarbons. Journal of Chemical Infor- mation and Modeling, 62(16):3704 37133 2022.
Xu, M., Yu, L., Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular con- formation generation. arXiv preprint arXiv:2203.085 2022.
Wallach, I., Dzamba, M., and Heifets, A. Atomnet: deep convolutional neural network for bioactivity predic- tion in trructurebassd drug discovery. arXiv preprint arXiv:1510.02885 2015.
Xu, Z., Wang. S., Zhu, F,, and Huang. J. Seq2seq fingerprint: An unsupervised deep molecular embedding for drug discovery. In Proceedings of the 8th ACM international conference on bioinformatics, computational biology, and health informatics, pp. 294, 2017.
Walters, P. We need better benchmarks for ma- chine learning in drug discovery, 2023. URL attps://ppee http://ppactiialcceemmffoooooo we -e--eedd-betee---encheeee html. Accessed: 2024-09-26.
Yang, J., Zheng, K., Long, S., Nie, Z., Zhang, M., Dai, X., Ma, W.-Y., and Zhou, H. Mol-ae: Auto-encoder based molecular representation learning with 3d cloze test objective. bioRxiv, pp. 2024 04, 2024.
for machine
Wang, L., Liu, Y., Lin, Y., Liu, H., and Ji, S. Comenet: Towards complete and efficient message passing for 3d molecular graphs. Advances in Neural Information Pro- cessing Systems, 35:650-664, 2022a.
Yin, T., Zhou, X., and Krähenbühl, P. Multimodal virtual point 3d detection. Advances Neural Information Pro- cessing Systems, 34:1 34:16494 14446464-66667,, 2021.
Wang, S., Guo, Y., Wang, Y., Sun, H., and Huang. J. Smiles- bert: large scale unsupervised pre- training for molecular property prediction. In Proceedings of the 10th ACM in- ternational conference on bioinformatics, computational biology and health informatics, pp. 429-436, 2019.
Yu, Q., Zhang, Y., Ni, Y., Feng, S., Lan, Y.. Zhou, H., and Liu, J. Multimodal molecular pretraining via modality blending. In The Twelfth International Conference on Learning Representations, 2024.
Zaidi, S., Schaarschmidt, Y. W., Sanchez-Gonzalez, A., Battaglia, P., Pascanu, R., and Godwin, Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.00133 2022.
M., Martens, J., Kim, H., Teh,
Wang, Y., Wang, J., Cao, Z., and Barati Farimani, A. Molec- ular contrastive learning of representations via graph neu- ral networks. Nature Machine Intelligence, 4(3):279-287, 2022b.
Wang, Y., Xu, C., Li, 7.,, and Barati Farimani, Denoise pretraining on nonequilibrium molecules for accurate and transferable neural potentials. Journal of Chemical Theory and Computation, 19(15)557775508, 2023.
Zee, A. Quantum theory in a nutshell, volume 7. Prince- ton university press, 2010.
Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z., Zhang, L., and Ke, G. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh Inter- national Conference on Learning Representations, ICLR 2023, Kigali. Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https::/ /ttp:r////p/eeeeeee... ftw.:rruumeeetem i Ld=622RM6WVqKu.
Wang, Y., Cheng, C., Li, S., Ren, Y., Shao, B., Liu, G., Heng, P.-A., and Zheng, N. Neural p²m: A long-range interaction modeling enhancer for geometric gnns. arXiv preprint arXiv:2409. 17622, 2024.
Weinberg, S. The quantum theory of fields, volume 2. Cam- bridge university press, 1995.
Wu, H., Wen, C., Shi, S., Li, x., and Wang, C. Virtual sparse convolution for multimodal 3d object detection. In Proceedings of the IEFE/CVF Conference on Computer Vision and Puttern Recognition, pp. 21653-21662. 2023.
Zhu, H., Deng, J., Zhang, Y., Ji, J., Mao. Q., Li, H., and Zhang, Y. Vpfnet: Improving 3d object detection with virtual point based lidar and stereo data fusion. JEEE Transactions on Multimedia, 1:5291-5330 2022a.

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Zhu, J., Xia, Y., Liu, C., Wu, L., Xie, S., Wang, Y., Wang. T., Qin, T., Zhou, W., Li, H., et al. Direct molecular con- formation generation. arXiv preprint 2022b.
12

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
A. Experiment Details
The pretraining settings are detailed Table the downstream finetuning settings in Table 6, and the downstream tasks in Table
Table Pretraining Settings
Hyper-parameters
Value
Peak learning rate
le-4
LR scheduler
Linear
Warmup ratio
0.01
Total updates
1M
Batch size
128
Residual dropout
0.1
Attention dropout
0.1
Embedding dropout
0.1
Encoder layers
Encoder attention heads
Encoder embedding dim
512
Encoder FFN dim
2048
MAE-Decoode layers
4
MAE-Decoder attention heads
4
MAE-Decoder embedding dim
256
MAB-Decoder FFN dim
1024
Adam (31, (2)
(0.9, 0.99)
Gradient clip
Mask ratio
0.3
Cell edge length df
0.49A
Om for in-cell position discretization
0.01A
Merge level
3
Table 6. Finc-tuning Settings
Hyper-parameten parameters
Value
Peak learning rate
[5e-5, le-4]
Batch size
[32, 64]
Epochs
200
Pooler dropout
[0.0, 0.1]
Warmup ratio
0.06
B. More Experiments
B.1. Benchmark on QM9 dataset
We have expanded our evaluation include the HOMO, LUMO, and GAP tasks the QM9 dataset, which is standard in molecular property prediction. Our model demonstrates better performance compared with the baselines, detailed in Table 8.
13

---

Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Table 7. Summary information of the downstream datasets
Category
Task
Task type
Metrics
Samples
Describe
The highest energy molecular orbital that is occupied 20,000
HOMO
Regression
MAE
by electrons
The lowest energy molecular orbital that is not occupied
LUMO
Regression
MAR
20,000
by electrons
The energy difference between the HOMO and LUMO
GAP
Regression
MAE
20,000
The energy of the first excited state computed using the CC2 method.
Computational Properties
E1-CC2
Regression
MAE
21,722
The energy the second excited state computed using the CC2 method.
E2-CC2
Regression
MAE
21,722
The free energy of the first excited state computed using the CC2 method.
f1-CC2
Regression
MAE
21,722
The free energy the second excited computed using the CC2 method.
F2-CC2
Regression
MAE
21,722
Dipmom
Regression
MAE
8,678
Dipole moment
Regression
MAE
8,678
Adiabatic ionization potential
D3 doid
Regression
MAE
8,678
D3 dispersion corrections
HLM
Regression
MAE
3,087
Human liver microsome stability
MME
Regression
MAE
2,642
efflux ratio
Experimental Properties
Solu
Regression
MAE
2,713
Aqueous solubility
BBBP
Classification
AUC
1,965
Blood brain barrier penetration
BACE
Classification
AUC
1,513
Binding results of human BACE-1 inhibitors
Table 8. Benchmark on QM9 dataset
Model
HOMO
LUMO
GAP 1
3D Informax
0,000952
0.000794
0.00155
Uni-Mol
0.000857
0.000763
0.00151
Spacefoormrr
0.000594
0.000544
0.00106
14


Reviews:
[{"review_id": "vomRpeCuJI", "reviewer": "Reviewer_sLbT", "summary": "This paper tackles molecular pretrained representation (MPR) by arguing that previous methods which focused solely on atom positions and types miss crucial physical context in the surrounding 3D space. Motivated by physical principles (e.g., the presence of electron densities and electromagnetic fields), the authors first demonstrate that simply adding randomly sampled virtual points can boost performance. Building on this observation, they propose SpaceFormer, a Transformer-based framework that discretizes the entire 3D space into grid cells, employs grid sampling and adaptive grid merging to manage computational cost, and integrates efficient 3D positional encoding (extending RoPE and using random Fourier features). The model is pre-trained via a Masked Autoencoder (MAE) strategy and is shown to outperform several baseline MPR models on diverse downstream tasks covering both computational and experimental molecular properties.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Regarding \"The experimental downstream tasks should also include widely used benchmarks such as QM9 and MD17, which are standard in molecular property prediction. Moreover, since grid/mesh-based methods are particularly adept at capturing long-range interactions [1, 2], the inclusion of the MD22 dataset would further validate the model's ability.\"**\n\nThank you for your insightful feedback and for highlighting the importance of these benchmarks. Following your comments, we have expanded our evaluation to include the HOMO, LUMO, and GAP tasks on the QM9 dataset, where our model demonstrates better performance, as detailed in the following table. \n|      Model    |      HOMO     |    LUMO       |      GAP      |\n|---------------|---------------|---------------|---------------|\n|  3D Infomax   |   0.000952    |    0.000794   |    0.00155    |\n|     Uni-Mol   |   0.000857    |    0.000763   |    0.00151    |\n|   SpaceFormer |   0.000594    |    0.000544   |    0.00106    |\n\nWe also attempted to evaluate our approach on the stachyose molecule from the MD22 dataset, following your suggestion and the references provided. This would further validate our model’s capability in capturing long-range interactions. However, the MD22 dataset primarily focuses on the prediction of energy and force, while our current model framework is geared towards property prediction. Adapting our model for force prediction requires additional development time. We kindly ask for your understanding and additional time to complete these experiments. We are committed to providing the results in the next discussion phase.\n\nWe greatly appreciate your constructive critique, which has strengthened our analysis. Thank you for your understanding and support.\n\n**2.  Regarding \"Since the grid-based methods are very time-consuming, what is the time consumption of Spaceformer and UniMol-AE (w/o grids)?\"**\n\nThanks very much for your question. We would like to clarify that the grid-based methods employed in this paper are optimized to be not very time-consuming. We have implemented two key strategies to enhance efficiency.\n\nFirstly, by merging or sampling grid cells, we significantly reduce the number of grid cells without compromising the model's performance. This is detailed in Table 2, which presents the number of grid cells after merging or sampling alongside the corresponding model performance.\n\nSecondly, our model framework is designed to handle long sequences efficiently. We employ FlashAttention and two effective 3D relative positional encodings, ensuring that the computational time increases almost linearly with sequence length, as illustrated in Figure 4.\n\nAdditionally, Mol-AE builds upon Uni-Mol with atom-based MAE pretraining, where only 15% of atoms are dropped for reconstruction during pretraining, resulting in negligible speed differences between Mol-AE and Uni-Mol. We compare the time consumption of Spaceformer and Uni-Mol in Figure 4, which demonstrates our model framework's superior efficiency for processing long sequences.\n\n**3. Regarding \"More grid-based approaches should be discussed:\n[1] Ewald-based long-range message passing for molecular graphs.\" International Conference on Machine Learning.\n[2] Neural P3M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs\"**\n\nThank you for your valuable feedback. We appreciate your insightful suggestions and will incorporate a discussion on these methods in our paper. \nThis addition will help provide a more comprehensive analysis of long-range interaction modeling in molecular graphs. We sincerely appreciate your guidance in improving our work.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "The experimental downstream tasks should also include widely used benchmarks such as QM9 and MD17, which are standard in molecular property prediction. Moreover, since grid/mesh-based methods are particularly adept at capturing long-range interactions [1, 2], the inclusion of the MD22 dataset would further validate the model's ability.\n\nSince the grid-based methods are very time-consuming, what is the time consumption of Spaceformer and UniMol-AE (w/o grids)?\n\n[1] Kosmala, Arthur, et al. \"Ewald-based long-range message passing for molecular graphs.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2] Wang, Yusong, et al. \"Neural P $^ 3$ M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs.\" The Thirty-eighth Annual Conference on Neural Information Processing Systems.", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Introduce a grid-based molecule pretraining method", "essential_references_not_discussed": "More grid-based approaches should be discussed:\n\n[1] Kosmala, Arthur, et al. \"Ewald-based long-range message passing for molecular graphs.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2] Wang, Yusong, et al. \"Neural P $^ 3$ M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs.\" The Thirty-eighth Annual Conference on Neural Information Processing Systems.", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "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": "ELYDaoOKvy", "reviewer": "Reviewer_f27K", "summary": "This paper introduces SpaceFormer, a Transformer-based molecular pretrained representation (MPR) model that explicitly encodes the 3D space surrounding molecules, going beyond traditional atom-only representations. The method discretizes the 3D space into grid cells, applying adaptive grid merging for efficiency, and employs masked autoencoder (MAE) pretraining. Extensive experiments demonstrate significant performance improvements over baselines (e.g., Uni-Mol [Zhou, 2023]) across 15 molecular property prediction tasks (Table 1). Ablations confirm the effectiveness of each component (Tables 2, 3, 4).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Regarding \"Minor weaknesses include no detailed runtime analysis and limited discussion on tasks where improvements were small (e.g., BBBP).\" and \"Why does SpaceFormer slightly underperform on BBBP compared to simpler baselines? Could it indicate an overfitting or a limitation in your method?\"**\n\nThank you for your insightful feedback. Regarding the runtime analysis on time cost, in Section 4.4, we provide a detailed analysis comparing SpaceFormer with Uni-Mol. The pretraining cost of Uni-Mol scales quadratically with the number of sampled cells, whereas SpaceFormer exhibits a more linear scaling pattern. \n\nFor runtime analysis on smaller improvements in BBBP,  one reason can be  the complexity and variability of the measurement environments for experimental properties, which can be easily influenced by external factors, like temperature and blood-brain barrier biofilm status for BBBP. Such variability can affect the stability of measurement results and may not accurately.\nAdditionally, as mentioned in the footnote of Section 4.1, the MoleculeNet dataset for the experimental task (including BBBP) has several limitations,  including invalid structures and inconsistent chemical representations, which affect its ability to differentiate molecular pretraining models. We will incorporate this analysis into the main text of our paper.\n\n**2. Regarding \"it omits discussion of earlier 3D grid-based methods like AtomNet or equivariant models such as SchNet, which could provide additional context.\"**\n\nThank you for your insightful feedback. We appreciate your suggestion to discuss earlier 3D grid-based methods. AtomNet and SchNet both use 3D convolutional neural networks, with SchNet focusing on atom positions. In contrast, SpaceFormer employs a transformer-based architecture to capture interactions across global grid cells, extending beyond just atoms positions. We will incorporate this discussion in our paper.\n\n**3. Regarding \"Computationally intensive due to large input size (though mitigated by merging)\"**\n\nThank you for raising this point. Firstly, the input size is smaller than expected due to grid merging or sampling strategies, which reduce the number of grid cells effectively while preserving model performance, as detailed in Table 2. Secondly, SpaceFormer efficiently handles computational complexity using FlashAttention and two 3D relative positional encodings, allowing our model's time to scale nearly linearly with sequence length, as shown in Figure 4.\n\n**4. Regarding \"does not explicitly enforce equivariance and have you considered equivariant architectures as baselines? Would these be complementary to your grid-based approach?\"**\n\nThank you for your insightful suggestion. We would like to clarify that both Uni-Mol and Mol-AE, which are included in our baselines, consider equivariance. Initially, SpaceFormer was developed with reference to Uni-Mol. However, we observed that explicitly enforcing equivariance, as Uni-Mol does, leads to increased computational complexity. For instance, Uni-Mol's pair-wise Gaussian distance  results in a complexity of $\\mathcal{O}(n^2)$. Consequently, SpaceFormer encodes this information more efficiently through  3D distance PE using RFF to approximate Gaussian distance, and 3D directional PE using RoPE to capture direction, although strict equivariance is not guaranteed. Additionally,  SpaceFormer also employs random rotation data augmentation to enhance robustness. Figure 4 illustrates SpaceFormer's superior efficiency compared to Uni-Mol, while Table 1 shows SpaceFormer outperforming other equivariant baselines, highlighting our model's robustness.\n\n**5. Regarding \"Clarify average token count per molecule after merging and positional encoding details for reproducibility\" and \"Minor typos should be corrected\"**\n\nThank you for your insightful feedback. In Table 2, we present the average token count per molecule (avg. cells) after the merging/sampling process. We recognize the importance of clarity and reproducibility, and we will include additional details about the positional encoding in the next version. Additionally, we plan to open source our code to facilitate reproducibility.  We will correct the typo  in the main text, as well as any other minor errors.\n\n**6. Regarding \"Could you briefly explain why random virtual points improved performance initially (Figure 2)? What might the model be capturing from these points?\"**\n\nThank you for your insightful question. The inclusion of virtual points with physical positions  serves to provide additional spatial context, allowing the model to develop intermediate representations that enhance its understanding of the overall spatial information. In addition, during pre-training, the task of identifying atom locations in 3D grid cells becomes more challenging, which ultimately improves the model's performance in downstream tasks. This approach helps the model capture spatial and positional information more effectively.", "claims_and_evidence": "The main claim—that modeling space beyond atoms enhances molecular representations—is strongly supported by experiments. Empirical results (e.g., ~20% error reduction on key quantum property tasks like HOMO/LUMO) convincingly validate this claim (Table 1). Ablation studies (Table 4) show improvements aren't just due to increased capacity or random noise points, but genuinely from modeling spatial context.", "methods_and_evaluation_criteria": "Methods and evaluations are appropriate and rigorous. The chosen tasks (quantum and experimental properties) test generalization in limited-data scenarios using OOD splits, effectively evaluating pretraining value. Baselines (Uni-Mol, GEM, 3D-Infomax) are strong, ensuring a fair assessment.", "theoretical_claims": "No major theoretical issues. Their choice of grid resolution (0.49 Å) ensuring one atom per cell is logical. Positional encoding is reasonable and supported by experimental gains (Table 3). The theoretical arguments (e.g., empty space containing meaningful fields) are scientifically sound though qualitative.", "experimental_designs_or_analyses": "Experiments and analyses are well-executed. Results averaged over multiple seeds and ablation studies thoroughly validate the method (Tables 2, 4). Minor weaknesses include no detailed runtime analysis and limited discussion on tasks where improvements were small (e.g., BBBP).", "relation_to_broader_scientific_literature": "The paper clearly positions itself relative to prior 3D MPR models (Uni-Mol, 3D-Infomax) and makes a novel contribution by explicitly modeling empty space. However, it omits discussion of earlier 3D grid-based methods like AtomNet [Wallach, 2015] or equivariant models such as SchNet [Schütt, 2017], which could provide additional context.", "essential_references_not_discussed": "See the answer for \"Relation To Broader Scientific Literature\"", "other_strengths_and_weaknesses": "Strengths: Highly original idea; consistent improvements across diverse tasks; careful experimental design and thorough ablation studies.\nWeaknesses: Computationally intensive due to large input size (though mitigated by merging); does not explicitly enforce equivariance; unclear why method underperformed on BBBP (Table 1).", "other_comments_or_suggestions": "- Minor typos (e.g., “concergence” → convergence in second column of page 7) should be corrected.\n- Clarify average token count per molecule after merging and positional encoding details for reproducibility.", "questions_for_authors": "1. Why does SpaceFormer slightly underperform on BBBP compared to simpler baselines? Could it indicate an overfitting or a limitation in your method?\n2. Have you considered equivariant architectures as baselines? Would these be complementary to your grid-based approach?\n3. Could you briefly explain why random virtual points improved performance initially (Figure 2)? What might the model be capturing from these points?", "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": "BcAXX4MKHD", "reviewer": "Reviewer_9MWX", "summary": "The paper introduces a novel framework called *SpaceFormer*, which models molecular as three-dimensional “images” by discretizing 3D space into grid units. This approach captures the spatial context of molecular more comprehensively than traditional methods that treat them as discrete collections of atoms. The authors demonstrate that *SpaceFormer* significantly outperforms existing 3D molecular pre-trained representation models across various tasks, especially in data-limited scenarios, by effectively leveraging additional spatial information. Their findings also reveal that randomly sampled virtual points can enhance performance, motivating further exploration of systematic 3D modeling in molecular representation.\n\n## update after rebuttal\nI thank the authors for their efforts in rebuttal. After reading all comments and rebuttals, I think the author has roughly solved the doubts about motivation, and this is crucial. So I'm willing to update the score 2->3", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Due to word limit, we will not quote the original text but refer to the title of the reviewer's feedback content.\n\n**1. Regarding Claims And Evidence**\n\nThank you for your insightful feedback. In our paper, we have provided substantial experimental evidence to validate our claims. For instance, our results demonstrate that incorporating virtual points enhances model performance and techniques like grid sampling or merging improve efficiency without sacrificing effectiveness. Moreover,  various components, like position encoding, contribute significantly to the model's performance improvements. \n\n**2. Regarding Methods And Evaluation Criteria and Weaknesses 2**\n\nThank you for your valuable feedback. We will revise the manuscript to enhance the readability, clarify our methodology  and add more details to the main text. We will also make our code open source to enhance understanding.\n\n**3. Regarding Experimental Designs Or Analyses  and Weaknesses 3**\n\n- Thank you for your feedback. In Section 4.1, we detailed the dataset's origin, significance, processing, and splitting methods, and included units and sample numbers in Table 1. Additional data size, description, task types, and evaluation metrics are found in Table 5 in the Appendix. We will release the data processing script with the code. If any part of this explanation remains unclear, we would appreciate further clarification from the reviewer. Regarding fairness in comparative experiments, we conducted additional tests on the QM9 dataset. The results, which show SpaceFomer's superior performance, are included in the first response to reviewer sLbT.\n\n- Thanks for your suggestion. We would like to clarify that  our results in experimental property-related tasks remain superior (with 4 out of 5 tasks ranked in the top two). We acknowledge that the enhancements are less pronounced than in computational property tasks. This may be due to the complexity and variability of the measurement environments for experimental properties, affected by factors like temperature and pH, which can impact measurement stability and accuracy.\n\n**4. Regarding Essential References Not Discussed**\n\nThank you for your feedback. We apologize for any confusion and appreciate the chance to clarify. In our paper, we briefly mention the MoleculeNet dataset in the footnote of Section 4.1, focusing on challenges with pre-trained models, as discussed in previous work. We did not explore the data generation process.\n\nIf specific statements cause misunderstandings, please let us know so we can address them thoroughly. We value your insights and welcome continued dialogue to improve our manuscript.\n\n**6. Regarding Reviewer's Weaknesses 1  and  Question 3**\n\nThanks for your thoughtful feedback. We appreciate the opportunity to clarify our motivation. We acknowledge that atom coordinates encapsulate significant information, but relying solely on them necessitates extensive computation to derive electron densities or potential fields. Our approach introduces grid cells beyond atoms, offering a computationally efficient means to incorporate additional spatial information. Furthermore, many computational simulation methods in physics like electronic density distributions and potential fields are functions of the entire 3D space, not just atom positions. This underscores our motivation to explore modeling beyond atoms.  The interaction between atoms and virtual points with physical positions can be seen as a way for the model to efficiently learn some intermediate representations of the space beyond atoms, approximate spatial field effects through learnable interactions, thereby enhancing its understanding of the overall spatial information. We acknowledge this is an intuitive understanding and plan to pursue more theoretical validation in the future.\n\n**7. Regarding Reviewer's Other Comments Or Suggestions**\n\nThank you for your suggestion. We use bold fonts for vectors and text fonts for functions. Specifically, in Section 3.2, we denote the embedding function as \"\\text{Embed}_t(\\cdot)\", following LaTeX conventions.  According to ICML guidelines,  tables can span two columns if within page margins. We've ensured compliance and appreciate your feedback, hoping you'll reconsider.\n\n**9. Regarding Question 1**\n\nThank you for your question. In our model, each atom type is assigned a unique identifier mapped to an embedding, independent of atomic mass. For the virtual point type, a 'NULL' type is introduced with a distinct embedding to differentiate it from other atom types, without any mass-related encoding.\n\n**10. Regarding Question 2**\n\nThanks for your question. Since we divide the 3D space where the molecules are located into grids and obtain the final grid cells/virtual points by merging, this method is deterministic and therefore does not affect the SE-(3) equivariance. Regarding your question on dimensions, we have detailed these specifications for each layer in Table 5, located in  Appendix A.", "claims_and_evidence": "Several of the paper’s claims are incorrect or not well-supported.", "methods_and_evaluation_criteria": "The description of the method is not sufficiently detailed, with some parts being confusing.", "theoretical_claims": "I double-checked the full text, including the appendix, and the author did not provide a statement of the theory.", "experimental_designs_or_analyses": "- The authors used the processed dataset but did not provide dataset details. The fairness of the comparative experiment needs to be further verified. \n- The performance of the molecular experimental properties dataset in Table 1 is not optimal, and the authors did not analyze it further.", "relation_to_broader_scientific_literature": "Not involved.", "essential_references_not_discussed": "The main text cites many literature related to quantum mechanics, but the authors do not understand the data generation of the MoleculeNet dataset, that is, they have a wrong understanding of the physical meaning of the data. Please refer to [1].\n\n\n\n[1] Montavon G, Rupp M, Gobre V, et al. Machine learning of molecular electronic properties in chemical compound space[J]. New Journal of Physics, 2013, 15(9): 095003.", "other_strengths_and_weaknesses": "**Strengths**:\n- SpaceFormer improves performance by making effective use of additional spatial information and utilizing randomly sampled virtual points. \n- The authors perform a comprehensive evaluation against baseline models across a variety of datasets. The experiment section is solid.\n- The method in this paper is innovative.\n\n**Weaknesses**:\n- The motivation of the paper is questionable. The authors suggest that existing molecular representations focus only on atomic coordinates and overlook electron density information, leading to the introduction of virtual points for electromagnetic field representation. However, the electromagnetic field is inherently linked to the atoms themselves, and the coordinates in the benchmark are the result of the optimization of the atomic structure and contain the atomic interaction information. Thus, the motivation is not well-founded.\n- The writing requires improvement; the methodology is difficult to follow, and many details are not clear. \n- The description of the dataset is not sufficiently detailed, with some parts being confusing. The authors should provide statistical details of the benchmark. At the same time, authors should compare models on existing benchmarks.", "other_comments_or_suggestions": "1. Functions or mathematical symbols are distinguished by other fonts, e.g. Embed_t (·)\n2. Table 2 extends beyond the boundaries of the double columns.", "questions_for_authors": "- The atomic type is generally the corresponding atomic mass, for example, O is 16, how to quantify the virtual point type as NULL?\n- After adding the virtual point, is the SE-(3) equivariance still valid? What are the specific dimensions of each encoding molecule on the transformer?\n- Positional relative encoding is generally used to represent the relative interaction between atoms, what is the specific physical meaning of the relative position representation of VP and atoms in this paper?", "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:

Meta-review: The paper presents a new representation learning framework for molecules based on explicitly modeling the surrounding 3D space. All the reviewers appreciated the idea and the provided experiments. Although some reviewers initially raised concerns regarding computational efficiency and specific dataset details, the authors provided thorough rebuttals and additional experiments that addressed these issues. Overall, I recommend acceptance for the paper.
Decision: Accept (poster)