3D-Transformer: Molecular Representation with Transformer in 3D Space 

Abstract: Spatial structures in the 3D space are important to determine molecular properties. Recent papers use geometric deep learning to represent molecules and predict properties. These papers, however, are computationally expensive in capturing long-range dependencies of input atoms; and more importantly, they have not considered the non-uniformity of interatomic distances, thus failing to learn context-dependent representations at different scales. To deal with such issues, we introduce 3D-Transformer, a variant of the Transformer for molecular representations that incorporates 3D spatial information. 3D-Transformer operates on a fully-connected graph with direct connections between atoms. To cope with the non-uniformity of interatomic distances, we develop a multi-scale self-attention module that exploits local fine-grained patterns with increasing contextual scales. As molecules of different sizes rely on different kinds of spatial features, we design an adaptive position encoding module that adopts different position encoding methods for small and large molecules. Finally, to attain the molecular representation from atom embeddings, we propose an attentive farthest point sampling algorithm that selects a portion of atoms with the assistance of attention scores, overcoming handicaps of the virtual node and previous distance-dominant downsampling methods. We validate 3D-Transformer across three important scientific domains: quantum chemistry, material science, and proteomics. Our experiments show significant improvements over state-of-the-art models on the crystal property prediction task and the protein-ligand binding affinity prediction task, and show better or competitive performance in quantum chemistry molecular datasets. This work provides clear evidence that biochemical tasks can gain consistent benefits from 3D molecular representations and different tasks require different position encoding methods.  

Reviews and author responses:

 We appreciate your valuable comments and polish our work as follows:


1. Hard to read.

We apologize for our poor and inconsistent arrangement of the paper. Due to the limited paper length, we put aside many parts and important information in the appendix. However, as you suggested, we should not concentrate too much on the abstract, introduction, and contributions. Therefore, we shorten these three sections, leaving only three major contributions of our work, and move essential information ahead to make it more consistent. Also, we delete experiments in material science, because we believe it is better to demonstrate the superiority of our approach in some mature datasets and domains. 

2. Contribution is not properly claimed. 

Thanks for pointing out some work similar to ours. However, we believe our method is significantly different from [1][2]
in the following aspects. (1) We propose a multi-scale self-attention to deal with non-uniformity in atomic distances, and introduce a new subsampling method to accumulate the molecular representations. However, both [1] and [2] focus on how to utilize the distance matrix, adjacency matrix, and shortest paths. Namely, they are concentrating more on coding position, which is most similar to our convolutional position encoding. (2) We design trainable motif-based embeddings to store semantic meanings of important substructures that are defined by functional groups. This is our innovation to first incorporate motifs to a Transformer model. 

Reference: [1] Molecular Attention Transformer [2] Relative Molecular Attention Transformer

3. Title.

Thanks for your suggestion on how to change the name. Three of four reviewers hold the opinion that our title is misleading. We correct it and if you still have any advice, we are pleased to know.

4. Ablation study.

Thanks for pointing out our contributions of applying 3D coordinates, but we believe other improvements like motif-based embeddings, AFPS are also novel and interesting. As you suggest, we compare our methods with the naive methods which directly add such coordinates to the graphs. The ablation study shows that our convolutional position encoding significantly reduces MAE in QM7 from 63.2 to 17.6, and in QM8 from 0.0205 to 0.104. 

5. Performance in QM7.

We re-exam our experiments in QM7 and confirm that our results are correct. I believe different models might perform differently in different tasks/targets. For instance,  MPNN performs well in classification problems but achieves higher MAE in regression tasks. Our method is good at capturing the local patterns with a multi-scale mechanism, this might be the most important factor in predicting the only property in QM7 but not other properties (QM8 has 12 properties). 

6. Classification tasks.

We value your instructions and add more experiments of classification tasks, including BBBP, ClinTox, and BACE. Our results show that our model outperforms [1] and [2], and is the best of all non-pretrained models. However, pretrained GROVE is better than ours, which indicates the power of pretraining on large-scale datasets. In those classification tasks with fewer samples, pretraining is of great effectiveness and we believe if our model is pretrained its performance will be significantly improved. 

 1. Novelty.

