大语言模型是语境中的分子学习者

DRUGAI

今天为介绍一篇来自于香港理工大学，上海交通大学和上海人工智能实验室的文章，也是MolReGPT[1]的续作。

论文题目：Large Language Models are In-Context Molecule Learners

论文链接：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分子指纹）去提升分子召回的质量。分子图召回首先采用一个图神经网络去对分子图进行编码，得到分子图向量:

具体来说，我们采用了Mole-BERT[4]作为图神经网络编码器。

这样，我们就可以计算余弦相似度来衡量当前的query和其他分子图之间的相似度：

对于描述到分子生成任务，我们沿用了先前表现良好的BM25进行相似度的计算：

召回后重排

先前的上下文召回存在一定的问题，例如最相似的例子反而是距离当前query最远的，以及排名靠前的分子中可能过于相似导致多样性缺失。因此，我们提出了随机游走和序列反转去重新组织上下文例子的排布。

其中，随机游走会给每个例子一定的概率被跳过，例如在N个粗排结果中，第j个被跳过的概率可以表示为

，这样最后就至少能有一个例子被选中。具体来说，N如果比精排结果n大很多，召回的质量就可能被影响，而n如果太小（接近于N），则可能会损害结果多样性，因此在本文中，考虑到上下文长度的限制，n是个不超过5的数字，而N被设置为10.

这之后，序列反转会将原本最相似的例子靠在最接近于当前query的位置，以增强其在上下文中的地位。

上下文分子微调

MolReGPT[1]已经表明，上下文学习对于提示大语言模型进行分子-描述翻译任务是有效的，但是为了达到这种效果，需要依赖于上下文学习和推理能力，而这种能力通常都对模型的参数量和性能有着极高的要求。因此，我们提出上下文分子微调，去训练大语言模型，使得他们具有从上下文中提取有用信息和学习分子片段和描述片段之间的对应的能力，不同于直接微调的损失函数：

上下文分子微调同时学习了上下文例子所给出的对应关系

，整个上下文例子中具有的对应可以以集合形式表述

，因而这块额外的学习其对应的损失可以表述为：

最终，整个的损失函数还需要加上纯文本信息

（包括prompt和例子）：

结果

本文的实验在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.