作者 | 香港理工大学李佳潼
编辑 | ScienceAI
去年,香港理工大学研究团队开发了一个基于检索的提示范式MolReGPT,利用大型语言模型探索分子发现,在分子和分子文本描述之间进行翻译。
近日,香港理工大学、上海交通大学和上海人工智能实验室联合发表了题为《Large Language Models are In-Context Molecule Learners》的文章,也是MolReGPT[1]的续作。
论文链接:https://arxiv.org/abs/2403.04197
模型链接:https://huggingface.co/phenixace/
介绍
近来,大语言模型在生物化学领域展现出了优异的性能,尤其是分子-描述翻译任务[1],这减小了自然语言描述与分子空间的差异。
但是,之前的方法要么是需要进行大量的领域预训练,要么是在分子与描述的对齐上存在过于粗糙的问题[2],要么是对大语言模型的能力和参数规模有着严苛的要求[1]。
为了解决这些问题,我们提出了上下文分子微调(ICMA),作为一个让大语言模型学习分子-描述对齐的新方法。
具体来说,上下文分子微调由以下三个部分组成:混合模态上下文召回,召回后重排,和上下文分子微调。
- 首先,混合模态上下文召回沿用了BM25和分子图召回分别对分子描述和分子进行召回,以增强模型的输入。
- 此外,为了解决召回中存在的不准确、样本间过度重复等问题,我们设计了两个召回后处理方法:序列反转和随机游走,以提升召回结果的质量。
- 最后,上下文分子微调借助于大语言模型的上下文学习能力,生成最终的分子/分子描述。
我们在ChEBI-20[3]和PubChem324K[2]两个分子-分子描述数据集上展开实验,实验结果表明,上下文分子微调可以使得Mistral-7B在分子-分子描述上取得SOTA或接近的结果,无需领域预训练和复杂的模型结构。
我们的贡献在于:上下文分子微调不需要额外的领域预训练就可以发挥作用,为最新最先进的十亿甚至百亿参数大语言模型在化学任务上的运用提供了新思路。与此同时,上下文分子微调通过学习上下文例子中分子描述片段与分子SMILES结构的对应关系,精细化了分子描述和分子的对齐。此外,上下文分子微调不那么依赖于模型的上下文学习和推理能力,即便是稍小的模型也能获得良好的性能增益。
方法
分子的相似性原理指出,相似的分子一般会具有相似的性质,如图1所示。借助于分子SMILES和分子描述片段的对应,我们可以基于这些共现信息大致推理出,如果有另一个相似的分子,它的SMILES片段会对应哪些相应分子描述片段,即具有什么样的结构和化学性质。
图1: 三个相似分子以及他们的分子描述。分子可以被表示为SMILES表达式或者分子图,而分子描述说明了分子的特征。这里三个分子在图结构上是相似的,分子描述中重叠的地方被蓝色和粉色标出。
因此,在这篇文章中,我们提出了上下文分子微调(ICMA),其流程如图2所示,总共三个阶段:混合模态上下文召回,召回后重排,和上下文分子微调。
图2: 上下文分子微调(ICMA)框架。
混合模态上下文召回
召回的质量对上下文例子的信息量至关重要,如果我们召回到的分子足够相似,那么他们在描述中的重叠就很可能会更多,这样模型就可以更好地学习到分子SMILES片段与描述片段的对应。
本文提出了混合模态上下文召回,在分子描述召回过程中,沿用了BM25对文本模态进行召回。而在分子召回过程中,我们提出了分子图召回,从图模态来召回相似的分子。
在分子到描述任务中,ICMA首先采用分子图召回(取代了MolReGPT[1]的Morgan分子指纹)去提升分子召回的质量。分子图召回首先采用一个图神经网络去对分子图进行编码,得到分子图向量:
召回后重排
上下文分子微调
研究结果
本文的实验在ChEBI-20和PubChem324K两个数据集上展开。表2、3列出了ChEBI-20数据集上,上下文分子微调和直接微调、基线模型的效果比较。首先,在两个子任务上,上下文分子微调相较于直接微调都可以带来可观的性能增益,而相较于先前的基线模型,上下文分子微调也可以取得SOTA或者相接近的性能。
表4、5列出了PubChem324K数据集上,上下文分子微调与直接微调的效果比较,其结果与ChEBI-20的结果相类似,上下文分子微调相较于直接微调都带来了可观的表现增益。
与此同时,本文也展开了大量的实验去探究影响上下文分子微调效果的因素,集中探究了召回算法对效果的影响、上下文设置的影响、不同基座模型的影响(模型类型和模型参数量)
表6、7展示了使用不同召回算法进行上下文分子微调的效果区别,显示出BM25和分子图召回所带来的效果增益。
图3则说明了上下文设置(例子数量和截断长度)对上下文分子微调的影响,表明截断长度必须适应上下文长度,以及例子增多所可能带来的负面影响(如过多的噪声)。
图3: 模型性能随着上下文设置(例子数量和截断长度)变化的比较。
而图4则展现了上下文分子微调所具备的scaling law,更大的模型因其具有更强的上下文学习和推理能力,因而会获得更好的上下文分子微调效果。
图4: 上下文分子微调(ICMA)的scaling law,这里选取了Galactica-125M,Galactica-1.3B,和Mistral-7B进行比较。
表11、12进一步说明了上下文分子微调的通用性,可以使得不同的大模型都能够在分子-描述翻译任务上取得性能增益,而不仅仅局限于某一个或者某一类模型。
最后,如表8、9所示,消融实验也论证了召回后重排策略所带来的性能增益。
结论
本文提出了上下文分子微调,作为一个新的范式,来将大语言模型运用到分子-描述翻译任务上,不需要额外的领域预训练和复杂的模型结构,上下文分子微调使得大语言模型能够运用他们的上下文学习和推理能力学习到分子和描述之间的对齐,提升了大语言模型在分子-描述翻译任务上的能力,也表明了大语言模型是天然的上下文分子学习者。我们的方法为将十亿以上的大语言模型运用到科学领域提供了新方法和新思路。
参考文献:
[1] Li, J., Liu, Y., Fan, W., Wei, X. Y., Liu, H., Tang, J., & Li, Q. (2024). Empowering molecule discovery for molecule-caption translation with large language models: A chatgpt perspective. IEEE Transactions on Knowledge and Data Engineering.
[2] Liu, Z., Li, S., Luo, Y., Fei, H., Cao, Y., Kawaguchi, K., ... & Chua, T. S. (2023). Molca: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. arXiv preprint arXiv:2310.12798.
[3] Edwards, C., Lai, T., Ros, K., Honke, G., Cho, K., & Ji, H. (2022). Translation between molecules and natural language. arXiv preprint arXiv:2204.11817.
[4] Xia, J., Zhao, C., Hu, B., Gao, Z., Tan, C., Liu, Y., ... & Li, S. Z. (2022, September). Mole-bert: Rethinking pre-training graph neural networks for molecules. In The Eleventh International Conference on Learning Representations.