Paper: Neural Atoms: Propagating Long-range Interaction in Molecular Graphs through Efficient Communication Channel (ICLR-2024)

Reviewer comment (Experiments):
does not utilize 3D coordinate information available for some molecules.

Author response:
Thanks for the valuable review! We would like to clarify the reason for not adopting 3D coordinate information and the consideration of inter-molecular interactions. Due to the time limitation, we are unable to incorporate the 3D information into the neural Atom mechanism, **as the 3D coordinate position of the neural atom needs to be carefully initialized and handled to avoid potential bias from the 3D position.** Here, we directly adopt the neural atoms to the 3D scenario without 3D coordinate information and provide the comparison with Ewald-based GNN [1] (denoted as Ewald Block) on the OE62 [2] dataset. We follow the same setting from [1] and use SchNet [3] as the backbone GNN. Results show that **our method achieves competitive performance even with the absence of 3D coordinate information**, showing the generalization and effectiveness of neural atoms.

|                   | Energy MAE $\downarrow$ | Energy MSE $\downarrow$ | Using 3D information |
| :---------------: | :---------------------: | :---------------------: | :------------------: |
| Baseline (SchNet) |         0.1351          |         0.0658          |     $\checkmark$     |
|   + Ewald Block   |       **0.0811**        |       **0.0301**        |     $\checkmark$     |
|   + Neural Atom   |       **0.0834**        |       **0.0309**        |       $\times$       |

We have added this discussion to our draft in Appendix D.

[1] Ewald-based Long-Range Message Passing for Molecular Graphs. In ICML, 2023

[2] Atomic structures and orbital energies of 61,489 crystal-forming organic molecules. In Scientific Data, 2020.

[3] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In NeurIPS, 2017.

Reviewer comment (Theory):
mainly considers intra-molecular, not inter-molecular interactions.

Author response:
Thanks for the valuable review! We would like to clarify the reason for not adopting 3D coordinate information and the consideration of inter-molecular interactions. Due to the time limitation, we are unable to incorporate the 3D information into the neural Atom mechanism, **as the 3D coordinate position of the neural atom needs to be carefully initialized and handled to avoid potential bias from the 3D position.** Here, we directly adopt the neural atoms to the 3D scenario without 3D coordinate information and provide the comparison with Ewald-based GNN [1] (denoted as Ewald Block) on the OE62 [2] dataset. We follow the same setting from [1] and use SchNet [3] as the backbone GNN. Results show that **our method achieves competitive performance even with the absence of 3D coordinate information**, showing the generalization and effectiveness of neural atoms.

|                   | Energy MAE $\downarrow$ | Energy MSE $\downarrow$ | Using 3D information |
| :---------------: | :---------------------: | :---------------------: | :------------------: |
| Baseline (SchNet) |         0.1351          |         0.0658          |     $\checkmark$     |
|   + Ewald Block   |       **0.0811**        |       **0.0301**        |     $\checkmark$     |
|   + Neural Atom   |       **0.0834**        |       **0.0309**        |       $\times$       |

We have added this discussion to our draft in Appendix D.

[1] Ewald-based Long-Range Message Passing for Molecular Graphs. In ICML, 2023

[2] Atomic structures and orbital energies of 61,489 crystal-forming organic molecules. In Scientific Data, 2020.

[3] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In NeurIPS, 2017.

Reviewer comment (Experiments):
hyperparameter tuning needed for number of neural atoms, as the problem with pooling operation known earlier in graph learning literature.

Author response:
Thanks for the valuable review. We would like to clarify further the setting of the proposed neural atoms. We do not directly finetune the number of neural atoms but define the number as the proportion of the average number of nodes for the specific dataset. For example, we use only 22 neural atoms for GCN in both Peptides-func and Peptides-struct datasets, which is far less than the original number of atoms (around 150 per graph). We update our hyperparameters with the specific number of neural Atoms used. We kindly refer the reviewer to Table 4 and Table 7 to Table 9 for more detailed hyperparameter settings, which show that our method can exceed the baseline with even one-tenth number of nodes w.r.t. the average number of nodes in the original graph.

Reviewer comment (Theory):
lacks strong theory on optimal atom groupings.

Author response:
Thank you for the question. We have to admit the (1) unavailability of the theory of the atom groupings strategy and (2) potential useful discussion and future works.

