[Paper Review] MolReFlect: Towards In-Context Fine-grained Alignments between Molecules and Texts

Jiatong Li, Yunqing Liu+7 more|arXiv (Cornell University)|Nov 22, 2024

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

MolReFlect proposes a teacher-student framework that enables in-context, fine-grained alignments between molecular sub-structures and textual descriptions in molecule-caption translation. By leveraging a large teacher LLM for zero-shot alignment extraction and reflection-based refinement, followed by Chain-of-Thought In-Context Molecule Tuning, the student LLM achieves state-of-the-art performance on ChEBI-20, significantly improving both generation quality and model explainability.

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.

Motivation & Objective

• Address the lack of fine-grained alignments between molecular sub-structures and descriptive phrases in molecule-caption translation tasks.
• Overcome the limitations of treating molecules as holistic SMILES strings or graphs, which hinders explainability and accuracy.
• Develop a cost-effective, parameter-efficient method to refine low-quality alignments without fine-tuning the teacher LLM.
• Enable smaller student LLMs to learn from high-quality, detailed alignments generated by a larger teacher model.
• Improve model interpretability by explicitly linking sub-structures (e.g., functional groups) to corresponding natural language phrases.

Proposed method

• Use a large teacher LLM to perform zero-shot alignment extraction by identifying key phrases in captions and mapping them to corresponding sub-structures in SMILES or molecular graphs.
• Implement In-Context Selective Reflection, where the teacher LLM retrieves similar past alignment examples and reflects on them to refine its output before passing to the student.
• Employ a smaller student LLM to selectively choose from the refined reflections and prior extractions, reducing noise and improving alignment quality.
• Apply Chain-of-Thought In-Context Molecule Tuning (CoT-ICMT), integrating fine-grained alignments into a reasoning-aware prompting format to enhance student LLM performance.
• Train the student LLM via instruction tuning using the refined alignments, enabling it to learn from the teacher’s reasoning and filtering process.
• Use retrieval-based few-shot examples during reflection to improve the consistency and accuracy of alignment refinement without modifying the teacher’s parameters.

Experimental results

Research questions

• RQ1 Can in-context, fine-grained alignments between molecular sub-structures and textual phrases improve molecule-caption translation performance?
• RQ2 How effective is In-Context Selective Reflection in refining alignment quality using retrieved examples and teacher reflection?
• RQ3 What is the impact of Chain-of-Thought In-Context Molecule Tuning on student LLM performance and alignment reasoning?
• RQ4 Is the teacher-student framework necessary, or can a single large LLM achieve comparable results via prompt engineering?
• RQ5 To what extent do fine-grained alignments enhance model explainability in molecule-text generation tasks?

Key findings

• MolReFlect achieves state-of-the-art performance on the ChEBI-20 benchmark, outperforming all prior baselines in both Mol2Cap and Cap2Mol tasks.
• The student LLM’s performance improves by 9.94% in Mol2Cap and 14.22% in Cap2Mol when using instruction tuning with in-context alignments versus naive supervised fine-tuning.
• Chain-of-Thought prompting alone on the teacher LLM (Llama-3-70B) improves Mol2Cap performance by 41.80% but reduces Cap2Mol performance by 1.05%, indicating noise in raw alignments.
• Few-shot prompting with the teacher LLM does not significantly improve performance, underscoring the need for a student model to filter and refine noisy alignments.
• The teacher-student framework is indispensable: the student model effectively learns to discern and mitigate noise in teacher-generated alignments, leading to superior generalization.
• Comprehensive case studies confirm that MolReFlect enables more interpretable and accurate alignments, explicitly linking functional groups like anhydrides to their textual descriptions.