We cannot deny the fact that we use some classic methods in 3D cloud points for reference. However, our approach has the following significant difference and contributions: (1) Multi-scale mechanism has been proposed for many years, but it has been barely used after the explosion of data and the invention of the Transformer. This is due to the fact that the pretrained Transformer does not need this mechanism to catch local patterns. We exam this idea in our ablation study and find that in larger datasets like QM8, its improvements are much smaller. We re-adopt this mechanism in the molecular representation problems for the sake of alleviating the problem of small datasets. (2) As for AFPS, it is based on the traditional FPS, but its motivation is unique. In a 3D point cloud, the points have no type information. Nevertheless, in molecules, each atom (point) has its class such as H, O, C, N, F. Their contributions to the property are different. Therefore, we design an AFPS to consider both their importance by using the attention matrix and the distances between each other. Therefore, from our point of view, it's not a simple FPS with a threshold. (3) Position encoding is a very essential part of the Transformer. Without it, Transformer cannot distinguish its locations. We adopt the pairwise distance matrix to encode their position and propose a convolutional PE in a linear time of complexity.

Additionally, we update our paper and add a trainable motif-based embedding module to incorporate the knowledge of momentous substructures defined by functional groups. Our experiments show it leads to better performance. We also conduct ablation studies to show that different kinds of motifs contribute differently to the downstream tasks. In sum, we offer a promising way to inject domain knowledge into Transformer by constructing a motif token vocabulary and hope to shed light on how to make better usage of those informative motifs.  


2. Suggestions about the title.

As pointed out by another reviewer, we are certainly aware of the overblown of our title. We reconsider and choose 'Motif-based Roto-translation Invariant Transformer for Molecular Property Prediction in 3D Space'. Do you think this one better describes the proposed method in our research domain more elegantly?

3. Reply to other minor concerns.

(1) Two 3DGNNs.

We recheck our baselines in Table 5, and believe it is due to our writing that causes your confusion. To be specific, in table 5, Townshend et al. (2020) propose two different methods: 3DGNN and 3DCNN. One is based on the convolutional neural networks, while another is based on the graph convolutional neural networks. These 3DCNNs can encode translational and permutational symmetries, but incur the significant computational expense and cannot capture rotational symmetries without data augmentation. On the other hand, 3DGNNs encode their pairwise connectivity, capturing a structured representation of atomistic data. These two methods are very close in their names and can be easily incorrectly recognized with each other. 

(2) Performance.

As a matter of fact, the performance is really close between our model and MPNN in QM8: MAE of MPNN is 0.015, and our MAE is 0.009. There does exist a huge gap between our model and MPNN in QM7. We suppose this is due to the fact that the results of MPNN in QM7 are reproduced by [1] but its performance in QM9 is in the original paper. [1] exam MPNN in a set of molecular tasks from MoleculeNet. However, its performance is not stable. For instance, it achieves excellent results in classification tasks like BBBP, Tox21, but performs badly in regression tasks like QM7, FreeSolv. 

Reference: Self-Supervised Graph Transformer on Large-Scale Molecular Data


 Hi, we appreciate you provide some valuable suggestions, and our responses are listed below:

1. Statement referring to prior work.

You are absolutely correct that we overclaim in some parts of our paper. We delete the declaration of "To the best of our knowledge, 3D-Transformer is the first transformer variant..." and modify some statements in a more humble way. 

However, as for hierarchical learning [1, 2], we believe their method is significantly distinct from our subsampling approach. Since they aim at proteins, which have a very uniform sequence arrangement, they select alpha carbons as the intermediate set of points and aggregate information at the level of those carbons for the entire complex. However, the structures of small molecules have no such stable paradigm. We provide a novel universal methodology to adaptively subsample atoms without any prior assumptions on the atom arrangement. 

reference: [1] Hierarchical, rotation‐equivariant neural networks to select structural models of protein complexes
[2] Protein model quality assessment using rotation-equivariant, hierarchical neural networks


2. Methodological novelty.

The distance map representation for small molecules and sinusoidal positional encoding do exist in previous work. However, we combine these two methods together into a universal Transformer architecture for all molecular presentation learning problems. Besides, we design a threshold based on the complexity of different types of molecules to decide when to adopt different position encoding strategies. According to these facts, we believe our work can provide some insights into the biochemical community. 


3. How to encode the protein-ligand complexes. 
We are pleased that you like our adaptive representation. It seems to become a default setting that small molecules and large molecules require different kinds of model architecture in recent years. Therefore, we analyze their difference and establish a threshold to distinguish them clearly. As for the protein-ligand task, we combine each pair together into one and feed it directly to the model. There we did not separate these two molecules apart and this compound is a large molecule, which has a very complex 3D structure. Nevertheless, it might be worth a try if we separately learn representations of each small molecule and protein with different position encoding methods. But we are confident that this will be a bad choice instead of seeing them as a whole because the latter allows atoms inside small molecules to immediately connect with atoms inside proteins. 

4. Other advice.

