MolReFlect: Towards Fine-grained In-Context Alignment between Molecules and Texts

Authors: Jiatong LI, Yunqing LIU, Wei Liu, Jingdi Lei, Di Zhang, Wenqi Fan, Dongzhan Zhou, Yuqiang Li, Qing Li
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Large Language Models, In-Context Tuning, Reflection Tuning, Molecule Discovery, Molecule-Text Alignment

Abstract:
Molecule discovery is a pivotal research field, impacting everything from the medicines we take to the materials we use. Recently, Large Language Models (LLMs) have been widely adopted in molecule understanding and generation, yet the alignments between molecules and their corresponding captions remain a significant challenge. Previous endeavours often treat the molecule as a general SMILES string or molecular graph, neglecting the fine-grained alignments between the molecular sub-structures and the descriptive textual phrases, which are crucial for accurate and explainable predictions. In this case, we introduce MolReFlect, a novel teacher-student framework designed to contextually perform the molecule-caption alignments in a fine-grained way. Our approach initially leverages a larger teacher LLM to label the detailed alignments by directly extracting critical phrases from molecule captions or SMILES strings and implying them to corresponding sub-structures or characteristics. To refine these alignments, we propose In-Context Selective Reflection, which retrieves previous extraction results as context examples for teacher LLM to reflect and lets a smaller student LLM select from in-context reflection and previous extraction results. Finally, we enhance the learning process of the student LLM through Chain-of-Thought In-Context Molecule Tuning, integrating the fine-grained alignments and the reasoning processes within the Chain-of-Thought format. Our experimental results demonstrate that MolReFlect enables LLMs like Mistral-7B to significantly outperform the previous baselines, achieving SOTA performance on the ChEBI-20 dataset. This advancement not only enhances the generative capabilities of LLMs in the molecule-caption translation task, but also contributes to a more explainable framework.

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

Reviews and discussion:
<thread>
creation date:
1730206795148

thread_type:
official_review

<thread_object>
creator: reviewer_8viw

<content>
summary:
The paper presents MolReFlect, a teacher-student framework that enhances the alignment between molecular structures and their captions, addressing challenges in molecule understanding and generation using Large Language Models (LLMs). By leveraging fine-grained alignments and innovative In-Context Selective Reflection technique, the approach achieves state-of-the-art performance on the ChEBI-20 dataset.

strengths:
* The paper introduces MolReFlect teacher-student framework, a novel approach to molecular caption alignment, addressing a critical gap in existing methods that often overlook molecular substructures.

* The experimental evaluation against a limited baselines set showcases the effectiveness of the proposed approach for generative cross-modal tasks on molecules and texts. 

* A thorough description of the proposed methodology with the detailed description of the experimental setting.

weaknesses:
* The experimental evaluation of the proposed methodology in conducted on the single ChEBI-20 dataset and lacks comparison against more recent state-of-the-art molecule captioning/generation approaches.
* While SELFIES molecular string representations are mentioned in the paper's Introduction, the experiments are conducted on SMILES. This motivation for choosing SMILES is not explained. 
* The proposed MolReFlect method is implemented using only a single set-up with Llama-3-70B-Instruct and Mistral-7B as teacher and student model, respectively. No experiments with larger student LLM/smaller teacher LLM are provided.
* High computational complexity and memory footprint of the proposed method. The implementation of the proposed pipeline has at least 77B parameters. For comparison, BioT5 is reported to have 252M parameters only.
* The proposed pipeline includes additional tasks, for instance, Named Entity Recognition for molecule substructure extraction and entity linking during In-context Selective Reflection. From the paper, the performance of LLMs on this indermideate tasks is not clear.

questions:
* Add more experiments on other datasets, for instance, Mol-Instructions [5].
* Are the improvements over baseline models statistically significant?
* Add more experimental comparison against more chemical LMs: e.g., nach0 [1], Text+Chem T5 [2], SciFive [3], PRESTO [4], GitMol [6]. MoMu is mentioned in the Related Work but it is absent from experimental comparison.
* The usage of SELFIES is expected to give 100% SMILES validity due to the format specification. Is it possible to conduct additional experiments with SELFIES or explain the choice of SMILES in your experiments?
* Error Analysis/case study, revealing key challenges for the proposed framework/hard cases/interpretation of the cases where MolReFlect gains an improvement over baselines, would strengthen the paper.
* Line 220: the notation is not quite clear to me. What are alignments $K$ formally? A function? A tuple of pairs <text, molecule substructure>? A set?
* In section 4.2, more discussion on the adopted baselines would help: architecture, parameter count, hugging face checkpoint links, pretraining data, etc.

Typos:
* line 192: What is ICMA? The abbreviation is not introduced earlier.
* line 399: Mo2Cap -> Mol2Cap

References:

[1] Livne, Micha, et al. "nach0: Multimodal natural and chemical languages foundation model." Chemical Science 15.22 (2024): 8380-8389.

[2] Christofidellis, Dimitrios, et al. "Unifying molecular and textual representations via multi-task language modelling." International Conference on Machine Learning. PMLR, 2023.

[3] Phan, Long N., et al. "Scifive: a text-to-text transformer model for biomedical literature." arXiv preprint arXiv:2106.03598 (2021).

[4] Cao, He, et al. "PRESTO: Progressive Pretraining Enhances Synthetic Chemistry Outcomes." arXiv preprint arXiv:2406.13193 (2024).