(1) Since we mainly focus on the methodology and the experimental evaluation of our method, we have not derived a theoretical justification for the proposed method. However, we did explore most of the research about the long-range interaction and grouping theory on the molecular graphs, but the available information is limited. As such, **to the best of our knowledge, there might not be a strong theory about the atom grouping strategy since the functional groups of different molecular graphs vary [1].** We leave this theoretical justification for future works.
(2) **Alternatively, there is some research on the impact of the virtual node.** For example, [2] studies the benefit of introducing single or multiple virtual nodes for link prediction tasks and provides a theoretical claim from the graph random walk perspective; [3] investigates the relationship between GNN and Graph Transformers via the bridge of the virtual node, and provide the assumptions and theory that guarantee the approximation of GNN with virtual node to the Graph Transformer. Our neural atoms might establish a similar theoretical framework like [2], as we reduce the communication distance that could affect the result of random walk.

[1] Chemistry: the central science. Brown, Theodore (2002). ISBN 0130669970.

[2] An analysis of virtual nodes in graph neural networks for link prediction. In LoG, 2022.

[3] On the Connection Between MPNN and Graph Transformer, In ICML, 2023.

Reviewer comment (Reproducibility):
How are the neural atom groupings initialized? Are they randomly assigned?

Author response:
Thanks for your question. The allocation matrix is computed as ${\hat{A}} = \sigma(Q^{(\ell)}_{\text{NA}} K^\top) \in \mathbb{R}^{K \times N}$. Therefore, the $\hat{{A}}$ is dynamically calculated according to the embedding similarity between the neural atoms and the atom within the molecular graph.

Reviewer comment (Theory):
is there a theoretical justification for why this approach models long-range interactions well?

Author response:
Thank you for the question. We have to admit the (1) unavailability of the theory of the atom groupings strategy and (2) potential useful discussion and future works.

(1) Since we mainly focus on the methodology and the experimental evaluation of our method, we have not derived a theoretical justification for the proposed method. However, we did explore most of the research about the long-range interaction and grouping theory on the molecular graphs, but the available information is limited. As such, **to the best of our knowledge, there might not be a strong theory about the atom grouping strategy since the functional groups of different molecular graphs vary [1].** We leave this theoretical justification for future works.
(2) **Alternatively, there is some research on the impact of the virtual node.** For example, [2] studies the benefit of introducing single or multiple virtual nodes for link prediction tasks and provides a theoretical claim from the graph random walk perspective; [3] investigates the relationship between GNN and Graph Transformers via the bridge of the virtual node, and provide the assumptions and theory that guarantee the approximation of GNN with virtual node to the Graph Transformer. Our neural atoms might establish a similar theoretical framework like [2], as we reduce the communication distance that could affect the result of random walk.

[1] Chemistry: the central science. Brown, Theodore (2002). ISBN 0130669970.

[2] An analysis of virtual nodes in graph neural networks for link prediction. In LoG, 2022.

[3] On the Connection Between MPNN and Graph Transformer, In ICML, 2023.

Reviewer comment (Experiments):
how does this approach scale to very large molecular graphs?

Author response:
Thanks for the valuable question! To scale the neural atoms to the very large molecular graph, one could employ a linear Transformer like Performer [1] or BigBrid [2]. In addition, one could prune the connection between neural atoms and the atoms within the original molecular graph based on the attention score to reduce the potential computation burden. **We have highlighted this point in the revised submission for the potential extension to the broad applications.**

The scale of the dataset we adopted, i.e., the Long-Range Graph Benchmark, is considered to be **larger than the general molecular graph dataset w.r.t. the graph size.** We show the comparison of the dataset statistic in the following table:

|                | #Graphs | #Nodes per graph | #Edges per graph |            Task Type             | Metric  |
| :------------: | :-----: | :--------------: | :--------------: | :------------------------------: | :-----: |
|  ogbg-molhiv   | 41,127  |       25.5       |       27.5       |   Graph binary classification    | ROC-AUC |
|  ogbg-molpcba  | 437,929 |       26.0       |       28.1       |   Graph binary classification    |   AP    |
|  LRGB-pepfunc  | 15,535  |      150.94      |      307.3       | Graph multi-label classification |   AP    |
| LRGB-pepstruct | 15,535  |      150.94      |      307.3       |   Graph multi-label regression   |   MAE   |

