项目标题自监督学习搞定药物分子研发快一倍1. 先聊透为什么药物研发会卡在“数据不够”这个坎上做AI制药相关项目的人多半都有过这种经历模型结构选好了环境配好了跑起来也顺畅结果一上真实实验数据效果直接“见光死”。问题往往不落在算法本身而落在数据上——药物分子领域的标注数据太少、太贵了。1.1 一个靶点的活性数据可能只有几百条一个典型的药物分子研发流程里针对某个靶点的已知活性分子数据公开可用的往往只有几百到几千条。这点数据喂给常规的深度学习模型尤其是图神经网络、Transformer这类参数动辄百万级的结构几乎是“杯水车薪”。你让模型从零开始学分子结构里的隐藏规律它只能记住这几百个分子的“长相”根本谈不上泛化。而药物研发恰恰最看重泛化——你得预测那些从未合成过、但理论上可能与靶点结合的分子活性。我最初做这个项目时也试着硬碰硬用分子指纹做输入接一个随机森林效果能到0.6左右的AUC换成图神经网络反而是0.55不升反降。原因很简单数据量撑不起模型复杂度。后来把方案换成自监督预训练同样的图神经网络结构、同样的下游数据AUC涨到0.72。这个差距对虚拟筛选来说可能就意味着同一批候选分子里多筛出一倍以上的真阳性。1.2 自监督学习解决的核心矛盾自监督学习的思路说白了就是“没有标签自己造标签”。对于分子数据来说无标注的分子结构信息其实是海量的——公开的化合物库里有上亿个分子结构它们不像生物活性数据那样需要做湿实验才能获得。我们只需要分子本身的结构表达式就能构造出大量训练样本。这样一来整个研发流程被拆成两段先在海量无标注分子上做预训练让模型学会“分子结构”本身的语言和规律再在有标注的少量活性数据上做微调。预训练阶段花掉一个月微调阶段只需要一两天但最终模型在小样本任务上的表现比直接在有标注数据上训练高一截。整条研发链路的时间成本确实可以接近砍半。1.3 “快一倍”的计算逻辑这里补一个具体的算术账。常规做法下一个项目团队花3个月做实验、1个月做数据清洗和建模最终拿到可用模型。自监督方案里预训练虽然要花时间但完全可以沉淀下来复用——同一个预训练模型可以服务后续几十个靶点项目。平摊到每个新靶点上增量成本就只剩微调和评估整体时间能压缩到原来的50%甚至更低。这就是“快一倍”的实际来源。2. 分子表征方式与预训练任务的搭配直接决定效果上限很多做深度学习的人转到分子领域时最先问的问题就是分子到底该用什么形式喂给模型答案没有唯一解但不同选择对应的预训练任务是完全不同的。2.1 三种主流分子表征各自适配的模型结构先搞清楚基本盘。分子表征主要有三种主流形式序列字符串SMILES、图结构原子为节点、化学键为边、三维构象原子坐标加空间关系。它们的关系可以简单类比SMILES是“读一行化学式”图结构是“看一张关系图”三维构象是“摸一个立体模型”。在选择上我建议项目早期直接锚定SMILES和图结构这两个方向三维构象放在第二步。原因有两点一是数据可得性SMILES和图结构都能直接从公开库中批量获取三维构象需要额外计算成本翻倍二是模型生态这两个方向的预训练方法最成熟社区可复用的经验最多。2.2 预训练任务设计的三个坑位确定表征后预训练任务的选择也有讲究。分子领域常用的自监督任务主要有三类掩码重建Masked Reconstruction、对比学习Contrastive Learning和属性预测Property Prediction。掩码重建的思路很直观——把SMILES序列里部分原子符号遮住让模型根据上下文预测被遮住的内容。这模仿的就是语言模型里完形填空的逻辑。分子SMILES里原子的排布规律确实有内在语法模型在预测被遮掩的原子时必须学会原子间的依赖关系。对比学习这里要特别注意分子领域的正样本对构造远不像图像领域那么顺手。图像可以随机裁剪、旋转分子不能随便改——你改动一个原子分子可能就不是同一个分子了。常见做法是同一个分子的不同SMILES枚举形式同一个三维结构可以有多种合法的SMILES写法作为正样本对。这个方案我强烈建议试试实现简单、效果稳定。属性预测作为辅助任务也很有价值。比如把分子的水溶性、脂溶性等比较容易通过计算获得的属性作为自监督信号之一多任务同时训练这在业务上完全可行因为这些属性可以通过快速计算接口批量获得成本远低于活性实验。2.3 我最终采用的组合方案综合对比下来我最终跑通的方案是SMILES序列输入 Transformer双向编码器 掩码重建为骨干任务配合分子图视角做辅助对比学习头。图文双塔结构的好处是梯度信号更丰富——序列信息里有原子顺序关系图信息里有拓扑结构关系两者互补。多嘴提一句模型里的注意力层数不是越多越好。我试过6层和12层的配置在预训练阶段12层的loss更漂亮但微调阶段在几百条小样本数据上反而表现下滑。12层模型在微调时更容易过拟合最终选择6层参数量大约在同配置语言模型的四分之一性价比最高。3. 从零搭一套分子自监督预训练管线的全过程想直接复现这条管线的朋友接下来这部分是我最希望你细看的。整个流程分四段数据准备、预训练、微调、评估。每一段都有值得记下来的细节。3.1 数据准备清洗比想象中更消耗时间先说数据来源公开的分子库可以直接用但拿回来千万别直接喂给模型。我踩过的坑是数据里有相当比例的无效SMILES、盐类复合物、异常原子价态这些脏数据被掩码后模型根本预测不出来loss常年居高不下。数据清洗的几个基本规则用标准化工具统一分子表示比如去掉盐类部分、保留中性母体分子去除含有罕见元素的分子例如硒、碘等这类分子在公开库里占比极低模型见过几次也不足支撑学习同一个分子只保留一个SMILES形式避免预训练阶段同一结构的过拟合风险剔除分子量过小少于10个重原子和过大超过50个重原子的分子。清洗后我一般会再做一次数据增强即对每个分子枚举出多个SMILES变体。这一步能有效扩充样本多样性预训练阶段模型能见到同一个分子的不同视角对比学习任务因此自动获得正样本对。3.2 预训练实操关键参数和经验值预训练本身其实是标准的Transformer预训练流程但有几个分子领域特有的参数值得记录。Token字典分子的SMILES本质上是一串字符建议按字符级别做分词。原子符号用单字符C、N、O等括号、数字、环标记这类结构符号也作为独立Token。控制字典规模在256以内比很多语言模型的词表都小Embedding层参数占比不大好处是容易收敛。序列长度SMILES序列的平均长度在50-80字符之间超过128的属于极少数。我统一截断到128远超99%的分子。序列长度设太短会损失长链分子信息设太长会浪费训练时间这个阈值是实践折中出来的。掩码比例常规语言模型掩码一般取15%分子SMILES上我建议更高25%左右。原因不复杂SMILES字符串的冗余度远低于自然语言掩码比例太低预测任务就过于简单模型学不到深刻的结构规律。但也不建议超过35%否则上下文信息太少任务变得过难训练不稳定。训练时长和策略预训练数据量在1亿分子级时学习率建议用schedule先线性warmup到1e-4再逐步衰减。大约训练到每个分子过2-3个epoch时loss就会进入平台期。我在训练过程中发现每5万步保存一次checkpoint然后根据下游任务小样本上的表现选择最优checkpoint这个原则非常管用——预训练loss最低的点不一定对应下游任务最优的点。3.3 下游微调怎么把预训练模型迁移到具体靶点任务微调阶段是真正考验工程细节的地方。输入数据是靶点特定的一批已知活性分子正样本和非活性分子负样本规模常常只有几百到几千条。我采用的微调方案是在预训练模型顶部加一个多层感知机头输入层用模型的平均池化特征或[CLS]位输出。训练时冻结模型主干只微调顶层模块。当模型收敛后再解锁全部参数做小幅微调这在实践中比直接全量微调更稳定。关键参数给一个参考下游训练轮次控制在30-50个epoch批大小为16-32。学习率设置上冻结阶段用5e-4全量微调阶段降到2e-5。这里尤其要警惕学习率过大导致的“灾难性遗忘”——我曾在一个项目里把全量微调学习率设为1e-4结果模型在训练集上AUC直接破0.99测试集上跌到0.60典型过拟合信号。正负样本比例的坑也不小。真实的活性筛选数据里正样本往往远少于负样本。有些项目为了方便直接做成1:10结果模型全学成了“预测全负样本也能拿高分”。我建议分类任务至少保持1:5以内的比例如果正样本实在太少可以尝试整体投到对比学习框架里而不是简单的二分类。3.4 评估环节不能光看AUC很多人觉得评估这一步简单跑个AUC就交差。做药物分子项目时这份想法要尽快收起来。评估模型是否真正满足研发需求至少看三个维度命中率从预测分数最高的前5%分子中有多少能被后续实验验证。这是研发团队最关心的指标——他们不想合成100个分子只中2个。多样性预测结果里有多少不同骨架类型。如果模型给出的是同一个母核的一堆类似物即便活性预测分数很高实际参考价值也有限。排序稳定性重复多次训练更换随机种子后模型对分子的相对排名是否稳定。药物研发中需要跨实验批次比较排名不稳定会直接打乱筛选节奏。上述三个指标综合起来才能完整表达“快一倍”的真正含义不是单次命中率高一点而是整个筛选流程更稳定、更省资源。4. 实测中容易踩的坑以及对应的排查方案项目做到第五周的时候我遇到了一堆在论文里根本不会写的问题。这里挑四个最典型的记录一下帮后来的人省点时间。4.1 数据泄漏模型“看着答案做题”最隐蔽的坑之一是训练集和测试集之间有数据重叠。分子领域这个问题尤其严重——同一个分子可能有不同的SMILES表示形式或者测试集里的分子恰好是训练集分子的类似物同骨架、同母核只是侧链改了一个基团。模型在预训练阶段都已经见过这些结构测试时自然“高看一眼”。排查方法也比较直接算Tanimoto相似度分子指纹之间的相似度把训练集和测试集里相似度超过0.8的样本对找出来逐一检查。我的经验是宁可把边界划到0.75把相似分子从训练集中剔除确保测试集真正陌生。4.2 对比学习分支退化成“常量输出”第二个高频坑出现在对比学习分支上。模型学到后期表征向量看起来有区分度但对比学习的loss已经不再贡献梯度信息正样本对和负样本对的表征在表征空间里几乎一样整体信息只依靠掩码重建分支在支撑。我排查这个问题时先检查了正样本构造逻辑。如果枚举SMILES时同一个分子的多个变体彼此高度相似字符重叠率超过85%对比学习的任务区分度就是不够的模型很容易走捷径。改进方案是除了SMILES枚举额外引入原子掩码和键删除等图增强操作增加正样本对的差异度让模型学到更抽象的结构不变性特征。4.3 小样本微调时的过拟合问题小样本场景下过拟合几乎必然出现。最典型的表现是训练集AUC一路狂奔到接近1.0验证集AUC在20个epoch后掉头向下。这个情况的应对我在前面的“冻结全量微调”方案里提过。我再补一个更省事的技巧只用预训练模型的表征做分类不微调任何参数。把表征作为特征输入逻辑回归或者随机森林往往能拿到70%以上的表现——作为竞赛基线这一招甚至可以打败一些精心调参的神经网络。依赖预训练特征而不微调的好处很明显整个训练时间压缩到几秒钟而且几乎不引入额外的过拟合风险适用于快速验证新靶点项目的可行性。4.4 预测分数和真实活性可能“对不上号”最后一个常见困惑模型预测的高分分子在实验验证里活性并不理想。这里面的原因很多时候不在模型本身而在于训练数据里的正负样本标签本身带有较大噪声——分子活性测试实验本身的重复性就有限有些标注为“非活性”的分子纯属实验条件导致。模型不可能在噪声标签里学到绝对真理你要做的不是无限调参而是设置一个合理的“可接受区间”。实际操作中我一般指导研发团队将预测分数前20%的分子视为“优先候选区”而不是死盯着前1%。这样做的好处是把模型排序能力发挥到极致同时给实验误差留出缓冲空间。5. 这个方案后续还能怎么迭代项目做到这个阶段当前的成果已经把常规研发流程的时间压缩了不少。但这个方向远没到天花板从我个人的判断看还有三个后续值得继续投入的方向。第一是多模态预训练。目前方案只用了一维序列和二维拓扑图后续可以把三维构象信息、分子力场特征一并放入预训练框架。融合空间信息后模型对分子“如何与靶点结合”的理解能力有望再上一个台阶。第二是统一分子生成与筛选。预训练模型生成的候选分子往往能保持较好的结构有效性与合成合理性可以将这个生成模型与活性预测模型在同一预训练框架上整合实现“先生成、再筛选、再生成”的闭环。这个方向对做生成式分子设计的人吸引力很大。第三是跨靶点迁移。预训练模型已经学会了分子结构的一般规律但对于靶点之间的共性和差异理解还比较粗浅。后续可以尝试在多个靶点数据上做多任务微调让模型在不同靶点之间迁移知识。我粗试验过一次在有类似靶点结构数据支持的情况下新靶点的初始预测精度能提升3-5个百分点。最后再分享一个小技巧不论预训练模型用了多少数据不要直接把它当成生产环境的最终模型。先在一个中等规模的开源基准数据上复现出与参考文献一致的效果再引入内部数据这条验证路径能让项目少走很多弯路。用好自监督学习从来都不是把模型做得更大更复杂而是让数据做到物尽其用让模型在有限标签下泛化得足够远。