More Space Is All You Need: Revisiting Molecular Representation Learning

Thank you very much for the further clarification. As our model does not support force prediction for now, we may not be able to complete it within the discussion period. However, we will ensure the energy prediction benchmark is completed and will make our best effort to include force prediction as well.

We sincerely thank the reviewer for the kind suggestion. Based on our current understanding of neural potential networks, gradient-based force prediction typically requires a supervised loss on forces during training. This process involves higher-order gradients, which may present challenges in our current framework (e.g., the flash-attention kernel might not support it). Nevertheless, we will make every effort to implement it.

5. On PCA and SE(3) Invariance:

PCA is not intended to handle SE(3) invariance in our method. Its motivation is solely to reduce the grid size, not to address SE(3) invariance. Also, as reviewer suggested, in molecular data, the inherent symmetry often causes PCA to fail to produce a unique coordinate system, thereby making it unreliable for ensuring SE(3) invariance. As a result, even with PCA, we still need a stable method to encode SE(3)-invariance features, like pair-wise distances. To handle cases where PCA eigenvectors are degenerate, we use atom weights to break the symmetry by determining the axis directions based on the side with the larger mean atom weight. Additionally, we have a fallback mechanism: if PCA fails to produce a valid coordinate system (e.g., non-orthogonality or axis swapping), we revert to the original coordinate system to ensure robustness.

6. On adaptive gridding strategies:

Thank you for the suggestion. In the early stages of our work, we experimented with an adaptive gridding strategy based on an octree structure. Specifically, eight adjacent empty cells could be merged in a bottom-up, hierarchical manner, resulting in fewer but larger cells in coarse regions and more but smaller cells in fine regions with higher atomic density. This approach is conceptually similar to our proposed importance sampling strategy, which samples fewer cells in coarse regions and more cells in fine regions. However, the proposed sampling strategy offers better control over the number of sampled cells, providing greater flexibility and efficiency. For this reason, we adopted it in the final design.

7. On extending SpaceFormer to larger systems:

We are actively working on extending SpaceFormer to handle larger systems. With the proposed grid sampling strategy, the total number of cells can be effectively controlled, but two additional bottlenecks remain:

• Full-Length Attention: As the number of cells increases, the current full-length attention mechanism becomes a significant bottleneck. To address this, we plan to adopt locality-based optimizations such as sliding-window attention or Swin attention. Furthermore, incorporating local-to-global or multi-level hierarchical attention mechanisms could further enhance both scalability and performance.

• Extremely Large Systems: While the above optimizations allow SpaceFormer to handle most cases in protein databases, a small number of extremely large proteins still pose challenges. For these cases, we utilize a random spatial cropping strategy during training to reduce the computational load. Notably, cropping is disabled during inference to ensure accurate predictions.

8. On discussion about the computational cost of SpaceFormer and other models:

We discuss the computational cost of SpaceFormer in Section 4.4, Point 5. Specifically, we note that "SpaceFormer scales much more efficiently with empty cells. While Uni-Mol’s computational cost increases quadratically with the number of empty points (No. 1-7), SpaceFormer scales linearly (No. 8-14). For example, within 100 hours, SpaceFormer can process 1,500 cells, whereas Uni-Mol can only handle 400 points." If this discussion is not enough, we would appreciate more details on specific aspects of computational cost that we can further elaborate on.

We thank the reviewer for their thoughtful comments and concerns, particularly regarding the experimental results, including those on QM9. Below, we address these points in detail:

1. On the performance of QM9:

It is important to note that the QM9 dataset used in this paper is a subsampled version containing 40k samples, whereas the original QM9 dataset consists of 133k samples. Therefore, direct comparison of our results with those reported in previous papers using the full dataset would be unfair. The subsampling in our study is intentional, as the paper focuses on evaluating the performance of a pretrained MRL model on downstream tasks with limited data. In real-world applications, the availability of labeled data is often constrained due to the high costs of computation and wet-lab experiments. This setup, therefore, enables a more realistic assessment of the model’s effectiveness in practical scenarios. Regarding the units in QM9, we use the original units provided in the dataset. We will revise the paper to explicitly include this detail.