We will take care of the way to provide the code in the future. Instead of including a GitHub link, we will submit the code immediately as part of the manuscript. 


 1. Representation presentation improvement. 

Thanks for pointing out some unclear portions of our work! \boldsymbol{A} = [a_{ij}] \in \mathbb{R}_{N\times N} is the attention matrix. We described it in the preliminary, and as you hope, we remove it after the Equation so that readers can comprehend it more easily. Thus, there is no inconsistency between equations. You are correct that a "max" is missing in equation (4) and we modify this typo mistake. 

Here we provide a more clear explanation of our proposed AFPS. As you know, each row of the attention matrix \boldsymbol{A} sums up to 1 if we operate a softmax on its rows. Therefore, the accumulation of its columns can reflect the importance of each atom on the calculations of self-attention. Therefore, we sum up \boldsymbol{A} along columns and use it to select atoms based not only on distance but also on its attentive significance. We correct our writing on this part and believe it is more readable and friendly for understanding.

2. Motivations regarding the absolute position embeddings.

Absolute position encodings do have the capability to distinguish chiral molecules, but unfortunately, it fails to achieve roto-translation transformation invariance. After serious consideration, we delete all parts of the usage of absolute position encodings and adopt the convolutional position encoding for our model. Accordingly, we concentrate our experiments on small molecules instead of three domains (namely, excluding materials and proteins).  

As for why chirality is not that important for small molecules, we are pleased to offer some opinions. Actually, chirality exists in both small and large molecules, but not all properties differ according to chiral molecules. In other words, different properties are influenced by chirality to varying degrees. Most properties of small molecules are not chiral while some properties of large molecules are strongly influenced by chirality. For example, as stated in another paper submitted to ICLR 2022 (see reference below), they declare that chirality can strongly alter biochemical interactions, particularly protein-drug binding. Therefore, we guess that is why absolute position encoding succeeds in all representation problems in large molecules. 

Reference: Learning 3D Representations of Molecular Chirality with Invariance to Bond Rotations

3. Questions about the threshold.

It is a hyper-parameter tuned in a grid search. As you suggest, 'adapt' seems over-claim, and it might be a good choice to make it a trainable parameter or perhaps immediate using a linear combination of both position encoding methods. However, since we abandon the AEP due to the concerns of roto-translation invariance, we move our attention only on the small molecules and no longer establish a threshold to separate both methods. 

4. Comparison of AFPS with other approaches. 

We are happy that you like the idea of AFPS. Your worry is reasonable. Therefore, we add some more ablation studies to compare our AFPS with other pooling (readout) methods like set2set, max pooling, mean pooling. And the results prove that our AFPS outperforms other pooling mechanisms.

5. Response to other suggestions. 
(1) We add more details to the hyperparameters in Appendix for better reproducibility and a deeper understanding of our proposed methods. We adopt threshold\tilde{\phi} = , \epsilon = 0.1 in AFPS.There are three scales, and the distance threshold \tau = [\rho / 2, \rho, 2 * \rho], where \rho is the density of each dataset. 
(2) Thanks for your kind reminder. We agree with your suggestion to move APE to the main text. Do you think we need to move the equation of expanding and translating cell units to the main body? 
(3) We move the table from the appendix into the main text and will perform more ablations on other datasets to demonstrate the effectiveness of each individual component.  


The authors target the molecular property prediction problem and propose the 3D-transformer to utilize 3D geometric information of molecules and take advantage of the advanced (graph) transformer architecture. Specifically, the authors propose three key components of the 3D-transformer, namely, the adaptive position encoding, the multi-scale self-attention, and the attentive farthest point sampling.  The paper studies an emerging problem of learning representations from 3D structure of molecules. Several interesting methods taking advantage of the molecular graph transformers and molecule 3D information are proposed. The experimental results seem solid. However, there are several concerns and questions regarding the methods and presentation.

1. The current presentation needs improvement. The description of the entire computing pipeline is not clear, and many details are missing or confusing. For example, it’s unclear that what the notation \boldsymbol{A} stands for in Eq.(4) and the notation of input/output are inconsistent for Eqs (4) and (5). The description “A are accumulated along column … computation by maximizing \tilde{A}_i” and Eq. (6) (a “max” is missing?) are confusing. Also, according to the description, line 1 in Algorithm 1 seems incorrect (\tilde A <- sum of A along column?). I suggest the authors examine the writing again and make the presentation more clear.

2. The motivation of absolute position embedding for large molecules is reasonable to me, i.e., it aims to distinguish chiral molecules. However, it is not clear that whether the absolute position embedding still satisfy the roto-transformation invariance, as it is still an important property for even large molecules. Besides, could the authors explain why chirality is not that important for small molecules?

