Improving Molecule-Language Alignment with Hierarchical Graph Tokenization

Authors: Yongqiang Chen, Quanming Yao, Juzheng Zhang, James Cheng, Yatao Bian
Primary area: learning on graphs and other geometries & topologies
Keywords: molecular-language alignment, large language models, hierarchical graph neural networks, tokenization, biomolecular studies, molecule

Abstract:
Recently there has been a surge of interest in extending the success of large language models (LLMs) to graph modality, such as molecules. As LLMs are predominantly trained with 1D text data, most existing approaches adopt a graph neural network to represent a molecule as a series of node tokens and feed these tokens to LLMs for molecule-language alignment. Despite achieving some successes, existing approaches have overlooked the hierarchical structures that are inherent in molecules. Specifically, in molecular graphs, the high-order structural information contains rich semantics of molecular functional groups, which encode crucial biochemical functionalities of the molecules. We establish a simple benchmark showing that neglecting the hierarchical information in graph tokenization will lead to subpar molecule-language alignment and severe hallucination in generated outputs. To address this problem, we propose a novel strategy called HIerarchical GrapH Tokenization (HIGHT). HIGHT employs a hierarchical graph tokenizer that extracts and encodes the hierarchy of node, motif, and graph levels of informative tokens to improve the graph perception of LLMs. HIGHT also adopts an augmented molecule-language supervised fine-tuning dataset, enriched with the hierarchical graph information, to further enhance the molecule-language alignment. Extensive experiments on **14** molecule-centric benchmarks confirm the effectiveness of HIGHT in reducing hallucination by **40%**, as well as significant improvements in various molecule-language downstream tasks.

Decision: reject
Reviewer scores: 6, 3, 6, 3

Reviews and discussion:
<thread>
creation date:
1730084506899

thread_type:
official_review

<thread_object>
creator: reviewer_my8l

<content>
summary:
This paper proposes a framework named HIGHT to align molecular data with LLMs. It identifies a shortcoming of LLM on learning functional groups, and proposes to extend the graph tokenization to motif level. Specifically, its input to the LLM includes node/atom embeddings as well as motif embeddings. The model is fine-tuned with motif prediction tasks on a dataset constructed using RDKit. The model shows good performance on molecule properties prediction compared to language models.

strengths:
The overall presentation is clear and the paper is easy to follow. The work proposes a complete pipeline to build a model with stronger motif/functional group querying ability. Using motif tokens is a straight-forward solution to enhance such ability. Various experiments are conducted to validate the model.

weaknesses:
- From the novelty and contribution perspective, taking motif representations/tokens is not new. By simply searching on Google, I found several papers that extract motifs for graph modeling [1, 2] (as the author also mentioned in the paper). This work is a simple extension of these techniques to align the motifs to the LLM.

- If I understand correctly, the motif tokenization algorithm, BRICS, will break the molecule in a very chemistry-aligned way. For example, a "OH" functional group will be tokenized into a motif. The downstream task of identifying the functional group will be very easy (simply aligning a single motif token with the text description of the function group, and the task is like asking "does a -OH motif have -OH functional group"). The author should justify how this is a helpful task besides simply telling the LLM that "there is such a functional group." For example, the author should show that the method has better downstream performance than simply telling the LLM the existence of functional groups.

- The distinction between specialist model and generalist model is arbitrary to me. Methods like MolFM and Text+Chem T5-augm-base have the same functionality as the proposal, yet they achieved better performance than HIGHT. I think the HIGHT is more specialized, as it requires explicit and specialized atom and motif tokenization. Can you be more specific about the distinction, and what's the advantage of a generalist model?

- Even without the motif tokens, many models achieved stronger performance. Can you explain why a better motif prediction ability does not lead to better downstream performance? Link back to weakness 1, does this also mean that the proposed task is too easy for the motif tokenization, preventing the model from learning meaningful/molecule-property-related from the pretraining process?

[1] Zhang, Zaixi, et al. "Motif-based graph self-supervised learning for molecular property prediction." Advances in Neural Information Processing Systems 34 (2021): 15870-15882.
[2] Chen, Xuexin, et al. "Motif graph neural network." IEEE Transactions on Neural Networks and Learning Systems (2023).

