3D Molecular Pretraining via Localized Geometric Generation

Thank you for the thorough review and positive feedback - your suggestions will lead to a stronger analysis. To address your concern, we provide the following point-to-point responses.

Q1:...... I believe that the RL reconstruction of molecular geometries from 3D fragments in Flam-Shepherd et al (https://arxiv.org/abs/2202.00658) relates closely to the core inspiration of this method and might warrant discussion in the intro....

Response: We appreciate you highlighting the promising fragment-based reconstruction approach by Flam-Shepherd et al. As you suggested, we will cite this relevant work in the introduction and discuss extensions of our tetrahedron segmentation to incorporate larger functional fragments as building blocks. This could enable more complex biologically meaningful reconstruction objectives in our future work.

Q2: Can the authors comment on the disparity of the approach of their validation scores on moleculeNet vs those on the large-scale challenge?

Response: For PCQM4M-v2, we would like to argue that existing top approaches incorporate sophisticated graph encoding priors, while we apply a pure transformer. Nevertheless, the benefits of our pretraining approach can be demonstrated by the performance increase over the non-pretrained TokenGT model. The primary contribution of this work is to give a glimpse at how proper selection of semantic units impacts 3D molecular pretraining, and we believe a further introduction of graph inductive bias will further improve our result in table2.

Q3: The ablation study is not helpful: the perturbation of the model is too limited, and table 3 suggests strongly that the parameters are actually not optimal. The random perturbation pretraining of table 4 is not described in a clear enough fashion.

Response: Thank you for the feedback requesting more extensive ablation experiments and analysis, and we have revised the abaltion study section in our paper. In table 3, we show that excessive noise in 3D structure denoising can lead to training divergence and detrimental impacts and validate the setting of our perturbation strategy. In table 4, we validate the benefits of structural pretraining by comparing against a naive random atom-level denoising baseline with equivalent 36% masking.

We really appreciate the helpful comments and critiques that allow us to improve LEGO. Regarding your concern, here is our response.

Q1: The geometrical justification of tetrahedra as simplest polyhedron might make sense, but in chemistry it makes a lot less sense. The authors literally show benzene in Figure 1b, which has 120-degree bonding pattern (so called planar-trigonal in chemistry) that is NOT a tetrahedron at all, and has very different local symmetries

......

What happens for atoms with more than 4 bonds (sulfur, phosphorous, etc) ?

Response: Thank you for highlighting this important point. We have expanded the motivation section (page 4) to better address cases where local molecular structures deviate from standard tetrahedral geometry. While tetrahedra serve as an informative reference decomposition target due to their simplicity, our segmentation strategy adapts accordingly when this arrangement does not naturally arise.

For instance, center atoms with fewer than four one-hop neighbors, as commonly observed in functional groups (e.g. ketenes, amines, ethers in Fig. 1b), are handled by treating the local unit as a degraded tetrahedra. Meanwhile, atoms forming more than four bonds, like sulfur and phosphorous, we incorporate all its one-hop neighbors into the unit. For cyclic structures such as benzene, non-adjacent ring atoms are selected, effectively partitioning the system into triangular fragments. In summary, while tetrahedral local structures provide a useful illustration, our method maintains flexibility to accommodate chemical environments where alternative local symmetries dominate.

Q2: "We attribute this to the different 3D structures molecules exhibit in liposome compounds." What does this mean ? What are these structures different ?

Response: Thank you for raising this question. To clarify, when referring to the "different 3D structures exhibited by molecules in liposome compounds", we intended to highlight potential differences between conformations in a lipid environment versus the conformations in a water environment, or the equilibrated 3D conformations used during pre-training. As you pointed out, our original phrasing was unclear on this distinction.

We hypothesize that the weaker performance on the Lipo task may arise because the stable conformations adopted by molecules in a lipid-like environment can differ substantially from common isolated conformations. As our pre-training exclusively saw the latter, transferring representations to predict lipid-specific properties poses an additional challenge.

Q3: All the biochemistry prediction tasks are actually properties of the graph, not the 3D structure, What 3D structure is being used ?

Response: We would like to point out that we adopt only 2D graph inputs for the biochemistry prediction tasks to conform with established dataset conventions and existing baseline approaches on MoleculeNet, including GraphMVP, 3DInfomax, and other 3D pre-trained models - all of which use 2D inputs. As such, while not directly evaluating 3D structure modeling, these tasks probe valuable model capabilities: 1.Encoding functional semantics related to downstream biological effects into 2D topology representations. 2.Assessing transferability of features learned from 3D pre-training when adapted to 2D inputs.

Conversely, we intentionally select 3D-sensitive quantum property prediction as an additional testbed for explicit 3D molecular structure modeling, using true 3D conformations as inputs. This directly examines the model's proficiency at encoding geometric details within continuous predictions - complementing the more functional evaluation on biochemical properties.

We sincerely thank you for the valuable remarks that help us better examine our work. To address your concern, we present the following responses.

Q1: What exactly constitutes the input for TokenGT-3D? Is it the original molecules, or do you utilize each local structure segmentation after masking and perturbation? Or is it the masked and perturbed local structure segmentations of a single molecule, or something else entirely?

Response: During pretraining, TokenGT-3D still take node and edge tokens, i.e. the original molecules as the input, as is shown in Figure 3 on page 4. The attribute mask and coordinate noise at perturbation stage are implemented token-wise during perturbation. Afterwards, the model takes both clean and perturbed tokens of the original molecule as the input; the segmentation on molecular conformations does not function as a 3D tokenizer.

Q2: Given the proposed method centers on pre-trained representation learning, what form do the learned representation embeddings take for downstream tasks?

Response: A $[`CLS`]$ token is concatnated after the augmented node and edge embedding before the forward process of the Transformer encoder and serve as the graph representation for downstream tasks. This BERT-style method follows the implementation in the original TokenGT.

Q3: Could you elaborate on how the local structures are reconstructed? What serves as the input for this process: a single embedding from the TokenGT-3D output, or a collection of embeddings from local structure segmentations within a single molecule?

Response: Thanks for your feedback. To better illustrate the local structure reconstruction procedure, we have revised the paper to reposition Algorithm 1 within the methodology section. As is previously stated in Q1, the model takes both clean and perturbed tokens as input, producing an embedding $\mathbf{Emb}$ in shape $[n+m,dim]$. During pretraining, we utilize three separate mask indicators -- $M_\text{center}, M_\text{edge},M_\text{leaf}$ -- for center atoms, edges and leaf atoms within the whole molecule. For reconstruction, the central atom embeddings, $\mathbf{Emb}[M_\text{center}]$, are fed into a specialized head to predict global orientation. Likewise, edge embeddings $\mathbf{Emb}[M_\text{edge}]$ and leaf node embeddings $\mathbf{Emb}[M_\text{leaf}]$are input into dedicated modules for reconstructing edge lengths and angles, respectively.