2. On more benchmark datesets than QM9:

We would like to clarify that our benchmark datasets are not limited to QM9. In addition to QM9, we include tasks from MoleculeNet (2018) and Biogen ADME (2023). These datasets focus on wet-lab experimental properties, most of which have very limited data, making them well-suited for evaluating pretrained MRL models on real-world tasks. The reviewer’s suggested datasets -- MD17, MD22, and QM7-X -- are not ideal for evaluating pretrained MRL models in our scenario. Specifically:

• MD17 and MD22: These datasets focus on molecular dynamics and are primarily used to train neural potential models for predicting energy and forces, which differ from the property prediction tasks targeted by our method.
• QM7-X: While QM7-X provides a large number of labeled samples, it does not align with the primary scenario of our work, which focuses on downstream tasks with limited data.

Although these datasets are not suitable for evaluation in our current framework, we agree that they could serve as valuable datasets for pretraining MRL models. We plan to investigate their utility for this purpose in future work.

3. On the Claim That the Comparison with Atom-Based Models Is “Flawed”:

We regret that we do not fully understand the reviewer’s assertion that our comparison with atom-based models is “flawed.” As shown in Table 6, we explicitly incorporate additional empty points into the atom-based model. The table clearly demonstrates the increased computation times as more empty points are added, reflecting a significant increase in “internal FLOPS” for Uni-Mol. The primary limitation here is that Uni-Mol is not efficient when scaling with a large number of empty points, which is why we restricted our experiments to 400 empty points. However, the observed trend suggests that even with more empty points, the performance does not improve. This finding underscores the inherent limitations of atom-based models in leveraging empty space information effectively.

4. On the Quantitative Comparison with the Electron Density Map:

We thank the reviewer for pointing this out and conducted a further comparison with the electron density map. Specifically, we directly compared the normalized density map from the learned representation with the electron density by calculating the element-wise absolute error at grid cells and averaging the results over 50 data samples. The resulting error is approximately 7.4e-4. We also implemented the distance-based solution suggested by the reviewer, which yielded an error of about 7.9e-4. While the learned representation slightly outperforms the distance-based solution, the improvement is small. We agree that this comparison may not be particularly meaningful at this stage due to the small gain. As a result, we will remove this analysis from the paper and plan to conduct a deeper investigation into the learned representation in future work. We sincerely thank the reviewer for this suggestion, as it has helped us improve the quality of the paper.

We sincerely thank the reviewer for the additional comments.

As stated in our initial response, we respectfully disagreed that a comparison with neural potential models (NPMs) is necessary to evaluate our contribution, as it falls beyond the scope of this paper. However, since the reviewer insisted, we conducted additional experiments to address this concern. Despite the significant effort involved in obtaining these results, the reviewer appears to doubt their validity. While we considered sharing the running logs to provide further evidence, we felt that it might not fully address the reviewer’s concerns. Nevertheless, we assure the reviewer that we have made every effort to address these points comprehensively and transparently.

We understand that the reviewer may be highly familiar with NPM and has evaluated our work through that lens, expecting a comparison. However, in our opinion, such a comparison is not essential for evaluating the contributions of this work.

Scope of Our Paper

Our work focuses on molecular representation learning (MRL), where large-scale data is used to pretrain a model, which is then evaluated on downstream tasks with very limited labeled data. In contrast, NPMs focus on predicting energy and force, primarily for molecular dynamics simulations.

Moreover, our baselines—including Grover, GEM, Uni-Mol, and Mol-AE—do not include NPMs as baseline models. This demonstrates that NPMs are designed with fundamentally different goals compared to MRL models, and by default, our work does not aim to compare with NPMs. Extending MRL models to NPN tasks may be an interesting direction, but it is not the focus of our work.

Our Designed Benchmark

