Transformer 在分子中的应用：从 SMILES 到图、3D 和多模态

Transformer 已经渗透到分子建模的各个层面。但「用 Transformer」不是一个统一的做法——在 SMILES 上、在图上、在三维坐标上，它面临的问题与解决方式完全不同。

四种用法

各用法的关键问题

# ---- SMILES Transformer ----
#
# 核心问题：SMILES 不是自然语言（见 145）
#   - 语法严格，没有需要「理解」的歧义
#   - 同一分子多种写法
#   - 环闭合标记打破局部性
#
# 应对：
#   - SMILES 增强（随机化写法）
#   - 相对位置编码（RoPE，见 147）
#   - 用 SELFIES 替代 SMILES（保证生成的分子有效）
#
# 价值场景：
#   【生成任务】—— 自回归生成天然适合 Transformer
#   而在性质预测上，相对指纹基线优势有限

# ---- 图 Transformer ----
#
# 动机：解决 GNN 的两个问题（见 159）
#   1) 过平滑限制深度
#   2) 感受野受层数限制
#
# 做法：全局注意力，每个原子直接关注所有原子
#   → 一层就有全局感受野
#
# 关键问题：【如何注入图结构信息】
#   纯注意力丢失了「哪些原子相连」
#   → 必须用结构编码：
#     - 【最短路径距离】作为注意力偏置（Graphormer）
#     - 中心性编码（度数）
#     - 边特征编码
#     - 拉普拉斯特征向量作为位置编码
#
# 【实际效果】：
#   在大数据集（如 PCQM4M 千万级）上明显优于 GNN
#   在小数据集（几千个分子）上【常常不如 GNN】
#   → 参数多、归纳偏置弱 → 需要更多数据

# ---- 三维 Transformer ----
#
# 关键问题：如何保证对称性（见 160）
#
# 两种做法：
#   1) 用距离作为注意力偏置（Uni-Mol）
#      → 保证不变性，实现简单
#      → 但不能输出向量
#   2) 等变注意力（Equiformer、SE(3)-Transformer）
#      → 完整的等变性，可输出坐标
#      → 实现复杂，计算慢
#
# 选择：只预测标量用 1)，需要坐标用 2)

# ---- 多模态 ----
#
# 关键问题：【不同模态如何对齐】（见 164）
#   蛋白残基、配体原子、核苷酸
#   → AF3/Boltz 的做法：统一到【原子级】表示
#   → 用成对表示建模跨模态关系

注意力机制在分子上的局限

• 二次复杂度：分子通常不大（几十个原子），但蛋白可达数千残基——需要线性注意力或稀疏注意力；
• 归纳偏置弱。这是最根本的问题。Transformer 的强项是「从大量数据中学习任意关系」，而分子建模的特点是数据少但物理规律明确——此时把规律编码进架构（如 GNN 的图结构、等变网络的对称性）比让模型自己学更有效；
• 注意力不等于相互作用：不应把注意力权重直接解读为化学相互作用（见 141《MolTrans 论文精读》、168《可解释性 AI》）；
• 位置编码的选择很关键：对分子来说，「位置」的含义不像文本那样明确——用什么编码结构信息，往往比注意力机制本身更重要。

实践判断

# 什么时候用 Transformer：
#
# ✓ 【生成任务】
#   自回归生成、序列到序列转换
#   → Transformer 是自然选择
#
# ✓ 【大规模预训练】
#   有千万级以上的无标注数据
#   → Transformer 的规模化能力强
#
# ✓ 【多模态融合】
#   需要建模不同类型实体之间的关系
#   → 注意力天然适合
#
# ✓ 【需要长距离依赖】
#   如蛋白序列中远距离残基的关系
#
# ✗ 【小数据的性质预测】
#   几百到几千个样本
#   → 用 ECFP + 树模型 或 GNN
#
# ✗ 【只需要局部结构信息】
#   → GNN 更高效
#
# ✗ 【严格的对称性要求且需输出坐标】
#   → 专门的等变架构更可靠
#
# 【一个务实的观察】：
#   分子领域的 Transformer 成功案例，
#   几乎都在【生成】或【超大规模预训练】场景
#   而在「用几千个样本做性质预测」这个
#   最常见的实际任务上，
#   它并没有明显优势

结构编码：图 Transformer 的核心

# 纯注意力不知道分子的连接关系
# 必须显式注入 —— 这是图 Transformer 成败的关键

import torch
import numpy as np
from rdkit import Chem

def shortest_path_bias(smiles, max_dist=8):
    """计算原子间最短路径距离，用作注意力偏置"""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    # RDKit 直接给出拓扑距离矩阵
    dm = Chem.GetDistanceMatrix(mol)
    dm = np.clip(dm, 0, max_dist).astype(np.int64)
    return torch.tensor(dm)

def laplacian_pe(smiles, k=8):
    """拉普拉斯特征向量作为位置编码"""
    mol = Chem.MolFromSmiles(smiles)
    A = Chem.GetAdjacencyMatrix(mol).astype(float)
    deg = A.sum(axis=1)
    # 归一化拉普拉斯
    d_inv_sqrt = np.divide(1.0, np.sqrt(deg), where=deg > 0)
    L = np.eye(len(A)) - (d_inv_sqrt[:, None] * A * d_inv_sqrt[None, :])
    vals, vecs = np.linalg.eigh(L)
    # 取最小的 k 个非平凡特征向量
    pe = vecs[:, 1:k+1]
    # 【符号不确定性】：特征向量的符号是任意的
    #   → 训练时随机翻转符号做增强
    return torch.tensor(pe, dtype=torch.float)

# 在注意力中使用：
#   attn_score = Q @ K.T / sqrt(d) + bias[spd_matrix]
#   → spd_matrix 索引到可学习的偏置表
#
# 【这个偏置项让模型知道分子的拓扑结构】
# 没有它，图 Transformer 退化成「原子袋」模型

延伸资源

• ChemBERTa：145《ChemBERTa 论文精读》；MolFormer：147《MolFormer 论文精读》；Uni-Mol：142《Uni-Mol 论文精读》；
• GNN：159《图神经网络 GNN》；等变网络：160《等变神经网络》；多模态：164《多模态模型》。