[1] Rethinking attention with performers. In ICLR, 2021.

[2] Big Bird: Transformers for Longer Sequences. In NIPS, 2021.

Reviewer comment (Experiments):
The proposed approach has the potential to be employed across various types of graphs, as the concept of "virtual atoms" could be (and already have been) utilized in other datasets like the OGB benchmarks. However, the authors did not explore this possibility in their work.

Author response:
Thanks for your valuable feedback. Here, we would like to clarify the difference between Nerual Atoms and other supernode methods. We will elaborate our justification from five: the **number of supernode or neural atoms, the information aggregating strategy, the interaction among them, the method to project the information back to the original graph**

Both the pipeline of supernode and neural atoms for obtaining global graph information can be described as two steps. 

- **Information aggregating**: The information of atoms within the molecular graph is aggregated into either supernode or multiple neural atoms with pair-wise connection. 
- **Interaction among node/atoms**: The second step shows differences in interaction among node/atoms. Supernode commonly exists alone, which means there is **only one** super node and thus lacks the ability to interact with others like neural atoms. 
- **Backward projection:** The final step shows differences in terms of the interaction among nodes/atoms. As supernode commonly exists alone, it thereby lacks the ability to interact with others like neural atoms. 

We provide a detailed comparison of supernodes and neural atoms in the following table.