The evaluation of MRL models typically relies on limited labeled data due to the high cost of acquiring such data, especially for experimental properties. To reflect this, our benchmark includes subsampled QM9 to evaluate learning ability under limited labeled data conditions. While full QM9 is widely used in NPMs, the goal of our benchmark is very different from them.

In our paper (Line 320), we elaborated on the rationale for proposing a new benchmark, noting the limitations of the widely used MoleculeNet dataset. Specifically:

• Walters (2023) identified several issues with MoleculeNet.
• Sun et al. (2022) demonstrated that MoleculeNet fails to adequately distinguish the performance of different molecular pretraining models.

Our designed benchmark aims to address these shortcomings and provide a fair evaluation framework tailored to real-world MRL scenarios.

Fair Comparison with Previous MRL Models

To ensure an apple-to-apple comparison, we adopted the same pretraining data as the most recent works, Uni-Mol and Mol-AE. Additionally, we used the same model size as Uni-Mol and aligned all hyperparameters, including those for pretraining and fine-tuning, to ensure fairness in comparison.

While reviewer rJDR raised concerns about our data split, we provided evidence demonstrating that our splitting method is highly similar to the one they recommended, further validating the robustness of our evaluation methodology.

Summary

We understand that the reviewer places significant emphasis on comparisons with NPMs, likely due to their expertise in this area. While we hold differing views on the necessity of such comparisons for this work, we value the reviewer’s input and have clarified our perspective to the best of our ability. Our goal in this reply is not to convince the reviewer to change their opinion but rather to clearly articulate our rationale and the scope of our work.

We sincerely thank the reviewer again for their thoughtful effort and constructive feedback throughout the discussion.

Additional Comments on iso-FLOP Comparison

Regarding the reviewer’s statement, "Neither the number of empty points used, nor the time for pre-training are a proper measure for establishing that models were compared under an iso-FLOP setting," we remain unclear about what constitutes "a proper measure" as we thought that training cost or additional empty points is a direct measure. We are not challenging the reviewer but are genuinely curious to understand their perspective and would greatly appreciate further clarification.

updated:

We now may understand the reviewer's point regarding the "iso-FLOP setting." There seems to be a potential misunderstanding that SpaceFormer’s model capacity is significantly larger than Uni-Mol’s. However, SpaceFormer and Uni-Mol share almost identical transformer architectures and model capacities, differing only in their positional encoding mechanisms. As a result, when the number of input tokens is the same, the FLOPs required by SpaceFormer and Uni-Mol should be nearly identical. Based on this, we considered the "number of empty points" a reasonable proxy for ensuring an "iso-FLOP setting.

We sincerely thank the reviewer for the insightful comments. Below, we provide detailed answers to your questions:

1. On spatial continuity:

We appreciate the concern regarding the potential disruption of spatial continuity caused by gridding, especially for atoms near cell boundaries. As described in Section 3.1, we address this issue by preserving the precise atomic positions, regardless of which cells the atoms belong to. These precise, continuous atomic positions are used consistently throughout the model, ensuring that gridding does not affect atom spatial information.

Additionally, as an implementation detail, we introduce a small translation noise to all atomic positions before gridding. This softens the boundaries between grid cells and mitigates potential issues with spatial continuity, providing smoother transitions and improving overall robustness.

2. On handling irregularly shaped or large molecules:

Our grid sampling strategy (Section 3.2) is designed to handle this issue. Specifically, the strategy samples fewer cells in coarse regions and more cells in fine regions with higher atomic density. For regions without atoms, only a small number of cells are sampled, ensuring computational efficiency while capturing the relevant spatial structure.

3. On Reliance on Accurate Atomic Positions:

Both the pretraining and downstream tasks do not rely on highly accurate atomic positions. We adopt the approach from Uni-Mol, using RDKit -- a widely used chemical toolkit -- to generate atomic positions when accurate atomic positions are not provided. RDKit efficiently generates positions in just a few milliseconds per molecule, making this a practical solution. Therefore, the lack of high-precision atomic positions does not limit the applicability of our method.