questions:
- In tables 3, 4, 5, are all baselines also fine-tuned with the same dataset as HIGHT?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730102306811

thread_type:
official_review

<thread_object>
creator: reviewer_p2nz

<content>
summary:
This paper proposes HIerarchical GrapH Tokenization (HIGHT), which tries to improve how LLMs understand and process molecular data. The key idea of HIGHT is to introduce a hierarchical graph tokenizer that extracts and encodes information at multiple levels: atoms, motifs, and the overall molecule. The paper demonstrates that HIGHT can reduce hallucination and improve performance across various molecular tasks, including property prediction, molecular description generation, and chemical reaction prediction.

strengths:
1.  This paper proposes to incorporate hierarchical graph information into LGLMs, and the authors achieve this with new architecture and instruction tuning dataset HiPubChem.
2. To address hallucination issue, the paper creates MotifHallu, the first hallucination benchmark based on the existence of common functional groups.
3. The paper includes extensive experiments with 14 real-world molecular and reaction comprehension benchmarks. The results show that HIGHT significantly reduces the hallucination on MotifHallu and demonstrates significant improvement on a number of tasks.

weaknesses:
1. The hierarchical graph tokenization process, which involves the use of multiple adapters, is likely to be more computationally expensive than traditional node-centric tokenization. The paper does not discuss the computational complexity. Also, LLM is tuned using LORA, and the number of parameters tuned should be discussed.
2. One motivation of applying LLMs for graph data is to utillize the generalization capability of LLMs. However, this paper do not provide experimental results on zero-shot or few-shot scenarios of the proposed model. I think it will greatly strength the paper if HIGHT has good performance under such cases.
3. The performance of HIGHT will largely depend on the backbone LLMs, and only vicuna-v-1.3-7B is evaulated.

questions:
1. At line 273, the authors said "attach positional encodings $p$ to all of the tokens", How are position encodings of motifs obtained?
2. If the input graphs use the positional encodings, then, should the original positional encodings in the LLMs be diabled? e.g, the widely used ROPE for the graph input part? 
3. What is the papameter count to be tuned?
4. Besides the vicuna-v-1.3-7B, can the authors provide experimental resutls for other LLM backbones? Since different backbones may have  a big impact on the performance.
5. How is the proposed model performance for zero-shot or few-shot scenarios?
6. In table 2, llama 13b has wrose performance than llama 7b on most of datasets. Also, Galactica-120B has a sharp performance drop on BACE. Any explanations on these results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730656525347

thread_type:
official_review

<thread_object>
creator: reviewer_qitu

<content>
summary:
The authors study Large Graph Language Models (LGLM). Drawing inspiration from Multimodal LLMs, authors focus on the task of incorporating graph data as a separate modality with a GNN encoder and an adapter. Authors conclude that node-centric tokenization of molecules leads to LLM hallucinations when asked about the presence of specific fragments. To overcome this issue, the authors propose to enrich the molecule's description by adding the tokens corresponding to BRICKS-fragments that are present in the molecule. The experimental results demonstrate that such a tokenization scheme reduces the amount of motif-related hallucinations and improves performance on other tasks.

strengths:
An improved tokenization of molecular graphs that enriches molecule's description with motif tokens.

weaknesses:
Specifically, most existing LGLMs directly take the node tokens from GNNs as inputs to LLMs (Cao et al., 2023): 
The paper cites InstructMol as a previous approach that utilizes node-centric tokenization. However, if I understand correctly, InstructMol takes the embedding of the whole graph along with the SMILES representations of the molecule. Moreover, it is not clear which previous models use the node-centric tokenization and whether there are such models at all.

Section 4.3 describes the fine-tuning approach that involves two stages, where the second stage is the finetuning on MoleculeNet, CheBI-20 and Mol-instructions specialized datasets. In my opinion, this implies that the resulting model is specialized. Please, provide better explanation for specialist and generalist models.	 	 	 		