[5] Yin Fang, Xiaozhuan Liang, Ningyu Zhang, Kangwei Liu, Rui Huang, Zhuo Chen, Xiaohui Fan, Huajun Chen:
Mol-Instructions: A Large-Scale Biomolecular Instruction Dataset for Large Language Models. ICLR 2024

[6] Liu, P., et al. "GIT-Mol: A multi-modal large language model for molecular science with graph, image, and text." Computers in Biology and Medicine 171 (2024): 108073-108073.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730349382806

thread_type:
official_review

<thread_object>
creator: reviewer_hpy1

<content>
summary:
The paper intends to enable fine-grained alignments and introduce a large teacher LLM to label the detailed alignments. Then the paper proposes in-context selective reflection, to let a student LLM select from in-context reflection and previous extraction results. Finally, the paper introduces Chain-of-Thought, which also increases the performance.

strengths:
1. The performance of generation is impressive.

2. The paper is well-motivated to enable the fine-grained alignment.

weaknesses:
1. The main concern remains in potential data leakage. The paper widely uses the pretrained large model, such as teacher model (Llama-3-70B), Mole-BERT and student model (Mistral-7B). The statement of data leakage in the paper is not convincing enough. Even though the teacher model cannot perform well in generation, given that the teacher model can do zero-shot alignment labeling, there's still possible potential leakage. Additionally, the pre-trained Mole-BERT may also introduces leakage, which should be further discussed. Although there may leakage in student model, the authors have already presented the effectiveness of the proposed pipeline. However, the comparison of the baselines may not be fair. 

2. In Line 265, why an unsupervised metric can avoid leakage from Galactiva-125M?

3. Is this paper possible also evaluated on retrieval tasks?

4. The justification in related work about 2D and 3D graph is not fair. Note that the molecular data is limited compared with image-text pairs, and the structural information is important for molecules.

I am happy to increase the score if the problems are solved.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730651367028

thread_type:
official_review

<thread_object>
creator: reviewer_xb7s

<content>
summary:
This paper presents MolReFlect, a framework that refines molecule-caption alignments through a teacher-student model. It introduces an approach of zero-shot alignment extraction, in-context refinement, and Chain-of-Thought In-Context Molecule Tuning (CoT-ICMT), claiming state-of-the-art results on the ChEBI-20 dataset.

strengths:
1. The proposed methods work reasonably well.
2. The paper contains plenty of ablations to validate the effectiveness of different components.

weaknesses:
1. While the model performs well on generation tasks, its effectiveness on other molecular tasks, such as molecule retrieval and property prediction tasks, is not evaluated.

2. The model is only evaluated on the ChEBI-20 dataset, making it difficult to fully assess the model's generative capabilities.

3. For generation tasks, language models tend to memorize the training dataset. Presenting the scores of molecule novelty would be more comprehensive and show that MolReFlect can perform well in generation tasks.

4. Lack of human evaluation: For tasks like molecule captioning, incorporating human expert evaluation could provide additional validation of the model's performance and practical utility.

5. Training code and implementation for reproducing the paper are missing.

6. The training requires expensive computations. The required computational comparisons of the proposed method and other methods should be presented.

questions:
1. In Section 3.2 on molecule retrieval, you use a pre-trained Mole-BERT as the graph encoder to calculate the cosine similarities between molecule graph embeddings. Why is the graph encoder employed here when the SMILES string is taken as input, as stated in the abstract (Line 23)?
1. How does MolReFlect perform on more complex molecular structures or rare chemical compounds not well-represented in the training data?
2. How does the performance of MolReFlect change when dealing with highly technical or domain-specific textual descriptions that might require expert knowledge to interpret?
3. Have you explored the potential of MolReFlect for other cross-modal (text and corresponding data modalities) tasks beyond the molecular domain? If so, what were the results or challenges encountered?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730656396860

thread_type:
official_review

<thread_object>
creator: reviewer_zye2

<content>
summary:
The authors present MolReFlect, a novel teacher-student framework designed to refine the in-context alignments between molecular substructures and their corresponding textual descriptions. They show that within this framework the training of a small student network may benefit from the alignments, generated by a teacher model. The model trained with this framework shows a better performance on Mol2Cap and Cap2Mol tasks on the ChEBI-20 dataset compared to the baseline trained with SFT.

strengths:
A new framework for training Chemical LLMs, that allows to successfully  finetune relatively small models with limited resources. 
Good results for the molecular captioning tasks without domain pretraining.

weaknesses:
All the baselines except ICMA differ from MolReFlect by at least an order of magnitude by the number of parameters.  
The ablation studies are partially unconvincing: they are done only for one setup Cap2Mol or Mol2Cap
The results for w/o In-Context Reflection and w/o Selection setups are practically indistinguishable from the full model, especially in the table 4.

questions:
1. The idea of using the retrieval in  Chain-of-Thought In-Context Molecule Tuning seems to be similar to the RAG. It would be interesting to compare these frameworks.
2. The authors don’t site arxiv.org/abs/2408.11866 where the better performance is achieved with a similar framework, but probably lager models.
3. Please, fix the bold numbers in Table 4 for the Validity metric.
4. It is interesting to check how robust the trained model is. (i.e. in terms of Ganeeva, Veronika, et al. "Chemical Language Models Have Problems with Chemistry: A Case Study on Molecule Captioning Task." The Second Tiny Papers Track at ICLR 2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1731996005935

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all the reviewers for their valuable comments and contributions to the ICLR community. We have responded to the concerns and questions correspondingly below the reviews and updated our paper for you to review.

We sincerely hope that our responses could resolve your questions.
</content>
</thread_object>

</thread>


Metareview: