Hierarchical Graph Tokenization for Molecule-Language Alignment

Dear Reviewer VbFG,

Thank you for your time and suggestions to our work. Please find our detailed responses to your questions below:

Related work section

Thank you for acknowledging our contribution as the first to incorporate the hierarchical graph information. As for the referred works, we have revised Sec 2 as the related work section, and included the discussion of existing attempts:

Existing works also try to enrich the molecule-language alignment with additional modalities, such as 2D[2] and 3D[1] information. In contrast, we focus on the intrinsic hierarchical information of the molecules, such as motifs.

It is relatively more straightforward given the function groups in the inputs

We kindly refer Reviewer VbFG to Table 17 in the appendix, where we can find that, even with the functional groups in the inputs (i.e., HIGHT w/o HiPubChem), it remains challenging for LGLM to correctly recognize whether the functional groups exist in the input molecule.

In addition, it is not straightforward that the hierarchical alignment across atom-functional group-molecule is beneficial to broader downstream tasks.

Why choose VQVAE to obtain node embedding

We choose to use VQVAE because it is one of the simplest and standard architectures (e.g., Mole-Bert and other LGLMs as in here). We would like to kindly note that, when adding SMILES/SELFIES to LLMs, it does not directly output node embeddings. Nevertheless, in Table 16 and 17, we also consider adding SELFIES to LLMs, which slightly decreases the perception of functional groups, and increases the chemical reaction prediction performances.

Computational overhead

We have revised our manuscript to include the computational overhead report, including training and inference latency, as well as tunable parameters, which can be found here. Although HIGHT requires longer training time and relatively higher tunable parameters, the absolute values are not high. Moreover, during inference, as LLM latency consumes most of the computation, HIGHT can even reduce the inference latency by generating more concise answers.

Performance gain in Table 1

We kindly refer Reviewer VbFG to Table 17 in the appendix, where we can find that, merely using the HiPubChem for instruction tuning with InstructMol still suffers from severe hallucination. It means that both tokenizers and the instruction tuning in the alignment phase are necessary to reduce hallucination. In addition, in Table 16, we can find similar phenomena in downstream tasks of chemical reaction prediction, which further strengthens the necessity for the two components.

GIMLET in MotifHallu

We have evaluated two state-of-the-art SMILES based LGLMS, i.e., GIMLET and Galactica-6.7B in MotifHallu. The results are given here. Interestingly, those models demonstrate a high hallucination of negative classes. Therefore, we recommend reporting macro F1 scores (F1 scores averaged across classes). HIGHT demonstrates significant improvements in terms of macro F1 scores over all baselines, up to 14.

Comparison with the instruction data

We provided a comparison between the instruction tuning data used in HIGHT and other LGLMs here. It can be found that HIGHT uses significantly less information and data to achieve relatively good capabilities across multiple tasks.

code

We have provided AC an anonymous link to our code, and kindly asked AC to share to the reviewers according to the ICML policy.

Dear Reviewer zQSo,

Thank you for your time and insightful suggestions, as well as your acknowledgment of the value and convinceness of our work. Following your suggestions, we have revised our manuscript to include a discussion on future work regarding extending HIGHT to incorporate 3D information.

• Hierarchical 3D Tokenizer: To incorporate the 3D hierarchical molecular information, we need to design a new tokenizer based on GNN backbones such as EGNN and SE(3) Transformers compatible with 3D information. For example, we can extend the hierarchical VQVAE in HIGHT, to accommodate 3D properties of motifs, such as SE(3) equivariance to the motif representations.

• Scaling 3D Data for Alignment: When extending to 3D information, one could curate both the molecules, proteins and RNA data. The motifs in proteins then can be certain amino acids. In RNA, the motifs can be certain recurrent structures.

• 3D Tokenizer Training: To train the tokenizer to fully capture the 3D molecular information, we can design several self-supervised learning objectives such as 3D spatial position recovery, masked atom/motif prediction.

• 3D Alignment Tuning: To facilitate the alignment of molecule and language, we need to extend the instruction tuning dataset to include descriptions of 3D molecular properties, such as 3D positions. In addition to incorporating the 3D hierarchical information, we also need to incorporate language captions of the 3D hierarchical structures, such as properties of the amino acids. In addition, one could also incorporate pairwise hierarchical graph information, such as binding affinities of molecules to proteins.

Therefore, it is promising to align 3D graphs with languages for broader scentific tasks, for which HIGHT provides the foundation for incorporating 3D hierarchical graph information for better alignment of graphs and languages.

Dear Reviewer brfq,

Thank you for acknowledging our performance improvements and constructive suggestions. Please find our explanations to your questions below:

it is not entirely clear what fraction of the improvement comes from the hierarchical tokenization itself.

We kindly refer Reviewer brfq to Figure 3c and Table 16 in the appendix, where we conduct comprehensive ablation studies to analyze the contributions of each component in HIGHT. The differences between HIGHT and InstructMol lie in the positional encoding, hierarchical instruction tuning, and the hierarchical tokenizer. Thus, we study the performances of InstructMol plus those components, where each can also be considered as a variant of HIGHT without hierarchical tokenization and some components.

In experiments, one can find that, compared to HIGHT, merely incorporating positional encodings, using the hierarchical instruction tuning, or using a larger tokenizer in InstructMol can not bring improvements or even lead to decreased performance. It demonstrates that, while each component in HIGHT is critical to the performance, the hierarchical tokenizer contributes most to the alignment performance.