3. Is the threshold \tilde{\phi} a trainable parameter or hyper-parameter tuned in the grid search? If it is a hyper-parameter, then “Adapt” seems an over-claim. Also, the discrete condition separating large molecules and small molecules does not make sense to me, unless there is a pivot point on the molecule size and their properties. Are the thresholds the same for all datasets or tuned individually (I suppose the threshold would be the same according to the motivation)? Is there any insights on using the discrete condition instead of soft conditions such as a linear combination of CPE and APE? It would be better the authors also provide the learned/tuned threshold values for each dataset or other results to support the claim.

4. The idea of AFPS is interesting as a novel sample-based global pooling (readout) operation. However, there might be a potential issue with the operation during optimization. The selection of the subset is based on the representation of each node, and the optimization of node representations requires gradient back-propagated through the node. If a node is not initially selected, the optimization of the node may lack of supervision (even with the gradient from attention, which can be weak), preventing the node to be further selected. Could the author compare the effect of different pooling methods?


Other suggestions:
1. Could the author provide all tuned values of hyper-parameters or constant numbers (threshold \tilde{\phi}, \epsilon in AFPS, distance threshold \tau_s, etc.) in the appendix for better reproducibility and a deeper understanding of the proposed methods.
2. It would be better to move equations describing the computation of proposed methods to the main text (e.g., the computation of APE).
3. As the proposed 3D-transformer consists of multiple novel components/operators, it would be better to conduct more ablations like Table 8 in C.1 (on more datasets) to help understand the effectiveness of each component.
 The paper studies an emerging problem of learning representations from 3D structure of molecules. Several interesting methods/operations taking advantage of the molecular graph transformers and molecule 3D information are proposed. The experimental results seem solid. However, there are several concerns and questions regarding the methods and presentation.

The paper presents an approach to represent 3D atomic structure in a machine learning context. 

The paper distinguishes between representations for "small" and "large" molecules based on a metric that takes the spatial extension and number of atoms in the molecule into account. Structures of small molecules are represented by an interatomic distance map. Large molecules are represented by what the authors refer to as a "sinusoidal function-based absolute position encoding method". In both cases, a transformer architecture is used on top of this initial representation. The authors further introduce a subsampling procedure to select a subset of points/ atoms and aggregate information. 

The authors demonstrate quantitative results on datasets relating to 1) small molecule property prediction (QM datasets), 2) protein-ligand binding (PDBbind dataset), and 3) a dataset from material science on metalorganic compounds (CoRE-MOF dataset) 

 The question of how to efficiently represent and learn from 3D atomic structure has attracted a lot of interest in recent years due to a wide range of promising applications. Different representations might be adequate for small molecules (with the defining characteristic of the chemical formula) compared to structures of larger molecules (such as proteins). 

My main concern regarding this paper are statements referring to prior work that I think are incorrect. 
An example of this is the first bullet point listed unter contributions "To the best of our knowledge, 3D-Transformer is the first Transformer variant that incorporates 3D spatial information for molecular representation learning.". The authors of Fuchs et al (2020) and the recent Alpha Fold paper will likely disagree. There is also a body of work of hierarchical learning from 3D molecular structures that is not referenced in relation to the introduced subsampling operation (see for example Eismann et al, Proteins, 2020). 

The methodological novelty in terms of molecule representation is limited, the distance map representation for small molecules and the sinusoidal function-based absolute position encoding method (Li et al., 2019) for larger molecules have previously been established. 

The adaptive representation is in general an interesting idea. However, it would be great if the authors could clarify how they encoded the protein-ligand complexes in the PDBbind dataset. In section 5, the authors state that their proposed metric Phi is "a good indicator of small or large molecules". The protein-ligand complexes contain both a large molecule (or parts of it in terms of the binding pocket) and a small molecule (the ligand). 

The quantitative results for the three datasets indicate that the method is performing well.

In the future, I would not include a Github link as part of the manuscript but provide the code as part of the submission. 

 My recommendation is based on the 1) methods' representation with respect to prior work and 2) limited technical novelty. 