4. On the "fixed grid size":

If we correctly interpret the reviewer’s concern, it pertains to the fixed grid cell edge length. While we use a fixed edge length in our design, we initially experimented with an adaptive gridding strategy based on an octree structure. This approach merged adjacent grid cells hierarchically, creating fewer, larger cells in coarse regions and more, smaller cells in fine regions with higher atomic density. Conceptually, this is similar to our importance sampling strategy, which also allocates fewer cells to coarse regions and more to fine regions. However, our sampling strategy offers better control over the total number of cells, providing greater flexibility and efficiency. For this reason, we adopted it in the final design.

5. On the data split method:

Thank you for your insightful feedback regarding splitting strategies. We chose the Out-of-Distribution (OOD) split based on fingerprint similarity to investigate the model’s ability to generalize to chemically distinct compounds. Fingerprints, which encode molecular structures into bit vectors, are widely used in cheminformatics for tasks like compound similarity and diversity evaluation. As such, OOD splits based on fingerprint similarity offer a meaningful way to assess how well a model performs when encountering structurally dissimilar compounds. We agree that scaffold and temporal splits are highly relevant and frequently employed in drug discovery. While fingerprint-based splits provide an initial evaluation of robustness, scaffold splits capture structural novelty, and temporal splits simulate real-world scenarios where newer compounds must be predicted based on historical data. In the revised manuscript, we will expand the discussion to address these complementary approaches and propose their inclusion as future work.

6. On more baseline models:

Thank you for your thoughtful comment regarding the comparison of SpaceFormer with other recent works. In this study, we have already compared SpaceFormer with both graph-based and 3D representation molecular pretraining methods, ensuring that the baselines are representative of the state of the art in these paradigms. we acknowledge the value of extending comparisons to include recent advances like NoisyNodes, GraphMVP, and 3D-InfoMax. While additional experiments were not feasible due to time and resource constraints, we will expand the discussion and include further experiments in the revised manuscript to provide a broader evaluation.

Thank you for pointing out that I did not respond to your point about scaffold split 8 days ago. Please note that in my initial review, I unambiguously requested for "in real-world drug discovery applications, much more common splitting strategies are scaffold split or temporal split. I would be curious to see how the model performs under such settings." In their response, the authors simply acknowledged that these splits are sensible but no experiments were shown. However, if the authors believe I was not clear enough in my initial review, I apologize. Also if the authors believe that scaffold or temporal splits are not necessary for their work, I simply disagree. I gave my reasons and it's ok that the authors disagree with my reasoning.

As I said, I cannot champion the acceptance of this work in its current form even though I think it proposes an interesting idea, but I will not oppose its acceptance if AC and other reviewers wish to and I am happy for my reasons, reviews and ratings to be discounted and discarded from consideration. Good luck!

Thank you for your response. I do not mean to come across as asking for infeasible or less meaningful experiments and I am not sure why the authors have this impression. The authors could have simply run pre-training for one property or showed partial experiments with different data splitting strategies if time was an issue. Especially when I mentioned that in real-world drug discovery random splitting is not what practitioners do and that is a big issue with building ML models.

Regarding baselines, I understand that NoisyNodes, 3D-InfoMax etc. are different pre-training strategies. My question is if given a choice in real-world setting, why and when should practitioners use the presented method vs the other pre-training methods, especially when no comparison or discussion is presented about merits and demerits of each? I assume that the authors will agree that the goal of any such work would be to help the practitioners and help drug discovery applications. Hence, for a strategy to be useful in practice, we should know when to use which method and without a comparison with methods that are known to work well, it is hard to do so.

I understand that preparing rebuttals, revising the paper and conducting experiments takes time and I fully understand the authors. Anyways, I will respectfully suggest that in future authors may choose to do these experiments. I understand that reviewer jmPG who have been highly engaged in discussion and raised good points also agree with my concerns. Anyways, I am happy to be overruled and my suggestions and ratings discounted by AC and other reviewers if they consider them infeasible or less meaningful too. Good luck!

