NV-KERMT-70M-v2代码实现详解：从SMILES解析到分子嵌入生成

NV-KERMT-70M-v2代码实现详解：从SMILES解析到分子嵌入生成

【免费下载链接】NV-KERMT-70M-v2项目地址: https://ai.gitcode.com/hf_mirrors/nvidia/NV-KERMT-70M-v2

NV-KERMT-70M-v2是一款基于图Transformer架构的分子表示学习模型，专为药物发现中的ADMET（吸收、分布、代谢、排泄、毒性）属性预测任务设计。该模型通过创新的多任务预训练策略，将2D分子图转化为具有化学意义的潜在表示，为药物研发提供强大的分子特征提取能力。

模型架构概览：融合图Transformer与对比学习

核心组件解析

NV-KERMT-70M-v2采用扩展自GROVER的graph-transformer架构，主要包含以下关键模块：

• 编码器（Encoder）：6层消息传递+注意力层，隐藏维度800，4个注意力头，配备PReLU激活函数和0.1的dropout率
• 解码器（Decoder）：3层Transformer结构，8个注意力头，512隐藏/潜在维度，2048维前馈网络，采用旋转位置编码（RoPE）
• 多任务学习头：包含概率潜在头、批内对比辅助变量、SMILES重构Transformer解码器以及化学特异性词汇预测头

![模型架构示意图]

创新点：四合一联合概率目标

模型通过以下四种损失函数的联合优化实现化学知识的深度融合：

1. SMILES重构损失（通过Transformer解码器从采样潜变量生成SMILES）
2. 对称潜变量密度正则化（A-MIM项）
3. 批内对比判别损失（匹配/非匹配分类）
4. 化学特异性自监督损失（原子上下文、键上下文和官能团预测）

数据处理流程：从SMILES到分子图

SMILES解析与图构建

输入的SMILES字符串首先通过RDKit工具包进行解析，转化为2D原子-键图结构。这一过程包括：

1. 化学键类型识别（单键、双键、三键等）
2. 原子属性提取（元素类型、电荷、氢原子数量等）
3. 分子拓扑结构构建（原子连接关系）

注意：推荐的SMILES解码器最大序列长度为512 tokens，超过此长度的分子需要截断或忽略。

化学特异性词汇表

模型使用三个关键的词汇表文件实现分子特征的向量化：

• 原子词汇表：pretrain_atom_vocab.json
• 键词汇表：pretrain_bond_vocab.json
• SMILES词汇表：pretrain_smiles_vocab.pkl

这些词汇表包含了药物发现领域常见的化学实体，确保模型能够准确理解分子结构的化学含义。

分子嵌入生成：从图编码到潜在表示

图Transformer编码过程

分子图的编码通过以下步骤完成：

1. 局部消息传递：捕获原子间的局部化学环境信息
2. 全局自注意力：建模分子范围内的长距离相互作用
3. 多任务块处理：融合多种化学特征学习任务

这一过程将分子图转化为原子级和键级的表示，保留了分子的化学特性和结构信息。

潜在表示学习

编码器的输出通过一个学习投影层映射到高斯潜在分布，生成512维的分子嵌入向量。这一向量不仅包含分子的结构信息，还融合了化学性质的预测能力，可直接用于下游ADMET属性预测任务。

实际应用：模型部署与微调指南

环境要求

• PyTorch 2.x
• RDKit cheminformatics toolkit
• 模型 checkpoint：kermt_contrastive_v2.0.pt

快速开始步骤

1. 克隆仓库：

git clone https://gitcode.com/hf_mirrors/nvidia/NV-KERMT-70M-v2

2. 加载预训练模型和词汇表文件
3. 准备SMILES格式的分子数据
4. 生成分子嵌入或进行下游任务微调

应用场景

• ADMET属性预测模型的特征提取器
• 多端点ADMET数据集的微调基础模型
• 分子表示学习研究的基准模型

总结：分子AI的强大工具

NV-KERMT-70M-v2通过创新的图Transformer架构和多任务预训练策略，为药物发现领域提供了一个强大的分子表示学习工具。其核心优势在于：

• 化学感知能力：专为分子结构设计的词汇表和编码方式
• 多任务学习：融合多种自监督任务，提升表示质量
• 灵活性：可作为特征提取器或微调基础模型使用

无论是计算化学研究者还是药物发现科学家，都能从这一模型中获得有价值的分子表示，推动药物研发的效率和准确性。

完整技术细节请参考：Adrian et al. (2025). "Multitask finetuning and acceleration of chemical pretrained models for small molecule drug property prediction." arXiv:2510.12719

【免费下载链接】NV-KERMT-70M-v2项目地址: https://ai.gitcode.com/hf_mirrors/nvidia/NV-KERMT-70M-v2