This paper proposes 3D-Transformers, which is a variant of Transformer that incorporates the 3D spatial information.
The authors develop a multi-scale attention module, an adaptive position encoding module, and a farthest sampling scheme for the 3D-Transformer.
The authors validate the performance of 3D-Transformer on three different domains. My major concern regarding this paper is the contribution and the experiments.
The paper proposes a "multi-scale attention" module and a "farthest point sampling" module. These are all built upon the design of PointNet++(Qi et al., 2017b) without significant modifications. Another contribution, namely ATPE, is just thresholding based on existing techniques. Most importantly, the effectiveness of these designs is now properly validated in the experiments (quantitatively). Without these comparisons that clearly show the difference for each module, the proposed "3D-Transformer" is very similar to an attention-based PointNet++ that is specially designed for molecular data.

Also, I don't like the title 3D-Transformer. 3D deep learning is a very large research field. Because the paper only focuses on molecular representation, the current title would cause confusion to readers in the actual 3D computer vision area, such as point cloud and mesh. I suggest the author reconsider the title that could better describe the proposed method in the corresponding research domain to avoid overblown.

Other minor concerns are:
In table 5, why there are two 3DGNNs? what is their difference?
Why the performance of the 3D-Transformer is close against MPNN on QM9, but shows a very huge gap on QM7 and QM8?
There are many typos and grammar mistakes in the paper. I recommend the authors carefully polish the writing.


 Overall, there are several components proposed in this paper. But they are not properly evaluated, which makes it hard to verify the significance of these contributions.

In the paper, the authors proposed a new Transformer dedicated to working with molecular representation. It is a modification of classical Transformer using the information of 3D coordinates of atoms. 3D-Transformer operates on a fully-connected graph with direct connections between atoms. In the paper, the authors proposed 3D-Transformer. It is a modification of classical Transformer with using the information of 3D coordinates of atoms. 3D-Transformer operates on a fully-connected graph with direct connections between atoms.

The paper is hard to read since the authors concentrate on the abstract, introduction, and listing contribution. Some important information was moved to the Appendix (like intuition behind the measure of molecular structure complexity in the 3D space). Unfortunately, it makes it not consistent. 

The main contribution is problematics since in literature, existing papers uses 3D coordinates in transformer architecture dedicated to  molecular representation:
[1] Molecule Attention Transformer Ł. Mazurka at all.
[2] Relative Molecule Self-Attention Transformer Ł. Mazurka at all.
In my opinion, authors should add to related work comparison to [1,2] and numerical comparison to [1]. 
The methods seem to be different, but the main contribution is problematic. 

It is not trivial to understand Fig. 1. Therefore, the caption should better describe the image. 

3D-transformer is misleading name since it suggests that the authors use 3D point cloud representation. In the title, the authors should add some information that the graph was used. 

If the main contribution is the application of 3D coordinates, the author should compare such a solution with the naive method, which adds such coordinates to the label of the graph.
 
At the QM7 dataset, 3D-transformer outperforms other algorithms. Do you have some idea why it happens, what caused such difference? 

Competition between baselines is applied only on regression tasks. Therefore, I think that authors should use classification tasks from [1,2] or from (cited in the article) “Self-Supervised Graph Transformer on Large-Scale Molecular Data,” where authors use BBBP, SIDER, ClinTox, BACE, Tox21, ToxCast datasets. The paper proposes an interesting modification of transformer architecture dedicated to the chemical compound. Unfortunately, the method is not well situated in the existing literature. In my opinion, the Abstract, Introduction, and list of contributions are too extended. On the other hand, some critical pieces of information are moved to the appendix. Furthermore, experiments are applied only on classification tasks. 

Meta-review: The authors present a method to learn representations of 3D atomic structure. They consider two cases: "small" and "large" molecules based on a metric that takes the spatial extension and number of atoms in the molecule into account. Small molecules are represented by an interatomic distance map. Large molecules are represented by a "sinusoidal function-based absolute position encoding method". Both settings make use of a transformer architecture on top of the initial representation. The authors also introduce a subsampling step to select a subset of points/atoms and aggregate information from these.  Experimental results are shown for datasets relating to small molecule property prediction, protein-ligand binding and a dataset from material science on metalorganic compounds.

Strengths:

- interesting modification of transformer architecture dedicated to the chemical compound.

Weaknesses: 

- Poor presentation of methods with respect to prior work
- Limited technical novelty. The distance map representation for small molecules and the sinusoidal function-based absolute position encoding method for larger molecules have previously been proposed. Many components are built upon the design of PointNet++(Qi et al., 2017b) without significant modifications. The proposed "3D-Transformer" is very similar to an attention-based PointNet++ that is specially designed for molecular data.
- Experiments are applied only on classification tasks.

All reviewers voted for rejection. I recommend the authors to address the limitations listed above by improving the presentation with respect to prior work, clarifying more the novelty of the methods and including a more diverse range of experiments.

Decision: Reject