In addition, in Figure 3a, we also experiment with different base LLMs such as vicunna and Llama2, where HIGHT demonstrates consistent improvements.

Discussion between Hu et al. 2024.

First, we would like to kindly refer Reviewr brfq to the ICML reviewer instruction, that ``Authors cannot expect to discuss other papers that have only been made publicly available within four months of the submission deadline.''. Even Hu et al., came out in November 2024, it is within four months of the ICML deadline.

Nevertheless, we will revise our manuscript to include a discussion with Hu et al. Although both of our works incorporate the hierarchical information into LGLMs,

• Hu et al., does not consider the influence of the hierarchical information on the molecule-language alignment, i.e., hallucinations.
• They also do not consider incorporating the hierarchical information into the alignment tuning phase.
• They evaluate the usefulness of the hierarchical graph information in a limited number of tasks.

Questions related to hierarchical graph representation

Hierarchical graph representation refers to aggregating the low-level (i.e., node-level) into meaningful higher-level abstractions, usually via subgraph aggregation. For example, in a social network, the higher-level abstraction could be a small group or a community. Although it is not limited to a restricted number of abstraction layers, in molecules, the hierarchy of atom-motif-molecule is the usual schema of the abstraction[1]. Hence, we simply follow the common practice in the community to obtain the hierarchical molecular representations, i.e., using BRICS to identify common functional groups[2,3]. Nevertheless, one could also apply subgraph learning strategies to enable more flexible hierarchies.

In fact, when going beyond small molecules to biomolecules such as proteins, more abstract subgraphs, such as the 3D folded architecture, or combined 3D structure between small drugs and protein pockets, are also critical to the functionalities of proteins[4].

Other minor suggestions

We have revised our manuscript according to your suggestions.

For the caption of Fig 1: HIGHT brings improvements to all tasks. Due to the heterogeneity of the evaluation metrics, we transform the numerical values a bit for better visualization. The new Figure 1b and the corresponding details are given here.

References

[1] The art and practice of structure-based drug design: A molecular modeling perspective. Medicinal Research Reviews 1996.

[2] Motif-based graph self-supervised learning for molecular property prediction, NeurIPS'21.

[3] Molecular representation learning via heterogeneous motif graph neural networks, ICML'22.

[4] Independent SE (3)-Equivariant Models for End-to-End Rigid Protein Docking, ICLR'22.

Dear Reviewer eiUy,

Thank you for your time and insightful suggestions for our paper. Please find our responses to your concerns below.

In molecular property prediction and chemical reaction tasks, HIGHT does not demonstrate substantial advantages.

We need to clarify that the performance gaps between HIGHT and the baselines lie in the differences of the base models. For example, Uni-Mol is built on a T5 architecture pretrained on large-scale 3D molecular data. GIMLET is built on a T5 pretrained on molecular property prediction-focused data.

We present a detailed table listing the information and pretrained data used by HIGHT and baselines. It can be found that some "baselines" are pretrained onto significantly larger datasets with additional information, which are not fair baselines for comparison. We report their performances in order to have an overview of the overall progress.

Therefore, our direct baseline is InstructMol, which uses the same architecture and pretraining data. Compared to InstructMol, HIGHT demonstrates significant improvements.

how does the inclusion of motif tokens contribute meaningfully to other downstream tasks beyond functional group recognition?

The visualization of Figure 1b may introduce some perceptual biases about the improvements brought about by the different scales of evaluation metrics used in the method. Due to the heterogeneity of the evaluation metrics, we transform the numerical values a bit for better visualization. The new Figure 1b and the corresponding details are given here.

Comparing InstructMol and HIGHT, indeed, incorporating motif-related information improves other downstream tasks.

Limitatiions of tokenization by BRICS; Hierarchy in the tokenization process; Primary novelty of the work

We need to clarify that the primary novelty of this work lies in the introduction of the hierarchical tokenization. It is well known that motif information is critical to molecular understanding. Without that, LGLMs will exhibit high hallucination and perform subpar on molecule-language alignment.

We acknowledge the limitations of BRICS in tokenization, while it's sufficient to demonstrate the usefulness of hierarchical graph information. Despite the simplicity of our methodology, it brings significant improvements compared to LGLMs without hierarchical graph information. Meanwhile, our method possesses a high degree of extensibility and can be augmented by incorporating advanced motif extraction techniques (such as[1,2]), or learnable motif extraction, which could further boost the performance of HIGHT.

In addition, we would also kindly note a key technical contribution in HIGHT lies in not only the tokenization, but also in the alignment tuning process via HiPubChem. We feel the original title may introduce confusion, thus we would like to change our title to ``Hierarchical Molecule-Language Alignment'' in order to highlight the key focus of our work, i.e., the importance of hierarchical information for molecule-language alignment.

Metrics in MotifHallu

We have supplemented the details in our manuscript:

• To calculate the F1 scores for positive and negative classes, we first divide the motif hallu dataset according to the class label. For the positive F1 score, the calculation is normal. For the negative F1 score, we reverse the labels and predictions.
• Motif Hallu is quite imbalanced. There are 4,124 positive and 19,800 negative samples.
• As Acc will be biased due to the imbalance, we updated the table to report classwise averaged F1 score, i.e., macro F1. To obtain an overview of the results, we present a detailed table here with various metrics. It can be found that, HIGHT still obtains significant improvements up to 14 in terms of macro F1.

References

[1] Motif-based graph self-supervised learning for molecular property prediction, NeurIPS'21.

[2] Molecular representation learning via heterogeneous motif graph neural networks, ICML'22.