浙大侯廷军开发Token-Mol 1.0：基于LLM的token化药物设计

编辑 | 萝卜皮

众所周知，药物研发需要经历一段极其漫长且复杂的过程。好在大型语言模型 (LLM) 在药物设计中的应用，正逐渐给这一挑战带来曙光。不过，现有基于 LLM 的方法往往难以有效地整合三维分子结构。

在最新的研究中，浙江大学侯廷军团队提出了 Token-Mol，这是一个仅包含 token 的 3D 药物设计模型，它将二维和三维结构信息以及分子属性编码成离散的 token。

Token-Mol 建立在 Transformer 解码器之上，并经过因果掩蔽训练，引入了专为回归任务定制的高斯交叉熵损失函数，从而在多个下游应用中实现了卓越的性能。

能力方面，该模型超越了现有方法，在两个数据集上将分子构象生成性能分别提高了 10% 和 20% 以上，同时在性质预测方面比仅使用 token 的模型高出 30%。

在基于口袋的分子生成中，它分别将药物相似性和合成可及性提高了约 11% 和 14%。并且，Token-Mol 的运行速度比「专家」扩散模型快 35 倍。

该研究以「Token-Mol 1.0: tokenized drug design with large language models」为题，于 2025 年 5 月 13 日发布在《Nature Communications》。

️背景与挑战

在药物研发中，获取带注释数据集的高成本仍然是该领域发展的重大障碍。近年来，以 BERT 和 GPT 为代表的无监督学习框架的快速发展，为化学和生物等学科引入了无监督化学和生物预训练模型。

这些模型经过大规模无监督训练，学习小分子或蛋白质的表征，并随后针对特定应用进行微调。通过在大规模数据集上进行无监督学习，这些预训练模型有效地解决了与稀疏标记和分布外泛化能力欠佳相关的挑战。

大规模分子预训练模型大致可分为两大类：基于化学语言的模型和利用分子图的模型。

化学语言模型使用诸如简化分子输入系统 (SMILES) 或自引用嵌入字符串 (SELFIES) 等表示形式对分子结构进行编码。它们往往采用类似于 BERT 或 GPT 的训练方法。

基于图的分子预训练模型表现出更高的通用性。它们以图形形式表示分子，节点表示原子，边表示化学键。预训练方法包括各种技术。与基于语言的模型不同，基于图的分子预训练模型本质上包含几何信息。

但是这里有一个很尴尬的问题，基于化学语言的大规模分子预训练模型无法处理对分子特性至关重要的3D结构信息，所以它难以胜任相关下游任务；基于图结构的预训练模型虽能整合 3D 信息却仅聚焦性质预测，同时难以与通用 NLP 模型融合。

因此，当前亟需开发能克服这两类模型局限、适用于所有药物设计场景并可无缝对接通用大语言模型的新型预训练模型。

️Token-Mol

在本研究中，为了克服现有模型的局限性，浙江大学侯廷军、康玉、Chang-Yu Hsieh 团队提出了 Token-Mol，一个用于分子预训练的大规模语言模型。

为了增强与现有通用模型的兼容性，研究人员采用了仅使用 token 的训练范式，将所有回归任务重新定义为概率预测任务。Token-Mol 采用 Transformer 解码器架构构建，通过 SMILES 和扭转角 token 整合必要的二维和三维结构信息。

具体来讲，Token-Mol 的完整工作流程如下图所示。初始阶段通过随机因果掩蔽对数据集进行预训练。随后，模型在定制数据集上进行微调，以完成特定的下游任务，包括构象生成、基于口袋的分子生成以及多种属性的预测。对于回归任务，在微调过程中使用 GCE 损失函数。此外，还可以使用强化学习进一步优化特定下游任务的性能。

图示：Token-Mol 概述。（来源：论文）

在此背景下，预训练数据集的预处理至关重要。对标准 SMILES 格式的整个分子进行深度优先搜索（DFS）遍历，以提取分子结构中嵌入的扭转角。之后，每个提取出的扭转角都会被同化为一个 token，并附加到 SMILES 字符串中。