We sincerely thank the reviewer for the insightful comments. Below, we provide detailed responses to your questions:

1. On the Claim That “ANI and SchNet Have Already Definitively Demonstrated the Value of Considering Complete 3D Space”:

We respectfully disagree with this claim. Both ANI and SchNet are atom-based models that take only atoms as input. Additionally, the “Behler-Parrinello symmetry” feature vector used in these methods considers only local neighborhood atoms around each atom, rather than the entire 3D space. While methods like SchNet employ continuous kernels, such as the RBF kernel, these primarily encode pairwise atom distances and only apply to atom-pair interactions. As such, these models may implicitly capture limited spatial information beyond atoms but do not explicitly consider or demonstrate the value of the complete 3D space, particularly the regions beyond atoms.

In contrast, SpaceFormer explicitly incorporates non-atom cells into its input, allowing it to directly utilize information from the full 3D space. Our experimental results show that this approach consistently outperforms atom-only models, emphasizing the benefits of explicitly modeling non-atomic regions.

2. On the "absence of comparisons with models like ANI2 or SchNet":

We appreciate the reviewer’s suggestion to include comparisons with SchNet and ANI2, as these models indeed demonstrate strong performance in their respective domains. However, we believe their focus and application differ significantly from the scope of our work.

SchNet and ANI2 are neural potential models designed specifically for energy and force prediction tasks, such as molecular dynamics simulations. Their primary objective is to approximate potential energy surfaces with high accuracy for quantum chemistry and molecular mechanics applications. In contrast, this paper focuses on molecular representation learning (MRL), where the goal is to pretrain on large-scale unlabeled molecular data and transfer the learned representations to downstream tasks, such as property prediction, with limited labeled data. While there are conceptual overlaps, MRL models and neural potential models are designed for fundamentally different applications and are evaluated using distinct protocols.

Furthermore, most MRL studies, including the baselines in our manuscript, do not consider neural potential models as direct baselines because their training paradigms and evaluation metrics are not aligned with MRL objectives. Consequently, comparisons with SchNet and ANI2 would not provide meaningful insights into the effectiveness of our approach for MRL tasks.

3. On the claim that "SchNet frequently outperforms the proposed SpaceFormer model on the QM9 dataset":

It is important to note that the QM9 dataset used in this paper is a subsampled version containing 40k samples, whereas the original QM9 dataset consists of 133k samples. Therefore, direct comparison of our results with those reported in previous papers using the full dataset would be unfair. The subsampling in our study is intentional, as the paper focuses on evaluating the performance of a pretrained MRL model on downstream tasks with limited data. In real-world applications, the availability of labeled data is often constrained due to the high costs of computation and wet-lab experiments. This setup, therefore, enables a more realistic assessment of the model’s effectiveness in practical scenarios.

4. On the instability of PCA:

The primary instability of PCA arises from the arbitrary orientation of axes and the potential for axis swapping. To address these issues, we incorporate atom weights to determine axis directions. Specifically, atoms are projected onto each PCA axis, and the mean atom weight on either side is calculated. The side with the larger average atom weight is designated as the positive direction, reducing ambiguity and ensuring consistency. Additionally, as a safeguard, we verify the validity of the generated coordinate system by checking for orthogonality and right-handedness. If the system fails these checks, we fall back to the original coordinate system to maintain robustness.

It is also important to note that we use PCA solely to define an effective cuboid for grid size reduction. As such, axis flips have minimal impact on the grid structure and, consequently, only a negligible effect on our method. In fact, during early experiments, we applied PCA without the above stability optimizations and observed no difference in performance after addressing these issues.

5. Reproducibility:

Very sorry for missing Reproducibility Statement as it is not in ICLR template. We will open source all our code after acceptance.

Update to the Reply 2:

We have conducted additional comparisons with neural potential models, such as SchNet and PaiNN, for both property prediction and energy/force prediction. The results can be found in our replies: https://openreview.net/forum?id=LBsr2llHz0&noteId=AL0f3opV3m and https://openreview.net/forum?id=LBsr2llHz0&noteId=Yv5WQCVQ7H.

These experiments required substantial time and effort, demonstrating our commitment to thoroughly addressing the reviewer’s concerns. We sincerely hope the reviewer will recognize our dedication and reconsider their evaluation.

Belows are the responses to your other comments.

4. On the necessity of gridding:

Gridding becomes essential when considering both atoms and empty space, as SpaceFormer does. At a high level, gridding in SpaceFormer serves as an efficient strategy to sample empty points by constraining them to the centers of grid cells. The ablation study in Section 4.4 of the paper explicitly confirms the necessity of gridding. Without gridding, previous atom-based models sample empty points from a large, continuous 3D space, which fails to improve performance even with an increased number of sampled points. In contrast, SpaceFormer leverages gridding to define empty cells systematically, enabling it to utilize information from empty space more effectively.

5. On the comparison with additional virtul point used in point-could domain:

We acknowledge that we overlooked citing and discussing related works on point clouds, as our primary focus was on molecular representation learning (MRL). We thank the reviewer for pointing this out and have since conducted further investigation. In the point cloud domain, virtual or grid points are often introduced as intermediate representations to address challenges such as resolution inconsistencies (e.g., between sensors or distant vs. nearby objects), enhance feature aggregation, or increase sampling density. These intermediate points act as bridges to improve geometric consistency and feature representation.

In contrast, MRL tasks primarily focus on predicting molecular properties, where there is no natural need for intermediate representations like virtual points. Moreover, as demonstrated in Section 4.4, simply adding virtual points fails to enhance the performance of existing models. This distinction explains why prior MRL methods have not adopted virtual points.

Our contribution lies in recognizing this gap and proposing a framework that effectively leverages empty space information, addressing an overlooked aspect of MRL. The validity and utility of our approach are supported by our experiments. We will revise the manuscript to include a discussion of related works on point clouds to clarify this distinction.

6. On the instability of PCA:

The primary instability of PCA arises from the arbitrary orientation of axes and the potential for axis swapping. To address these issues, we incorporate atom weights to determine axis directions. Specifically, atoms are projected onto each PCA axis, and the mean atom weight on either side is calculated. The side with the larger average atom weight is designated as the positive direction, reducing ambiguity and ensuring consistency. Additionally, as a safeguard, we verify the validity of the generated coordinate system by checking for orthogonality and right-handedness. If the system fails these checks, we fall back to the original coordinate system to maintain robustness.

It is also important to note that we use PCA solely to define an effective cuboid for grid size reduction. As such, axis flips have minimal impact on the grid structure and, consequently, only a negligible effect on our method. In fact, during early experiments, we applied PCA without the above stability optimizations and observed no difference in performance after addressing these issues.

7. On the Sampling Strategy of Learned Representations in Figure 2:

As stated in the paper (Line 520 in the first submitted version), the learned representation shown corresponds to the Full-Grid SpaceFormer encoder (No. 2 in Table 4). No sampling strategy is applied to ensure a fair comparison with the electron density.

8. Reproducibility

Very sorry for missing Reproducibility Statement as it is not in ICLR template. We will open source all our code after acceptance.

We sincerely thank the reviewer for their detailed comments. However, we believe there are some misunderstandings regarding our method, particularly concerning PCA and the proposed positional encoding. We address these points below and will revise the paper to make these aspects clearer.

1. On PCA and SE(3) Invariance:

