最近在整理做分子几何表征项目的思路想把这个过程中的关键取舍和经验捋一遍。这个主题叫“不止生成迈向更好的分子几何表征”说白了就是我们不只是想让模型生成一个看起来合理的分子三维构象更重要的是让这些几何信息真正变成“表征”喂给下游性质预测、分子生成、虚拟筛选这些任务时能稳定地带来增益。如果你也在做AI制药、材料数据挖掘、构象分析这类事情那这篇东西应该能对得上你的胃口。很多人上手就奔向“生成”这个目标给定一个二维骨架去预测三维坐标。但项目推进到中后期你会发现生成质量只是地基真正的分水岭在于“几何表征”好不好。同样的坐标数据放进不同的编码结构用不同的预训练目标去学最后的任务性能可能差一大截。所以这篇文章我想从需求定义讲到方案选型再讲实操细节、评估方法和避坑经验争取让你读完能直接复制一套完整的流程去跑你自己的分子体系。1. 先搞清楚分子几何表征到底要解决什么问题1.1 从2D到3D为什么几何信息这么关键分子在二维层面只是一张拓扑图原子连成快照但现实中分子是三维的。同一个分子骨架可以扭转出无数种构象有些构象能量高有些能量低而生物活性和物理性质往往取决于低能量构象群甚至取决于分子在靶点口袋里的“变形能力”。这就是几何信息不能被忽略的直接原因一个简单的例子是手性分子互为镜像但空间结构完全不同化学性质却可以天差地别。如果没有三维坐标模型根本没法区分这种差异。我之前在处理一个药物小分子的性质预测任务时做过对比只用拓扑图特征模型的精度卡在一个平台期上不去后来引入3D构象和原子间距离张量误差直接降了一截。这就是几何表征的价值所在——它补全了拓扑图缺失的空间约束信息让模型有机会理解分子的实际形状、电子云排布、位阻效应。但这里必须警惕一个陷阱有坐标不等于有好表征。如果你只是把坐标堆成一个向量塞进MLP模型学到的其实是绝对坐标的某种记忆转动一下分子输出就变了。这不符合物理规律一个分子的能量和性质不该因你在电脑里把分子旋转90度而改变。所以几何表征的核心问题可以归纳成两句话怎么编码几何信息才能让模型尊重物理对称性怎么学习几何信息才能让下游任务真正受益而不只是头疼医头。1.2 几何表征的技术难点与核心约束分子几何表征有三个绕不开的硬约束我建议在做方案选型之前先背下来。第一个是等变性equivariance。想象你手里拿着一张椅子你把它原地转动一下椅子的外形、你该坐的位置都没变但每条腿的坐标全变了。分子的情况更苛刻合法操作不只是旋转还有平移。一个模型的预测目标如果是坐标、位移、速度向量这类量那它必须随输入一起旋转这叫SE(3)等变性。如果预测目标是标量性质比如能量、溶解度那模型必须做到旋转平移后输出不变这叫不变性。一个不具备等变性的几何表征模型在物理任务上总会有系统性偏差。第二个是连续空间的高维性。分子几何存在于一个连续的高维空间每个原子有三个空间自由度一个50个原子的分子就是150维空间。想在这个空间里做采样、做优化、做生成朴素方法会遇到维度灾难。这也是为什么纯靠坐标回归的模型往往抖得厉害因为高维连续空间里没有足够的数据密度来支撑稳定学习。第三个是构象分布的多样性。分子在室温下不是固定在一个构象里的它有一个Boltzmann分布不同构象之间只有微小的能量差异。生成任务里要求模型“找出一个合理构象”并不难难的是让模型理解整个构象集合的分布特征并把这种分布编码进表征里。好多项目在“生成”阶段玩得飞起一到下游任务就拉胯问题往往就出在这里模型没见过足够多样的几何变化学到的表征是“断断续续”的换一个构象就认不出来了。理解这三点之后你就能看懂为什么主流方案会往等变图神经网络和去噪预训练方向走这不是炫技是物理约束决定的。2. 方案拆解从构象生成到几何表征学习的完整链路2.1 第一环构象样本的质量决定表征天花板很多人觉得既然要学几何表征那输入坐标随便用RDKit生成一个就行。这个想法坦白说浪费了整个管线的一半潜力。构象样本的质量决定了模型能看到多少几何变化也决定了表征的天花板在哪儿。我一般把构象获取分成三个层次。第一个层次是快速力场法比如RDKit里面的ETKDG、MMFF优化速度和成本都很低一个分子几毫秒就能出一个构象。它的问题是只能覆盖低能量区域的一些局部极小值很多高能量或长程折叠的构象根本采不到。第二个层次是半经验或密度泛函级别的结构优化精度上去了但计算成本也上去了适合中小规模数据集不适合动不动几十万分子的场景。第三个层次是生成式方法比如基于扩散模型的GeoDiff、Torsional Diffusion这类方案它们直接从数据里学习扭转角分布出来的构象多样性和合理性都明显更好但训练和推理成本都不低。我的建议是如果你做的是大规模预训练第一层就可以但要做多样化的筛选如果你做的是小规模高精度表征尽量用带力场精修的第二层如果你有新分子骨架的专项任务第三层的生成模型能带来额外优势。我在实际项目里用的是“RDKit粗采样 能量排序 去重”的组合管线一个分子出20个候选构象按MMFF能量排序后保留低能量且RMSD差异明显的5个。这样可以兼顾质量和多样性模型在后面学到的几何变化集就比较完整。2.2 第二环选择等变表征模型而不是普通图网络处理好输入构象之后下一步就是选编码器。这个环节最常见的失误是把2D图神经网络原封不动搬到3D数据上只是额外加了距离特征。这么做不是完全不行但问题很多。普通GNN的节点更新是建立在“消息是标量”这个假设上的虽然距离是旋转不变的但你丢弃了方向信息模型其实是在用一个降维的、残缺的几何描述去做预测。更好的选择是等变图神经网络。当年处理这个项目时我在EGNN、PaiNN、GemNet这几个架构之间做了一番折腾最终主力用的是EGNN的等变消息传递框架。原因有三点第一它直接把坐标作为模型的一部分参与更新而不是当作额外的边特征几何信息没有被降维处理第二它在做坐标预测和力场学习的时候天然满足等变性不用额外加数据增强去“骗”模型学会旋转不变第三实现复杂度适中比GemNet那些动辄几层球谐函数张量积的架构好维护得多。等变网络的核心思想其实不复杂在每一层更新里既要更新原子特征标量又要更新坐标向量而且坐标更新的方式必须和旋转保持一致。你可以把这个过程理解成一个木偶戏人体标量特征在动线条向量坐标也跟着动但不管你怎么转舞台木偶和线条之间的相对关系是不变的。这么设计出来的模型它在物理上的归纳偏置是内生的不需要靠大量数据来领悟“旋转了分子标签还应该一样”这种底层规律。2.3 第三环预训练任务与下游适配架构选好之后最关键的灵魂问题是模型用什么目标去学几何表征这也是标题里“不止生成”的核心含义。如果你的项目只做构象生成那训练目标显然是最小化生成坐标和真实坐标的误差但如果要做更好的表征那生成本身只是辅助任务你更应该思考的是“如何让模型学到既能表达几何细节、又能泛化到不同任务目标的特征”。我在项目里用了三个训练阶段的组合。第一是无监督的几何去噪预训练做法很简单把构象坐标做微小随机扰动模型要预测出原始无噪声的坐标。这个任务逼着模型编码器理解分子内部坐标的可信度哪些原子位置是局部能量约束很强的哪些是松散可动的。第二是自监督的掩码原子预测随机遮盖一部分原子及其局部几何让模型从剩余原子环境去恢复被遮盖环境的几何类型。第三才是生成式预训练用扩散损失去学习构象空间的结构。实践证明前两个任务对下游性质预测的提升更直接生成任务更多是帮助模型在困难构象上建立泛化能力。下游适配时又要注意一件事纯拿预训练的embedding做聚类或线性探针并不总是最优。我在验证一个活性预测任务时发现冻结encder只训一个MLP头的效果反而不如在encoder后面加一个轻量的任务适配层用低学习率微调几轮来得好。原因是分子几何表征与化学性质的映射往往是非线性的线性探针会丢掉关系信息而全量微调又可能破坏预训练学到的几何规律所以折中路径是只微调最后两层加投影头。3. 实操手记构建一个分子几何表征项目的完整流程3.1 数据准备与几何特征化第一步是从数据源准备干净的分子。如果你用的是SMILES字符串建议先做标准化处理去盐、去重、选择主组分、质子化状态统一。这一步很多人偷懒结果后面全乱套。举个例子同一个分子的中性形态和酸式盐形态如果在数据里并存模型会被迫去学那些和物理性质无关的假变化。接下来就是加氢和构象生成。我用的是RDKit代码上大概是这样from rdkit import Chem from rdkit.Chem import AllChem def molecule_to_conformers(smiles, num_conf20, seed42): mol Chem.MolFromSmiles(smiles) if mol is None: return None mol Chem.AddHs(mol) params AllChem.ETKDGv3() params.randomSeed seed params.useRandomCoords True params.numThreads 0 cids AllChem.EmbedMultipleConfs(mol, numConfsnum_conf, paramsparams) results [] for cid in cids: try: converged AllChem.MMFFOptimizeMoleculeConfs(mol, maxIters200) # returns (cid, energy) except ValueError: continue conf mol.GetConformer(cid) positions conf.GetPositions() results.append((positions, converged[cid][1])) return mol, results这段代码里有两个容易被忽略的细节。第一个是AddHs一定要做不加氢的分子连键长信息都不完整很多机器学习的3D模型会把不饱和价电子状态搞混。第二个是MMFF优化那一行如果不做力场精修ETKDG直接吐出来的坐标在键长键角上会有明显的不自然感模型学了这种畸变几何后面去噪任务会不收敛。构象拿到后构建3D分子图。我建议以原子为节点、原子间空间距离为连边依据采用“半径截断前k近邻”混合策略默认截断半径取4.5埃再确保每个原子至少连到前8个近邻。边特征领域不要只用纯距离标量可以尝试加入单位方向向量的特征投影这在等变网络里是加分项。节点特征上除了原子序数、形式电荷、杂化类型我还会补充局部环境的平面性指标和手性标记这能帮助模型区分分子构型的细微差别。3.2 模型构建的关键实现细节等变图神经网络的实现里最容易写错的就是坐标更新分支。以EGNN为例一个等变层的更新步骤大致是这样import torch import torch.nn as nn class EGCLayer(nn.Module): def __init__(self, hidden_dim): super().__init__() self.phi_m nn.Sequential( nn.Linear(hidden_dim * 2 1, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim) ) self.phi_h nn.Sequential( nn.Linear(hidden_dim * 2, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim) ) self.phi_x nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1, biasFalse) ) def forward(self, h, x, edge_index): i, j edge_index d_ij (x[i] - x[j]).norm(dim-1, keepdimTrue) # (E, 1) d_ij torch.sqrt(torch.clamp(d_ij**2, min1e-6)) edge_input torch.cat([h[i], h[j], d_ij], dim-1) m_ij self.phi_m(edge_input) agg torch.zeros_like(h) agg.index_add_(0, i, m_ij) h_new self.phi_h(torch.cat([h, agg], dim-1)) coef self.phi_x(m_ij) / d_ij # 关键除以距离保证等变 delta_x torch.zeros_like(x) delta_x.index_add_(0, i, (x[i] - x[j]) * coef) x_new x delta_x return h_new, x_new这里最需要注意的细节是coef的计算。坐标更新里出现了(x_i - x_j) / d_ij这个表达式其实是一个单位方向的向量它本身是等变的。从旋转角度看只要坐标向量旋转这个差值向量也会跟着旋转从平移角度看坐标系原点移动时差值不变。所以这个设计天然保证了等变性。如果你只把坐标差值直接乘一个标量不除以距离那么更新幅度会严重依赖分子尺度的大小训练起来会特别不稳定。另外索引归约这里用index_add_而不是手动循环不然E条边的数量一上来内存和时间都扛不住。我在项目里用的隐藏维度和层数是版本化的底线上保证每层都有残差连接和层归一化防止等变更新的梯度在深层消失。3.3 训练策略、损失函数与超参数选择几何表征模型的训练目标可以多任务组合但损失之间的配比和权重需要谨慎。我用的是以下加权损失L w_noise * L_noise w_contrast * L_contrast w_gen * L_genL_noise是坐标去噪损失用预测噪声和真实噪声之间的均方误差L_contrast是对比损失把同一个分子不同构象的embedding拉近同时把不同分子的embedding推开L_gen则是扩散生成损失让模型能重构训练集中的构象分布。权重上我实际取w_noise1.0, w_contrast0.5, w_gen0.1。生成损失权重不敢给大因为一旦它主导模型会退化成纯生成模型特征空间里全是构象重建细节反而对下游任务不利。超参数方面我整理了一份通用参考配置你可以拿去做起点参数推荐值备注原子隐藏维度128分子量不太大时够用太大会导致显存爆炸消息传递层数5层数太浅学不到长程几何约束学习率1e-3用AdamW cosine schedule批量大小64和构象采样数量有关每个样本多构象时可以减半梯度裁剪10.0防止坐标更新偶发剧烈震荡坐标扰动噪声sigma0.05 Å去噪预训练专用太大变成模糊随机旋转增强on每个epoch对坐标做随机SE(3)旋转训练过程中有一个经验可以分享等变网络在用混合精度训练时坐标更新那个分支特别容易发鬼火。fp16下的坐标差值计算动态范围比标量特征大得多动不动就溢出。我的方案是坐标分支保持在fp32计算只对特征矩阵用混合精度实现对PyTorch来说就是把等变层的坐标更新部分放到torch.autocast(dtypetorch.float32)作用域里其他部分保持默认。这样既不损失太多训练速度又稳定得多。4. 质量评估如何判断几何表征“好”还是“不好”4.1 下游任务基准与定量结论判断几何表征质量我坚持只看下游任务上的实际提升不看预训练loss降到多低。跑完预训练之后固定encoder参数用线性探针和微调两种方式分别在公开数据集上做定量评测。常用的标准做法是在QM9数据上做分子性质预测比如极化率、HOMO-LUMO能隙、零点能这些目标。正常配置下一个5层等变编码器加上去噪预训练在HOMO-LUMO这种复杂性质上的MAE能降到0.05 eV量级甚至更低而不做几何预训练的基线模型通常要到0.08 eV以上。再一个是MD17这种分子动力学数据集能量预测MAE可以控制在0.5 kcal/mol以内力预测的MAE也有可观下降。这不是说我拿到的绝对数值有多厉害重点在于同一个实验条件下几何表征是否稳定地比拓扑基线更准。我还建议加一个等变性检验随机旋转一个测试集分子然后比较模型输出和原始输出。对坐标预测任务输出应该完美跟着旋转走对标量性质任务输出应该完全不变。如果模型对这个测试敏感说明你的等变模块里有泄漏或规范化不均匀得回头检查图构造阶段是否有非等变的特征混进去了。4.2 消融实验不是堆模型就有用很多人工程会上来就把生成、对比、去噪三个损失全堆上模型参数量也一路增加最后表现却不一定最好。你需要做消融实验来精确判定每个模块的真实贡献。我做的消融排序大概是这样的同时去掉对比和去噪只留下生成损失做训练下游任务的预测误差上升最明显只保留去噪对比去掉效果次之只有对比没有去噪效果垫底但也比完全没有几何预训练强。这说明在几何表征学习中坐标去噪任务贡献了最大比例的信息对比学习负责稳定特征空间结构而生成损失更像是一个正则或数据增广的角色。但是要注意消融结论和数据集规模强相关。在大规模非标注构象库上生成任务能更好地帮助模型建立构象流形在小规模精细数据集上对比和去噪反而更高效。我建议你们做项目时不要抄固定结论花两天时间跑消融矩阵每一列的改动都记录下游任务指标这比闭着眼睛堆模块有价值得多。提到稳定性我还会测一个东西表征对构象噪声的鲁棒性。把验证集的坐标加上0.2埃的随机扰动看看下游预测指标的波动幅度。一个好的几何表征应该具备天然的平滑性小扰动不改变预测结果如果发现波动很大说明模型有一部分在死记坐标位置没有真正学到几何常量这通常需要增强去噪任务权重或换更严格的构象预处理策略。5. 常见问题与排查技巧实录5.1 构象生成质量差导致训练崩溃这是我遇到最频繁的问题没有之一。现象是训练loss一直在下降但下游任务验证集误差高得离谱。排查下来八成是训练输入的构象空间里混杂了大量高能量不合理构象模型的去噪任务无法区分“合理的构象变化”和“物理上不可能的畸变”。解决办法分两条线走。一条线是回数据源头把距离窗口内的镜像故障构象过滤掉比如某些原子对的距离小于共价半径之和就该直接删除该样本。另一条线是在训练时给构象加一个RMSD分布筛选同一分子保留的多个构象互相之间的RMSD不能太小也不能太大太小了模型见不到多样性太大了说明中间肯定有坏构象。我用的是RMSD在0.5到2.0埃之间的过滤窗口实测效果很好。5.2 等变网络不收敛或发散等变网络训练发散时容易让人怀疑“等变结构是不是不行”。实际上很多时候只是细节处理有问题。先检查坐标中心化状况如果输入分子的质心不在原点模型的更新里会多出一部分不需要学习的平移成分导致隐藏维度里有一路神经元专门在补平移动态训练起来特别不稳。我都习惯在做预处理时就把所有构象对齐到质心为零。另一个是梯度裁剪加上去抖动技巧在等变更新里coef那个分支的梯度经常会产生瞬间的大数值因为当两个原子距离接近零时除以距离就会把更新幅度放大到离谱。所以我在距离计算里加了最小值钳制clamp(min1e-6)同时所有坐标分支的梯度只允许在[-10, 10]范围内流动。如果你在训练日志里看到loss曲线每隔几十步就跳一下十有八九是这里的问题。5.3 数据泄漏与评估陷阱做分子几何表征时数据泄漏比你想的更隐蔽。最常见的泄漏是把同一个分子的多个构象同时放在训练集和验证集里。虽然构象不同但它们来自同一分子骨架模型只要记住分子身份就能在验证集上“作弊”下游指标看起来很漂亮换到新分子上立刻原型毕露。正确做法是按分子骨架拆分子集确保验证集里的骨架在训练集中一次都没出现过。另外还有一个隐患是构象生成的随机种子泄漏。如果训练集和验证集的分子都是同一批RDKit种子生成的构象模型可能在隐空间里学到种子的bias。我后期调整方案把随机种子也要和数据拆分绑定每个分子拆分位置重新生成随机种子。这个坑很冷门但排查真会浪费一整周。5.4 常见问题速查表现象可能原因处理思路预训练loss很低下游精度差预训练任务和下游任务目标错位调整任务权重增加去噪比例等变性测试不过图构建阶段混入绝对坐标特征检查边特征保证只用距离和相对方向同一个分子不同构象embedding距离过大对比损失权重过低提高对比权重降低生成权重显存很快打满边数量爆炸缩小截断半径或k近邻数考虑子图采样坐标更新输出NaN距离接近零导致除零检查距离clamp降低学习率6. 一些我在实操中的个人体会搞完这一整套流程之后我的一个明显体会是分子几何表征的真正价值不是做出一个能生成漂亮构象的模型而是让几何信息以符合物理规律的方式进入模型的语言系统。等变网络解决的问题就是“坐标系变了语义不散架”这件事一定不要靠数据增广硬学要放在模型结构里。具体操作层面我最后还想分享一个很有用的细节在预训练阶段对输入构象做随机旋转加微小抖动这个做法被很多人以为是多此一举但实测下来它对于训练稳定性尤其是坐标更新分支的帮助非常大。原因在于即使等变结构理论上是旋转等变的数值计算里的浮点误差在大的旋转角度下还是会累积出差异数据增强等于给模型打了一针疫苗让它在很宽的输入范围内都保持一致的输出行为。另外有一个值得往下走的方向如果把几何表征和力场物理约束直接嵌在一个损失里比如要求模型预测的隐空间坐标变化方向接近分子动力学模拟的低频振动模式这类混合方案未来可能会比纯数据驱动的几何表征更接近“理解分子”。我还没完全跑通但在小规模体系上已经有了一些迹象有兴趣的人可以往这个方向试试看。这个项目的下一步我打算把几何表征模型和构象生成器融合成一个Encoder-Decoder框架共用一套等变encoder但把decoder从扩散模型换成更轻量化的规范流模型这样既能保持生成的多样性也能确保encoder学到的表征不会在反向传播中被生成损失带偏。路还长但每一步踩实了后面的成果就水到渠成。