|                                          |                          Super Node                          |                         Neural Atoms                         |
| :--------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: |
|             #Atoms / #Nodes              | 1 (In most cases, there is only one single supernode for aggregating the global graph information, as increasing the number of which might not lead to performance improvement.) | $K$ (Our neural atoms can be defined as the proportion of the average number of atoms of the molecular graph dataset, which is **significantly smaller than the original number of atoms. The more neural atoms, the better the performance.**) |
|         Information aggregating          | The global pooling method, e.g., global sum/mean/max pooling, **treats all the information from nodes the same**, thus lacking diversity among different nodes in the graph. | The multi-head Attention mechanism allows **aggregating information with different weights according to the similarity between specific neural atoms and the original atoms**, which allows diversity among different atoms. |
| Interaction among supernode/neural atoms | None. (None. (Since supernode usually exists alone, it thus lacks the ability to interact with others.) | A fully connected graph with an attention mechanism to bridge the neural atoms for information exchange. This allows the information located in a different part of the graph, even with a large hop distance, to share information based on their embedding similarities. |
|           Backward projection            | Direct element-wise adding. The information of the supernode is directly added to the representation of each node, which fuses the information from the single supernode, which might lead to the similarity among different atoms. | Weighted combination. The neural atoms can fuse the information within according to the similarity score between t neural atoms and the atoms in the original molecular graph. Such a mechanism allows the model to obtain diversity representation for further purposes. |

We have added this to our Appendix F; we kindly refer the reviewer to check our latest version of the draft.

Reviewer comment (Writing):
The quality of writing in the paper seems to get worse as you read further, especially in Section 4. There are a lot of grammar mistakes in Section 4.2. Additionally, there are some terms like "ratio" in Table 4 and "varying proportion" in the last paragraph on page 7 that are introduced without prior explanation. This gives the impression that the paper was completed hastily.

Author response:
Thanks for pointing out these problems. Following your suggestions, we have thoroughly revised the submission to improve the writing and presentation, which are highlighted in blue. We kindly refer the reviewer to assess the updated version. Besides, the concept of “proportion” is introduced at the beginning of Sec 4.2, which refers to the number of neural atoms, as it is set according to the average number of atoms in the dataset. In addition, we remove the term “ratio” to eliminate the potentially misleading information.

Reviewer comment (Novelty):
The novelty of the proposed method is limited. For example, from the perspective of general graph machine learning, supernodes have been widely used which are the same as the concept of neural atoms.

Author response:
**Reply:** Thanks for your valuable feedback. Here, we would like to clarify the difference between Nerual Atoms and other supernode methods. Analytically, the pipeline of both supernode and neural atoms for obtaining global graph information can be described as three steps. 

- **Information aggregating:** The information of atoms within the molecular graph is aggregated into either supernode or multiple neural atoms with pair-wise connection. 
- **Interaction among node/atoms:** The second step shows differences in interaction among node/atoms. Supernode commonly exists alone, which means there is **only one** super node and thus lacks the ability to interact with others like neural atoms. 
- **Backward projection:** The final step shows differences in terms of the interaction among nodes/atoms. As supernode commonly exists alone, it thereby lacks the ability to interact with others like neural atoms. 

Further, we provide a detailed comparison of supernodes and our proposed neural atoms in the following table from five: (1) the number of supernode or neural atoms, (2) the information aggregating strategy, (3) the interaction among them, (4) the method to project the information back to the original graph.

|                                          |                          Super Node                          |                         Neural Atoms                         |
| :--------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: |
|             #Atoms / #Nodes              | 1 (In most cases, there is only one single supernode for aggregating the global graph information, as increasing the number of which might not lead to performance improvement.) | $K$ (Our neural atoms can be defined as the proportion of the average number of atoms of the molecular graph dataset, which is **significantly smaller than the original number of atoms. The more neural atoms, the better the performance.**) |
|         Information aggregating          | The global pooling method, e.g., global sum/mean/max pooling, **treats all the information from nodes the same**, thus lacking diversity among different nodes in the graph. | The multi-head Attention mechanism allows **aggregating information with different weights according to the similarity between specific neural atoms and the original atoms**, which allows diversity among different atoms. |
| Interaction among supernode/neural atoms | None. (None. (Since supernode usually exists alone, it thus lacks the ability to interact with others.) | A fully connected graph with an attention mechanism to bridge the neural atoms for information exchange. This allows the information located in a different part of the graph, even with a large hop distance, to share information based on their embedding similarities. |
|           Backward projection            | Direct element-wise adding. The information of the supernode is directly added to the representation of each node, which fuses the information from the single supernode, which might lead to the similarity among different atoms. | Weighted combination. Our proposed neural atoms can fuse the information within according to the similarity score between neural atoms and the atoms in the original molecular graph. Such a mechanism allows the model to obtain diversity representation for further purposes. |

We have added the above comparison to Appendix F; we kindly refer the reviewer to check our latest version of the draft.

Reviewer comment (Evaluation):
Some claims are controversial: in the Introduction, the authors claimed the disadvantages of transformers, especially the self-attention mechanism, but the way to project atom representations to neural atom representations in the proposed method still uses multi-head attention. Similarly, the controversy also happens in the running time experiment.

Author response:
**Reply:** Thanks for this comment. Here, we would like to (1) clarify the difference between neural atoms and Graph Transformers w.r.t. the number of atom processing, and (2) the rationality for our neural atoms adopts multi-head attention mechanism as a way of atom grouping strategy. In addition, we provide the (3) explanation for the small running time gaps between GNN enhanced with neural atoms and Transformer-based methods.

(1) Due to the Graph Transformer operating message passing fully connected, it could involve irrelevant information during the message aggregation and update process. In contrast, the neural atom map the atoms in the molecular graph to $K$ neural atoms, which are **far less than the size of the atoms of the original molecular graph**. Such a mechanism allows us to **filter out** potentially irrelevant information, since the attention mechanism aggregates the information according to the similarity between the embedding of neural atoms and the atoms within the original graph.

(2) Rather than processing attention on the fully-connected graph with huge time consumption, the neural atom can reduce the computation burden and thus restrict the running time. Specifically, instead of directly modeling the atom interaction in a fully connected graph, we **map the potential interaction into the space constructed by neural atoms, which is more sparse compared to the original graph** (Table 4/7/8/9, which shows that our method can exceed the baseline with even one-tenth number of node w.r.t. the average number of node in the original graph).

(3) For the running time experiment, the size of the graph (150.94 nodes, 307.3 edges per graph for both LRGB-pepfunc and LRGB-pepstruct) is relatively smaller compared to the graph of product (the Amazon Product-Computer, which have 13,752 nodes with 491,7222 edges in a single graph) or citation networks (the ogbn-arXiv, which have 169,343 nodes with 1,166,243 edges in one single graph). **The inherent nature of the molecular graph we target implies that it would not have a large graph size.** Thus, the running time gap compared to the other Transformer-based method might not be noteworthy.

Alternatively, we provide a running time comparison on the two graph datasets mentioned above, under the same experimental environment mentioned in the draft, to demonstrate the gap between the Transformer-based method and our proposed Nerual Atoms. The statistical information of the two applied datasets and running times are listed in the following tables.

|                         | #Graphs | #Nodes per graph | #Edges per graph |            Task Type            |  Metric  |
| :---------------------: | :-----: | :--------------: | :--------------: | :-----------------------------: | :------: |
|       ogbn-arXiv        |    1    |     169,343      |    1,166,243     | Node multi-class classification | Accuracy |
| Amazon Product-Computer |    1    |      13,752      |     491,7222     | Node multi-class classification | Accuracy |

|       Type       |       Model       |                    ogbn-arXiv                     |               Amazon Product-Computer               |
| :--------------: | :---------------: | :-----------------------------------------------: | :-------------------------------------------------: |
|       GNN        |        GCN        | w/ 2.1s (Params. 5352) ; w/o  0.5s (Params. 3032) | w/ 0.2s (Params. 15082) ; w/o  0.2s (Params. 12746) |
|       GNN        |       GCNII       | w/ 2.2s (Params. 5352) ; w/o  0.4s (Params. 3016) | w/ 0.2s (Params. 15354) ; w/o  0.2s (Params. 12730) |
|       GNN        |       GINE        | w/ 2.2s (Params. 5624) ; w/o  0.4s (Params. 3304) | w/ 0.3s (Params. 15066) ; w/o  0.2s (Params. 13018) |
|       GNN        |     GatedGCN      | w/ 2.3s (Params. 6520) ; w/o  0.6s (Params. 4184) | w/ 0.2s (Params. 16234) ; w/o  0.2s (Params. 13898) |
| Transformer +GNN |        GPS        |                OOM (Params. 6440)                 |                0.4s (Params. 16154)                 |
|   Transformer    | Transformer+LapPE |                OOM (Params. 5016)                 |                0.4s (Params. 14730)                 |

Reviewer comment (Experiments):
Experiments have been conducted on only three relatively small-size benchmark datasets. In fact, there are some larger and more recent benchmark datasets such as OC20 and OE62.

Author response:
**Reply:** Thanks for the valuable question! Here, we would like to (1) further clarify the setting of the molecular graph we focus on and (2) the reason why we are not comparing our method with either the OE62 or OC60 datasets. In addition, we provide the (3) explanation for the unavailability of a larger benchmark. 

(1) **As we focus on 2D intra-molecule interaction, both OC20 and OE62 3D inter-molecular benchmarks are incompatible with our setting.** These two datasets require the model to have the ability to handle the 3D coordinate information, and the neural atom does not have the ability to handle the 3D coordinate information for now, as we target the 2D scenario. **We have highlighted this point in the revised submission to avoid misunderstanding.**

(2) For the comparison with more recent work like [1], the Ewald-based GNN shares a similar spirit to our paper, but it mainly focuses on the inter-molecule interaction for 3D molecules. **This is completely different from our setting, as we pay more attention to the 2D intra-molecule interaction.** 

(3) The scale of the dataset we adopted, i.e., the Long-Range Graph Benchmark, is considered to be larger than the general molecular graph dataset w.r.t. the graph size. We show the comparison of the dataset statistic in the following table:

|                | #Graphs | #Nodes per graph | #Edges per graph |            Task Type             | Metric  |
| :------------: | :-----: | :--------------: | :--------------: | :------------------------------: | :-----: |
|  ogbg-molhiv   | 41,127  |       25.5       |       27.5       |   Graph binary classification    | ROC-AUC |
|  ogbg-molpcba  | 437,929 |       26.0       |       28.1       |   Graph binary classification    |   AP    |
|  LRGB-pepfunc  | 15,535  |      150.94      |      307.3       | Graph multi-label classification |   AP    |
| LRGB-pepstruct | 15,535  |      150.94      |      307.3       |   Graph multi-label regression   |   MAE   |

As we are interested in the atom-level interaction in an inductive learning scenario, the OGB link prediction benchmark might not meet our demand as it only provides a single graph with a transductive setting.

Due to the time limitation, we are unable to incorporate the 3D information into the neural atom mechanism, as the 3D coordinate position of the neural atom needs to be carefully initialized and handled to avoid potential bias from the 3D position.

Alternatively, we directly adopt the neural atom to the 3D scenario without 3D coordinate information and provide the comparison with Ewald-based GNN [1] (denoted as Ewald Block) on the OE62[2] dataset. We follow the same setting from [1] and use SchNet[3] as the backbone GNN. Results show that our method achieves **competitive performance even with the absence of 3D coordinate information**, showing the generalization and effectiveness of our proposed neural atom.

|                   | Energy MAE $\downarrow$ | Energy MSE $\downarrow$ | Using 3D information |
| :---------------: | :---------------------: | :---------------------: | :------------------: |
| Baseline (SchNet) |         0.1351          |         0.0658          |     $\checkmark$     |
|   + Ewald Block   |       **0.0811**        |       **0.0301**        |     $\checkmark$     |
|   + Neural Atom   |       **0.0834**        |       **0.0309**        |       $\times$       |

We have added this discussion to our draft in Appendix D.

[1] Ewald-based Long-Range Message Passing for Molecular Graphs. In ICML, 2023

[2] Atomic structures and orbital energies of 61,489 crystal-forming organic molecules. In Scientific Data, 2020.

[3] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In NeurIPS, 2017.

Reviewer comment (Evaluation):
From the comparison of the running time, there are no clear advantages of the proposed method compared to the transformers which conflicts with the claim in the Introduction (GTs are more computationally expensive than GNNs). What are the reasons? The multi-head attention used in step 1?

Author response:
**Reply:** Thanks for this comment. Here, we would like to (1) clarify the difference between neural atoms and Graph Transformers w.r.t. the number of atom processing, and (2) the rationality for our neural atoms adopts multi-head attention mechanism as a way of atom grouping strategy. In addition, we provide the (3) explanation for the small running time gaps between GNN enhanced with neural atoms and Transformer-based methods.

(1) Due to the Graph Transformer operating message passing fully connected, it could involve irrelevant information during the message aggregation and update process. In contrast, the neural atom map the atoms in the molecular graph to $K$ neural atoms, which are **far less than the size of the atoms of the original molecular graph**. Such a mechanism allows us to **filter out** potentially irrelevant information, since the attention mechanism aggregates the information according to the similarity between the embedding of neural atoms and the atoms within the original graph.

(2) Rather than processing attention on the fully-connected graph with huge time consumption, the neural atom can reduce the computation burden and thus restrict the running time. Specifically, instead of directly modeling the atom interaction in a fully connected graph, we **map the potential interaction into the space constructed by neural atoms, which is more sparse compared to the original graph** (Table 4/7/8/9, which shows that our method can exceed the baseline with even one-tenth number of node w.r.t. the average number of node in the original graph).

(3) For the running time experiment, the size of the graph (150.94 nodes, 307.3 edges per graph for both LRGB-pepfunc and LRGB-pepstruct) is relatively smaller compared to the graph of product (the Amazon Product-Computer, which have 13,752 nodes with 491,7222 edges in a single graph) or citation networks (the ogbn-arXiv, which have 169,343 nodes with 1,166,243 edges in one single graph). **The inherent nature of the molecular graph we target implies that it would not have a large graph size.** Thus, the running time gap compared to the other Transformer-based method might not be noteworthy.

Alternatively, we provide a running time comparison on the two graph datasets mentioned above, under the same experimental environment mentioned in the draft, to demonstrate the gap between the Transformer-based method and our proposed Nerual Atoms. The statistical information of the two applied datasets and running times are listed in the following tables.

|                         | #Graphs | #Nodes per graph | #Edges per graph |            Task Type            |  Metric  |
| :---------------------: | :-----: | :--------------: | :--------------: | :-----------------------------: | :------: |
|       ogbn-arXiv        |    1    |     169,343      |    1,166,243     | Node multi-class classification | Accuracy |
| Amazon Product-Computer |    1    |      13,752      |     491,7222     | Node multi-class classification | Accuracy |

|       Type       |       Model       |                    ogbn-arXiv                     |               Amazon Product-Computer               |
| :--------------: | :---------------: | :-----------------------------------------------: | :-------------------------------------------------: |
|       GNN        |        GCN        | w/ 2.1s (Params. 5352) ; w/o  0.5s (Params. 3032) | w/ 0.2s (Params. 15082) ; w/o  0.2s (Params. 12746) |
|       GNN        |       GCNII       | w/ 2.2s (Params. 5352) ; w/o  0.4s (Params. 3016) | w/ 0.2s (Params. 15354) ; w/o  0.2s (Params. 12730) |
|       GNN        |       GINE        | w/ 2.2s (Params. 5624) ; w/o  0.4s (Params. 3304) | w/ 0.3s (Params. 15066) ; w/o  0.2s (Params. 13018) |
|       GNN        |     GatedGCN      | w/ 2.3s (Params. 6520) ; w/o  0.6s (Params. 4184) | w/ 0.2s (Params. 16234) ; w/o  0.2s (Params. 13898) |
| Transformer +GNN |        GPS        |                OOM (Params. 6440)                 |                0.4s (Params. 16154)                 |
|   Transformer    | Transformer+LapPE |                OOM (Params. 5016)                 |                0.4s (Params. 14730)                 |

Reviewer comment (Experiments):
What will be the performance on larger datasets? Can the proposed method beat more recent SOTA such as [1]?

Author response:
**Reply:** Thanks for the valuable question! Here, we would like to (1) further clarify the setting of the molecular graph we focus on and (2) the reason why we are not comparing our method with either the OE62 or OC60 datasets. In addition, we provide the (3) explanation for the unavailability of a larger benchmark. 

(1) **As we focus on 2D intra-molecule interaction, both OC20 and OE62 3D inter-molecular benchmarks are incompatible with our setting.** These two datasets require the model to have the ability to handle the 3D coordinate information, and the neural atom does not have the ability to handle the 3D coordinate information for now, as we target the 2D scenario. **We have highlighted this point in the revised submission to avoid misunderstanding.**

(2) For the comparison with more recent work like [1], the Ewald-based GNN shares a similar spirit to our paper, but it mainly focuses on the inter-molecule interaction for 3D molecules. **This is completely different from our setting, as we pay more attention to the 2D intra-molecule interaction.** 

(3) The scale of the dataset we adopted, i.e., the Long-Range Graph Benchmark, is considered to be larger than the general molecular graph dataset w.r.t. the graph size. We show the comparison of the dataset statistic in the following table:

|                | #Graphs | #Nodes per graph | #Edges per graph |            Task Type             | Metric  |
| :------------: | :-----: | :--------------: | :--------------: | :------------------------------: | :-----: |
|  ogbg-molhiv   | 41,127  |       25.5       |       27.5       |   Graph binary classification    | ROC-AUC |
|  ogbg-molpcba  | 437,929 |       26.0       |       28.1       |   Graph binary classification    |   AP    |
|  LRGB-pepfunc  | 15,535  |      150.94      |      307.3       | Graph multi-label classification |   AP    |
| LRGB-pepstruct | 15,535  |      150.94      |      307.3       |   Graph multi-label regression   |   MAE   |

As we are interested in the atom-level interaction in an inductive learning scenario, the OGB link prediction benchmark might not meet our demand as it only provides a single graph with a transductive setting.

Due to the time limitation, we are unable to incorporate the 3D information into the neural atom mechanism, as the 3D coordinate position of the neural atom needs to be carefully initialized and handled to avoid potential bias from the 3D position.

Alternatively, we directly adopt the neural atom to the 3D scenario without 3D coordinate information and provide the comparison with Ewald-based GNN [1] (denoted as Ewald Block) on the OE62[2] dataset. We follow the same setting from [1] and use SchNet[3] as the backbone GNN. Results show that our method achieves **competitive performance even with the absence of 3D coordinate information**, showing the generalization and effectiveness of our proposed neural atom.

|                   | Energy MAE $\downarrow$ | Energy MSE $\downarrow$ | Using 3D information |
| :---------------: | :---------------------: | :---------------------: | :------------------: |
| Baseline (SchNet) |         0.1351          |         0.0658          |     $\checkmark$     |
|   + Ewald Block   |       **0.0811**        |       **0.0301**        |     $\checkmark$     |
|   + Neural Atom   |       **0.0834**        |       **0.0309**        |       $\times$       |

We have added this discussion to our draft in Appendix D.

[1] Ewald-based Long-Range Message Passing for Molecular Graphs. In ICML, 2023

[2] Atomic structures and orbital energies of 61,489 crystal-forming organic molecules. In Scientific Data, 2020.

[3] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In NeurIPS, 2017.