Uni-Mol 论文精读：3D 预训练何时比 SMILES/GNN 更有意义

Uni-Mol（Zhou 等，ICLR 2023）主张分子的三维信息应在预训练阶段学习。这个方向有道理——但「三维预训练是否比 SMILES/GNN 更好」的答案高度依赖具体任务，理解这个依赖关系比记住结论重要。

方法概要

# 表示：
#   输入 = 原子类型 + 三维坐标
#   模型 = Transformer + 成对距离编码
#
#   关键：注意力中加入基于原子间距离的偏置
#         → 保证 SE(3) 不变性（旋转平移不改变输出）
#
# 两个预训练任务：
#   1) 【3D 位置去噪】
#      给坐标加噪声，让模型还原原始坐标
#      → 学习「合理的三维结构长什么样」
#
#   2) 【掩码原子预测】
#      遮住原子类型，根据空间环境预测
#      → 学习原子的化学环境
#
# 两个独立的预训练模型：
#   分子模型：约 2 亿个 RDKit 生成的构象
#   口袋模型：约 300 万个 PDB 蛋白口袋
#
# 下游任务：
#   分子性质预测：用分子模型
#   对接/结合预测：分子模型 + 口袋模型组合（见 104）

三维预训练什么时候有优势

任务类型	三维是否有帮助	原因
量子化学性质（QM9 等）	明显有	这些性质本来就由三维结构决定
构象相关的性质	有	如某些光谱性质
结合亲和力	可能有	结合依赖三维互补
logP、溶解度	有限	主要由官能团组成决定
毒性、代谢	有限	更依赖结构警示与官能团
分类任务（小数据）	常常没有	三维带来的额外参数容易过拟合

判断的经验法则：问「这个性质在多大程度上由分子的空间形状而非官能团组成决定？」如果答案是「主要看有什么官能团」，那么二维表示已经够了，三维模型只会增加成本与过拟合风险。

三维预训练的固有难题

# 难题一：【构象从哪来】
#
#   预训练用 RDKit 生成的低能构象，
#   但分子在实际环境中：
#     - 在溶液中是构象系综，不是单一构象
#     - 结合时采取的可能是高能构象（见 049）
#     - 结晶态构象又不同
#
#   → 「预训练时见的三维」与「任务需要的三维」不一致
#
#   应对：
#     - 用多个构象取平均（成本高）
#     - 用能量加权（需要能量计算）
#     - 接受这个近似（多数工作的做法）
#
#   【实践后果】：
#     结果对构象生成方式敏感
#     → 必须固定构象生成流程（工具、参数、随机种子）
#     → 否则结果不可复现

# 难题二：【计算成本】
#   三维模型比二维模型慢数倍
#   → 大规模筛选时是实际约束
#
# 难题三：【预训练数据的构象质量】
#   RDKit 生成的 2 亿个构象，质量参差
#   → 模型学到的「合理三维结构」可能带有
#     ETKDG 算法的偏好

# 难题四：【评估的公平性】
#   预训练模型见过大量分子，
#   下游测试集可能与预训练数据重叠
#   → 应检查测试集分子是否在预训练数据中

正确的评估方式

# 比较预训练模型与传统方法时，必须控制的变量：
#
# 1) 【相同的数据划分】
#    骨架划分或时间划分，不用随机划分
#
# 2) 【基线获得同等调优】
#    ECFP + LightGBM 也要认真调参（见 131）
#
# 3) 【多个随机种子】
#    报告均值 ± 标准差
#
# 4) 【报告微调成本】
#    预训练模型的微调时间、显存需求
#    → 如果提升 0.02 AUC 但成本高 100 倍，值得吗？
#
# 5) 【消融预训练】
#    对照实验：同样架构但【不做预训练】，直接训练
#    → 差距才是预训练的真实贡献
#    → 【这个对照常常缺失】
#
# 6) 【不同数据量下的表现】
#    预训练的价值通常在小数据时最明显
#    → 画出「训练数据量 vs 性能」曲线
#    → 比单点比较有信息量得多

import numpy as np

def learning_curve_comparison(X, y, models, sizes=(100, 300, 1000, 3000)):
    """比较不同数据量下各方法的表现"""
    results = {name: [] for name in models}
    for n in sizes:
        for name, fit_predict in models.items():
            scores = []
            for seed in range(5):
                rng = np.random.default_rng(seed)
                idx = rng.choice(len(X), size=min(n, len(X)), replace=False)
                scores.append(fit_predict(X[idx], y[idx]))
            results[name].append((np.mean(scores), np.std(scores)))
    return results

# 典型的发现：
#   小数据（< 500）：预训练模型优势明显
#   大数据（> 5000）：差距缩小甚至反转
#   → 【这个规律比单个数字有用得多】

使用

pip install unimol_tools

from unimol_tools import MolTrain, MolPredict

# 微调做回归
clf = MolTrain(
    task="regression",
    data_type="molecule",
    epochs=20,
    learning_rate=1e-4,
    batch_size=16,
    early_stopping=5,
    metrics="rmse",
    split="scaffold",          # 【用骨架划分】
    save_path="./exp",
    smiles_col="SMILES",
    target_cols=["TARGET"],
)
clf.fit(data="train.csv")

# 预测
predictor = MolPredict(load_model="./exp")
y_pred = predictor.predict(data="test.csv")

# 只取表示做下游任务
from unimol_tools import UniMolRepr
repr_model = UniMolRepr(data_type="molecule", remove_hs=False)
reprs = repr_model.get_repr(["CCO", "c1ccccc1"], return_atomic_reprs=True)
X = reprs["cls_repr"]        # 分子级表示，可喂给任意 sklearn 模型

# 【实用做法】：
#   把 Uni-Mol 表示与 ECFP 拼接，一起喂给 LightGBM
#   → 常常比单独用任一种都好
#   → 成本低，值得先试

与其它表示的实际比较

• 没有普遍最优的表示：不同任务、不同数据量下最优选择不同；
• 拼接常常有效：ECFP + 描述符 + 预训练表示一起用，交给梯度提升模型——这个组合在实践中很稳健；
• 先跑基线：ECFP + LightGBM 训练几分钟，Uni-Mol 微调可能要几小时。先知道基线在哪，再决定是否值得；
• 部署成本要计入：预训练模型的推理需要 GPU，运维复杂度更高。

延伸资源

• Uni-Mol Docking：104《Uni-Mol Docking》；ChemBERTa：145《ChemBERTa 论文精读》；MolFormer：147《MolFormer 论文精读》；
• GNN：159《图神经网络 GNN》；Chemprop：131《Chemprop / D-MPNN 论文精读》；构象生成：049《构象生成入门》。