PCA is not intended to handle SE(3) invariance in our method. Its motivation is solely to reduce the grid size, not to (and cannot) address SE(3) invariance. In molecular data, the inherent symmetry often causes PCA to fail to produce a unique coordinate system, thereby making it unreliable for ensuring SE(3) invariance. Furthermore, if PCA could perfectly handle SE(3) invariance, there would not be hundreds of papers addressing SE(3)-invariance/equivariance challenges. The reviewer’s comemnt “PCA vectors have arbitrary signs and uncertainties in direction when eigenvalues are close” further supports this. Consequently, since PCA cannot ensure SE(3) invariance, our proposed positional encoding, which captures pairwise distances using Random Fourier Features, is essential for providing SE(3)-invariant input features.

2. On Random Fourier Features and Gaussian Kernels:

Approximating Gaussian kernels with Random Fourier Features (RFF) is a well-established and validated approach. Its original paper [1], which received the Test-of-Time Award at NeurIPS 2017, demonstrates the effectiveness of this technique. The reviewer's comment that “For the Gaussian kernel, k(δ) is not convex, so k is not everywhere positive and δk(δ) is not a probability distribution…” refers to the Random Binning Features approximation, another approach in that paper, not the Random Fourier Features method we use. For further details on Gaussian kernels approximated with RFF, we kindly direct you to Section 3 of paper [1].

3. On 3D RoPE and Pairwise Relative Position Encoding:

We apologize if the term “3D Angular Positional Encoding with RoPE” caused any confusion. The term “Angular” might indeed be misleading; “Directional” is a more accurate description in this context. Below, we provide a clearer explanation, starting from the concept of RoPE. Rotary Positional Embedding (RoPE) is a widely used method for encoding pairwise relative positions. Originally proposed for 1D sequences in language models, RoPE has since been successfully adapted for 2D image tasks [2]. In our work, 3D RoPE extends this concept to encode the pairwise relative position of two points in 3D space. For points A and B, with coordinates (x1, y1, z1)and (x2, y2, z2), their relative position is represented as $\vec{AB} = B - A = (x2-x1, y2-y1, z2-z1)$. Since vector $\vec{AB}$ contains directional information, we refer to it as “Directional Positional Encoding.” This direction is inherently dependent on the coordinate system and varies with rotation, which is precisely our goal. By doing so, the model can capture coordinate system information and better predict point positions within it. This approach is consistent with existing SE(3)-equivariance models, such as EGNN [3], which also leverage relative positional information for prediction.

To summary, we believe that our proposed method is rigorously correct, and our ablation study on positional encoding (Sec. 4.3) also demonstrated the effectiveness of the proposed method. We hope these clarifications resolve the misunderstandings and please let us know if you have any other questions.

Reference: [1] Rahimi, Ali, and Benjamin Recht. "Random features for large-scale kernel machines." Advances in neural information processing systems 20 (2007).

[2] Heo, Byeongho, Song Park, Dongyoon Han, and Sangdoo Yun. "Rotary Position Embedding for Vision Transformer." arXiv e-prints (2024): arXiv-2403.

[3] Satorras, Vıctor Garcia, Emiel Hoogeboom, and Max Welling. "E (n) equivariant graph neural networks." In International conference on machine learning, pp. 9323-9332. PMLR, 2021.

We believe the reviewer may still have some misunderstandings regarding the mechanism of 3D RoPE.

3D RoPE does not attempt to reconstruct 3D geometry from 2D projections, nor does it rely on combining in-plane rotations as in the Euler convention. Instead, 3D RoPE directly operates on 3D relative positional differences. For instance, for two 3D points with coordinates $(x_i, y_i, z_i)$ and $(x_j, y_j, z_j)$, 3D RoPE directly encodes the positional deltas $(x_j - x_i, y_j - y_i, z_j - z_i)$ along the three axes. This is achieved without involving any 2D projections or 2D angles.

We are uncertain how the reviewer interprets RoPE, so allow us to provide a clearer explanation below.

Let us start with the 1D case. In natural language processing, given two tokens located at positions $x_i$ and $x_j$, the original RoPE mechanism is designed to capture their relative position $x_j - x_i$. This concept is straightforward and widely accepted.

