ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

FraGAT+:基于分子片段的多尺度图注意力机制提升分子性质预测性能

FraGAT+:基于分子片段的多尺度图注意力机制提升分子性质预测性能 开头做分子性质预测的这几年我最大的感受是分子表征方式直接决定了模型性能的天花板。早期大家用Morgan指纹配合机器学习算法虽然简单有效但指纹是离散的、固定的本质上是一种词袋模型后来图神经网络GNN流行起来把分子当成原子节点和化学键边组成的图确实把结构信息用起来了但大多数GNN模型都停留在原子级别——每一个原子作为一个节点消息传递在原子之间进行。这个粒度够吗对于很多性质来说不够。尤其是当性质取决于某个官能团、某个环系的整体行为时模型不得不在多跳消息传递中反复堆叠层数才能看到更大的结构单元不仅训练更慢还容易出现过度平滑。FraGAT是发表在Bioinformatics 2021上的一个工作全称是基于分子原始片段的多尺度图注意力机制的分子性质预测模型。它最核心的思想是不把分子当成一张原子级拓扑图而是先把分子切分成有化学语义的片段fragment再在片段之间做图注意力。这个多尺度体现在两个层面——原子级和片段级同时建模信息在两个粒度之间交互。这个思路本身并不复杂但它在分子性质预测任务上确实带来了显著的性能提升。这篇文章我会尽可能完整地拆解FraGAT的设计动机、构图策略、注意力机制、信息融合方式也会聊聊我在类似任务上复现和迁移使用时的经验与坑。如果你正在做药物分子性质预测、ADMET筛选、分子优化这类工作并且觉得原子级GNN已经到头了这篇内容应该足够给你新的思路。1. 分子表征的进化脉络为什么原子级GNN还有短板要理解FraGAT为什么值得关注得先把分子表征这条路从头捋一遍。不是浪费时间因为很多人在模型跑不动的时候第一反应是换网络结构、调超参但实际上问题可能出在表征粒度这个根上。1.1 从指纹到图解决了有没有的问题传统的分子指纹ECFP/Morgan指纹把分子的局部结构哈希成固定长度的二进制向量。它的优点是计算快、可解释性强、和树模型配合得很好PySpark、XGBoost、随机森林那一套都能直接用。缺点是指纹是预定义的、不可学习的它无法针对特定任务自动提取最优的化学信息。图神经网络解决了一个根本性的问题让分子表征变成可学习的。分子天然就是一张图原子是节点、化学键是边。GNN通过消息传递和聚合操作让每个原子节点的隐向量逐步融合邻居信息。理论上堆叠k层GNN中心原子就能感知k跳范围内的结构。1.2 原子级GNN的结构盲区但是仅仅把分子看成原子图其实丢掉了化学中非常重要的尺度概念。举一个很简单的例子要判断一个分子是否具有某个毒性警报比如PAINS警报你真正关心的并不是单个原子而是某几个原子连接在一起形成的一个特定的亚结构——比如β-内酰胺环、Michael受体这类基团。在原子级GNN中模型要捕捉这样的片段级信息必须依赖层数堆叠第一层看到相邻原子第二层看到两跳半径的亚结构第三层才能看到更大的环系。问题在于层数越深信息越模糊。图神经网络存在过度平滑问题堆到三四层之后所有节点表征趋同区分度反而下降。计算开销大。深层的消息传递对显存和时间都不友好。化学语义丢失。分子中的环、官能团、手性中心这些天然具有明确化学含义的结构在原子级图里被拆得七零八落模型需要自己悟出它们的存在。我开始做分子性质预测时也踩过这个坑。当时在一个hERG抑制活性数据集上跑了GAT发现不管怎么调dropout和学习率AUC都卡在0.78左右。后来分析了一下数据里的活性分子发现它们几乎都含有带正电荷的叔胺基团——但模型的注意力权重并没有集中在这些基团上说明它根本没看到这个尺度上的结构模式。1.3 片段化把化学知识重新注入模型片段化fragmentation就是把分子按照一定规则拆成具有化学意义的子结构。常见的拆分方式有拆分方式代表性工具特点逆合成规则BRICS按可逆合成断键位置拆分片段更接近合成子官能团识别RDKit官能团库按已知官能团边界拆分保留化学语义环侧链自定义把环单独拎出来连接的侧链作为独立片段数据驱动图分解算法无监督地从数据集中学拆分位置FraGAT使用的是基于原始片段的拆分方案拆出来的片段保留了片段内部的原子连接关系因而叫原始并不是把一个环完全拍平成超节点。这个设计非常关键——它让后续的编码器依然能利用片段内的细粒度结构信息而不是像有些片段级模型那样把整个片段压缩成一个节点的向量丢失内部细节。从整体来看片段化解决的不只是信息粒度不够的问题更是一种把化学先验知识引入模型架构的方式。你不需要让模型从零开始学会一个苯环是什么你直接告诉它这个苯环是一个整体。模型省下来的学习成本自然就花在了更重要的地方——理解片段之间的交互关系上。2. 片段从哪里来构图策略决定信息质量FraGAT里第一个值得展开的模块就是分子的构图策略。不同模型对图的定义完全不同而图的定义其实预设了模型能够捕捉什么样的结构信息。2.1 双层级图结构节点可以是分子内部的所有原子也可以是片段FraGAT的图结构设计是分层的这也是它多尺度的第一层含义。具体来说原子级子图内部图每个片段内部原子和原子之间的连接关系完整保留。片段内部的原子不是孤立节点而是由化学键连成的连通子图。这个子图上的特征和连接方式都会输入到编码器中进行消息传递。片段级图片段的连接拓扑把片段看作超节点片段之间通过原先的化学键拆分点处断裂的键连接构成一个更高层次的图。这个图上的节点数远小于原子数计算成本低得多但保留了分子骨架的粗粒度拓扑。举个例子一个萘环如果按某种规则拆成两个苯环片段那么内部子图就是两个苯环各自的六元碳环结构而片段级图上这两个片段之间有一条边代表它们共用一个键或者原先是稠合关系。2.2 为什么要原始片段而不是简化片段我看到很多方法喜欢把片段直接抽象成一个节点这样构建的分子图非常简洁几十个片段就能描述一个大分子。但FraGAT刻意保留了片段内部的原子结构原因在于片段内部的细节往往决定性质。一个官能团是羧酸还是酯就差在一个原子上但两者对溶解度、渗透性、代谢稳定性的影响可能完全不同。如果简化成节点这种区分度就丢了。注意力机制需要原始信息作为输入。片段级图上做注意力聚合时片段节点的初始特征来自片段内部原子子图的编码结果。如果原子信息已经在片段化时被丢弃后续怎么聚合都是无米之炊。用预构建图可学习编码的组合还带来一个实际好处整个模型是端到端可训练的。片段化本身虽然是规则定义的不可学习但片段内部特征编码、片段间聚合这些环节都是可学习的任务信号可以通过梯度反传直接影响特征提取器的参数。2.3 片段化工具选型的实践建议在实际落地时最常用的是RDKit的BRICS拆分以及基于RECAP规则的拆分。我用过两者的区别很明显BRICS拆分出来的片段一般偏大会保留较完整的功能团和环系适合做骨架分析RECAP按照组合化学的逆向合成逻辑拆片段更小、更模块化比较适合虚拟组合库设计。FraGAT这种模型训练场景下我没有看到论文对某个具体拆分工具做死板限定实践上选择哪套规则要根据你的分子集合特征来定如果分子普遍很大比如超过50个重原子片断化应该更激进反之小分子库最好采用保守的拆分策略避免拆出来的片段就是单个原子退化成原子级图。3. 多尺度注意力机制两个级别的信息如何协同工作光有图结构还不够。FraGAT真正有意思的部分是它的注意力机制设计。多尺度注意力不是说我有两个注意力模块这么简单关键是这两个模块分别在什么粒度上提取信息以及它们之间如何互相参考、互相增强。3.1 片段内部原子级注意力负责提取局域化学特征片段内部的子图规模通常不大几个到十几个原子在这里做GAT层消息传递计算量几乎可以忽略。原子级注意力的作用是把片段内部的电子效应、空间效应、杂原子位置等微观信息编码成片段的初始嵌入。打个比方你在看一张照片时第一步不是立刻判断整张照片的内容而是先看清照片里每个局部区域这是一棵树那是一辆车。原子级编码做的就是这件事——先把每个片段内部看清楚。对于每个原子节点i它的邻居节点j对它的贡献权重通过注意力系数计算$$\alpha_{ij} \frac{\exp(\text{LeakyReLU}(a^T [W h_i | W h_j]))}{\sum_{k \in N(i)} \exp(\text{LeakyReLU}(a^T [W h_i | W h_k]))}$$然后节点表征更新为邻居消息按注意力权重的加权求和。这一步和标准的GAT完全一致区别只在于图是从分子拆分得到的片段内部子图而不是整分子的原子连接图。因为子图小这里可以用多头注意力来增强表达力不用太担心显存。3.2 片段之间片段级注意力捕捉全局功能协同片段级图上的注意力是整个模型的灵魂。在这个尺度上每个节点代表一个片段节点特征是片段内部编码器输出的汇总向量。片段之间通过注意力机制互相交流学习哪些片段的组合对目标性质真正起作用。这个设计的化学直觉是分子性质往往取决于多个官能团的协同作用。一个分子同时含有氢键供体和氢键受体它可能很容易穿透血脑屏障一个分子同时含有亲脂芳环和可旋转键数目多它可能会表现出较高的代谢清除率。这些都不是某个片段单独决定的而是片段之间的相互作用。片段级注意力的实现方式和原子级类似但参与权重计算的节点更少、更抽象。由于节点数是分子被拆分后的片段数通常只有几个到二十几个计算非常轻量这也让整个模型的推理速度比深层的原子级GNN更快。3.3 多尺度的融合方式从局部到整体再反向修正这里要特别说一下多尺度不是简单地把原子级特征和片段级特征拼在一起。FraGAT的设计是分层递进的先完成片段内部编码再把片段表征放到片段级图上做信息传递得到更新的片段表征最后在全局readout阶段聚合得到分子级表征。我个人的理解是这相当于先看局部再看整体得到一个考虑了上下文信息的片段表征。打个比方你先看到一个人手上的戒指片段内部特征然后再看到他在婚礼现场片段间上下文你这时对戒指这个片段的理解就跟单独看戒指完全不一样了。片段在分子整体环境中的语义角色正是通过片段级注意力来调整的。这种设计还带来一个额外的好处经过片段级注意力更新后的片段表征本身已经包含了整分子的上下文信息因此可以作为分子的局部结构描述子输出用于分析哪些片段对预测结果贡献最大。这在药物化学的骨架跃迁、片段生长等场景里特别有用——你可以直接可视化注意力的权重分布让模型不再是一个黑盒。4. 从片段到分子Readout层怎么把散落的信息拧成一股绳前面所有步骤产出的都是片段级表征序列接下来最关键的一步是Readout——将一组不定长的片段表征聚合成一个固定维度的分子级向量供下游MLP做性质预测。这一层看似简单但在无数模型里怎么聚合决定了最终预测性能的上限。4.1 注意力池化每个片段对性质的贡献是不同的简单的池化方式——Sum池化、Mean池化——虽然计算简单但都有一个隐藏假设所有片段对分子的最终性质同等重要。这在真实化学场景中几乎不成立。某个分子对hERG的抑制活性可能主要由其芳环骨架上的疏水片段贡献而旁边的亲水侧链片段几乎不参与。如果使用Mean池化占多数的非关键片段会把关键片段的信号稀释掉Sum池化则会让分子大小成为主导因素。FraGAT使用注意力池化作为Readout层让模型为每个片段学一个重要性权重。形式上可以理解为$$z_molecule \sum_{k1}^{K} w_k \cdot h_k$$其中$w_k$是第k个片段的注意力池化权重由片段特征经过一个小型映射网络归一化得到。权重值直接反映了该片段对当前预测任务的贡献强度。在实际实现中这个权重还可以与分子规模做归一化避免大分子因为片段多天然占优。4.2 多层特征融合除了最终层中间层特征也别浪费一个容易被人忽略的细节是模型经过多轮信息传递后中间层产生的特征其实也包含不同感受野的信息。FraGAT这类模型通常会在多个层级收集特征做拼接或加权融合作为最终的分子表征。这种方式类似计算机视觉里的FPN特征金字塔网络区别是读出的特征来自不同的图层级原子、片段、分子而不是不同的图像分辨率。我调模型的经验是特征融合策略不需要太复杂一般把片段内部编码更新前后的特征拼接再把片段级注意力更新前后的特征也拼进去效果就比只取最后一层好很多。原因在于更新前的特征保留了局部原始信息更新后的特征携带了全局上下文两者是互补的单独看哪一个都有信息损失。4.3 任务细节回归还是分类Readout之后的分支设计Readout之后连的是什么头Head取决于任务类型。回归任务如logP、溶解度通常接一个单输出的MLP损失函数用MSE或Smooth L1。分类任务如毒性二分类、血脑屏障渗透性二分类则在输出层用sigmoid配合BCE loss。多任务学习也是一个实际中经常碰到的需求——比如同时预测一个分子的多种ADMET属性。多任务场景下我建议Readout输出的共享分子表征保持不变每个任务单独接一组MLP头。共享表征让不同任务之间互相正则化提高泛化性独立的任务头避免任务之间的梯度冲突过于严重。FraGAT的骨干架构完全支持这种改装你只需要在接口处做一点调整即可。5. 实验分析与效果对比FraGAT凭什么比基线好评估一个模型不能只看它自己的指标得放在同一批基线模型、同一批数据集上看相对提升。这个部分结合我在类似任务上的实验结果一并说明。5.1 常用基准数据集与评估指标分子性质预测的常用公开数据集包括数据集任务类型规模难点ESOL水溶性回归~1100个分子化学多样性大需良好结构表征Lipophilicity亲脂性回归~4200个分子需要区分异构体BACEβ-分泌酶抑制剂分类~1500个分子活性类似物多需捕捉细微差异BBBP血脑屏障通透性分类~2000个分子正负样本不均衡Tox21多任务毒性分类~8000个分子多任务、标签稀疏hERG心脏毒性回归/分类视版本而定工业界极关注模型易过拟合评估指标回归任务通常看RMSE和R²分类任务看ROC-AUC和PR-AUC。多任务分类还要关注每个任务单独的结果不能只看平均。5.2 和原子级GNN的对比结果差异FraGAT的论文和后续复现工作里一个反复出现的现象是在回归任务ESOL、Lipophilicity上的提升比分类任务BBBP、BACE更大。我自己的推测是回归任务对分子表征的数值精度要求更高片段级表征保留了明确的化学语义预测出的数值更稳定而分类任务往往存在一个明显的决定性亚结构原子级GNN经过多层注意力也能找出那个结构所以提升空间小一些。在Tox21这种多任务稀疏标签数据集上FraGAT的优势也很明显——多个毒性终点同时预测时共享的片段级表征能把不同任务的有效特征整合起来比每个任务单独学一套原子级特征更省数据、更抗过拟合。5.3 与指纹基线方法对比这一点对于工业应用特别重要。很多人觉得深度学习模型一定全面碾压指纹GBM的组合实测并非如此。指纹方法在小数据集、数据分布简单的情况下依然非常能打而且可解释性极强——你可以直接定位哪个子结构导致预测出毒性。FraGAT这类片段级GNN在中小数据集上的提升往往只有3~5个百分点的AUC但如果数据量充足、分子结构多样这个差距会进一步扩大。我的建议是如果你的数据量不足2000条先用指纹SVM/GBM做个快速baseline别直接上GNN数据量超过5000条并且结构多样性高这时候用FraGAT这类片段级模型才划算。6. 复现与落地过程中的几个坑最后聊聊我在实现和迁移这个模型时踩过的坑。FraGAT的公开细节还算充分但真正把整个流程跑通、跑准没你想的那么顺利。6.1 片段化工具选择与分子集合特征的匹配片段的粒度选择对模型性能影响非常大。我在一个类药小分子库上测试过不同BRICS拆分阈值阈值设太松常见的芳香环会被拆碎成单个原子分子图的片段粒度退化成原子级别多尺度优势完全消失阈值设太紧大分子只拆出两三块片段图信息量不足模型基本退化成分子指纹方法。比较实用的调节思路是先统计你的数据集分子中可旋转键数量的分布按中位数来设置拆分时允许断键的类型和位置。如果你的分子普遍柔性较大可旋转键多可以大方地使用BRICS按可逆合成规则断开如果分子是刚性稠环居多我不知道就应该保守处理只在末端侧链处断开。6.2 原子初始特征与片段边界条件的构造原子级特征通常包括原子类型、度、形式电荷、杂化方式、芳香性、氢原子数、手性类型等。这些特征用RDKit的GetSymbols、GetDegree、GetFormalCharge、GetHybridization、GetIsAromatic、GetTotalNumHs等函数就可以直接取到。但要注意的是在片段边界处被切断的化学键涉及的原子需要做封端处理。如果不封端片段边缘碳原子会处于悬空的非正常价态RDKit在计算某些描述符时会报警甚至报错而且靠近断键的原子特征也会失真。我的做法是在所有断裂处补一个[ * ]虚拟原子它不参与消息传递只用于满足原子的化合价要求。这个虚拟原子不会进入片段内部子图的节点集合只影响特征构造过程。6.3 数值稳定性注意力权重退化问题如果你直接在片段级图上跑GAT一个容易遇到的问题是多轮更新后注意力权重分布变得过于平滑——每个片段对每个片段的重要性都差不多等于是告诉模型所有的片段同样重要这实际上丢了注意力机制的全部意义。我在实验中加了一个温度系数和一个熵正则项来对抗这个退化。温度系数让注意力分布更尖锐熵正则项设置一个下界来防止单个片段一家独大。实测在这两个技巧的配合下模型的片段注意力分布能从摊大饼变成有突出峰不仅在uc精度上有微小提升更重要的是可视化时能合理找到关键片段。6.4 显存与训练速度的调优建议片段级模型的显著优势是计算量比深层原子级GNN小得多。训练的时候一个标准的FraGAT比如6层注意力编码在我的单张4090上batch size开到512都没有压力。相比之下同样效果的原子级GNN可能需要显存翻倍以上。不过有个容易忽略的瓶颈在DataLoader这边分子片段化过程如果放在训练循环里实时执行会成为明显的性能瓶颈。我的建议是预处理阶段就做好片段切分和索引构建存成分子图对象或对应的稀疏矩阵格式训练时直接加载不在线计算。药企实际落地时面对百万级分子的库用这种离线预处理的策略可以把整个训练时间缩短一个数量级。最后分享一个我实际用FraGAT思路做项目的体会。当时我们参加一个内部ADMET预测竞赛大部分队伍都在用原子级GNN或者Transformer。我把市面上顶级的指纹方法、图方法、以及基于片段的预训练Transformer都跑了一遍baseline最后在片段级图注意力架构上做改进加入多任务头和注意力正则拿到了总排名前5的成绩。最重要的不是排名而是这个架构让我第一次清晰地理解了哪个片段决定了这个分子的哪个属性——这种可解释性带来的业务价值是单纯堆精度的人很难体会的。如果你也想在自己的项目里引入这个思路建议从最简单的版本开始一个BRICS拆分器、一个内部原子GAT、一个片段级GAT、一个注意力池化不到300行PyTorch代码就能跑通基线。先把框架跑通再根据你的特定任务打磨细节这条路比直接追求复现完整论文要快得多。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进