Uni-Mol Docking：预训练分子模型如何进入对接任务

Uni-Mol（深势科技，2023）的核心主张是：分子的三维信息应该在预训练阶段就被学习，而不是在下游任务中临时处理。它把这个思路应用到包括对接在内的多个任务上。

预训练设计

# 与二维预训练的区别（见 142）
#
# 二维路线（如 MolCLR、GROVER）：
#   输入分子图（原子 + 键）
#   预训练任务：掩码原子/键预测、子图对比
#   → 学到的是拓扑规律
#
# Uni-Mol：
#   输入三维构象（原子 + 坐标）
#   → 直接建模空间关系
#
# 两个预训练任务：
#   1) 3D 位置恢复
#      给原子坐标加噪声，让模型还原
#      → 学习「合理的三维结构长什么样」
#
#   2) 掩码原子预测
#      遮住某些原子的类型，根据周围环境预测
#      → 学习原子的化学环境
#
# 架构：
#   基于 Transformer，但注意力中加入了
#   成对距离的编码（pair representation）
#   → 保证 SE(3) 不变性
#
# 预训练数据：
#   约 2 亿个分子构象（用 RDKit 生成）
#   + 300 万个蛋白口袋（来自 PDB）
#
# 关键点：
#   分子构象与蛋白口袋分别预训练，
#   得到两个编码器，下游任务中组合使用

Uni-Mol Docking 的做法

• 输入：蛋白口袋 + 配体（含 RDKit 生成的初始构象）；
• 预测：配体原子与口袋原子之间的距离矩阵（与 TankBind 思路类似，见 103《TankBind 论文精读》）；
• 坐标重建：从预测的距离约束出发，用优化算法生成坐标；
• Uni-Mol Docking V2（2024）：加入了对物理合理性的改进，显著提高了 PoseBusters 通过率——这是针对早期深度学习对接方法共同短板的直接回应。

V2 的关键改进

问题	V2 的应对
物理有效性差	训练中加入物理约束损失；推理时做几何精修
键长键角异常	约束配体内部几何来自合理构象
空间冲突	显式惩罚与蛋白的重叠
手性错误	加入手性约束

V2 报告的 PoseBusters 有效性通过率大幅提升，这说明物理有效性问题是可以通过工程手段解决的——不是深度学习范式的固有缺陷，而是早期工作没有重视。

使用

# 安装
pip install unimol_tools
# 或从源码：git clone https://github.com/deepmodeling/Uni-Mol

# 用作分子表示（性质预测任务）
from unimol_tools import UniMolRepr

clf = UniMolRepr(data_type="molecule", remove_hs=False)
smiles = ["CC(=O)Nc1ccc(O)cc1", "CCO"]
repr_dict = clf.get_repr(smiles, return_atomic_reprs=True)

print(repr_dict["cls_repr"][0].shape)      # 分子级表示
print(repr_dict["atomic_reprs"][0].shape)  # 原子级表示

# 这些表示可作为下游模型的输入特征

# 微调做性质预测
from unimol_tools import MolTrain

clf = MolTrain(
    task="regression",
    data_type="molecule",
    epochs=20,
    learning_rate=1e-4,
    batch_size=16,
    metrics="rmse",
    save_path="./exp",
)
clf.fit(data="train.csv")   # 需要 SMILES + TARGET 列

# 注意：三维预训练模型对构象敏感
#   → 构象生成方式会影响结果
#   → 应固定构象生成流程（见 279）

实际表现的判读

• 在多个分子性质基准上表现好：但与所有预训练分子模型一样，相对精心调优的传统方法（指纹 + 梯度提升）的优势并不总是显著（见 142《Uni-Mol 论文精读》）；
• 三维预训练的收益取决于任务：对依赖三维结构的性质（如构象相关的性质）收益明显；对主要由二维拓扑决定的性质，收益有限；
• 对接任务上 V2 有竞争力：在 PoseBusters 基准上的表现与传统对接方法可比，且速度快；
• 与传统方法的正确比较：必须在相同的数据划分（尤其是骨架划分或时间划分）下比较，随机划分会高估所有方法（见 238《PDBbind》）。

三维预训练的普遍问题

# 问题一：构象从哪来
#   预训练用 RDKit 生成的构象，
#   而真实的结合构象可能与之差别很大
#   → 「预训练时见的三维」与「应用时需要的三维」不一致
#
# 问题二：单构象 vs 构象集合
#   分子在溶液中是构象集合，
#   用单一构象表示丢失了信息
#   → 有些工作用多构象平均，但成本高
#
# 问题三：评估的可比性
#   预训练模型通常在大量数据上训练，
#   下游测试集可能与预训练数据有重叠
#   → 需要仔细检查
#
# 问题四：计算成本
#   三维模型比二维模型慢得多
#   → 大规模筛选时是实际约束
#
# 实践建议：
#   把 Uni-Mol 当作「候选方法之一」，
#   在自家数据上与指纹+LightGBM 基线对比后再决定
#   （见 159）

生态定位

• Uni-Mol 是国内 AI 制药开源生态中影响较大的工作，配套的 Uni-Mol Tools 让使用门槛降低；
• 与深势科技的商业平台联动：Hermite 等产品使用了相关技术（见 316《DP Technology 深势科技》）；
• 持续迭代：从 Uni-Mol 到 Docking V2，再到后续版本，是少数持续维护的中文社区分子预训练项目。

延伸资源

• 预训练分子模型：142《Uni-Mol 论文精读》；分子表示对比：159《图神经网络 GNN》；
• DiffDock：101《DiffDock 论文精读》；深势科技：316《DP Technology 深势科技》。