Next, we extend this concept to 2D, as described in [1]. Given two points in 2D space with positions $(x_i, y_i)$ and $(x_j, y_j)$, the goal is to encode their positional differences $x_j - x_i$ and $y_j - y_i$. 2D RoPE achieves this by encoding each positional difference independently. Specifically, in the context of multi-head attention, half of the attention heads are assigned to encode $x_j - x_i$, and the other half to encode $y_j - y_i$.

Similarly, this concept extends naturally to 3D: we encode the relative positional differences along all three axes $(x_j - x_i, y_j - y_i, z_j - z_i)$ by dividing the attention heads into three sets. Each set is dedicated to encoding the positional difference along one axis. This ensures that 3D RoPE directly encodes relative positions in 3D space, without involving any 2D projections or rotations.

From the above explanation, it is clear that the key to 3D RoPE is using three independent sets of 1D RoPE to encode relative positions along the three axes. At no point does it rely on projecting 3D information into 2D space.

We suspect that the reviewer may have misunderstood the role of the 2D rotation matrix used in 1D RoPE. In 1D RoPE, the 2D rotation matrix is used to represent a rotation with an angle corresponding to the token position (e.g., $x_i$ and $x_j$). During the Query-Key dot product operation, the relative position $x_j - x_i$ is encoded through the composition of the two 2D rotation matrices. Thus, in all cases—whether 1D, 2D, or 3D—the 2D rotation matrices are solely used for encoding the positional difference along a single axis and are not related to projecting 3D information into 2D space. For further details, please refer to Equation (2) in the paper or the original RoPE paper [2].

Reference:

[1] Rotary Position Embedding for Vision Transformer.

[2] Enhanced Transformer with Rotary Position Embedding.

Thank you for your response and valuable feedback. We hope our previous responses have addressed your other concerns.

Regarding your request for an additional ablation study, due to the limited time before the discussion ends, we are afraid we may not be able to complete them in time. However, we believe our current results are sufficient to demonstrate the differences between PCA and RoPE in terms of their contributions to the final performance. Specifically, by combining the results from Table 3 and Table 5, we present the following observations:

From these results, it is clear that PCA does not significantly contribute to the final performance, while the proposed positional encoding techniques (RoPE and RFF) play a much larger role in improving the final performance .

Regarding your second point: If we only encode $x_i - x_j$, the direction $\frac{x_i - x_j}{\||x_i - x_j\||}$ is indeed not explicitly represented. However, in our approach, we also encode the distance $\||x_i - x_j\||$ using RFF. This allows the model to implicitly capture the direction $\frac{x_i - x_j}{\||x_i - x_j\||}$ through the combined encoding of $x_i - x_j$ (by 3D RoPE) and $\||x_i - x_j\||$ (by RFF).

We hope our response above addresses your concern regarding PCA. Please let us know if you have any additional questions or require further clarification.

We are uncertain about the comment, "...3D orientation from the combination of 2D rotations (rotations do not commute)." We suspect there may be a misunderstanding regarding how RoPE works. While we may be mistaken, we provide a more detailed explanation of RoPE below for clarification.

In RoPE, the 2D rotations are not used to encode rotations in a 2D space. Instead, they serve to encode 1D positions $i$ (and $j$) by converting these positions into 2D rotation matrices with angles corresponding to the positions. During the dot product operation between the Query and Key, the relative position $j-i$ is encoded through the composition of two 2D rotation matrices.

For the 3D scenario, since attention mechanisms typically involve multiple heads, we can divide the heads into three groups, with each group encoding the positional information along one of the axes. In this way, 3D RoPE directly encodes $x_{i} - x_{j}$ effectively.

Thank you for adding the additional results. The comment "The results clearly demonstrate that SpaceFormer outperforms SchNet and PaiNN" is based on what exactly? I unfortunately cant see any statistical tests that would underlay that claim - also no uncertainty estimations on the results. Could you please provide them as well - otherwise its hard to compare single numbers with each other, not talking any variability into consideration.

Thank you