目前用于分子和文本联合表示的深度学习模型主要依赖于一维或二维分子格式，而忽略了提供有价值的物理见解的重要三维结构信息，因此限制了模型在广泛的模态范围内的多功能性和适应性。相反，专注于明确的三维表示的有限研究往往忽略了生物医学领域内的文本数据。因此，将二维和三维分子信息以及生物医学文本相结合的分子表示学习模型仍十分缺乏。
2024年6月28日，美国耶鲁大学Xiangru Tang等人在Bioinformatics上发表文章MolLM: a unified language model for integrating biomedical text with 2D and 3D molecular representations。
作者提出了一个统一的预训练语言模型，MolLM，可以同时捕获二维和三维分子信息以及生物医学文本信息。MolLM由一个文本Transformer编码器和一个分子Transformer编码器组成，设计用于编码二维和三维分子结构。为了支持MolLM的自监督预训练，模型采用对比学习作为学习的监督信号。实验结果表明，MolLM在分子表示学习方面超越了现有的方法。
为了便于将分子和文本编码到MolLM的联合潜在空间中，MolLM对分子的文本描述使用了文本编码器，并分别对分子的二维/三维表示使用了分子编码器，如图1所示。
图1 MolLM结构图
为了便于分子的自监督预训练，需要一个大规模的文本-分子配对数据集。因此，作者从两个来源收集数据集。首先，为了收集分子数据和直接相关的描述，作者访问了PubChem分子数据库，该数据库包含超过1亿个分子化合物的信息。从这个存储库中，作者使用PubChem分类浏览器来识别具有大量物理描述的160K个分子。具体来说，作者的标准是在“Record Description”（记录描述）字段中选择至少有150个字符的条目。例如“acetylcarnitine”。然后，使用PubChem REST API，作者查询二维和三维分子图数据，以及PubChem提供的基本物理描述。
为了收集相关的文本数据，作者遵循MoMu采用的匹配方法，从S2ORC语料数据库中的论文中提取与某些分子相关的句子。然后，作者将这些文本描述与每个分子的二维和三维图形表示配对，以创建分子图形-文本对。具体来说，使用前面提到的例子，作者在S2ORC数据库中搜索与“acetylcarnitine”相关的论文，从每篇论文的摘要、介绍和结论部分提取包含该分子或其同义词的文本。由于文本编码器的长度限制，作者使用256个标记的截止限制来构建文本数据。据作者所知，MolLM的数据集是这个领域内最大的，包括学术文本和全面的多模态分子信息。
在完成数据集构建后，首先使用KV-PLM的检查点初始化文本编码器的模型权重，旨在为作者数据集中学术文献中提取的句子生成嵌入。KV-PLM通过微调与分子相关的学术文本(包括SMILES字符串)的性能来增强BERT，这是一种广泛使用的通用文本编码器。作者还将标记器输出长度限制增加了一倍，从128个标记增加到256个标记。这种调整有助于在处理文本序列时避免不必要的截断和细节丢失，从而允许包含更详细的文本描述。因此，作者的方法能够有效地捕获完整句子和段落的核心意义和上下文。
MolLM旨在通过处理二维和三维结构信息的途径获得分子数据的表示。二维路径利用从分子的二维图结构中提取的信息，包括度、最短路径距离和边。这些数据元素表示分子中原子之间的空间关系，使模型能够理解其二维结构。相比之下，三维路径侧重于分子的三维几何结构，计算三维结构内原子之间的空间距离。这种几何视角使模型能够深入了解原子的三维排列及其复杂的相互作用，这对于预测分子性质和相互作用通常至关重要。总之，这三种途径使MolLM对分子有了透彻的理解，跨越了它的序列和它的二维和三维结构，从而产生了丰富的，多模态的表示。
受Transformer-M的启发，MolLM采用类似的策略在统一模型中处理分子表示中的二维图形结构和三维几何结构。MolLM使用图Transformer作为基础，将结构数据(如边缘特征、键类型和三维空间关系)编码为注意力机制中的偏差值。Transformer使用这些值来衡量和组合输入的不同部分。作者在注意力机制中实现了这种修改，因为它允许模型中原子之间的关系直接受到它们的二维和三维空间关系的影响，与这些几何关系的物理意义保持一致。
在图的结构方面，顶点代表原子，边代表原子之间的键。对于边特征的编码，作者通过考虑边特征行向量来利用二维图结构，也就是沿着任意两个顶点之间的最短路径。这些边缘特征向量可以表示化学键的特征，包括键类型和立体化学特征。此外，作者的目标是在作者的结构编码中包括二维距离，特别是通过遵循它们的键来确定任意两个原子之间的最短路径。因此，对于二维空间中原子之间的关系，作者定义了一个编码来表示两个原子之间的最短路径距离。
类似地，对三维空间关系，作者没有使用标准欧几里得距离，而是使用高斯基核函数，因为它提供了距离的非线性函数，提供了三维空间关系的更细致的表示。对于多层感知器(MLP)，通过处理高斯基核给出的距离表示来获得最终编码。MLP允许更丰富的表示，包含可以从三维距离推断出的高级特征，以及有选择地强调某些方面的能力，例如三维距离内的特定方向。在应用高斯基核后，将顶点之间的三维距离编码进行处理，找到原子之间的最短路径距离。
然后，这些来自分子图结构和三维几何结构的二维和三维位置编码被线性组合。作者将这种组合编码合并到图Transformer层的注意力机制中，也就是直接将该矩阵作为偏差添加到标准Transformer注意力计算中。将这个矩阵添加到Transformer注意力机制背后的动机是捕获图的顶点和原子之间的上下文关系，同时在最终嵌入中权衡每个顶点和原子的重要性。通过纳入基于二维和三维空间数据的偏差，该模型在注意力计算中直接考虑了这些丰富的空间数据，从而考虑了分子结构中原子之间的空间关系。
对于MolLM的预训练，数据增强策略使作者能够增强处理分子相关任务的鲁棒性和效率。作者利用对比学习作为目标，旨在通过比较成对的正样本和负样本来获得有意义的表征。作者的方法涉及两种类型的损失:(1)对比损失，它比较相同数据样本中不同模态的表示;(2)自对比损失，它对比相同数据样本中相同模态的不同增强。
首先，作者对MoMu中提出的数据增强进行了扩展，引入了两种增强，改变了每个分子的化学特征，如分子基序和官能团。这些附加功能补充了MoMu中使用的两种增强功能，这两种增强功能可以改变图的特征，如节点和边。总的来说，作者对每个二维分子图应用了四种数据增强:(1)节点下降，(2)子图采样，(3)化学转化（添加子结构），(4)子结构去除。然后，作者计算并将三维原子位置插入到这些增广图中，旨在构建适合预训练的语义一致的分子结构。最近的工作已经证明了这种增强在预训练中的有效性，提高了模型的鲁棒性、泛化和对官能团和化学结构组成的理解。
给定一小批N个分子图，作者为每个图生成四种不同的增广。对比学习的目标是最小化这一批中每个图的对比损失，同时考虑跨模态和自对比损失。跨模态对比损失的目的是在潜在空间中对齐分子和文本的表征，而自对比损失的目的是在同一模态内对齐增强表征。换言之，跨模态对比损失使同一分子的不同模态(即分子图和相关文本描述)之间的距离最小化，而使不同分子之间的距离最大化。自对比损耗使同一分子的不同增强之间的距离最小化，而使不同分子的增强之间的距离最大化。因此，该模型学习在联合潜在空间中为分子和文本生成更鲁棒和语义上有意义的表示。
作者将MolLM与一些具有代表性的方法进行了比较，在文本句子级别的零样本学习（zero-shot）跨模态识别任务上，将准确度与前20的召回率（R@20）作为对比指标（越高越好）。如表1所示，MolLM在分子到文本（Molecule-to-Text, M-T）和文本到分子（Text-to-Molecule, T-M）上均超过了其他方法。
表1 与其他方法对比
作者设计了消融实验来验证模型设计的有效性。消融实验通过分子标注任务来进行。在这个任务中，模型根据分子的SMILES字符串、二维图和三维图生成与分子相关的文本。对于基线比较，作者选择MolT5和MoMu。为了说明加入三维信息的优势，作者分别实现了基于二维信息的模型和加入三维信息的模型，同时考虑了使用更小的数据集与作者所构建的较大的基础数据集的对比，如图2所示。总之，在BLEU-2, BLEU-4, ROUGE-L, METEOR这四种语言模型评价指标上，MolLM的性能超过基线模型，三维信息的加入和数据集的扩充都提升了模型性能。
图2 消融实验
作者还进行了案例分析。在预训练的数据增强（分子编辑）过程中，作者发现许多编辑导致分子与原始分子大不相同。因此，作者研究了一个相似度的度量，定义为原始分子及其修改版本的RDKit分子指纹之间计算的平均谷本系数，范围从0.0到1.0，如表2所示。对于四种不同情况（1）有更多的氢键供体/受体（2）药物相似性更高/更低（3）渗透性更高/更低（4）水溶性更高/更低的分子，相似度的度量始终徘徊在0.51左右，表明该模型对分子增强达到了相当鲁棒的修改水平。
表2 案例分析
本文提出了一种新的框架MolLM，强调结合生物医学文本与分子的二维和三维表示构建分子语言模型。这种方法体现了结合更丰富的分子表征的重要性，突出了结构和空间特征之间微妙的相互作用，同时表明了通过文本输入整合生物医学知识的重要作用。在未来，这种方法可以使用GPT进一步增强，以丰富可用的数据池进行预训练。总之，该模型展示了利用计算方法结合现有生物医学文献来解决复杂任务的前景，例如需要超越多种模式的复杂理解的药物发现。