在整个预训练阶段，研究人员会实施基于因果回归的随机因果掩蔽。预训练完成后，研究人员会对下游任务进行微调。值得注意的是，任务提示是专门为构建对话系统而设计的。

这一特性显示了仅使用 token 的模型相较于其他大规模模型的一个关键优势：它们能够促进实时交互。

对于基于口袋的分子生成任务，研究人员引入了口袋编码器和融合模块，从而更好地将蛋白质口袋信息整合到模型中。同时，他们利用多头条件注意力机制，将每个自回归步骤中生成的信息完整地整合到后续迭代中。该机制将自回归过程中生成的每个 token 视为后续生成的先决条件，从而确保整个查询、键和值矩阵均源自原始序列。

需要注意的是，在实际应用中，先导化合物不仅要对靶标表现出高亲和力，还要满足一系列标准，包括高生物活性和多种良好的药理特性。这对基于口袋的分子设计任务提出了更高的要求，而训练数据集中受体-配体分子对的整合存在固有的局限性。

该模型主要利用来自蛋白质口袋的信息来生成配体分子。因此，这些生成分子的性质受到训练数据的严重影响，限制了对其生物物理和化学性质的明确控制。当需要精确调节分子特性时，这些限制尤为明显。

Token-Mol 建立在自回归语言模型架构之上，其中 token 的生成与强化学习框架中的动作相一致，从而促进了强化学习的无缝优化，从而保证获得量身定制的结果。

️性能验证

为了验证 Token-Mol 的性能，研究人员将它应用于分子构象生成、性质预测和基于口袋的分子生成等任务。

图示：Token-Mol 和基线模型之间的分子特性分布。（来源：论文）

在分子构象生成任务中，Token-Mol 的表现优于现有的最佳方法，在两个数据集上的各项指标上分别实现了超过 10% 和 20% 的准确率提升。

在分子属性预测任务中，Token-Mol 与仅使用 token 的模型相比，在回归任务中平均提升了 30%。

在基于 Pocket 的分子生成任务中，Token-Mol 不仅生成了与当前最佳模型生成的分子 Vina 得分相当的分子，而且药物相似性 (QED) 和合成可及性 (SA) 分别提升了约 11% 和 14%。

图示：真实药物设计场景评估。（来源：论文）

为了进一步验证 Token-Mol 的泛化能力，团队还在真实的药物设计场景中进行了测试。值得一提的是，Token-Mol 生成的类药分子在 8 个不同靶点的评估中，平均成功率提高了 1 倍。

此外，将强化学习集成到 Token-Mol 中，能够有效提升特定下游任务在更现实场景下的性能，使其在 3D 任务上优于基于几何图神经网络的大型模型。

️与 Token-Mol 聊天

Token-Mol 的 token-only 框架相较于传统回归模型具有显著优势，因为它能够无缝集成前沿的大模型技术，包括快速学习、混合专家（MoE）和检索增强生成（RAG）。

为了说明此功能，该团队在论文里展示了几个简单的对话用例。通过运用提示学习，研究人员可以控制诸如本研究中提到的属性预测等任务的执行。

首先，他们插入特定的提示，例如「预测 ESOL」从而微调模型。这使得用户能够在提示后与模型直接交互，从而请求预测不同的分子属性。

图示：与 Token-Mol 聊天。（来源：论文）

使用时，研究人员查询了不同分子的各种属性，Token-Mol 成功地给出了相应的预测。这展现了 Token-Mol 拥有与化学家进行有意义对话的潜力。

用户可以提供分子构象，但由于 Token-Mol 可以生成相应的构象，因此最终输出将仅包含预测的目标属性。

研究人员表示，未来该模型可以整合 RAG。

当向 Token-Mol 查询分子的特定属性时，系统会使用基于嵌入的向量搜索，将查询转换为向量。然后，该向量会与数据库中高度相关的向量描述进行匹配，以提供上下文信息。查询以及检索到的上下文信息（例如空间结构信息和其他相关属性）随后被输入到 Token-Mol，最终生成答案。

总而言之，Token-Mol 的内在潜力巨大。它的新颖架构也为药物设计 AI 模型的标准化提供了参考。

论文链接：https://www.nature.com/articles/s41467-025-59628-y