[论文解读] 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 Science被引用 3

一句话总结

MolReFlect 提出了一种教师-学生框架，实现了分子子结构与文本描述在分子-图像字幕翻译任务中的上下文式细粒度对齐。通过利用大型教师 LLM 进行零样本对齐提取与基于反思的优化，随后采用思维链上下文分子微调（Chain-of-Thought In-Context Molecule Tuning），学生 LLM 在 ChEBI-20 上实现了最先进性能，显著提升了生成质量与模型可解释性。

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.

研究动机与目标

• 解决分子-图像字幕翻译任务中分子子结构与描述性短语之间缺乏细粒度对齐的问题。
• 克服将分子视为整体 SMILES 字符串或图结构所带来的局限，此类方法会损害可解释性与准确性。
• 开发一种成本低、参数高效的策略，以在不微调教师 LLM 的前提下优化低质量对齐。
• 使较小的学生 LLM 能够从大型教师模型生成的高质量、详细对齐中学习。
• 通过显式关联子结构（如官能团）与对应自然语言短语，提升模型的可解释性。

提出的方法

• 使用大型教师 LLM 通过识别字幕中的关键短语，并将其映射到 SMILES 或分子图中的对应子结构，实现零样本对齐提取。
• 实施上下文选择性反思（In-Context Selective Reflection），即教师 LLM 检索类似的历史对齐示例，并对其进行反思以优化输出，再传递给学生模型。
• 采用较小的学生 LLM 从优化后的反思结果和先前提取结果中进行选择性筛选，以降低噪声并提升对齐质量。
• 应用思维链上下文分子微调（Chain-of-Thought In-Context Molecule Tuning, CoT-ICMT），将细粒度对齐整合到推理感知的提示格式中，以增强学生 LLM 的性能。
• 通过指令微调方式训练学生 LLM，利用优化后的对齐结果，使其能够学习教师的推理与过滤过程。
• 在反思过程中使用基于检索的少样本示例，以在不修改教师参数的前提下提升对齐优化的一致性与准确性。

实验结果

研究问题

• RQ1 上下文式、细粒度的分子子结构与文本短语对齐是否能提升分子-图像字幕翻译性能？
• RQ2 在使用检索示例与教师反思的前提下，上下文选择性反思在优化对齐质量方面的有效性如何？
• RQ3 思维链上下文分子微调（CoT-ICMT）对学生 LLM 性能与对齐推理能力的影响如何？
• RQ4 教师-学生框架是否必不可少？单个大型 LLM 是否可通过提示工程实现相当的结果？
• RQ5 细粒度对齐在多大程度上提升了分子-文本生成任务中的模型可解释性？

主要发现

• MolReFlect 在 ChEBI-20 基准上实现了最先进性能，在 Mol2Cap 与 Cap2Mol 任务中均优于所有先前基线模型。
• 与朴素的监督微调相比，使用指令微调结合上下文对齐后，学生 LLM 在 Mol2Cap 任务上性能提升 9.94%，在 Cap2Mol 任务上提升 14.22%。
• 仅在教师 LLM（Llama-3-70B）上应用思维链提示，可使 Mol2Cap 性能提升 41.80%，但 Cap2Mol 性能下降 1.05%，表明原始对齐存在噪声。
• 使用教师 LLM 的少样本提示并未显著提升性能，凸显了学生模型在过滤与优化噪声对齐方面的重要性。
• 教师-学生框架不可或缺：学生模型能有效学习识别并缓解教师生成对齐中的噪声，从而实现更优的泛化能力。
• 全面的案例研究证实，MolReFlect 实现了更可解释、更准确的对齐，能够显式地将如酸酐等官能团与对应文本描述关联。