Token-Mol 1.0 深度解析:将三维分子结构“翻译”为语言模型的离散语言
摘要
随着大型语言模型(LLM)在药物设计领域的应用日益增多,如何有效融合分子的三维(3D)结构信息成为了一大核心挑战 1。传统的化学语言模型(如基于SMILES)本质上无法处理3D信息 2,而基于图的方法虽然可以包含几何信息,却难以与通用的NLP模型集成 3。Token-Mol 1.0 是一篇发表于 Nature Communications 的研究,它提出了一种创新的“纯词元化”(token-only)范式,旨在构建一个统一的AI药物设计基础模型,弥合二维化学语言与三维物理结构之间的鸿沟。
本解析将重点阐述Token-Mol的核心方法论,特别是其分子表征策略、模型架构与关键创新模块,并探讨其与通用大模型技术(如RAG)的潜在整合,为理解和借鉴其设计哲学提供深度视角。
核心方法:Token-Mol的分子表征哲学
Token-Mol的基石在于其独特的输入构建方式,它将复杂的分子信息完全转化为一个离散的词元(token)序列,使得标准的语言模型可以直接处理。
输入构建:融合2D与3D信息的“分子语言”
模型的核心思想是将一个带有三维构象的分子,编码为一个包含二维拓扑和三维几何信息的单一文本序列。这个过程如 图1a 所示,具体步骤如下:
-
获取二维拓扑信息 (SMILES):
-
输入:分子的二维连接性图。
处理
:首先,将分子结构转换为化学领域广泛应用的
SMILES(简化分子线性输入规范)字符串
4
。SMILES是一种用ASCII字符串明确描述分子结构的规范。
输出
:一串描述分子图的字符序列,例如
C1=CC(=CC(=C1)O)CN...
5
。这是最终序列的基础部分。
-
-
提取三维几何信息 (Torsion Angles):
挑战
:直接将原子的三维笛卡尔坐标(XYZ)作为输入,会导致序列过长且难以处理分子的旋转/平移不变性 6
。
Token-Mol的解决方案
:通过在分子的SMILES表示上进行
深度优先搜索(DFS)遍历
,来提取决定其三维构象的
关键可旋转键的扭转角(Torsion Angles)
7
。扭转角是描述分子构象的核心内部坐标,具有旋转不变性。
输出
:一系列代表扭转角度数的连续数值,例如
[-0.20, 3.14, 2.18, ...]
8
。
-
整合为最终的“Token-only”表征:
处理
:将提取出的扭转角数值
也作为独立的词元
,直接追加到SMILES字符串的末尾 9
。同样,分子的其他理化性质(如在性质预测任务中)也被处理成词元 1010
。
最终输入序列
:一个结合了SMILES和扭转角词元的长序列,能够同时表征分子的2D化学结构和3D空间构象 1111
。
设计哲学
:这种表征方式非常精妙,它
将决定分子3D构象的核心自由度(扭转角)从连续空间映射到了离散的词元空间
,同时保留了描述2D化学结构的SMILES语言。这使得一个基于Transformer的标准语言模型架构,能够在一个统一的框架内同时“阅读”和“理解”分子的2D和3D信息 1212
。
Token-Mol的核心机制在于其创新的数据表示方式。它使用广泛接受的SMILES(简化分子线性输入规范)字符串来表示分子的2D连接性,即原子类型和化学键排布 。SMILES本身是一种成熟的化学语言,但它本质上缺乏3D空间信息 。为了弥补这一缺陷,Token-Mol引入了扭转角(torsion angles)作为3D构象的关键描述符 。扭转角描述了沿化学键旋转的构象自由度,是决定分子三维形状的核心内部坐标之一 。
该模型的实现流程是,首先通过深度优先搜索(DFS)遍历分子图,提取出所有可旋转键的扭转角。然后,将这些连续的扭转角数值进行离散化处理,并作为特殊的“扭转角令牌”附加到SMILES字符串的相应位置。最终形成一个混合序列,例如 C(C<120.5>)C
,其中 <120.5>
就是一个代表特定扭转角度的令牌。这种方式巧妙地将2D拓扑(SMILES骨架)和3D几何(扭转角)编织成一种“3D注释的化学语言”,可以直接输入到Transformer解码器中进行自回归式学习 。
图1: Token-Mol总览。(a) 数据预处理流程,将分子的SMILES字符串与扭转角结合成单一的词元化表示 13。(b) 模型的预训练与微调工作流 14。(c) GCE损失函数的权重分配示意图 15。(d) 用于口袋生成任务的编码器与融合模块 16。
输入表征的优势与劣势
这种将SMILES与扭转角结合的“分子语言”是一种创新的折衷,具有独特的优缺点。
优势:
统一2D与3D信息
:最核心的优势在于,它成功地将2D拓扑信息(SMILES)和3D几何信息(扭转角)编码到一个单一的、离散的词元序列中,从而能够被标准的语言模型架构直接处理 17171717
。
兼容性与速度
:作为一种“纯词元化”模型,它与GPT等通用大语言模型的架构高度兼容,易于集成 18181818
。其推理速度极快,例如在分子生成任务中比基于几何的扩散模型快约35倍 191919191919191919
。
规避XYZ坐标的难题
:该方法避免了直接使用笛卡尔坐标(XYZ)带来的序列过长和缺乏旋转等变性的问题 20
。
-
** bridging a gap**:Token-Mol的表征为分子表示范式提供了
第三条路径
,有效连接了传统的2D序列模型(无法处理3D信息)和3D图模型(难以集成到通用LLM中) 21
。
生成更灵活的分子
:通过引入扭转角信息,模型能够生成比单纯基于2D信息的模型更柔性、更多样化的分子,以更好地适应不同形状的口袋 22
。
劣势与挑战:
对低频信息的学习不足
:模型在学习和准确预测那些
出现频率较低的扭转角
的分布时会遇到困难 23
。
数值敏感性有限
:尽管引入了GCE损失函数,但与基于图神经网络(GNN)的模型相比,Token-Mol对连续数值的敏感度仍然存在局限 24
。
结构有效性风险
:由于模型是自回归地生成序列,对扭转角数量或数值的预测不准确可能会导致最终生成的分子结构无效 25
。
对柔性分子的挑战
:分析表明,随着分子中
可旋转键数量的增加,所有评估指标的性能都呈下降趋势
26
。尽管Token-Mol在这种情况下依然表现出相对优势,但这仍然是一个固有的挑战 27
。
模型架构与训练策略
Token-Mol采用了一系列精心设计的策略来训练模型,以确保其能够从“分子语言”中学习到有用的知识。
模型骨干:Transformer解码器
架构
:模型基于一个包含
12层Transformer解码器
的架构,每层配备8个注意力头 28
。
自回归方法
:采用自回归(Autoregressive)方式进行训练和生成 29
。在训练时,通过掩码矩阵防止信息泄露 30
;在生成时,模型逐个预测下一个词元,从而构建出完整的分子序列 31
。
关键创新1:随机因果掩码 (Random Causal Masking)
挑战
:传统的从左到右的因果掩码不适合“完形填空”式的任务,限制了模型的灵活性 32
。
Token-Mol的策略
:在预训练阶段,采用
随机因果掩码
策略 33333333
。它会从泊松分布中采样要掩盖的片段数量(1到6个),然后在序列中随机选择位置进行掩码 34
。被掩盖的内容会附加在序列末尾,由特殊词元引导模型进行预测 35
。
目的
:这种策略极大地增强了模型“填空”的能力,使其能适应更多样化的下游任务,例如在分子的特定位置进行修饰或补全 36
。
关键创新2:高斯交叉熵损失函数 (Gaussian Cross-Entropy Loss)
挑战
:传统的交叉熵损失函数主要用于离散分类任务,它对数值大小不敏感 37
。例如,在预测扭转角时,如果真实值是2°,那么预测成3°和80°所产生的损失是完全相同的,这显然不合理 38
。
Token-Mol的解决方案
:针对回归任务(如预测扭转角和分子属性),作者提出了
高斯交叉熵(GCE)损失函数
39393939
。
工作原理
(如图1c):对于每一个要预测的数值标签,GCE会构建一个
以该标签值为中心的高斯分布
40
。这样,离真实标签值越近的词元会被赋予越高的概率权重,而离得远的词元权重则较低 41
。
效果
:这种加权方式使得模型在训练过程中能够
学习到数值之间的相对关系
42
,显著提升了其在回归任务上的表现。消融实验表明,缺少GCE会导致模型在回归任务上的平均RMSE增加约12% 43
。
下游任务:分子性质预测
分子性质预测是检验模型表征学习能力的关键。在Token-Mol的框架中,这不是预训练阶段的一部分,而是一个下游微调任务。
-
流程:模型首先在大型无标签分子数据集(GEOM)上进行预训练,学习通用的分子表征 44444444。然后,针对具体的性质预测任务,使用带有标签的特定数据集对模型进行微调(Fine-tuning) 45454545。
-
预测的性质:研究中评估了一系列分类和回归任务,数据集来源于MoleculeNet和TDC等基准平台 46。
-
分类任务 (Classification)
:
BACE
:预测分子是否为β-分泌酶1(BACE1)抑制剂 47
。
BBBP
(Blood-Brain Barrier Penetration):预测分子是否能穿透血脑屏障 48
。
ClinTox
:预测药物是否因毒性而在临床试验中失败 49
。
SIDER
:预测药物的副作用 50
。
Tox21
:预测化合物在12条毒性信号通路上的活性 51
。
ToxCast
:预测化合物在数百种高通量筛选实验中的毒性 52
。
-
回归任务 (Regression)
:
ESOL
:预测有机物在水中的溶解度 53
。
FreeSolv
:预测小分子在水中的水合自由能 54
。
Lipophilicity
:预测分子的亲脂性(油水分配系数) 55
。
Caco-2
:预测药物通过Caco-2细胞的渗透率,用以模拟肠道吸收 56
。
AqSolDB (Aqueous Solubility)
:预测分子的水溶性 57
。
Acute Toxicity LD50
:预测化学物质的急性口服毒性(半数致死剂量) 58
。
-
-
注意力分析:为了验证模型的可解释性,研究者分析了模型在进行溶解度(ESOL)和毒性(LD50)预测时的注意力权重 59。结果显示,在预测溶解度时,模型会高度关注极性基团(如羟基、氨基)和疏水基团(如氯苯) 60;在预测毒性时,模型会准确地将高注意力分配给已知的毒性基团(toxicophores),如亚硝酰胺和磷酸三酯 61616161。这证明模型不仅能做出准确预测,其决策过程也与化学直觉相符。
核心模块深度解析:基于口袋的分子生成
Token-Mol通过引入特定模块来处理复杂的下游任务,尤其是基于口袋的分子生成。
口袋编码器与融合模块 (Pocket Encoder and Fusion Block)
-
任务:在给定蛋白质口袋信息的前提下,生成能与之结合的配体分子 62。
-
输入:
蛋白质口袋信息
:使用一个
预训练好的蛋白质口袋编码器
来提取口袋的3D结构和理化性质特征 63636363
。该编码器在Token-Mol微调阶段其参数被
冻结
,仅作为特征提取器 64646464
。
- 部分生成的配体序列(在自回归过程中)。
-
模型设计与融合机制 (如图1d):
融合机制
:为了将口袋信息与正在生成的配体分子信息相融合,模型采用了一种
多头条件注意力(multi-head condition-attention)机制
6565656565656565
。
工作原理
:这与传统的交叉注意力有所不同。在这里,蛋白质口袋信息被视为一个
静态的“提示(Prompt)”
66666666
。在自回归生成配体的每一步中,注意力机制的
查询(Query)、键(Key)和值(Value)矩阵完全来源于已经生成的配体序列本身
67
。口袋的“提示”信息则被用来调节(condition)这个自注意力计算过程。
- 通俗解释:可以想象成,在写一个故事时,有一个固定的主题(口袋信息)放在旁边。在写每个新句子(生成新原子/键)时,你不仅要回头看自己已经写过的内容(已生成的配体部分),还要时刻瞟一眼那个主题,确保新写的内容与主题是相关的。这使得生成的分子在每一步都受到口袋环境的约束和引导,从而保证了其结构与目标口袋的高度匹配。
与强化学习(RL)的结合
目的
:针对特定目标(如最大化与某个靶点的结合亲和力),进一步优化生成的分子 68
。
可行性
:Token-Mol的自回归架构与强化学习框架天然契合,因为“生成一个词元”这个动作可以被看作是RL中的一个“行动(action)” 69
。
实施
:研究中使用了
REINVENT算法
对模型进行优化 70
。通过设计一个包含亲和力(Vina Score)和类药性(QED)的奖励函数,模型可以在满足约束条件(如类药性)的同时,逐步生成亲和力更高的分子 717171717171717171
。实验证明,通过RL优化,分子的Vina score得到了显著提升,平均值从-8左右优化到了约-9.5 72
。
未来展望:与通用大语言模型的融合
Token-Mol的“纯词元化”框架使其与前沿的通用大模型技术具有极佳的兼容性,为未来构建更强大的药物设计助手铺平了道路 73737373。
-
即时交互与提示学习 (Prompt Learning)
-
Token-Mol可以像与ChatGPT对话一样进行交互 74
。通过设计特定的任务提示(如
Predict ESOL
)并对模型进行微调后,用户可以直接向模型查询特定分子的性质,模型也能成功返回预测结果 75
。这展示了其作为
化学家对话伙伴
的巨大潜力 76
。
-
-
检索增强生成 (Retrieval-Augmented Generation, RAG)
-
论文明确指出了整合RAG的未来方向 77777777
。
工作流程设想
:当用户向Token-Mol查询一个分子的性质时,系统可以首先将查询(包含分子信息)转换为一个向量,然后在一个包含海量分子数据(如3D结构、其他理化性质、实验数据等)的
向量数据库
中进行检索 78
。
优势
:检索出的最相关的上下文信息,会与原始查询一起被输入到Token-Mol中 79
。这使得模型在生成答案时,不仅依赖其内部学到的知识,还能
参考外部数据库提供的、更丰富和准确的实时信息
,从而极大地增强其预测的可靠性和信息的全面性。
-
总结
Token-Mol 1.0 是首个专为药物设计定制的“纯词元化”大型预训练语言模型 80。通过将分子的2D(SMILES)和3D(扭转角)信息统一编码为离散词元序列,并引入高斯交叉熵(GCE)损失函数来处理连续数值,它成功地让一个标准的Transformer架构能够同时理解和生成复杂的分子结构 81。其模块化的设计,以及与强化学习、RAG等前沿AI技术的天然兼容性,为开发“单一基础模型驱动的全面药物设计”提供了一条极具前景的道路 82。