Taking into consideration that the difference between specialist and generalist models is not clear, the resulting model does not demonstrate performance superior to baselines in most of the experiments.

There is no comparison with [1] in Table 4. The results in [1] are superior to all the models from Table 4.

In Table 5, the Mol-instruction has the highest MACCS FTS for the retrosynthesis task. However, a smaller number is balded.

The comparison on MotifHallu is not complete. Please provide comparison with SMILES-based approaches. Moreover, the improvement on the MotifHally benchmark is expected, as the proposed approach was explicitly designed to better solve this task.

[1] Srinivas, S. S., & Runkana, V. (2024). Crossing New Frontiers: Knowledge-Augmented Large Language Model Prompting for Zero-Shot Text-Based De Novo Molecule Design. arXiv preprint arXiv:2408.11866.

questions:
Listed in Cons.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691342951

thread_type:
official_review

<thread_object>
creator: reviewer_mtje

<content>
summary:
This paper presents a new approach to aligning molecular graph representations with language using a method called Hierarchical Graph Tokenization (HIGHT). Traditional graph-language alignment models primarily focus on node-level information, often neglecting the inherent hierarchical structure of molecules, which leads to alignment issues and hallucination in large language models (LLMs).

The authors introduce HIGHT, which utilizes a hierarchical graph tokenizer to capture information at the node, motif, and entire molecule levels. This tokenizer incorporates both atom-level and motif-level tokens, which are then used to improve alignment with language models. To address the alignment of hierarchical molecular data with textual descriptions, the authors also develop an enhanced molecular instruction tuning dataset called HiPubChem, which provides detailed motif information.

strengths:
1. The Hierarchical Graph Tokenization (HIGHT) technique is a major advancement. By incorporating hierarchical structure at multiple levels (node, motif, and graph), the paper addresses a crucial gap in previous molecule-language alignment methods, which typically rely only on node-level information. This hierarchical approach captures the functional groups and structural motifs inherent in molecules, improving the model’s ability to represent complex biochemical properties accurately.

2. The introduction of HiPubChem, an augmented molecular instruction tuning dataset enriched with motif and functional group information, enhances model training by aligning molecular structural details with language descriptions. This contribution is valuable for future work in molecular and biochemical language model alignment.

3. The effectiveness of each of the two methods was verified through simple ablation studies.

weaknesses:
1. The introduction of the hierarchical graph tokenizer seems to make the tokenizer larger compared with the ordinary node-level tokenizer. It should be discussed that whether the performance gain comes from the larger tokenizer.

2. There should be more detail descriptions and discussions about the evaluation tasks.

questions:
1. How many parameters do those two tokenizers have respectively?
2. What are the ablation study results on other tasks such as property prediction and chemical reaction prediction? 
3. What are the input and output of the molecular property prediction task and other tasks? The performance gain mainly comes from hierarchical graph tokenization, and it has nothing to do with the new tuning dataset, right?
</content>
</thread_object>

</thread>


Metareview:
metareview:
While this paper proposes a novel hierarchical graph tokenization method (HIGHT) to address shortcomings in molecular-language alignment, it ultimately fails to meet the bar for acceptance at ICLR due to some concerns as follows. 
1. Inadequate Justification of Impact: The manuscript does not convincingly demonstrate that the proposed method's improvements on downstream tasks arise from meaningful advancements in model design, as opposed to task-specific tuning or dataset augmentation. Some experimental results suggest that simply adding motifs as input offers similar benefits without architectural changes.

2. Lack of Robustness Across Models and Tasks: Results indicate that the improvements are highly contingent on specific experimental setups (e.g., Vicuna 7B backbone) and may not generalize across alternative LLM architectures or unseen tasks. This limits the broader applicability of the approach.

3. Evaluation: Despite extensive experimentation, key benchmarks and comparisons are missing. For instance, there is insufficient evaluation against SMILES-based or contemporary molecule-text alignment methods. Additionally, some baselines are tuned differently, complicating fair comparisons. The method's practical impact is hindered by increased complexity and resource demands (e.g., additional preprocessing, longer training times) relative to the modest performance gains reported.

Given these concerns, a rejection recommendation is made.